Dump modules passed for linking

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] ");