wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
linux
d4r/patches/zluda/0014-skip-compiling-natively-served-modules.patch
T
woutermans 4064117244 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.
2026-10-18 23:03:58 +02:00

222 lines
10 KiB
Diff
Raw Permalink Blame History

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(())
}
Reference in new issue View Git Blame Copy Permalink