Add tests for heap alloc

3 files changed, 109 insertions, 13 deletions

diff --git a/zluda/src/impl/dark_api.rs b/zluda/src/impl/dark_api.rs
index c3b596c..e7eb4b1 100644
--- a/zluda/src/impl/dark_api.rs
+++ b/zluda/src/impl/dark_api.rs
@@ -188,16 +188,16 @@ impl CudaDarkApi for CudaDarkApiZluda {
     }
 
     unsafe extern "system" fn heap_alloc(
-        _halloc_ptr: *mut *mut zluda_dark_api::HeapAllocRecord,
-        _param1: usize,
-        _param2: usize,
+        alloc_ptr: *mut *mut zluda_dark_api::HeapAllocRecord,
+        destructor: Option<unsafe extern "system" fn(u32, usize)>,
+        value: usize,
     ) -> CUresult {
         super::unimplemented()
     }
 
     unsafe extern "system" fn heap_free(
-        _halloc: *mut zluda_dark_api::HeapAllocRecord,
-        _param2: *mut usize,
+        alloc_ptr: *mut zluda_dark_api::HeapAllocRecord,
+        value: *mut usize,
     ) -> CUresult {
         super::unimplemented()
     }
diff --git a/zluda/tests/dark_api.rs b/zluda/tests/dark_api.rs
new file mode 100644
index 0000000..3e5b3e0
--- /dev/null
+++ b/zluda/tests/dark_api.rs
@@ -0,0 +1,92 @@
+use crate::common::CudaDriverFns;

+use cuda_types::*;

+use std::{mem, ptr};

+mod common;

+

+cuda_driver_test!(heap_alloc_chain);

+

+unsafe fn heap_alloc_chain<T: CudaDriverFns>(cuda: T) {

+    assert_eq!(cuda.cuInit(0), CUresult::CUDA_SUCCESS);

+    let mut export_table = mem::zeroed();

+    let guid = zluda_dark_api::HeapAccess::GUID;

+    assert_eq!(

+        cuda.cuGetExportTable(&mut export_table, &guid),

+        CUresult::CUDA_SUCCESS

+    );

+    assert_eq!(3 * mem::size_of::<usize>(), *export_table.cast::<usize>());

+    let heap_access = zluda_dark_api::HeapAccess::new(export_table);

+    let mut record1 = ptr::null_mut();

+    assert_eq!(

+        CUresult::CUDA_SUCCESS,

+        heap_access.heap_alloc(&mut record1, None, 1)

+    );

+    let mut record2 = ptr::null_mut();

+    assert_eq!(

+        CUresult::CUDA_SUCCESS,

+        heap_access.heap_alloc(&mut record2, None, 2)

+    );

+    let mut record3 = ptr::null_mut();

+    assert_eq!(

+        CUresult::CUDA_SUCCESS,

+        heap_access.heap_alloc(&mut record3, None, 3)

+    );

+    assert_eq!((&*record1).destructor, None);

+    assert_eq!((&*record1).value, 1);

+    assert_eq!((&*record1).prev_alloc, record2);

+    assert_eq!((&*record1).next_alloc, ptr::null());

+    assert_eq!((&*record2).destructor, None);

+    assert_eq!((&*record2).value, 2);

+    assert_eq!((&*record2).prev_alloc, record3);

+    assert_eq!((&*record2).next_alloc, record1);

+    assert_eq!((&*record3).destructor, None);

+    assert_eq!((&*record3).value, 3);

+    assert_eq!((&*record3).prev_alloc, ptr::null());

+    assert_eq!((&*record3).next_alloc, record2);

+}

+

+cuda_driver_test!(heap_free);

+

+unsafe fn heap_free<T: CudaDriverFns>(cuda: T) {

+    assert_eq!(cuda.cuInit(0), CUresult::CUDA_SUCCESS);

+    let mut export_table = mem::zeroed();

+    let guid = zluda_dark_api::HeapAccess::GUID;

+    assert_eq!(

+        cuda.cuGetExportTable(&mut export_table, &guid),

+        CUresult::CUDA_SUCCESS

+    );

+    let heap_access = zluda_dark_api::HeapAccess::new(export_table);

+    let mut record1 = ptr::null_mut();

+    assert_eq!(

+        CUresult::CUDA_SUCCESS,

+        heap_access.heap_alloc(&mut record1, Some(shutdown), 11)

+    );

+    let mut record2 = ptr::null_mut();

+    assert_eq!(

+        CUresult::CUDA_SUCCESS,

+        heap_access.heap_alloc(&mut record2, Some(shutdown), 12)

+    );

+    let mut record3 = ptr::null_mut();

+    assert_eq!(

+        CUresult::CUDA_SUCCESS,

+        heap_access.heap_alloc(&mut record3, None, 13)

+    );

+    let mut value = 0usize;

+    assert_eq!(

+        CUresult::CUDA_SUCCESS,

+        heap_access.heap_free(record2, &mut value),

+    );

+    assert_eq!(value, 12);

+    assert_eq!(

+        mem::transmute::<_, usize>((&*record1).destructor),

+        shutdown as usize

+    );

+    assert_eq!((&*record1).value, 11);

+    assert_eq!((&*record1).prev_alloc, record3);

+    assert_eq!((&*record1).next_alloc, ptr::null());

+    assert_eq!((&*record3).destructor, None);

+    assert_eq!((&*record3).value, 13);

+    assert_eq!((&*record3).prev_alloc, ptr::null());

+    assert_eq!((&*record3).next_alloc, record1);

+}

+

+unsafe extern "system" fn shutdown(_unknown: u32, _value: usize) {}

diff --git a/zluda_dark_api/src/lib.rs b/zluda_dark_api/src/lib.rs
index 6849e0e..31e0fea 100644
--- a/zluda_dark_api/src/lib.rs
+++ b/zluda_dark_api/src/lib.rs
@@ -286,11 +286,13 @@ dark_api_table!(
     => HEAP_ACCESS [3] {
         0 => SIZE_OF,
         1 => heap_alloc(
-            halloc_ptr: *mut *mut HeapAllocRecord,
-            param1: usize,
-            param2: usize
+            alloc_ptr: *mut *mut HeapAllocRecord,
+            // destructor is called only on CUDA exit, on Windows
+            // that is DLL unload from DllMain
+            destructor: Option<unsafe extern "system" fn(u32, usize)>,
+            value: usize
         ) -> CUresult,
-        2 => heap_free(halloc: *mut HeapAllocRecord, param2: *mut usize) -> CUresult
+        2 => heap_free(halloc: *mut HeapAllocRecord, value: *mut usize) -> CUresult
     },
     // This fn table is used by OptiX
     [0xB1u8, 0x05, 0x41, 0xE1, 0xF7, 0xC7, 0xC7, 0x4A, 0x9F, 0x64, 0xF2, 0x23, 0xBE, 0x99, 0xF1, 0xE2]
@@ -487,10 +489,12 @@ pub enum ContextStateManager {}
 
 #[repr(C)]
 pub struct HeapAllocRecord {
-    param1: usize,
-    param2: usize,
-    _unknown: usize,
-    global_heap: *mut c_void,
+    pub destructor: Option<unsafe extern "system" fn(u32, usize)>,
+    pub value: usize,
+    // The two fields below are mainatined by the driver,
+    // they form a a doubly-linked list
+    pub prev_alloc: *const HeapAllocRecord,
+    pub next_alloc: *const HeapAllocRecord,
 }
 
 #[derive(Clone, Copy)]