diff --git a/iris-mpc-gpu/src/dot/share_db.rs b/iris-mpc-gpu/src/dot/share_db.rs index b0b5889c2..e3f3e653d 100644 --- 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 index 07d19936c..5dd0d0d61 100644 --- 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,