zluda: skip compiling modules whose kernels are all native (0014)

DLSS loads every module of a model at feature creation and ZLUDA compiled
each one, even when the bridge had already linked a native kernel for every
visible entry and the compiled code never ran. On RDNA2 the transformer
layers emulate MMA and each took minutes in LLVM's register allocator, so
install.sh's K test frame looked hung for about half an hour on a Steam Deck.

Such a module (no globals, 64-bit) now loads the first native code object as
a stand-in and keeps its PTX; it is compiled only if a function without a
native replacement is requested. The K frame takes ~40 s on the Deck.

patches/zluda/0014-skip-compiling-natively-served-modules.patch

@@ -0,0 +1,221 @@
+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<CString, Box<Function>>,
++    // 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<String>,
++    compiled: Option<hipModule_t>,
+ }
+ 
+ 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<hipModule_t> {
++    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>, 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<Metadata32Bit>), CUerror> {
++fn get_best_ptx_and_compile(image: CodeLibraryRef<'_>) -> Result<Module, CUerror> {
+     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<Metadata32Bit>), 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<Cow<'a, str>> {
+     }
+ }
+ 
+-pub(crate) fn load_hip_module(
+-    library: CodeLibraryRef,
+-) -> Result<(hipModule_t, u32, Option<Metadata32Bit>), CUerror> {
+-    let global_state = driver::global_state()?;
+-    get_best_ptx_and_compile(global_state, library)
++pub(crate) fn load_hip_module(library: CodeLibraryRef) -> Result<Module, CUerror> {
++    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(())
+ }
+