diff --git a/zluda/src/impl/library.rs b/zluda/src/impl/library.rs index 98be33e6..f089f5a7 100644 --- a/zluda/src/impl/library.rs +++ b/zluda/src/impl/library.rs @@ -21,8 +21,7 @@ impl Library { let module_lock = self.modules.get(device).ok_or(CUerror::INVALID_DEVICE)?; *module_lock.get_or_init(|| match self.data { LibraryData::Lazy(lib) => { - let (module, sm_version, meta32) = module::load_hip_module(lib)?; - Ok(Module::new(module, sm_version, meta32).wrap()) + Ok(module::load_hip_module(lib)?.wrap()) } LibraryData::Eager(()) => Err(CUerror::NOT_SUPPORTED), }) diff --git a/zluda/src/impl/module.rs b/zluda/src/impl/module.rs index 905b7d44..16deb701 100644 --- a/zluda/src/impl/module.rs +++ b/zluda/src/impl/module.rs @@ -1,6 +1,6 @@ use super::driver; use crate::r#impl::function; -use crate::r#impl::{driver::GlobalState, function::Function}; +use crate::r#impl::function::Function; use cuda_types::{cuda::*, dark_api::FatbinFileHeader}; use dark_api::FunctionArgInfo; use hip_runtime_sys::*; @@ -24,6 +24,10 @@ pub(crate) struct Module { struct ModuleMutable { functions: FxHashMap>, + // d4r: the PTX of a module loaded without compiling (native_placeholder), and the module compiled from + // it once a function without a native replacement is requested + deferred: Option, + compiled: Option, } impl ModuleMutable { @@ -40,6 +44,12 @@ impl ModuleMutable { let (func_handle, prep, block_z, grid_x, prep_key, prep_slots, post) = match native_override(entry.key()) { Some(native) => native, None => { + if let Some(text) = self.deferred.as_deref() { + let (compiled, _, _) = compile_ptx(Some(text), ptx_parser::parse_module_unchecked(text))?; + self.compiled = Some(compiled); + self.deferred = None; + } + let module = self.compiled.unwrap_or(module); let mut func_handle = unsafe { std::mem::zeroed() }; unsafe { hipModuleGetFunction(&mut func_handle, module, entry.key().as_ptr()) }?; (func_handle, None, 0, 0, None, 0, Vec::new()) @@ -64,6 +74,57 @@ impl ModuleMutable { } } +fn read_native_image(path: &std::path::Path) -> Option<&'static [u8]> { + let image = fs::read(path).ok()?; + // HIP can crash on a short or truncated code object: an ELF's section headers (e_shoff, e_shentsize, e_shnum) must fit + let le = |at: usize, size: usize| image[at..at + size].iter().rev().fold(0u64, |v, &b| (v << 8) | u64::from(b)); + if image.len() < 64 + || (image.starts_with(b"\x7fELF") && le(40, 8).saturating_add(le(58, 2) * le(60, 2)) > image.len() as u64) + { + eprintln!("[d4r-native] {} is truncated ({} bytes)", path.display(), image.len()); + return None; + } + // HIP may reference the image after loading; the override lives for the whole process. + Some(Box::leak(image.into_boxed_slice())) +} + +/// d4r: a module whose every visible kernel has a native replacement (and that has no globals to look up) +/// is not compiled when it is loaded: its kernels are served by native_override, and the PTX is compiled +/// only if another function is requested. On RDNA2 the DLSS transformer layers each took minutes to +/// compile (emulated MMA) for code that never ran. The returned module, the first kernel's native code +/// object, stands in for the compiled one (unload, lookups). +fn native_placeholder(ast: &ptx_parser::Module) -> Option { + let dir = std::env::var_os("D4R_ZLUDA_NATIVE_DIR")?; + if ast.invalid_directives != 0 || ast.address_size != 64 { + return None; + } + let mut first = None; + for directive in &ast.directives { + match directive { + ptx_parser::Directive::Variable(_, variable) => { + if variable.info.state_space != ptx_parser::StateSpace::Shared { + return None; + } + } + ptx_parser::Directive::Method(linking, function) => { + let name = &function.func_directive.name; + if !name.is_kernel() || !linking.contains(ptx_parser::LinkingDirective::VISIBLE) { + continue; + } + let path = std::path::Path::new(&dir).join(format!("{}.hsaco", name.text())); + if !path.is_file() { + return None; + } + first.get_or_insert(path); + } + } + } + let image = read_native_image(&first?)?; + let mut module = unsafe { mem::zeroed() }; + unsafe { hipModuleLoadData(&mut module, image.as_ptr().cast()) }.ok()?; + Some(module) +} + /// d4r: hand-written native replacement for a PTX kernel. With D4R_ZLUDA_NATIVE_DIR set, a kernel /// NAME is served from DIR/NAME.hsaco when that file exists (same name, launch shape and parameter /// block as the PTX kernel). An optional NAME_prep kernel in the same code object runs first on the @@ -76,17 +137,7 @@ fn native_override( ) -> Option<(hipFunction_t, Option<(hipFunction_t, u32)>, u32, u32, Option, u32, Vec<(hipFunction_t, u32, u32, u32, u32, u32)>)> { let dir = std::env::var_os("D4R_ZLUDA_NATIVE_DIR")?; let path = std::path::Path::new(&dir).join(format!("{}.hsaco", name.to_str().ok()?)); - let image = fs::read(&path).ok()?; - // HIP can crash on a short or truncated code object: an ELF's section headers (e_shoff, e_shentsize, e_shnum) must fit - let le = |at: usize, size: usize| image[at..at + size].iter().rev().fold(0u64, |v, &b| (v << 8) | u64::from(b)); - if image.len() < 64 - || (image.starts_with(b"\x7fELF") && le(40, 8).saturating_add(le(58, 2) * le(60, 2)) > image.len() as u64) - { - eprintln!("[d4r-native] {} is truncated ({} bytes)", path.display(), image.len()); - return None; - } - // HIP may reference the image after loading; the override lives for the whole process. - let image: &'static [u8] = Box::leak(image.into_boxed_slice()); + let image = read_native_image(&path)?; let mut native_module = unsafe { mem::zeroed() }; if unsafe { hipModuleLoadData(&mut native_module, image.as_ptr().cast()) }.is_err() { eprintln!("[d4r-native] failed to load {}", path.display()); @@ -176,6 +227,9 @@ impl ZludaObject for Module { type CudaHandle = CUmodule; fn drop_checked(&mut self) -> CUresult { + if let Some(compiled) = self.mutable.get_mut().ok().and_then(|m| m.compiled.take()) { + unsafe { hipModuleUnload(compiled) }?; + } unsafe { hipModuleUnload(self.base) }?; Ok(()) } @@ -263,6 +317,8 @@ impl Module { sm_version, mutable: Mutex::new(ModuleMutable { functions: FxHashMap::default(), + deferred: None, + compiled: None, }), bit32, } @@ -281,10 +337,7 @@ impl Module { } } -fn get_best_ptx_and_compile( - global_state: &GlobalState, - image: CodeLibraryRef<'_>, -) -> Result<(hipModule_t, u32, Option), CUerror> { +fn get_best_ptx_and_compile(image: CodeLibraryRef<'_>) -> Result { let mut ptx_modules = Vec::new(); unsafe { CodeLibraryRef::iterate_modules(image, |_, module| match module { @@ -373,6 +426,20 @@ fn get_best_ptx_and_compile( ) } }; + if let (Some(text), Some(placeholder)) = (text, native_placeholder(&module)) { + let loaded = Module::new(placeholder, module.sm_version, None); + loaded.mutable.lock().map_err(|_| CUerror::UNKNOWN)?.deferred = Some(text.to_string()); + return Ok(loaded); + } + let (hip_module, sm_version, zluda32) = compile_ptx(text.map(|text| &**text), module)?; + Ok(Module::new(hip_module, sm_version, zluda32)) +} + +fn compile_ptx( + text: Option<&str>, + module: ptx_parser::Module<'_>, +) -> Result<(hipModule_t, u32, Option), CUerror> { + let global_state = driver::global_state()?; // TODO: get this information on initialization let hip_properties = get_hip_properties()?; let gcn_arch = get_gcn_arch(&hip_properties)?; @@ -389,7 +456,7 @@ fn get_best_ptx_and_compile( let mut cache_with_key = match (text, global_state.cache_path.as_ref()) { (Some(text), Some(p)) => (|| { let cache = zluda_cache::ModuleCache::open(p)?; - let key = get_cache_key(gcn_arch, &text, &attributes)?; + let key = get_cache_key(gcn_arch, text, &attributes)?; Some((cache, key)) })(), _ => None, @@ -410,11 +477,8 @@ fn cow_bytes_to_str<'a>(data: Cow<'a, [u8]>) -> Option> { } } -pub(crate) fn load_hip_module( - library: CodeLibraryRef, -) -> Result<(hipModule_t, u32, Option), CUerror> { - let global_state = driver::global_state()?; - get_best_ptx_and_compile(global_state, library) +pub(crate) fn load_hip_module(library: CodeLibraryRef) -> Result { + get_best_ptx_and_compile(library) } #[derive(serde::Serialize)] @@ -605,16 +669,14 @@ pub(crate) fn load(module: &mut CUmodule, fname: &CStr) -> CUresult { image.push(0); let library = unsafe { CodeLibraryRef::try_load(image.as_ptr() as *const std::ffi::c_void) } .map_err(|_| CUerror::NO_BINARY_FOR_GPU)?; - let (hip_module, sm_version, meta32) = load_hip_module(library)?; - *module = Module::new(hip_module, sm_version, meta32).wrap(); + *module = load_hip_module(library)?.wrap(); Ok(()) } pub(crate) fn load_data(module: &mut CUmodule, image: &std::ffi::c_void) -> CUresult { let library = unsafe { CodeLibraryRef::try_load(image) }.map_err(|_| CUerror::NO_BINARY_FOR_GPU)?; - let (hip_module, sm_version, meta32) = load_hip_module(library)?; - *module = Module::new(hip_module, sm_version, meta32).wrap(); + *module = load_hip_module(library)?.wrap(); Ok(()) }