[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::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,56 @@ impl DistanceComparator {
         }
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub fn join_db_matches_with_bitmaps(
+        &self,
+        matches_bitmap_left: &[CudaSlice<u64>],
+        matches_bitmap_right: &[CudaSlice<u64>],
+        final_results: &[CudaSlice<u32>],
+        db_sizes: &[usize],
+        streams: &[CudaStream],
+        kernels: &[CudaFunction],
+        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],
+            );
+
+            unsafe {
+                kernels[i]
+                    .clone()
+                    .launch_on_stream(
+                        &streams[i],
+                        cfg,
+                        (
+                            &matches_bitmap_left[i],
+                            &matches_bitmap_right[i],
+                            &final_results[i],
+                            (self.query_length / ROTATIONS) as u64,
+                            db_sizes[i] as u64,
+                            &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>],
@@ -347,4 +405,16 @@ impl DistanceComparator {
             })
             .collect::<Vec<_>>()
     }
+
+    pub fn prepare_luc_bitmap(&self, db_size: usize) -> Vec<CudaSlice<u64>> {
+        let row_stride64 = (db_size + 63) / 64;
+        (0..self.device_manager.device_count())
+            .map(|i| {
+                self.device_manager
+                    .device(i)
+                    .alloc_zeros(row_stride64 * self.query_length / ROTATIONS)
+                    .unwrap()
+            })
+            .collect::<Vec<_>>()
+    }
 }

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

@@ -93,6 +93,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, 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 = (dbLength + 63) / 64;
+
+    size_t totalBits   = dbLength * queryLength;
+    size_t numElements = (totalBits + 63) / 64; //  div_ceil
+
+    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;
+            bool useOr = (orPolicyBitmap[rowIndex + (dbIdx / 64)]
+                          & (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,
         },
@@ -31,7 +32,7 @@ use iris_mpc_common::{
     iris_db::iris::IrisCode,
     IrisCodeDbSlice,
 };
-use itertools::Itertools;
+use itertools::{izip, Itertools};
 use rand::{rngs::StdRng, SeedableRng};
 use ring::hkdf::{Algorithm, Okm, Salt, HKDF_SHA256};
 use std::{collections::HashMap, mem, sync::Arc, time::Instant};
@@ -762,6 +763,22 @@ impl ServerActor {
         ///////////////////////////////////////////////////////////////////
         // MERGE LEFT & RIGHT results
         ///////////////////////////////////////////////////////////////////
+
+        // Initialize bitmap with OR rule, if exists
+        if let Some(or_rule_bitmap) = batch.or_rule_serial_ids {
+            // Populate the pre-allocated OR policy bitmap with the serial ids
+            let or_policy_bitmap = self.prepare_or_policy_bitmap(or_rule_bitmap, batch_size);
+            self.distance_comparator.join_db_matches_with_bitmaps(
+                &self.db_match_list_left,
+                &self.db_match_list_right,
+                &self.final_results,
+                &self.current_db_sizes,
+                &self.streams[0],
+                &self.distance_comparator.merge_batch_with_bitmap_kernels,
+                &or_policy_bitmap,
+            );
+        }
+
         tracing::info!("Joining both sides");
         // Merge results and fetch matching indices
         // Format: host_results[device_index][query_index]
@@ -1565,6 +1582,44 @@ impl ServerActor {
 
         Ok((compact_device_queries, compact_device_sums))
     }
+
+    fn prepare_or_policy_bitmap(
+        &mut self,
+        or_rule_serial_ids: Vec<Vec<u32>>,
+        batch_size: usize,
+    ) -> Vec<CudaSlice<u64>> {
+        let streams = self
+            .device_manager
+            .devices()
+            .iter()
+            .map(|dev| dev.fork_default_stream().unwrap())
+            .collect::<Vec<_>>();
+
+        let devices = self.device_manager.devices();
+
+        let mut or_policy_bitmap = Vec::with_capacity(devices.len());
+
+        for (dev, stream) in izip!(devices, streams) {
+            let row_stride64 = (self.max_db_size + 63) / 64;
+            let total_size = row_stride64 * batch_size / ROTATIONS;
+
+            // 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;
+                }
+            }
+            // Transfer the bitmap to the device
+            let _bitmap = htod_on_stream_sync(&bitmap, dev, &stream).unwrap();
+            or_policy_bitmap.push(_bitmap);
+        }
+        or_policy_bitmap
+    }
 }
 
 /// Internal helper function to log the timers of measured cuda streams.

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

@@ -79,6 +79,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:         Option<Vec<Vec<u32>>>,
     pub valid_entries:              Vec<bool>,
 }