Feat: even odd import from s3 (#897)

* import bytes in even-odd orders from s3

* trigger image push

* deploy to stage

* fix code len assert

* remove leftover change on loading from db

* change even/odd order

* bring back iris-mpc-py to workspace

* remove unused method

* simplify size

* move invalid id check to pre-load

* bump stage version

Author: eaypek-tfh

.github/workflows/temp-branch-build-and-push.yaml

@@ -3,7 +3,7 @@ name: Branch - Build and push docker image
 on:
   push:
     branches:
-      - "ps/potential-phantom-match"
+      - "feat/even-odd-import-from-s3"
 
 concurrency:
   group: '${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}'

deploy/stage/common-values-iris-mpc.yaml

@@ -1,4 +1,4 @@
-image: "ghcr.io/worldcoin/iris-mpc:v0.13.10"
+image: "ghcr.io/worldcoin/iris-mpc:v0.13.11"
 
 environment: stage
 replicaCount: 1

deploy/stage/smpcv2-0-stage/values-iris-mpc.yaml

@@ -75,7 +75,7 @@ env:
     value: "iris-mpc-db-exporter-store-node-0-stage-eu-north-1"
 
   - name: SMPC__DB_CHUNKS_FOLDER_NAME
-    value: "binary_output_16k"
+    value: "even_odd_binary_output_16k"
 
   - name: SMPC__LOAD_CHUNKS_PARALLELISM
     value: "32"

deploy/stage/smpcv2-1-stage/values-iris-mpc.yaml

@@ -75,7 +75,7 @@ env:
     value: "iris-mpc-db-exporter-store-node-1-stage-eu-north-1"
 
   - name: SMPC__DB_CHUNKS_FOLDER_NAME
-    value: "binary_output_16k"
+    value: "even_odd_binary_output_16k"
 
   - name: SMPC__LOAD_CHUNKS_PARALLELISM
     value: "32"

deploy/stage/smpcv2-2-stage/values-iris-mpc.yaml

@@ -75,7 +75,7 @@ env:
     value: "iris-mpc-db-exporter-store-node-2-stage-eu-north-1"
 
   - name: SMPC__DB_CHUNKS_FOLDER_NAME
-    value: "binary_output_16k"
+    value: "even_odd_binary_output_16k"
 
   - name: SMPC__LOAD_CHUNKS_PARALLELISM
     value: "32"

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

@@ -302,7 +302,7 @@ impl ShareDB {
         }
     }
 
-    pub fn load_single_record(
+    pub fn load_single_record_from_db(
         index: usize,
         db: &CudaVec2DSlicerRawPointer,
         record: &[u16],
@@ -339,6 +339,35 @@ impl ShareDB {
         };
     }
 
+    pub fn load_single_record_from_s3(
+        index: usize,
+        db: &CudaVec2DSlicerRawPointer,
+        a0_host: &[u8],
+        a1_host: &[u8],
+        n_shards: usize,
+        code_length: usize,
+    ) {
+        assert_eq!(a0_host.len(), code_length);
+        assert_eq!(a1_host.len(), code_length);
+
+        let device_index = index % n_shards;
+        let device_db_index = index / n_shards;
+
+        unsafe {
+            std::ptr::copy(
+                a0_host.as_ptr() as *const _,
+                (db.limb_0[device_index] + (device_db_index * code_length) as u64) as *mut _,
+                code_length,
+            );
+
+            std::ptr::copy(
+                a1_host.as_ptr() as *const _,
+                (db.limb_1[device_index] + (device_db_index * code_length) as u64) as *mut _,
+                code_length,
+            );
+        };
+    }
+
     pub fn preprocess_db(&self, db: &mut SlicedProcessedDatabase, db_lens: &[usize]) {
         let code_len = self.code_length;
         for device_index in 0..self.device_manager.device_count() {
@@ -381,7 +410,7 @@ impl ShareDB {
             .par_chunks(self.code_length)
             .enumerate()
             .for_each(|(idx, chunk)| {
-                Self::load_single_record(idx, &db.code_gr, chunk, n_shards, code_length);
+                Self::load_single_record_from_db(idx, &db.code_gr, chunk, n_shards, code_length);
             });
 
         // Calculate the number of entries per shard

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

@@ -461,36 +461,36 @@ impl ServerActor {
         self.current_db_sizes = db_lens1;
     }
 
-    pub fn load_single_record(
+    pub fn load_single_record_from_db(
         &mut self,
         index: usize,
         left_code: &[u16],
         left_mask: &[u16],
         right_code: &[u16],
         right_mask: &[u16],
     ) {
-        ShareDB::load_single_record(
+        ShareDB::load_single_record_from_db(
             index,
             &self.left_code_db_slices.code_gr,
             left_code,
             self.device_manager.device_count(),
             IRIS_CODE_LENGTH,
         );
-        ShareDB::load_single_record(
+        ShareDB::load_single_record_from_db(
             index,
             &self.left_mask_db_slices.code_gr,
             left_mask,
             self.device_manager.device_count(),
             MASK_CODE_LENGTH,
         );
-        ShareDB::load_single_record(
+        ShareDB::load_single_record_from_db(
             index,
             &self.right_code_db_slices.code_gr,
             right_code,
             self.device_manager.device_count(),
             IRIS_CODE_LENGTH,
         );
-        ShareDB::load_single_record(
+        ShareDB::load_single_record_from_db(
             index,
             &self.right_mask_db_slices.code_gr,
             right_mask,
@@ -499,6 +499,53 @@ impl ServerActor {
         );
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub fn load_single_record_from_s3(
+        &mut self,
+        index: usize,
+        left_code_odd: &[u8],
+        left_code_even: &[u8],
+        right_code_odd: &[u8],
+        right_code_even: &[u8],
+        left_mask_odd: &[u8],
+        left_mask_even: &[u8],
+        right_mask_odd: &[u8],
+        right_mask_even: &[u8],
+    ) {
+        ShareDB::load_single_record_from_s3(
+            index,
+            &self.left_code_db_slices.code_gr,
+            left_code_odd,
+            left_code_even,
+            self.device_manager.device_count(),
+            IRIS_CODE_LENGTH,
+        );
+        ShareDB::load_single_record_from_s3(
+            index,
+            &self.left_mask_db_slices.code_gr,
+            left_mask_odd,
+            left_mask_even,
+            self.device_manager.device_count(),
+            MASK_CODE_LENGTH,
+        );
+        ShareDB::load_single_record_from_s3(
+            index,
+            &self.right_code_db_slices.code_gr,
+            right_code_odd,
+            right_code_even,
+            self.device_manager.device_count(),
+            IRIS_CODE_LENGTH,
+        );
+        ShareDB::load_single_record_from_s3(
+            index,
+            &self.right_mask_db_slices.code_gr,
+            right_mask_odd,
+            right_mask_even,
+            self.device_manager.device_count(),
+            MASK_CODE_LENGTH,
+        );
+    }
+
     pub fn increment_db_size(&mut self, index: usize) {
         self.current_db_sizes[index % self.device_manager.device_count()] += 1;
     }