Anonymized stats through buckets

iris-mpc-gpu/Cargo.toml

@@ -9,7 +9,7 @@ repository.workspace = true
 
 [dependencies]
 bincode = "1.3.3"
-cudarc = { version = "0.12", features = ["cuda-12020", "nccl"] }
+cudarc = { git = "https://github.com/worldcoin/cudarc-fork.git", features = ["cuda-12020", "nccl"] }
 eyre.workspace = true
 tracing.workspace = true
 bytemuck.workspace = true

iris-mpc-gpu/benches/matmul.rs

@@ -58,7 +58,14 @@ fn bench_memcpy(c: &mut Criterion) {
                 &streams,
                 &blass,
             );
-            engine.dot_reduce(&query_sums, &db_slices.code_sums_gr, &db_sizes, 0, &streams);
+            engine.dot_reduce(
+                &query_sums,
+                &db_slices.code_sums_gr,
+                &db_sizes,
+                0,
+                &streams,
+                0,
+            );
             device_manager.await_streams(&streams);
         });
     });

iris-mpc-gpu/src/dot/distance_comparator.rs

@@ -1,23 +1,31 @@
 use super::ROTATIONS;
-use crate::helpers::{
-    device_manager::DeviceManager, launch_config_from_elements_and_threads,
-    DEFAULT_LAUNCH_CONFIG_THREADS,
+use crate::{
+    helpers::{
+        device_manager::DeviceManager, launch_config_from_elements_and_threads,
+        DEFAULT_LAUNCH_CONFIG_THREADS,
+    },
+    threshold_ring::protocol::{ChunkShare, ChunkShareView},
 };
 use cudarc::{
-    driver::{CudaFunction, CudaSlice, CudaStream, CudaView, LaunchAsync},
+    driver::{
+        result::{launch_kernel, memset_d8_sync},
+        sys, CudaFunction, CudaSlice, CudaStream, CudaView, DevicePtr, DeviceSlice, LaunchAsync,
+    },
     nvrtc::compile_ptx,
 };
-use std::{cmp::min, sync::Arc};
+use std::{cmp::min, ffi::c_void, sync::Arc};
 
 const PTX_SRC: &str = include_str!("kernel.cu");
 const OPEN_RESULTS_FUNCTION: &str = "openResults";
+const OPEN_RESULTS_BATCH_FUNCTION: &str = "openResultsBatch";
 const MERGE_DB_RESULTS_FUNCTION: &str = "mergeDbResults";
 const MERGE_BATCH_RESULTS_FUNCTION: &str = "mergeBatchResults";
 const ALL_MATCHES_LEN: usize = 256;
 
 pub struct DistanceComparator {
     pub device_manager:          Arc<DeviceManager>,
     pub open_kernels:            Vec<CudaFunction>,
+    pub open_batch_kernels:      Vec<CudaFunction>,
     pub merge_db_kernels:        Vec<CudaFunction>,
     pub merge_batch_kernels:     Vec<CudaFunction>,
     pub query_length:            usize,
@@ -37,6 +45,7 @@ impl DistanceComparator {
     pub fn init(query_length: usize, device_manager: Arc<DeviceManager>) -> Self {
         let ptx = compile_ptx(PTX_SRC).unwrap();
         let mut open_kernels: Vec<CudaFunction> = Vec::new();
+        let mut open_batch_kernels: Vec<CudaFunction> = Vec::new();
         let mut merge_db_kernels = Vec::new();
         let mut merge_batch_kernels = Vec::new();
         let mut opened_results = vec![];
@@ -58,12 +67,15 @@ impl DistanceComparator {
             device
                 .load_ptx(ptx.clone(), "", &[
                     OPEN_RESULTS_FUNCTION,
+                    OPEN_RESULTS_BATCH_FUNCTION,
                     MERGE_DB_RESULTS_FUNCTION,
                     MERGE_BATCH_RESULTS_FUNCTION,
                 ])
                 .unwrap();
 
             let open_results_function = device.get_func("", OPEN_RESULTS_FUNCTION).unwrap();
+            let open_results_batch_function =
+                device.get_func("", OPEN_RESULTS_BATCH_FUNCTION).unwrap();
             let merge_db_results_function = device.get_func("", MERGE_DB_RESULTS_FUNCTION).unwrap();
             let merge_batch_results_function =
                 device.get_func("", MERGE_BATCH_RESULTS_FUNCTION).unwrap();
@@ -90,13 +102,15 @@ impl DistanceComparator {
             );
 
             open_kernels.push(open_results_function);
+            open_batch_kernels.push(open_results_batch_function);
             merge_db_kernels.push(merge_db_results_function);
             merge_batch_kernels.push(merge_batch_results_function);
         }
 
         Self {
             device_manager,
             open_kernels,
+            open_batch_kernels,
             merge_db_kernels,
             merge_batch_kernels,
             query_length,
@@ -115,6 +129,81 @@ impl DistanceComparator {
 
     #[allow(clippy::too_many_arguments)]
     pub fn open_results(
+        &self,
+        results1: &[CudaView<u64>],
+        results2: &[CudaView<u64>],
+        results3: &[CudaView<u64>],
+        matches_bitmap: &[CudaSlice<u64>],
+        db_sizes: &[usize],
+        real_db_sizes: &[usize],
+        offset: usize,
+        total_db_sizes: &[usize],
+        ignore_db_results: &[bool],
+        match_distances_buffers_codes: &[ChunkShare<u16>],
+        match_distances_buffers_masks: &[ChunkShare<u16>],
+        match_distances_counters: &[CudaSlice<u32>],
+        code_dots: &[ChunkShareView<u16>],
+        mask_dots: &[ChunkShareView<u16>],
+        batch_size: usize,
+        streams: &[CudaStream],
+    ) {
+        for i in 0..self.device_manager.device_count() {
+            // Those correspond to 0 length dbs, which were just artificially increased to
+            // length 1 to avoid division by zero in the kernel
+            if ignore_db_results[i] {
+                continue;
+            }
+            let num_elements = (db_sizes[i] * self.query_length).div_ceil(64);
+            let threads_per_block = DEFAULT_LAUNCH_CONFIG_THREADS; // ON CHANGE: sync with kernel
+            let cfg = launch_config_from_elements_and_threads(
+                num_elements as u32,
+                threads_per_block,
+                &self.device_manager.devices()[i],
+            );
+            self.device_manager.device(i).bind_to_thread().unwrap();
+
+            let ptr_param = |ptr: *const sys::CUdeviceptr| ptr as *mut c_void;
+            let usize_param = |val: &usize| val as *const usize as *mut _;
+
+            let params = &mut [
+                // Results arrays
+                ptr_param(results1[i].device_ptr()),
+                ptr_param(results2[i].device_ptr()),
+                ptr_param(results3[i].device_ptr()),
+                ptr_param(matches_bitmap[i].device_ptr()),
+                usize_param(&db_sizes[i]),
+                usize_param(&(batch_size * ROTATIONS)),
+                usize_param(&offset),
+                usize_param(&num_elements),
+                usize_param(&real_db_sizes[i]),
+                usize_param(&total_db_sizes[i]),
+                ptr_param(match_distances_buffers_codes[i].a.device_ptr()),
+                ptr_param(match_distances_buffers_codes[i].b.device_ptr()),
+                ptr_param(match_distances_buffers_masks[i].a.device_ptr()),
+                ptr_param(match_distances_buffers_masks[i].b.device_ptr()),
+                ptr_param(match_distances_counters[i].device_ptr()),
+                ptr_param(code_dots[i].a.device_ptr()),
+                ptr_param(code_dots[i].b.device_ptr()),
+                ptr_param(mask_dots[i].a.device_ptr()),
+                ptr_param(mask_dots[i].b.device_ptr()),
+            ];
+
+            unsafe {
+                launch_kernel(
+                    self.open_kernels[i].cu_function(),
+                    cfg.grid_dim,
+                    cfg.block_dim,
+                    0,
+                    streams[i].stream,
+                    params,
+                )
+                .unwrap();
+            }
+        }
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn open_batch_results(
         &self,
         results1: &[CudaView<u64>],
         results2: &[CudaView<u64>],
@@ -143,7 +232,7 @@ impl DistanceComparator {
             self.device_manager.device(i).bind_to_thread().unwrap();
 
             unsafe {
-                self.open_kernels[i]
+                self.open_batch_kernels[i]
                     .clone()
                     .launch_on_stream(
                         &streams[i],
@@ -347,4 +436,41 @@ impl DistanceComparator {
             })
             .collect::<Vec<_>>()
     }
+
+    pub fn prepare_match_distances_buffer(&self, max_size: usize) -> Vec<ChunkShare<u16>> {
+        (0..self.device_manager.device_count())
+            .map(|i| {
+                let a = self.device_manager.device(i).alloc_zeros(max_size).unwrap();
+                let b = self.device_manager.device(i).alloc_zeros(max_size).unwrap();
+
+                self.device_manager.device(i).bind_to_thread().unwrap();
+                unsafe {
+                    memset_d8_sync(*a.device_ptr(), 0xff, a.num_bytes()).unwrap();
+                    memset_d8_sync(*b.device_ptr(), 0xff, b.num_bytes()).unwrap();
+                }
+
+                ChunkShare::new(a, b)
+            })
+            .collect::<Vec<_>>()
+    }
+
+    pub fn prepare_match_distances_counter(&self) -> Vec<CudaSlice<u32>> {
+        (0..self.device_manager.device_count())
+            .map(|i| self.device_manager.device(i).alloc_zeros(1).unwrap())
+            .collect::<Vec<_>>()
+    }
+
+    pub fn prepare_match_distances_buckets(&self, n_buckets: usize) -> ChunkShare<u32> {
+        let a = self
+            .device_manager
+            .device(0)
+            .alloc_zeros(n_buckets)
+            .unwrap();
+        let b = self
+            .device_manager
+            .device(0)
+            .alloc_zeros(n_buckets)
+            .unwrap();
+        ChunkShare::new(a, b)
+    }
 }

iris-mpc-gpu/src/dot/kernel.cu

@@ -27,7 +27,7 @@ extern "C" __global__ void matmul_correct_and_reduce(int *c, unsigned short *out
     }
 }
 
-extern "C" __global__ void openResults(unsigned long long *result1, unsigned long long *result2, unsigned long long *result3, unsigned long long *output, size_t chunkLength, size_t queryLength, size_t offset, size_t numElements, size_t realChunkLen, size_t totalDbLen)
+extern "C" __global__ void openResultsBatch(unsigned long long *result1, unsigned long long *result2, unsigned long long *result3, unsigned long long *output, size_t chunkLength, size_t queryLength, size_t offset, size_t numElements, size_t realChunkLen, size_t totalDbLen)
 {
     size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (idx < numElements)
@@ -51,6 +51,41 @@ extern "C" __global__ void openResults(unsigned long long *result1, unsigned lon
     }
 }
 
+extern "C" __global__ void openResults(unsigned long long *result1, unsigned long long *result2, unsigned long long *result3, unsigned long long *output, size_t chunkLength, size_t queryLength, size_t offset, size_t numElements, size_t realChunkLen, size_t totalDbLen, unsigned short *match_distances_buffer_codes_a, unsigned short *match_distances_buffer_codes_b, unsigned short *match_distances_buffer_masks_a, unsigned short *match_distances_buffer_masks_b, unsigned int *match_distances_counter, unsigned short *code_dots_a, unsigned short *code_dots_b, unsigned short *mask_dots_a, unsigned short *mask_dots_b)
+{
+    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < numElements)
+    {
+        unsigned long long result = result1[idx] ^ result2[idx] ^ result3[idx];
+        for (int i = 0; i < 64; i++)
+        {
+            unsigned int queryIdx = (idx * 64 + i) / chunkLength;
+            unsigned int dbIdx = (idx * 64 + i) % chunkLength;
+            bool match = (result & (1ULL << i));
+
+            // Check if we are out of bounds for the query or db
+            if (queryIdx >= queryLength || dbIdx >= realChunkLen || !match)
+            {
+                continue;
+            }
+
+            // Save the corresponding code and mask dots for later (match distributions)
+            unsigned int match_distances_counter_idx = atomicAdd(&match_distances_counter[0], 1);
+            if (match_distances_counter[0] < 1000)
+            {
+                match_distances_buffer_codes_a[match_distances_counter_idx] = code_dots_a[idx * 64 + i];
+                match_distances_buffer_codes_b[match_distances_counter_idx] = code_dots_b[idx * 64 + i];
+                match_distances_buffer_masks_a[match_distances_counter_idx] = mask_dots_a[idx * 64 + i];
+                match_distances_buffer_masks_b[match_distances_counter_idx] = mask_dots_b[idx * 64 + i];
+            }
+
+            // Mark which results are matches with a bit in the output
+            unsigned int outputIdx = totalDbLen * (queryIdx / ALL_ROTATIONS) + dbIdx + offset;
+            atomicOr(&output[outputIdx / 64], (1ULL << (outputIdx % 64)));
+        }
+    }
+}
+
 extern "C" __global__ void mergeDbResults(unsigned long long *matchResultsLeft, unsigned long long *matchResultsRight, unsigned int *finalResults, size_t queryLength, size_t dbLength, size_t numElements, unsigned int *matchCounter, unsigned int *allMatches, unsigned int *matchCounterLeft, unsigned int *matchCounterRight, unsigned int *partialResultsLeft, unsigned int *partialResultsRight)
 {
     size_t idx = blockIdx.x * blockDim.x + threadIdx.x;