Add common CUDA binary visitor

zluda_common/Cargo.toml

@@ -8,3 +8,4 @@ edition = "2021"
 cuda_types = { path = "../cuda_types" }
 hip_runtime-sys = { path = "../ext/hip_runtime-sys" }
 rocblas-sys = { path = "../ext/rocblas-sys" }
+dark_api = { path = "../dark_api" }

zluda_common/src/lib.rs

@@ -1,9 +1,20 @@
-use cuda_types::{cublas::*, cublaslt::cublasLtHandle_t, cuda::*, nvml::*};
+use cuda_types::{
+    cublas::*,
+    cublaslt::cublasLtHandle_t,
+    cuda::*,
+    dark_api::{FatbinFileHeaderFlags, FatbinHeader, FatbincWrapper},
+    nvml::*,
+};
+use dark_api::fatbin::{
+    Fatbin, FatbinError, FatbinFile, FatbinFileIterator, FatbinIter, FatbinSubmodule, ParseError,
+};
 use hip_runtime_sys::*;
 use rocblas_sys::*;
 use std::{
-    ffi::CStr,
+    borrow::Cow,
+    ffi::{c_void, CStr},
     mem::{self, ManuallyDrop, MaybeUninit},
+    ops::ControlFlow,
     ptr,
 };
 
@@ -412,3 +423,212 @@ pub fn drop_checked<T: ZludaObject>(handle: T::CudaHandle) -> Result<(), T::Erro
     unsafe { ManuallyDrop::drop(&mut wrapped_object) };
     underlying_error
 }
+
+/*
+pub struct CodeModuleRef<'a> {
+    pub kind: CodeModuleKind,
+    pub data: &'a [u8],
+}
+
+impl<'a> CodeModule<'a> {
+    /// Interprets `data` as a code module of some kind.
+    ///
+    /// This does not validate the contents of `data`, it only looks at the headers to determine
+    /// what kind of data it is.
+    pub fn parse(data: *mut c_void) -> Result<Self, CUerror> {
+        if data.len() >= 4 {
+            let kind = match &data[0..4] {
+                FatbincWrapper::MAGIC => CodeModuleKind::FatbincWrapper,
+                FatbinHeader::MAGIC => CodeModuleKind::FatbinHeader,
+                elf64::header::ELFMAG => CodeModuleKind::Elf,
+                _ => {
+                    if data.ends_with(&[0]) && data.iter().all(|&c| c != 0) {
+                        CodeModuleKind::Ptx
+                    } else {
+                        CodeModuleKind::ForeignElf
+                    }
+                }
+            };
+            Ok(CodeModule { kind, data })
+        } else {
+            Err(CUerror::INVALID_VALUE)
+        }
+    }
+}
+     */
+
+// We receive module as an opaque pointer. We want to handle three different
+// lifetime-related scenarios:
+// * The module has a 'static lifetime, but we don't want to use it just yet
+//   (`cuLibraryLoadData` with CU_LIBRARY_BINARY_IS_PRESERVED = 1), we might
+//   never use it.
+//       In this case we just keep the void pointer, we can pass it to
+//       the consuming function later
+// * The module has a non-'static lifetime, and we will use it in the future
+//   (`cuLibraryLoadData` with CU_LIBRARY_BINARY_IS_PRESERVED = 0)
+//       In this case we need to copy the data into its own buffers
+// * The module lifetime is scoped to the current function. E.g. zluda_trace
+//   might to parse a module to inspect and save it or it's cuModuleLoadData
+//        In this case we need to return either the compatible ELF or the
+//        iterator over `Cow` with decompressed PTX strings
+//            Even here there are two cases:
+//            * The consumer is cuModuleLoadData, if it's our ELF then it wants
+//              to load it directly from the pointer
+//            * The consumer is zluda_trace, it wants to compute the length of
+//              the ELF and save it to a file
+pub enum CodeLibaryRef<'a> {
+    FatbincWrapper(Fatbin<'a>),
+    FatbinHeader(FatbinSubmodule<'a>),
+    Text(&'a str),
+    Elf(*mut c_void),
+}
+
+impl<'a> CodeLibaryRef<'a> {
+    const ELFMAG: [u8; 4] = *b"\x7FELF";
+
+    unsafe fn try_load(ptr: *mut c_void) -> Option<Self> {
+        Some(match *ptr.cast::<[u8; 4]>() {
+            FatbincWrapper::MAGIC => Self::FatbincWrapper(Fatbin {
+                wrapper: &*(ptr.cast()),
+            }),
+            FatbinHeader::MAGIC => Self::FatbinHeader(FatbinSubmodule {
+                header: &*(ptr.cast()),
+            }),
+            Self::ELFMAG => Self::Elf(ptr),
+            _ => CStr::from_ptr(ptr.cast()).to_str().ok().map(Self::Text)?,
+        })
+    }
+
+    unsafe fn iterate_modules(
+        &self,
+        fn_: &mut impl FnMut((usize, usize), Result<CodeModule, FatbinError>),
+    ) {
+        match self {
+            CodeLibaryRef::FatbincWrapper(fatbin) => {
+                let module_iter = fatbin.get_submodules();
+                match module_iter {
+                    Ok(mut iter) => {
+                        let mut module_index = 0;
+                        while let Some(maybe_submodule) = iter.next() {
+                            match maybe_submodule {
+                                Ok(submodule) => Self::iterate_modules(
+                                    &CodeLibaryRef::FatbinHeader(submodule),
+                                    &mut |(_, subindex), module| {
+                                        fn_((module_index, subindex), module)
+                                    },
+                                ),
+                                Err(err) => {
+                                    fn_((module_index, 0), Err(FatbinError::ParseFailure(err)))
+                                }
+                            }
+                            module_index += 1;
+                        }
+                    }
+                    Err(err) => fn_((0, 0), Err(err)),
+                }
+            }
+            CodeLibaryRef::FatbinHeader(submodule) => {
+                let mut iter = submodule.get_files();
+                let mut index = 0;
+                while let Some(file) = iter.next() {
+                    fn_(
+                        (0, index),
+                        file.map(CodeModule::File)
+                            .map_err(FatbinError::ParseFailure),
+                    );
+                    index += 1;
+                }
+            }
+            CodeLibaryRef::Text(text) => fn_((0, 0), Ok(CodeModule::Text(*text))),
+            CodeLibaryRef::Elf(elf) => fn_((0, 0), Ok(CodeModule::Elf(*elf))),
+        }
+    }
+}
+
+enum CodeModule<'a> {
+    File(FatbinFile<'a>),
+    Text(&'a str),
+    Elf(*mut c_void),
+}
+
+/*
+enum PtxIterator<'a> {
+    FatbincWrapper(FatbinIter<'a>),
+    FatbinHeader(FatbinFileIterator<'a>),
+    Text(std::iter::Once<&'a str>),
+}
+
+impl<'a> PtxIterator<'a> {
+    fn next(&mut self) -> Option<Result<&'a str, ParseError>> {
+        match self {
+            PtxIterator::FatbincWrapper(iter) => {
+                while let Ok(Some(submodule)) = iter.next() {
+                    let mut files = submodule.get_files();
+                    while let Some(file) = unsafe { files.next().ok()? } {
+                        if file.header.kind == FatbinFileHeader::HEADER_KIND_PTX {
+                            return Some(
+                                unsafe { file.decompress().ok()? }
+                                    .as_slice()
+                                    .strip_suffix(&[0])?
+                                    .as_ref()
+                                    .and_then(|s| std::str::from_utf8(s).ok()),
+                            );
+                        }
+                    }
+                }
+                None
+            }
+            PtxIterator::FatbinHeader(iter) => unsafe {
+                iter.next().map(|file| file.map(|file| file.decompress()))
+            },
+            PtxIterator::Text(iter) => iter.next(),
+        }
+    }
+}
+
+fn decompress_payload<'a>(file: &'a FatbinFile) -> Result<Cow<'a, str>, FatbinError> {
+    let mut payload = if file
+        .header
+        .flags
+        .contains(FatbinFileHeaderFlags::CompressedLz4)
+    {
+        Cow::Owned(unwrap_some_or!(
+            fn_logger.try_return(|| decompress_lz4(&file).map_err(|e| e.into())),
+            continue
+        ))
+    } else if file
+        .header
+        .flags
+        .contains(FatbinFileHeaderFlags::CompressedZstd)
+    {
+        Cow::Owned(unwrap_some_or!(
+            fn_logger.try_return(|| decompress_zstd(&file).map_err(|e| e.into())),
+            continue
+        ))
+    } else {
+        Cow::Borrowed(file.get_payload())
+    };
+    match file.header.kind {
+        FatbinFileHeader::HEADER_KIND_PTX => {
+            while payload.last() == Some(&0) {
+                // remove trailing zeros
+                payload.to_mut().pop();
+            }
+            state.record_new_submodule(module, &*payload, fn_logger, "ptx")
+        }
+        FatbinFileHeader::HEADER_KIND_ELF => {
+            state.record_new_submodule(module, &*payload, fn_logger, "elf")
+        }
+        _ => {
+            fn_logger.log(log::ErrorEntry::UnexpectedBinaryField {
+                field_name: "FATBIN_FILE_HEADER_KIND",
+                expected: vec![
+                    UInt::U16(FatbinFileHeader::HEADER_KIND_PTX),
+                    UInt::U16(FatbinFileHeader::HEADER_KIND_ELF),
+                ],
+                observed: UInt::U16(file.header.kind),
+            });
+        }
+    }
+}
+*/