Support more export table functions

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(
+unsafe extern "system" fn device_get_something(result: *mut c_uchar, _dev: CUdevice) -> CUresult {
-    result: *mut c_uchar,
-    _dev: CUdevice,
-) -> CUresult {
     if result == ptr::null_mut() {
         return CUresult::CUDA_ERROR_INVALID_VALUE;
     }

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);
+            ORIGINAL_GET_MODULE_FROM_CUBIN_EXT1 = mem::transmute(original_fn);
-            get_module_from_cubin_ext as *const _
+            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)
     })
 }
 

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"] }
+winapi = { version = "0.3", features = ["processthreadsapi", "winbase", "winnt", "winerror", "libloaderapi", "tlhelp32", "handleapi", "std"] }