don't free mem

worldcoin · Dec 27, 2024 · e5c8c51 · e5c8c51
1 parent fc0b726
commit e5c8c51
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/iris-mpc-gpu/src/dot/share_db.rs b/iris-mpc-gpu/src/dot/share_db.rs
@@ -28,10 +28,11 @@ use cudarc::{
     nvrtc::compile_ptx,
 };
 use itertools::{izip, Itertools};
+use memmap2::MmapMut;
 use rayon::prelude::*;
 use std::{
     ffi::{c_void, CStr},
-    mem,
+    mem::{self, forget},
     sync::Arc,
 };
 
@@ -243,13 +244,23 @@ impl ShareDB {
             .devices()
             .iter()
             .map(|device| unsafe {
-                let host_mem0 = memmap2::MmapMut::map_anon(max_size * self.code_length).unwrap();
-                let host_mem1 = memmap2::MmapMut::map_anon(max_size * self.code_length).unwrap();
+                let host_mem0 = MmapMut::map_anon(max_size * self.code_length).unwrap();
+                let host_mem1 = MmapMut::map_anon(max_size * self.code_length).unwrap();
+
+                let host_mem0_ptr = host_mem0.as_ptr() as u64;
+                let host_mem1_ptr = host_mem1.as_ptr() as u64;
+
+                // Make sure to not drop the memory, even though we only use the pointers
+                // afterwards. This also has the effect that this memory is never freed, which
+                // is fine for the db.
+                forget(host_mem0);
+                forget(host_mem1);
+
                 (
                     StreamAwareCudaSlice::from(device.alloc(max_size).unwrap()),
                     (
                         StreamAwareCudaSlice::from(device.alloc(max_size).unwrap()),
-                        (host_mem0.as_ptr() as u64, host_mem1.as_ptr() as u64),
+                        (host_mem0_ptr, host_mem1_ptr),
                     ),
                 )
             })

diff --git a/iris-mpc-gpu/src/server/actor.rs b/iris-mpc-gpu/src/server/actor.rs
@@ -238,11 +238,15 @@ impl ServerActor {
             comms.clone(),
         );
 
+        let now = Instant::now();
+
         let left_code_db_slices = codes_engine.alloc_db(max_db_size);
         let left_mask_db_slices = masks_engine.alloc_db(max_db_size);
         let right_code_db_slices = codes_engine.alloc_db(max_db_size);
         let right_mask_db_slices = masks_engine.alloc_db(max_db_size);
 
+        tracing::info!("Allocated db in {:?}", now.elapsed());
+
         // Engines for inflight queries
         let batch_codes_engine = ShareDB::init(
             party_id,