cudaFree

rapidsai · Feb 12, 2024 · 291c140 · 291c140
1 parent 346f6cf
commit 291c140
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/rust/cuvs/src/dlpack.rs b/rust/cuvs/src/dlpack.rs
@@ -122,14 +122,19 @@ impl ManagedTensor {
 
             let mut ret = self.0.clone();
             ret.dl_tensor.data = device_data;
+            // call cudaFree automatically to clean up data
+            ret.deleter = Some(cuda_free_tensor);
             ret.dl_tensor.device.device_type = ffi::DLDeviceType::kDLCUDA;
-            // TODO: do we need to set the device id here too?
-            // TODO: set deleter here to call cudaFree
+
             Ok(ManagedTensor(ret))
         }
     }
 }
 
+unsafe extern "C" fn cuda_free_tensor(self_: *mut ffi::DLManagedTensor) {
+    let _ = ffi::cudaFree((*self_).dl_tensor.data);
+}
+
 impl Drop for ManagedTensor {
     fn drop(&mut self) {
         unsafe {