Implement static typing for dynamically-loaded CUDA DLLs

1 files changed, 87 insertions, 48 deletions

diff --git a/zluda_dump/src/lib.rs b/zluda_dump/src/lib.rs
index d79c391..04fc36e 100644
--- a/zluda_dump/src/lib.rs
+++ b/zluda_dump/src/lib.rs
@@ -2,6 +2,7 @@ use cuda_types::{
     CUdevice, CUdevice_attribute, CUfunction, CUjit_option, CUmodule, CUresult, CUuuid,
 };
 use paste::paste;
+use side_by_side::CudaDynamicFns;
 use std::io;
 use std::{
     collections::HashMap, env, error::Error, ffi::c_void, fs, path::PathBuf, ptr::NonNull, rc::Rc,
@@ -10,47 +11,50 @@ use std::{
 
 #[macro_use]
 extern crate lazy_static;
+extern crate cuda_types;
 
 macro_rules! extern_redirect {
-    ($abi:literal fn $fn_name:ident( $($arg_id:ident : $arg_type:ty),* ) -> $ret_type:path) => {
-        #[no_mangle]
-        pub extern $abi fn $fn_name ( $( $arg_id : $arg_type),* ) -> $ret_type {
-            let original_fn = |fn_ptr| {
-                let typed_fn = unsafe { std::mem::transmute::<_, extern "system" fn( $( $arg_id : $arg_type),* ) -> $ret_type>(fn_ptr) };
-                typed_fn($( $arg_id ),*)
-            };
-            let get_formatted_args = Box::new(move |writer: &mut dyn std::io::Write| {
-                (paste! { format :: [<write_ $fn_name>] }) (
-                    writer
-                    $(,$arg_id)*
-                )
-            });
-            crate::handle_cuda_function_call(stringify!($fn_name), original_fn, get_formatted_args)
-        }
+    ($($abi:literal fn $fn_name:ident( $($arg_id:ident : $arg_type:ty),* ) -> $ret_type:path);*) => {
+        $(
+            #[no_mangle]
+            pub extern $abi fn $fn_name ( $( $arg_id : $arg_type),* ) -> $ret_type {
+                let original_fn = |dynamic_fns: &mut crate::side_by_side::CudaDynamicFns| {
+                    dynamic_fns.$fn_name($( $arg_id ),*)
+                };
+                let get_formatted_args = Box::new(move |writer: &mut dyn std::io::Write| {
+                    (paste! { format :: [<write_ $fn_name>] }) (
+                        writer
+                        $(,$arg_id)*
+                    )
+                });
+                crate::handle_cuda_function_call(stringify!($fn_name), original_fn, get_formatted_args)
+            }
+        )*
     };
 }
 
 macro_rules! extern_redirect_with_post {
-    ($abi:literal fn $fn_name:ident( $($arg_id:ident : $arg_type:ty),* ) -> $ret_type:path) => {
-        #[no_mangle]
-        pub extern "system" fn $fn_name ( $( $arg_id : $arg_type),* ) -> $ret_type {
-            let original_fn = |fn_ptr| {
-                let typed_fn = unsafe { std::mem::transmute::<_, extern "system" fn( $( $arg_id : $arg_type),* ) -> $ret_type>(fn_ptr) };
-                typed_fn($( $arg_id ),*)
-            };
-            let get_formatted_args = Box::new(move |writer: &mut dyn std::io::Write| {
-                (paste! { format :: [<write_ $fn_name>] }) (
-                    writer
-                    $(,$arg_id)*
+    ($($abi:literal fn $fn_name:ident( $($arg_id:ident : $arg_type:ty),* ) -> $ret_type:path);*) => {
+        $(
+            #[no_mangle]
+            pub extern "system" fn $fn_name ( $( $arg_id : $arg_type),* ) -> $ret_type {
+                let original_fn = |dynamic_fns: &mut crate::side_by_side::CudaDynamicFns| {
+                    dynamic_fns.$fn_name($( $arg_id ),*)
+                };
+                let get_formatted_args = Box::new(move |writer: &mut dyn std::io::Write| {
+                    (paste! { format :: [<write_ $fn_name>] }) (
+                        writer
+                        $(,$arg_id)*
+                    )
+                });
+                crate::handle_cuda_function_call_with_probes(
+                    stringify!($fn_name),
+                    || (), original_fn,
+                    get_formatted_args,
+                    move |logger, state, _, cuda_result| paste! { [<$fn_name _Post>] } ( $( $arg_id ),* , logger, state, cuda_result )
                 )
-            });
-            crate::handle_cuda_function_call_with_probes(
-                stringify!($fn_name),
-                || (), original_fn,
-                get_formatted_args,
-                move |logger, state, _, cuda_result| paste! { [<$fn_name _Post>] } ( $( $arg_id ),* , logger, state, cuda_result )
-            )
-        }
+            }
+        )*
     };
 }
 
@@ -77,6 +81,7 @@ mod log;
 #[cfg_attr(windows, path = "os_win.rs")]
 #[cfg_attr(not(windows), path = "os_unix.rs")]
 mod os;
+mod side_by_side;
 mod trace;
 
 lazy_static! {
@@ -127,7 +132,8 @@ impl<T> LateInit<T> {
 
 struct GlobalDelayedState {
     settings: Settings,
-    libcuda_handle: NonNull<c_void>,
+    libcuda: CudaDynamicFns,
+    side_by_side_lib: Option<CudaDynamicFns>,
     cuda_state: trace::StateTracker,
 }
 
@@ -139,9 +145,8 @@ impl GlobalDelayedState {
     ) -> (LateInit<Self>, log::FunctionLogger<'a>) {
         let (mut fn_logger, settings) =
             factory.get_first_logger_and_init_settings(func, arguments_writer);
-        let maybe_libcuda_handle = unsafe { os::load_cuda_library(&settings.libcuda_path) };
-        let libcuda_handle = match NonNull::new(maybe_libcuda_handle) {
-            Some(h) => h,
+        let libcuda = match unsafe { CudaDynamicFns::load_library(&settings.libcuda_path) } {
+            Some(libcuda) => libcuda,
             None => {
                 fn_logger.log(log::LogEntry::ErrorBox(
                     format!("Invalid CUDA library at path {}", &settings.libcuda_path).into(),
@@ -149,11 +154,30 @@ impl GlobalDelayedState {
                 return (LateInit::Error, fn_logger);
             }
         };
+        let side_by_side_lib = settings
+            .side_by_side_path
+            .as_ref()
+            .and_then(|side_by_side_path| {
+                match unsafe { CudaDynamicFns::load_library(&*side_by_side_path) } {
+                    Some(fns) => Some(fns),
+                    None => {
+                        fn_logger.log(log::LogEntry::ErrorBox(
+                            format!(
+                                "Invalid side-by-side CUDA library at path {}",
+                                &side_by_side_path
+                            )
+                            .into(),
+                        ));
+                        None
+                    }
+                }
+            });
         let cuda_state = trace::StateTracker::new(&settings);
         let delayed_state = GlobalDelayedState {
             settings,
-            libcuda_handle,
+            libcuda,
             cuda_state,
+            side_by_side_lib,
         };
         (LateInit::Success(delayed_state), fn_logger)
     }
@@ -163,6 +187,7 @@ struct Settings {
     dump_dir: Option<PathBuf>,
     libcuda_path: String,
     override_cc_major: Option<u32>,
+    side_by_side_path: Option<String>,
 }
 
 impl Settings {
@@ -179,7 +204,7 @@ impl Settings {
                 None
             }
         };
-        let libcuda_path = match env::var("ZLUDA_DUMP_LIBCUDA_FILE") {
+        let libcuda_path = match env::var("ZLUDA_CUDA_LIB") {
             Err(env::VarError::NotPresent) => os::LIBCUDA_DEFAULT_PATH.to_owned(),
             Err(e) => {
                 logger.log(log::LogEntry::ErrorBox(Box::new(e) as _));
@@ -201,10 +226,19 @@ impl Settings {
                 Ok(cc) => Some(cc),
             },
         };
+        let side_by_side_path = match env::var("ZLUDA_SIDE_BY_SIDE_LIB") {
+            Err(env::VarError::NotPresent) => None,
+            Err(e) => {
+                logger.log(log::LogEntry::ErrorBox(Box::new(e) as _));
+                None
+            }
+            Ok(env_string) => Some(env_string),
+        };
         Settings {
             dump_dir,
             libcuda_path,
             override_cc_major,
+            side_by_side_path,
         }
     }
 
@@ -241,7 +275,7 @@ pub struct ModuleDump {
 
 fn handle_cuda_function_call(
     func: &'static str,
-    original_cuda_fn: impl FnOnce(NonNull<c_void>) -> CUresult,
+    original_cuda_fn: impl FnOnce(&mut CudaDynamicFns) -> Option<CUresult>,
     arguments_writer: Box<dyn FnMut(&mut dyn std::io::Write) -> std::io::Result<()>>,
 ) -> CUresult {
     handle_cuda_function_call_with_probes(
@@ -256,7 +290,7 @@ fn handle_cuda_function_call(
 fn handle_cuda_function_call_with_probes<T, PostFn>(
     func: &'static str,
     pre_probe: impl FnOnce() -> T,
-    original_cuda_fn: impl FnOnce(NonNull<c_void>) -> CUresult,
+    original_cuda_fn: impl FnOnce(&mut CudaDynamicFns) -> Option<CUresult>,
     arguments_writer: Box<dyn FnMut(&mut dyn std::io::Write) -> std::io::Result<()>>,
     post_probe: PostFn,
 ) -> CUresult
@@ -283,13 +317,18 @@ where
             (logger, global_state.delayed_state.as_mut().unwrap())
         }
     };
-    let name = std::ffi::CString::new(func).unwrap();
-    let fn_ptr =
-        unsafe { os::get_proc_address(delayed_state.libcuda_handle.as_ptr(), name.as_c_str()) };
-    let fn_ptr = NonNull::new(fn_ptr).unwrap();
     let pre_result = pre_probe();
-    let cu_result = original_cuda_fn(fn_ptr);
-    logger.result = Some(cu_result);
+    let maybe_cu_result = original_cuda_fn(&mut delayed_state.libcuda);
+    let cu_result = match maybe_cu_result {
+        Some(result) => result,
+        None => {
+            logger.log(log::LogEntry::ErrorBox(
+                format!("No function {} in the underlying CUDA library", func).into(),
+            ));
+            CUresult::CUDA_ERROR_UNKNOWN
+        }
+    };
+    logger.result = maybe_cu_result;
     post_probe(
         &mut logger,
         &mut delayed_state.cuda_state,