Unified fatbin versions behind a single iterator. (#398)

dark_api/src/fatbin.rs

@@ -75,21 +75,24 @@ impl<'a> Fatbin<'a> {
         Ok(Fatbin { wrapper })
     }
 
-    pub fn get_first(&self) -> Result<FatbinSubmodule, FatbinError> {
+    pub fn get_submodules(&self) -> Result<FatbinIter<'a>, FatbinError> {
-        let header: &FatbinHeader =
+        match self.wrapper.version {
-            parse_fatbin_header(&self.wrapper.data).map_err(|e| FatbinError::ParseFailure(e))?;
+            FatbincWrapper::VERSION_V2 => 
-        Ok(FatbinSubmodule::new(header))
+                Ok(FatbinIter::V2(FatbinSubmoduleIterator {
-    }
-
-    pub fn get_submodules(&self) -> Option<FatbinSubmoduleIterator> {
-        let is_version_2 = self.wrapper.version == FatbincWrapper::VERSION_V2;
-        if !is_version_2 {
-            return None;
-        }
-
-        Some(FatbinSubmoduleIterator {
                     fatbins: self.wrapper.filename_or_fatbins as *const *const std::ffi::c_void,
-        })
+                    _phantom: std::marker::PhantomData,
+                })),
+            FatbincWrapper::VERSION_V1 => {
+                let header = parse_fatbin_header(&self.wrapper.data)
+                    .map_err(FatbinError::ParseFailure)?;
+                Ok(FatbinIter::V1(Some(FatbinSubmodule::new(header))))
+            }
+            version => Err(FatbinError::ParseFailure(ParseError::UnexpectedBinaryField{
+                field_name: "FATBINC_VERSION",
+                observed: version,
+                expected: [FatbincWrapper::VERSION_V1, FatbincWrapper::VERSION_V2].into(),
+            })),
+        }
     }
 }
 
@@ -107,12 +110,27 @@ impl<'a> FatbinSubmodule<'a> {
     }
 }
 
-pub struct FatbinSubmoduleIterator {
+pub enum FatbinIter<'a> {
-    fatbins: *const *const std::ffi::c_void,
+    V1(Option<FatbinSubmodule<'a>>),
+    V2(FatbinSubmoduleIterator<'a>),
 }
 
-impl FatbinSubmoduleIterator {
+impl<'a> FatbinIter<'a> {
-    pub unsafe fn next(&mut self) -> Result<Option<FatbinSubmodule>, ParseError> {
+    pub fn next(&mut self) -> Result<Option<FatbinSubmodule<'a>>, ParseError> {
+        match self {
+            FatbinIter::V1(opt) => Ok(opt.take()),
+            FatbinIter::V2(iter) => unsafe { iter.next() },
+        }
+    }
+}
+
+pub struct FatbinSubmoduleIterator<'a> {
+    fatbins: *const *const std::ffi::c_void,
+    _phantom: std::marker::PhantomData<&'a FatbinHeader>,
+}
+
+impl<'a> FatbinSubmoduleIterator<'a> {
+    pub unsafe fn next(&mut self) -> Result<Option<FatbinSubmodule<'a>>, ParseError> {
         if *self.fatbins != ptr::null() {
             let header = *self.fatbins as *const FatbinHeader;
             self.fatbins = self.fatbins.add(1);

zluda/src/impl/module.rs

@@ -24,16 +24,17 @@ impl ZludaObject for Module {
 
 fn get_ptx_from_wrapped_fatbin(image: *const ::core::ffi::c_void) -> Result<Vec<u8>, CUerror> {
     let fatbin = Fatbin::new(&image).map_err(|_| CUerror::UNKNOWN)?;
-    let first = fatbin.get_first().map_err(|_| CUerror::UNKNOWN)?;
+    let mut submodules = fatbin.get_submodules().map_err(|_| CUerror::UNKNOWN)?;
-    let mut files = first.get_files();
 
+    while let Some(current) = unsafe { submodules.next().map_err(|_| CUerror::UNKNOWN)? } { 
+        let mut files = current.get_files();
         while let Some(file) = unsafe { files.next().map_err(|_| CUerror::UNKNOWN)? } {
-        // Eventually we will want to get the PTX for the highest hardware version that we can get to compile. But for now we just get the first PTX we can find.
             if file.header.kind == FatbinFileHeader::HEADER_KIND_PTX {
                 let decompressed = unsafe { file.decompress() }.map_err(|_| CUerror::UNKNOWN)?;
                 return Ok(decompressed);
             }
         }
+    }
 
     Err(CUerror::NO_BINARY_FOR_GPU)
 }

zluda_dump/src/trace.rs

@@ -262,13 +262,10 @@ pub(crate) unsafe fn record_submodules_from_wrapped_fatbin(
     state: &mut StateTracker,
 ) -> Result<(), ErrorEntry> {
     let fatbin = Fatbin::new(&fatbinc_wrapper).map_err(ErrorEntry::from)?;
-    let first = fatbin.get_first().map_err(ErrorEntry::from)?;
+    let mut submodules = fatbin.get_submodules()?;
-    record_submodules_from_fatbin(module, first, fn_logger, state)?;
-    if let Some(mut submodules) = fatbin.get_submodules() {
     while let Some(current) = submodules.next()? {
         record_submodules_from_fatbin(module, current, fn_logger, state)?;
     }
-    }
     Ok(())
 }