use the open_bucket function in the testcase

worldcoin · Jan 14, 2025 · 6dcc548 · 6dcc548
1 parent d9c4c7e
commit 6dcc548
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 61 deletions.
diff --git a/iris-mpc-gpu/src/server/actor.rs b/iris-mpc-gpu/src/server/actor.rs
@@ -1152,7 +1152,7 @@ impl ServerActor {
 
             let buckets = self
                 .phase2_buckets
-                .open_buckets(&mut self.buckets, batch_streams);
+                .open_buckets(&self.buckets, batch_streams);
 
             tracing::info!("BUCKETRESULT: {:?}", buckets);
 

diff --git a/iris-mpc-gpu/src/threshold_ring/protocol.rs b/iris-mpc-gpu/src/threshold_ring/protocol.rs
@@ -2525,11 +2525,7 @@ impl Circuits {
         self.buffers.check_buffers();
     }
 
-    pub fn open_buckets(
-        &mut self,
-        buckets: &mut ChunkShare<u32>,
-        streams: &[CudaStream],
-    ) -> Vec<u32> {
+    pub fn open_buckets(&mut self, buckets: &ChunkShare<u32>, streams: &[CudaStream]) -> Vec<u32> {
         let a = dtoh_on_stream_sync(&buckets.a, &self.devs[0], &streams[0]).unwrap();
         let b = dtoh_on_stream_sync(&buckets.b, &self.devs[0], &streams[0]).unwrap();
         let mut res = buckets.as_view();

diff --git a/iris-mpc-gpu/tests/buckets.rs b/iris-mpc-gpu/tests/buckets.rs
@@ -6,7 +6,7 @@ mod buckets_test {
     };
     use iris_mpc_common::iris_db::iris::IrisCodeArray;
     use iris_mpc_gpu::{
-        helpers::{device_manager::DeviceManager, dtoh_on_stream_sync, htod_on_stream_sync},
+        helpers::{device_manager::DeviceManager, htod_on_stream_sync},
         threshold_ring::protocol::{ChunkShare, Circuits},
     };
     use itertools::{izip, Itertools};
@@ -110,31 +110,6 @@ mod buckets_test {
         result
     }
 
-    fn open(party: &mut Circuits, x: &ChunkShare<u32>, streams: &[CudaStream]) -> Vec<u32> {
-        let mut view = x.as_view();
-        let dev = party.get_devices()[0].clone();
-
-        let mut a = dtoh_on_stream_sync(&x.a, &dev, &streams[0]).unwrap();
-        cudarc::nccl::result::group_start().unwrap();
-        // Result is in bit 0
-        party.comms()[0]
-            .send_view(&view.b, party.next_id(), &streams[0])
-            .unwrap();
-
-        party.comms()[0]
-            .receive_view(&mut view.a, party.prev_id(), &streams[0])
-            .unwrap();
-        cudarc::nccl::result::group_end().unwrap();
-        let b = dtoh_on_stream_sync(&x.b, &dev, &streams[0]).unwrap();
-        let c = dtoh_on_stream_sync(&x.a, &dev, &streams[0]).unwrap();
-
-        for (a, b, c) in izip!(a.iter_mut(), b, c) {
-            *a += b + c;
-        }
-
-        a
-    }
-
     fn install_tracing() {
         tracing_subscriber::registry()
             .with(
@@ -195,7 +170,7 @@ mod buckets_test {
             tracing::info!("id: {}, compute time: {:?}", id, now.elapsed());
 
             let now = Instant::now();
-            let result = open(&mut party, &bucket, &streams);
+            let result = party.open_buckets(&bucket, &streams);
             party.synchronize_streams(&streams);
             tracing::info!("id: {}, Starting tests...", id);
             tracing::info!(

diff --git a/iris-mpc-gpu/tests/one_bucket.rs b/iris-mpc-gpu/tests/one_bucket.rs
@@ -1,12 +1,12 @@
-#[cfg(feature = "gpu_dependent")]
+// #[cfg(feature = "gpu_dependent")]
 mod one_bucket_test {
     use cudarc::{
         driver::{CudaDevice, CudaStream},
         nccl::Id,
     };
     use iris_mpc_common::iris_db::iris::{IrisCodeArray, MATCH_THRESHOLD_RATIO};
     use iris_mpc_gpu::{
-        helpers::{device_manager::DeviceManager, dtoh_on_stream_sync, htod_on_stream_sync},
+        helpers::{device_manager::DeviceManager, htod_on_stream_sync},
         threshold_ring::protocol::{ChunkShare, Circuits},
     };
     use itertools::{izip, Itertools};
@@ -101,31 +101,6 @@ mod one_bucket_test {
         count
     }
 
-    fn open(party: &mut Circuits, x: &ChunkShare<u32>, streams: &[CudaStream]) -> u32 {
-        let mut view = x.as_view();
-        let dev = party.get_devices()[0].clone();
-
-        let mut a = dtoh_on_stream_sync(&x.a, &dev, &streams[0]).unwrap();
-        cudarc::nccl::result::group_start().unwrap();
-        // Result is in bit 0
-        party.comms()[0]
-            .send_view(&view.b, party.next_id(), &streams[0])
-            .unwrap();
-
-        party.comms()[0]
-            .receive_view(&mut view.a, party.prev_id(), &streams[0])
-            .unwrap();
-        cudarc::nccl::result::group_end().unwrap();
-        let b = dtoh_on_stream_sync(&x.b, &dev, &streams[0]).unwrap();
-        let c = dtoh_on_stream_sync(&x.a, &dev, &streams[0]).unwrap();
-
-        for (a, b, c) in izip!(a.iter_mut(), b, c) {
-            *a += b + c;
-        }
-
-        a[0]
-    }
-
     fn install_tracing() {
         tracing_subscriber::registry()
             .with(
@@ -180,7 +155,7 @@ mod one_bucket_test {
             tracing::info!("id: {}, compute time: {:?}", id, now.elapsed());
 
             let now = Instant::now();
-            let result = open(&mut party, &bucket, &streams);
+            let result = party.open_buckets(&bucket, &streams)[0];
             party.synchronize_streams(&streams);
             tracing::info!("id: {}, Starting tests...", id);
             tracing::info!(