[POP-2034] LUC implementation on iric-mpc

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

@@ -13,24 +13,26 @@ const PTX_SRC: &str = include_str!("kernel.cu");
 const OPEN_RESULTS_FUNCTION: &str = "openResults";
 const MERGE_DB_RESULTS_FUNCTION: &str = "mergeDbResults";
 const MERGE_BATCH_RESULTS_FUNCTION: &str = "mergeBatchResults";
+const MERGE_BATCH_RESULTS_WITH_OR_POLICY_BITMAP_FUNCTION: &str = "mergeDbResultsWithOrPolicyBitmap";
 const ALL_MATCHES_LEN: usize = 256;
 
 pub struct DistanceComparator {
-    pub device_manager:          Arc<DeviceManager>,
-    pub open_kernels:            Vec<CudaFunction>,
-    pub merge_db_kernels:        Vec<CudaFunction>,
-    pub merge_batch_kernels:     Vec<CudaFunction>,
-    pub query_length:            usize,
-    pub opened_results:          Vec<CudaSlice<u32>>,
-    pub final_results:           Vec<CudaSlice<u32>>,
-    pub results_init_host:       Vec<u32>,
-    pub final_results_init_host: Vec<u32>,
-    pub match_counters:          Vec<CudaSlice<u32>>,
-    pub all_matches:             Vec<CudaSlice<u32>>,
-    pub match_counters_left:     Vec<CudaSlice<u32>>,
-    pub match_counters_right:    Vec<CudaSlice<u32>>,
-    pub partial_results_left:    Vec<CudaSlice<u32>>,
-    pub partial_results_right:   Vec<CudaSlice<u32>>,
+    pub device_manager:                  Arc<DeviceManager>,
+    pub open_kernels:                    Vec<CudaFunction>,
+    pub merge_db_kernels:                Vec<CudaFunction>,
+    pub merge_batch_kernels:             Vec<CudaFunction>,
+    pub merge_batch_with_bitmap_kernels: Vec<CudaFunction>,
+    pub query_length:                    usize,
+    pub opened_results:                  Vec<CudaSlice<u32>>,
+    pub final_results:                   Vec<CudaSlice<u32>>,
+    pub results_init_host:               Vec<u32>,
+    pub final_results_init_host:         Vec<u32>,
+    pub match_counters:                  Vec<CudaSlice<u32>>,
+    pub all_matches:                     Vec<CudaSlice<u32>>,
+    pub match_counters_left:             Vec<CudaSlice<u32>>,
+    pub match_counters_right:            Vec<CudaSlice<u32>>,
+    pub partial_results_left:            Vec<CudaSlice<u32>>,
+    pub partial_results_right:           Vec<CudaSlice<u32>>,
 }
 
 impl DistanceComparator {
@@ -39,6 +41,7 @@ impl DistanceComparator {
         let mut open_kernels: Vec<CudaFunction> = Vec::new();
         let mut merge_db_kernels = Vec::new();
         let mut merge_batch_kernels = Vec::new();
+        let mut merge_batch_with_bitmap_kernels: Vec<CudaFunction> = Vec::new();
         let mut opened_results = vec![];
         let mut final_results = vec![];
         let mut match_counters = vec![];
@@ -60,13 +63,17 @@ impl DistanceComparator {
                     OPEN_RESULTS_FUNCTION,
                     MERGE_DB_RESULTS_FUNCTION,
                     MERGE_BATCH_RESULTS_FUNCTION,
+                    MERGE_BATCH_RESULTS_WITH_OR_POLICY_BITMAP_FUNCTION,
                 ])
                 .unwrap();
 
             let open_results_function = device.get_func("", OPEN_RESULTS_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();
+            let merge_batch_results_with_bitmap_function = device
+                .get_func("", MERGE_BATCH_RESULTS_WITH_OR_POLICY_BITMAP_FUNCTION)
+                .unwrap();
 
             opened_results.push(device.htod_copy(results_init_host.clone()).unwrap());
             final_results.push(device.htod_copy(final_results_init_host.clone()).unwrap());
@@ -88,17 +95,18 @@ impl DistanceComparator {
                     .alloc_zeros(ALL_MATCHES_LEN * query_length / ROTATIONS)
                     .unwrap(),
             );
-
             open_kernels.push(open_results_function);
             merge_db_kernels.push(merge_db_results_function);
             merge_batch_kernels.push(merge_batch_results_function);
+            merge_batch_with_bitmap_kernels.push(merge_batch_results_with_bitmap_function);
         }
 
         Self {
             device_manager,
             open_kernels,
             merge_db_kernels,
             merge_batch_kernels,
+            merge_batch_with_bitmap_kernels,
             query_length,
             opened_results,
             final_results,
@@ -166,6 +174,62 @@ impl DistanceComparator {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub fn join_db_matches_with_bitmaps(
+        &self,
+        max_db_size: usize,
+        matches_bitmap_left: &[CudaSlice<u64>],
+        matches_bitmap_right: &[CudaSlice<u64>],
+        final_results: &[CudaSlice<u32>],
+        db_sizes: &[usize],
+        streams: &[CudaStream],
+        or_policies_bitmap: &[CudaSlice<u64>],
+    ) {
+        for i in 0..self.device_manager.device_count() {
+            if db_sizes[i] == 0 {
+                continue;
+            }
+
+            let num_elements = (db_sizes[i] * self.query_length / ROTATIONS).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();
+
+            unsafe {
+                self.merge_batch_with_bitmap_kernels[i]
+                    .clone()
+                    .launch_on_stream(
+                        &streams[i],
+                        cfg,
+                        (
+                            &matches_bitmap_left[i],
+                            &matches_bitmap_right[i],
+                            &final_results[i],
+                            self.query_length,
+                            db_sizes[i] as u64,
+                            num_elements,
+                            max_db_size,
+                            &self.match_counters[i],
+                            &self.all_matches[i],
+                            &self.match_counters_left[i],
+                            &self.match_counters_right[i],
+                            &self.partial_results_left[i],
+                            &self.partial_results_right[i],
+                            // Additional args
+                            &or_policies_bitmap[i],
+                        ),
+                    )
+                    .unwrap();
+            }
+        }
+    }
+
     pub fn join_db_matches(
         &self,
         matches_bitmap_left: &[CudaSlice<u64>],

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

@@ -53,6 +53,7 @@ extern "C" __global__ void openResults(unsigned long long *result1, unsigned lon
 
 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;
     if (idx < numElements)
     {
@@ -93,6 +94,67 @@ extern "C" __global__ void mergeDbResults(unsigned long long *matchResultsLeft,
     }
 }
 
+extern "C" __global__ void mergeDbResultsWithOrPolicyBitmap(unsigned long long *matchResultsLeft, unsigned long long *matchResultsRight, unsigned int *finalResults, size_t queryLength, size_t dbLength, size_t numElements, size_t maxDbLength, unsigned int *matchCounter, unsigned int *allMatches, unsigned int *matchCounterLeft, unsigned int *matchCounterRight, unsigned int *partialResultsLeft, unsigned int *partialResultsRight, const unsigned long long *orPolicyBitmap) // 2D bitmap stored as 1D
+{
+
+    size_t rowStride64 = (maxDbLength + 63) / 64;
+    size_t totalBits   = maxDbLength * queryLength;
+
+    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < numElements)
+    {
+        for (int i = 0; i < 64; i++)
+        {
+
+            size_t globalBit = idx * 64 + i;
+            // Protect against any leftover bits if totalBits not multiple of 64
+            if (globalBit >= totalBits) break;
+
+            unsigned int queryIdx = globalBit / dbLength;
+            unsigned int dbIdx = globalBit % dbLength;
+            bool matchLeft = (matchResultsLeft[idx] & (1ULL << i));
+            bool matchRight = (matchResultsRight[idx] & (1ULL << i));
+
+            // Check bounds
+            if (queryIdx >= queryLength || dbIdx >= dbLength)
+                continue;
+
+            // Check for partial results (only used for debugging)
+            if (matchLeft)
+            {
+                unsigned int qmcL = atomicAdd(&matchCounterLeft[queryIdx], 1);
+                if (qmcL < MAX_MATCHES_LEN)
+                    partialResultsLeft[MAX_MATCHES_LEN * queryIdx + qmcL] = dbIdx;
+            }
+            if (matchRight)
+            {
+                unsigned int qmcR = atomicAdd(&matchCounterRight[queryIdx], 1);
+                if (qmcR < MAX_MATCHES_LEN)
+                    partialResultsRight[MAX_MATCHES_LEN * queryIdx + qmcR] = dbIdx;
+            }
+            size_t rowIndex = queryIdx * rowStride64;
+            size_t orPolicyBitmapIdx = rowIndex + (dbIdx / 64);
+
+            bool useOr = (orPolicyBitmap[orPolicyBitmapIdx]
+                          & (1ULL << (dbIdx % 64))) != 0ULL;
+
+            // If useOr is true => (matchLeft || matchRight),
+            // else => (matchLeft && matchRight).
+            bool finalMatch = useOr ? (matchLeft || matchRight)
+                                    : (matchLeft && matchRight);
+
+            if (finalMatch)
+            {
+                atomicMin(&finalResults[queryIdx], dbIdx);
+                unsigned int qmc = atomicAdd(&matchCounter[queryIdx], 1);
+                if (qmc < MAX_MATCHES_LEN)
+                    allMatches[MAX_MATCHES_LEN * queryIdx + qmc] = dbIdx;
+            }
+        }
+    }
+}
+
+
 extern "C" __global__ void mergeBatchResults(unsigned long long *matchResultsSelfLeft, unsigned long long *matchResultsSelfRight, 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;

iris-mpc-gpu/src/server/actor.rs

@@ -9,6 +9,7 @@ use crate::{
         self,
         comm::NcclComm,
         device_manager::DeviceManager,
+        htod_on_stream_sync,
         query_processor::{
             CompactQuery, CudaVec2DSlicerRawPointer, DeviceCompactQuery, DeviceCompactSums,
         },
@@ -606,6 +607,9 @@ impl ServerActor {
                 && batch_size * ROTATIONS == batch.db_right_preprocessed.len(),
             "Query batch sizes mismatch"
         );
+        if !batch.or_rule_serial_ids.is_empty() {
+            assert_eq!(batch.or_rule_serial_ids.len(), batch_size);
+        };
 
         ///////////////////////////////////////////////////////////////////
         // PERFORM DELETIONS (IF ANY)
@@ -762,16 +766,48 @@ impl ServerActor {
         ///////////////////////////////////////////////////////////////////
         // MERGE LEFT & RIGHT results
         ///////////////////////////////////////////////////////////////////
+
         tracing::info!("Joining both sides");
         // Merge results and fetch matching indices
         // Format: host_results[device_index][query_index]
-        self.distance_comparator.join_db_matches(
-            &self.db_match_list_left,
-            &self.db_match_list_right,
-            &self.final_results,
-            &self.current_db_sizes,
-            &self.streams[0],
-        );
+
+        // Initialize bitmap with OR rule, if exists
+        if !batch.or_rule_serial_ids.is_empty()
+            && !batch
+                .or_rule_serial_ids
+                .iter()
+                .all(|inner| inner.is_empty())
+        {
+            assert_eq!(batch.or_rule_serial_ids.len(), batch_size);
+
+            // Populate the pre-allocated OR policy bitmap with the serial ids
+            let host_or_policy_bitmap = prepare_or_policy_bitmap(
+                self.max_db_size,
+                batch.or_rule_serial_ids.clone(),
+                batch_size,
+            );
+
+            let device_or_policy_bitmap =
+                self.allocate_or_policy_bitmap(host_or_policy_bitmap.clone());
+
+            self.distance_comparator.join_db_matches_with_bitmaps(
+                self.max_db_size,
+                &self.db_match_list_left,
+                &self.db_match_list_right,
+                &self.final_results,
+                &self.current_db_sizes,
+                &self.streams[0],
+                &device_or_policy_bitmap,
+            );
+        } else {
+            self.distance_comparator.join_db_matches(
+                &self.db_match_list_left,
+                &self.db_match_list_right,
+                &self.final_results,
+                &self.current_db_sizes,
+                &self.streams[0],
+            );
+        }
 
         self.distance_comparator.join_batch_matches(
             &self.batch_match_list_left,
@@ -1565,6 +1601,20 @@ impl ServerActor {
 
         Ok((compact_device_queries, compact_device_sums))
     }
+
+    fn allocate_or_policy_bitmap(&mut self, bitmap: Vec<u64>) -> Vec<CudaSlice<u64>> {
+        let devices = self.device_manager.devices();
+
+        let mut or_policy_bitmap = Vec::with_capacity(devices.len());
+
+        for (device_idx, dev) in devices.iter().enumerate() {
+            // Transfer the bitmap to the device. It will be the same for each of the
+            // devices
+            let _bitmap = htod_on_stream_sync(&bitmap, dev, &self.streams[0][device_idx]).unwrap();
+            or_policy_bitmap.push(_bitmap);
+        }
+        or_policy_bitmap
+    }
 }
 
 /// Internal helper function to log the timers of measured cuda streams.
@@ -1844,3 +1894,25 @@ pub fn get_dummy_shares_for_deletion(
             .into();
     (iris_share, mask_share)
 }
+
+pub fn prepare_or_policy_bitmap(
+    max_db_size: usize,
+    or_rule_serial_ids: Vec<Vec<u32>>,
+    batch_size: usize,
+) -> Vec<u64> {
+    let row_stride64 = (max_db_size + 63) / 64;
+    let total_size = row_stride64 * batch_size;
+
+    // Create the bitmap on the host
+    let mut bitmap = vec![0u64; total_size];
+
+    for (query_idx, db_indices) in or_rule_serial_ids.iter().enumerate() {
+        for &db_idx in db_indices {
+            let row_start = query_idx * row_stride64;
+            let word_idx = row_start + (db_idx as usize / 64);
+            let bit_offset = db_idx as usize % 64;
+            bitmap[word_idx] |= 1 << bit_offset;
+        }
+    }
+    bitmap
+}

iris-mpc-gpu/src/server/mod.rs

@@ -2,7 +2,9 @@ mod actor;
 pub mod sync_nccl;
 
 use crate::dot::{share_db::preprocess_query, IRIS_CODE_LENGTH, MASK_CODE_LENGTH, ROTATIONS};
-pub use actor::{get_dummy_shares_for_deletion, ServerActor, ServerActorHandle};
+pub use actor::{
+    get_dummy_shares_for_deletion, prepare_or_policy_bitmap, ServerActor, ServerActorHandle,
+};
 use iris_mpc_common::galois_engine::degree4::{
     GaloisRingIrisCodeShare, GaloisRingTrimmedMaskCodeShare,
 };
@@ -79,6 +81,7 @@ pub struct BatchQuery {
     pub db_right_preprocessed:      BatchQueryEntriesPreprocessed,
     pub deletion_requests_indices:  Vec<u32>, // 0-indexed indicies in of entries to be deleted
     pub deletion_requests_metadata: Vec<BatchMetadata>,
+    pub or_rule_serial_ids:         Vec<Vec<u32>>,
     pub valid_entries:              Vec<bool>,
 }
 
@@ -124,6 +127,7 @@ impl BatchQuery {
         filter_by_indices!(self.store_left.mask, indices_set);
         filter_by_indices!(self.store_right.code, indices_set);
         filter_by_indices!(self.store_right.mask, indices_set);
+        filter_by_indices!(self.or_rule_serial_ids, indices_set);
         filter_by_indices_with_rotations!(self.query_left.code, indices_set);
         filter_by_indices_with_rotations!(self.query_left.mask, indices_set);
         filter_by_indices_with_rotations!(self.db_left.code, indices_set);