Support more export table functions

3 files changed, 79 insertions, 18 deletions

diff --git a/zluda/src/impl/export_table.rs b/zluda/src/impl/export_table.rs
index aa46788..1d778af 100644
--- a/zluda/src/impl/export_table.rs
+++ b/zluda/src/impl/export_table.rs
@@ -128,12 +128,14 @@ static CUDART_INTERFACE_VTABLE: [VTableEntry; CUDART_INTERFACE_LENGTH] = [
     VTableEntry { ptr: ptr::null() },

     VTableEntry { ptr: ptr::null() },

     VTableEntry {

-        ptr: get_module_from_cubin_ext as *const (),

+        ptr: get_module_from_cubin_ext1 as *const (),

     },

     VTableEntry {

         ptr: cudart_interface_fn6 as *const (),

     },

-    VTableEntry { ptr: ptr::null() },

+    VTableEntry {

+        ptr: get_module_from_cubin_ext2 as _,

+    },

     VTableEntry { ptr: ptr::null() },

 ];

 

@@ -237,8 +239,14 @@ unsafe extern "system" fn get_module_from_cubin(
     {

         return CUresult::CUDA_ERROR_INVALID_VALUE;

     }

-    //let result = result.decuda();

     let fatbin_header = (*fatbinc_wrapper).data;

+    get_module_from_cubin_unwrapped(fatbin_header, result)

+}

+

+unsafe extern "system" fn get_module_from_cubin_unwrapped(

+    fatbin_header: *const FatbinHeader,

+    result: *mut CUmodule,

+) -> CUresult {

     if (*fatbin_header).magic != FATBIN_MAGIC || (*fatbin_header).version != FATBIN_VERSION {

         return CUresult::CUDA_ERROR_INVALID_VALUE;

     }

@@ -273,11 +281,12 @@ unsafe extern "system" fn get_module_from_cubin(
     CUresult::CUDA_ERROR_COMPAT_NOT_SUPPORTED_ON_DEVICE

 }

 

-unsafe extern "system" fn get_module_from_cubin_ext(

+unsafe extern "system" fn get_module_from_cubin_ext1(

     result: *mut CUmodule,

     fatbinc_wrapper: *const FatbincWrapper,

     ptr1: *mut c_void,

     ptr2: *mut c_void,

+    _unknown: usize,

 ) -> CUresult {

     // Not sure what those two parameters are actually used for,

     // they are somehow involved in __cudaRegisterHostVar

@@ -333,7 +342,24 @@ unsafe fn decompress_kernel_module(file: *const FatbinFileHeader) -> Option<Vec<
     }

 }

 

-unsafe extern "system" fn cudart_interface_fn6(_: u64) {}

+unsafe extern "system" fn cudart_interface_fn6(_: usize) -> CUresult {

+    CUresult::CUDA_SUCCESS

+}

+

+// From the assembly looks like arg5 is a count of something in arg3 and arg4

+unsafe extern "system" fn get_module_from_cubin_ext2(

+    fatbinc_wrapper: *const FatbinHeader,

+    result: *mut CUmodule,

+    _arg3: usize,

+    _arg4: usize,

+    arg5: c_uint,

+) -> CUresult {

+    if arg5 != 0 {

+        CUresult::CUDA_ERROR_NOT_SUPPORTED

+    } else {

+        get_module_from_cubin_unwrapped(fatbinc_wrapper, result)

+    }

+}

 

 const TOOLS_TLS_GUID: CUuuid = CUuuid {

     bytes: [

@@ -535,7 +561,9 @@ static DEVICE_EXTENDED_RT_VTABLE: [VTableEntry; DEVICE_EXTENDED_RT_LENGTH] = [
     VTableEntry { ptr: ptr::null() },

     VTableEntry { ptr: ptr::null() },

     VTableEntry { ptr: ptr::null() },

-    VTableEntry { ptr: device_get_something as _ },

+    VTableEntry {

+        ptr: device_get_something as _,

+    },

     VTableEntry { ptr: ptr::null() },

     VTableEntry { ptr: ptr::null() },

     VTableEntry { ptr: ptr::null() },

@@ -573,10 +601,7 @@ unsafe extern "system" fn device_get_attribute_ext(
 

 // I don't know is this function return,

 // but on my GTX 1060 it returns 0

-unsafe extern "system" fn device_get_something(

-    result: *mut c_uchar,

-    _dev: CUdevice,

-) -> CUresult {

+unsafe extern "system" fn device_get_something(result: *mut c_uchar, _dev: CUdevice) -> CUresult {

     if result == ptr::null_mut() {

         return CUresult::CUDA_ERROR_INVALID_VALUE;

     }

diff --git a/zluda_dump/src/lib.rs b/zluda_dump/src/lib.rs
index f3dd44a..e334151 100644
--- a/zluda_dump/src/lib.rs
+++ b/zluda_dump/src/lib.rs
@@ -163,7 +163,7 @@ pub unsafe fn cuModuleLoadData(
 
 unsafe fn record_module_image_raw(module: CUmodule, raw_image: *const ::std::os::raw::c_void) {
     if *(raw_image as *const u32) == 0x464c457f {
-        os_log!("Unsupported ELF module: {:?}", raw_image);
+        os_log!("Unsupported ELF module image: {:?}", raw_image);
         return;
     }
     let image = to_str(raw_image);
@@ -708,12 +708,22 @@ static mut ORIGINAL_GET_MODULE_FROM_CUBIN: Option<
         fatbinc_wrapper: *const FatbincWrapper,
     ) -> CUresult,
 > = None;
-static mut ORIGINAL_GET_MODULE_FROM_CUBIN_EXT: Option<
+static mut ORIGINAL_GET_MODULE_FROM_CUBIN_EXT1: Option<
     unsafe extern "system" fn(
         result: *mut CUmodule,
         fatbinc_wrapper: *const FatbincWrapper,
         ptr1: *mut c_void,
         ptr2: *mut c_void,
+        _unknown: usize,
+    ) -> CUresult,
+> = None;
+static mut ORIGINAL_GET_MODULE_FROM_CUBIN_EXT2: Option<
+    unsafe extern "system" fn(
+        fatbinc_wrapper: *const FatbinHeader,
+        result: *mut CUmodule,
+        ptr1: *mut c_void,
+        ptr2: *mut c_void,
+        _unknown: usize,
     ) -> CUresult,
 > = None;
 
@@ -843,8 +853,12 @@ unsafe fn get_export_override_fn(
             get_module_from_cubin as *const _
         }
         (CUDART_INTERFACE_GUID, 6) => {
-            ORIGINAL_GET_MODULE_FROM_CUBIN_EXT = mem::transmute(original_fn);
-            get_module_from_cubin_ext as *const _
+            ORIGINAL_GET_MODULE_FROM_CUBIN_EXT1 = mem::transmute(original_fn);
+            get_module_from_cubin_ext1 as *const _
+        }
+        (CUDART_INTERFACE_GUID, 8) => {
+            ORIGINAL_GET_MODULE_FROM_CUBIN_EXT2 = mem::transmute(original_fn);
+            get_module_from_cubin_ext2 as *const _
         }
         _ => os::get_thunk(original_fn, report_unknown_export_table_call, guid, idx),
     }
@@ -906,7 +920,14 @@ unsafe fn get_module_from_cubin_impl(
     {
         return CUresult::CUDA_ERROR_INVALID_VALUE;
     }
-    let fatbin_header = (*fatbinc_wrapper).data;
+    get_module_from_cubin_unwrapped(module, (*fatbinc_wrapper).data, get_module_base)
+}
+
+unsafe fn get_module_from_cubin_unwrapped(
+    module: *mut CUmodule,
+    fatbin_header: *const FatbinHeader,
+    get_module_base: impl FnOnce() -> CUresult,
+) -> CUresult {
     if (*fatbin_header).magic != FATBIN_MAGIC || (*fatbin_header).version != FATBIN_VERSION {
         return CUresult::CUDA_ERROR_INVALID_VALUE;
     }
@@ -936,6 +957,8 @@ unsafe fn get_module_from_cubin_impl(
             },
             Err(_) => {}
         }
+    } else {
+        os_log!("Unsupported runtime module: {:?}", *module);
     }
     result
 }
@@ -949,14 +972,27 @@ unsafe extern "system" fn get_module_from_cubin(
     })
 }
 
-unsafe extern "system" fn get_module_from_cubin_ext(
+unsafe extern "system" fn get_module_from_cubin_ext1(
     module: *mut CUmodule,
     fatbinc_wrapper: *const FatbincWrapper,
     ptr1: *mut c_void,
     ptr2: *mut c_void,
+    _unknown: usize,
 ) -> CUresult {
     get_module_from_cubin_impl(module, fatbinc_wrapper, || {
-        ORIGINAL_GET_MODULE_FROM_CUBIN_EXT.unwrap()(module, fatbinc_wrapper, ptr1, ptr2)
+        ORIGINAL_GET_MODULE_FROM_CUBIN_EXT1.unwrap()(module, fatbinc_wrapper, ptr1, ptr2, _unknown)
+    })
+}
+
+unsafe extern "system" fn get_module_from_cubin_ext2(
+    fatbin_header: *const FatbinHeader,
+    module: *mut CUmodule,
+    ptr1: *mut c_void,
+    ptr2: *mut c_void,
+    _unknown: usize,
+) -> CUresult {
+    get_module_from_cubin_unwrapped(module, fatbin_header, || {
+        ORIGINAL_GET_MODULE_FROM_CUBIN_EXT2.unwrap()(fatbin_header, module, ptr1, ptr2, _unknown)
     })
 }
 
diff --git a/zluda_redirect/Cargo.toml b/zluda_redirect/Cargo.toml
index 130e244..85c2976 100644
--- a/zluda_redirect/Cargo.toml
+++ b/zluda_redirect/Cargo.toml
@@ -10,4 +10,4 @@ crate-type = ["cdylib"]
 [target.'cfg(windows)'.dependencies]
 detours-sys = { path = "../detours-sys" }
 wchar = "0.6"
-winapi = { version = "0.3", features = ["processthreadsapi", "winbase", "winnt", "winerror", "libloaderapi", "tlhelp32", "std"] }
\ No newline at end of file
+winapi = { version = "0.3", features = ["processthreadsapi", "winbase", "winnt", "winerror", "libloaderapi", "tlhelp32", "handleapi", "std"] }
\ No newline at end of file