diff options
| -rw-r--r-- | zluda_dump/src/cuda.rs | 12 | ||||
| -rw-r--r-- | zluda_dump/src/lib.rs | 187 | ||||
| -rw-r--r-- | zluda_dump/src/os_unix.rs | 4 | ||||
| -rw-r--r-- | zluda_dump/src/os_win.rs | 2 |
4 files changed, 161 insertions, 44 deletions
diff --git a/zluda_dump/src/cuda.rs b/zluda_dump/src/cuda.rs index 50082d1..35ec52c 100644 --- a/zluda_dump/src/cuda.rs +++ b/zluda_dump/src/cuda.rs @@ -2574,7 +2574,7 @@ extern_redirect! { stateOut: *mut CUlinkState, ) -> CUresult; } -extern_redirect! { +extern_redirect_with! { pub fn cuLinkAddData( state: CUlinkState, type_: CUjitInputType, @@ -2585,8 +2585,9 @@ extern_redirect! { options: *mut CUjit_option, optionValues: *mut *mut ::std::os::raw::c_void, ) -> CUresult; + super::cuLinkAddData; } -extern_redirect! { +extern_redirect_with! { pub fn cuLinkAddData_v2( state: CUlinkState, type_: CUjitInputType, @@ -2597,8 +2598,9 @@ extern_redirect! { options: *mut CUjit_option, optionValues: *mut *mut ::std::os::raw::c_void, ) -> CUresult; + super::cuLinkAddData; } -extern_redirect! { +extern_redirect_with! { pub fn cuLinkAddFile( state: CUlinkState, type_: CUjitInputType, @@ -2607,8 +2609,9 @@ extern_redirect! { options: *mut CUjit_option, optionValues: *mut *mut ::std::os::raw::c_void, ) -> CUresult; + super::cuLinkAddFile; } -extern_redirect! { +extern_redirect_with! { pub fn cuLinkAddFile_v2( state: CUlinkState, type_: CUjitInputType, @@ -2617,6 +2620,7 @@ extern_redirect! { options: *mut CUjit_option, optionValues: *mut *mut ::std::os::raw::c_void, ) -> CUresult; + super::cuLinkAddFile; } extern_redirect! { pub fn cuLinkComplete( diff --git a/zluda_dump/src/lib.rs b/zluda_dump/src/lib.rs index e334151..6425768 100644 --- a/zluda_dump/src/lib.rs +++ b/zluda_dump/src/lib.rs @@ -1,5 +1,5 @@ use std::{ - collections::{BTreeMap, HashMap}, + collections::{hash_map, BTreeMap, HashMap}, env, error::Error, ffi::{c_void, CStr}, @@ -14,8 +14,8 @@ use std::{ use std::{fs::File, ptr}; use cuda::{ - CUdevice, CUdevice_attribute, CUdeviceptr, CUfunction, CUjit_option, CUmodule, CUresult, - CUstream, CUuuid, + CUdevice, CUdevice_attribute, CUdeviceptr, CUfunction, CUjitInputType, CUjit_option, + CUlinkState, CUmodule, CUresult, CUstream, CUuuid, }; use ptx::ast; use regex::Regex; @@ -70,7 +70,10 @@ macro_rules! extern_redirect_with { mod cuda; pub static mut LIBCUDA_HANDLE: *mut c_void = ptr::null_mut(); +pub static mut PENDING_LINKING: Option<HashMap<CUlinkState, Vec<ModuleDump>>> = None; +pub static mut LINKED_CUBINS: Option<HashMap<*mut c_void, ModuleDump>> = None; pub static mut MODULES: Option<HashMap<CUmodule, ModuleDump>> = None; +pub static mut MODULE_DUMP_COUNTER: usize = 0; pub static mut KERNELS: Option<HashMap<CUfunction, KernelDump>> = None; static mut BUFFERS: Option<BTreeMap<usize, (usize, AllocLocation)>> = None; pub static mut LAUNCH_COUNTER: usize = 0; @@ -169,50 +172,100 @@ unsafe fn record_module_image_raw(module: CUmodule, raw_image: *const ::std::os: let image = to_str(raw_image); match image { None => os_log!("Malformed module image: {:?}", raw_image), - Some(image) => record_module_image(module, image), + Some(image) => record_module_image_with_module(module, raw_image, image), }; } -unsafe fn record_module_image(module: CUmodule, image: &str) { +unsafe fn record_module_image_with_module( + module: CUmodule, + raw_image: *const ::std::os::raw::c_void, + image: &str, +) { + match record_module_image_impl(raw_image, image) { + Ok(dump) => { + MODULES + .get_or_insert_with(|| HashMap::new()) + .insert(module, dump); + } + Err(e) => { + os_log!("{}", e); + } + } +} + +unsafe fn record_module_image_with_linker( + link_obj: CUlinkState, + raw_image: *const ::std::os::raw::c_void, + image: &str, +) { + match record_module_image_impl(raw_image, image) { + Ok(dump) => { + match PENDING_LINKING + .get_or_insert_with(|| HashMap::new()) + .entry(link_obj) + { + hash_map::Entry::Occupied(mut vec) => { + vec.get_mut().push(dump); + } + hash_map::Entry::Vacant(e) => { + e.insert(vec![dump]); + } + }; + } + Err(e) => { + os_log!("{}", e); + } + } +} + +unsafe fn record_module_image_impl( + raw_image: *const ::std::os::raw::c_void, + image: &str, +) -> Result<ModuleDump, Box<dyn Error>> { if !image.contains(&".version") { - os_log!("Malformed module image: {:?}", module); - } else { - let mut errors = Vec::new(); - let ast = ptx::ModuleParser::new().parse(&mut errors, image); - let kernels_args = match (&*errors, ast) { - (&[], Ok(ast)) => { - let kernels_args = ast - .directives - .iter() - .filter_map(directive_to_kernel) - .collect::<HashMap<_, _>>(); - Some(kernels_args) - } - (_, _) => { - // Don't print errors - it's usually too verbose to be useful - os_log!("Errors when parsing module: {:?}", module); - None - } - }; - MODULES.as_mut().unwrap().insert( - module, - ModuleDump { - content: Rc::new(image.to_string()), - kernels_args, - }, - ); + return Err(format!( + "Malformed module image (no `.version`): {:?}", + raw_image + ))?; } + let mut errors = Vec::new(); + let ast = ptx::ModuleParser::new().parse(&mut errors, image); + let kernels_args = match (&*errors, ast) { + (&[], Ok(ast)) => { + let kernels_args = ast + .directives + .iter() + .filter_map(directive_to_kernel) + .collect::<HashMap<_, _>>(); + Some(kernels_args) + } + (err_vec, res) => { + // Don't print errors - it's usually too verbose to be useful + os_log!( + "{} errors when parsing module image: {:?}", + err_vec.len() + res.iter().len(), + raw_image + ); + None + } + }; + let dump = ModuleDump { + content: Rc::new(image.to_string()), + kernels_args, + }; if let Err(e) = try_dump_module_image(image) { - os_log!("Errors when saving module: {:?}, {}", module, e); + return Err(format!( + "Errors when saving module image: {:?}, {}", + raw_image, e + ))?; } + Ok(dump) } unsafe fn try_dump_module_image(image: &str) -> Result<(), Box<dyn Error>> { let mut dump_path = get_dump_dir()?; - dump_path.push(format!( - "module_{:04}.ptx", - MODULES.as_ref().unwrap().len() - 1 - )); + dump_path.push(format!("module_{:04}.ptx", MODULE_DUMP_COUNTER)); + MODULE_DUMP_COUNTER += 1; let mut file = File::create(dump_path)?; file.write_all(image.as_bytes())?; Ok(()) @@ -952,7 +1005,9 @@ unsafe fn get_module_from_cubin_unwrapped( if let Some(text) = maybe_kernel_text { match CStr::from_bytes_with_nul(&text) { Ok(cstr) => match cstr.to_str() { - Ok(utf8_str) => record_module_image(*module, utf8_str), + Ok(utf8_str) => { + record_module_image_with_module(*module, text.as_ptr() as _, utf8_str) + } Err(_) => {} }, Err(_) => {} @@ -1056,3 +1111,61 @@ pub unsafe fn cuDeviceGetAttribute( } cont(pi, attrib, dev) } + +#[allow(non_snake_case)] +pub unsafe fn cuLinkAddData( + state: CUlinkState, + type_: CUjitInputType, + data: *mut ::std::os::raw::c_void, + size: usize, + name: *const ::std::os::raw::c_char, + numOptions: ::std::os::raw::c_uint, + options: *mut CUjit_option, + optionValues: *mut *mut ::std::os::raw::c_void, + cont: impl FnOnce( + CUlinkState, + CUjitInputType, + *mut ::std::os::raw::c_void, + usize, + *const ::std::os::raw::c_char, + ::std::os::raw::c_uint, + *mut CUjit_option, + *mut *mut ::std::os::raw::c_void, + ) -> CUresult, +) -> CUresult { + if let Some(image) = to_str(data) { + record_module_image_with_linker(state, data, image) + } else { + os_log!("PTX module not a string: {:?}", data); + } + cont( + state, + type_, + data, + size, + name, + numOptions, + options, + optionValues, + ) +} + +#[allow(non_snake_case)] +pub unsafe fn cuLinkAddFile( + state: CUlinkState, + type_: CUjitInputType, + path: *const ::std::os::raw::c_char, + numOptions: ::std::os::raw::c_uint, + options: *mut CUjit_option, + optionValues: *mut *mut ::std::os::raw::c_void, + cont: impl FnOnce( + CUlinkState, + CUjitInputType, + *const ::std::os::raw::c_char, + ::std::os::raw::c_uint, + *mut CUjit_option, + *mut *mut ::std::os::raw::c_void, + ) -> CUresult, +) -> CUresult { + cont(state, type_, path, numOptions, options, optionValues) +} diff --git a/zluda_dump/src/os_unix.rs b/zluda_dump/src/os_unix.rs index 49e1825..0d9db04 100644 --- a/zluda_dump/src/os_unix.rs +++ b/zluda_dump/src/os_unix.rs @@ -19,12 +19,12 @@ pub unsafe fn get_proc_address(handle: *mut c_void, func: &CStr) -> *mut c_void macro_rules! os_log {
($format:tt) => {
{
- eprintln!($format);
+ eprintln!("[ZLUDA_DUMP] {}", format!($format));
}
};
($format:tt, $($obj: expr),+) => {
{
- eprintln!($format, $($obj,)+);
+ eprintln!("[ZLUDA_DUMP] {}", format!($format, $($obj,)+));
}
};
}
diff --git a/zluda_dump/src/os_win.rs b/zluda_dump/src/os_win.rs index 0cd8f3d..e99b653 100644 --- a/zluda_dump/src/os_win.rs +++ b/zluda_dump/src/os_win.rs @@ -90,7 +90,7 @@ macro_rules! os_log { pub fn __log_impl(s: String) {
let log_to_stderr = std::io::stderr().as_raw_handle() != ptr::null_mut();
if log_to_stderr {
- eprintln!("[ZLUDA_DUMP] {}\n", s);
+ eprintln!("[ZLUDA_DUMP] {}", s);
} else {
let mut win_str = String::with_capacity("[ZLUDA_DUMP] ".len() + s.len() + 2);
win_str.push_str("[ZLUDA_DUMP] ");
|