From 6be29ac9c1299edc09851f83e34e1c9fdf76c1bf Mon Sep 17 00:00:00 2001 From: Andrzej Janik Date: Fri, 5 Sep 2025 17:21:35 +0000 Subject: [PATCH] Add common CUDA binary visitor --- zluda_common/Cargo.toml | 1 + zluda_common/src/lib.rs | 224 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 223 insertions(+), 2 deletions(-) diff --git a/zluda_common/Cargo.toml b/zluda_common/Cargo.toml index 4c528e5..ca70ab8 100644 --- a/zluda_common/Cargo.toml +++ b/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" } diff --git a/zluda_common/src/lib.rs b/zluda_common/src/lib.rs index 94c795b..c611f9f 100644 --- a/zluda_common/src/lib.rs +++ b/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(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 { + 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 { + 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), + ) { + 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> { + 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, 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), + }); + } + } +} +*/