Au/iris workers NUMA-awareness (POP-2820)

iris-mpc-cpu/src/execution/hawk_main.rs

@@ -21,7 +21,7 @@ use crate::{
 };
 use clap::Parser;
 use eyre::{eyre, Report, Result};
-use futures::try_join;
+use futures::{future::try_join_all, try_join};
 use intra_batch::intra_batch_is_match;
 use iris_mpc_common::job::Eye;
 use iris_mpc_common::{
@@ -630,6 +630,7 @@ impl HawkActor {
                     self.iris_store[0].write().await,
                     self.iris_store[1].write().await,
                 ],
+                iris_pools: self.workers_handle.clone(),
             },
             GraphLoader([
                 self.graph_store[0].write().await,
@@ -652,6 +653,7 @@ pub struct IrisLoader<'a> {
     party_id: usize,
     db_size: &'a mut usize,
     irises: BothEyes<Aby3SharedIrisesMut<'a>>,
+    iris_pools: BothEyes<IrisPoolHandle>,
 }
 
 #[allow(clippy::needless_lifetimes)]
@@ -665,14 +667,16 @@ impl<'a> InMemoryStore for IrisLoader<'a> {
         right_code: &[u16],
         right_mask: &[u16],
     ) {
-        for (side, code, mask) in izip!(
+        for (store, pool, code, mask) in izip!(
             &mut self.irises,
+            self.iris_pools.clone(),
             [left_code, right_code],
             [left_mask, right_mask]
         ) {
             let iris = GaloisRingSharedIris::try_from_buffers(self.party_id, code, mask)
                 .expect("Wrong code or mask size");
-            side.insert(vector_id, iris);
+            let iris = pool.numa_realloc_blocking(iris).unwrap();
+            store.insert(vector_id, iris);
         }
     }
 
@@ -846,6 +850,70 @@ impl From<BatchQuery> for HawkRequest {
 }
 
 impl HawkRequest {
+    async fn numa_realloc(self, workers: BothEyes<IrisPoolHandle>) -> Self {
+        // TODO: Result<Self>
+        let start = Instant::now();
+
+        let (queries, queries_mirror) = join!(
+            Self::numa_realloc_orient(self.queries, &workers),
+            Self::numa_realloc_orient(self.queries_mirror, &workers)
+        );
+
+        metrics::histogram!("numa_realloc_duration").record(start.elapsed().as_secs_f64());
+        Self {
+            batch: self.batch,
+            queries,
+            queries_mirror,
+            ids: self.ids,
+        }
+    }
+
+    async fn numa_realloc_orient(
+        queries: SearchQueries,
+        workers: &BothEyes<IrisPoolHandle>,
+    ) -> SearchQueries {
+        let (left, right) = join!(
+            Self::numa_realloc_side(&queries[LEFT], &workers[LEFT]),
+            Self::numa_realloc_side(&queries[RIGHT], &workers[RIGHT])
+        );
+        Arc::new([left, right])
+    }
+
+    async fn numa_realloc_side(
+        requests: &VecRequests<VecRots<Aby3Query>>,
+        worker: &IrisPoolHandle,
+    ) -> VecRequests<VecRots<Aby3Query>> {
+        // Iterate over all the irises.
+        let all_irises_iter = requests.iter().flat_map(|rots| {
+            rots.iter()
+                .flat_map(|query| [&query.iris, &query.iris_proc])
+        });
+
+        // Go realloc the irises in parallel.
+        let tasks = all_irises_iter.map(|iris| worker.numa_realloc(iris.clone()).unwrap());
+
+        // Iterate over the results in the same order.
+        let mut new_irises_iter = try_join_all(tasks).await.unwrap().into_iter();
+
+        // Rebuild the same structure with the new irises.
+        let new_requests = requests
+            .iter()
+            .map(|rots| {
+                rots.iter()
+                    .map(|_old_query| {
+                        let iris = new_irises_iter.next().unwrap();
+                        let iris_proc = new_irises_iter.next().unwrap();
+                        Aby3Query { iris, iris_proc }
+                    })
+                    .collect_vec()
+                    .into()
+            })
+            .collect_vec();
+
+        assert!(new_irises_iter.next().is_none());
+        new_requests
+    }
+
     fn request_types(
         &self,
         iris_store: &Aby3SharedIrises,
@@ -1314,6 +1382,10 @@ impl HawkHandle {
         tracing::info!("Processing an Hawk job…");
         let now = Instant::now();
 
+        let request = request
+            .numa_realloc(hawk_actor.workers_handle.clone())
+            .await;
+
         // Deletions.
         apply_deletions(hawk_actor, &request).await?;
 

iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs

@@ -23,6 +23,16 @@ use tracing::info;
 
 #[derive(Debug)]
 enum IrisTask {
+    /// Move an iris code to memory closer to the pool (NUMA-awareness).
+    Realloc {
+        iris: ArcIris,
+        rsp: oneshot::Sender<ArcIris>,
+    },
+    Insert {
+        vector_id: VectorId,
+        iris: ArcIris,
+        rsp: oneshot::Sender<VectorId>,
+    },
     DotProductPairs {
         pairs: Vec<(ArcIris, VectorId)>,
         rsp: oneshot::Sender<Vec<RingElement<u16>>>,
@@ -40,12 +50,34 @@ enum IrisTask {
 
 #[derive(Clone, Debug)]
 pub struct IrisPoolHandle {
-    workers: Arc<Vec<Sender<IrisTask>>>,
+    workers: Arc<[Sender<IrisTask>]>,
     next_counter: Arc<AtomicU64>,
     metric_latency: FastHistogram,
 }
 
 impl IrisPoolHandle {
+    pub fn numa_realloc(&self, iris: ArcIris) -> Result<oneshot::Receiver<ArcIris>> {
+        let (tx, rx) = oneshot::channel();
+        let task = IrisTask::Realloc { iris, rsp: tx };
+        self.get_next_worker().send(task)?;
+        Ok(rx)
+    }
+
+    pub fn numa_realloc_blocking(&self, iris: ArcIris) -> Result<ArcIris> {
+        Ok(self.numa_realloc(iris)?.blocking_recv()?)
+    }
+
+    pub async fn insert(&self, vector_id: VectorId, iris: ArcIris) -> Result<VectorId> {
+        let (tx, rx) = oneshot::channel();
+        let task = IrisTask::Insert {
+            vector_id,
+            iris,
+            rsp: tx,
+        };
+        self.get_next_worker().send(task)?;
+        Ok(rx.await?)
+    }
+
     pub async fn dot_product_pairs(
         &mut self,
         pairs: Vec<(ArcIris, VectorId)>,
@@ -85,7 +117,7 @@ impl IrisPoolHandle {
     ) -> Result<Vec<RingElement<u16>>> {
         let start = Instant::now();
 
-        let _ = self.get_next_worker().send(task);
+        self.get_next_worker().send(task)?;
         let res = rx.await?;
 
         self.metric_latency.record(start.elapsed().as_secs_f64());
@@ -121,7 +153,7 @@ pub fn init_workers(shard_index: usize, iris_store: SharedIrisesRef<ArcIris>) ->
     }
 
     IrisPoolHandle {
-        workers: Arc::new(channels),
+        workers: channels.into(),
         next_counter: Arc::new(AtomicU64::new(0)),
         metric_latency: FastHistogram::new("iris_worker.latency"),
     }
@@ -130,6 +162,23 @@ pub fn init_workers(shard_index: usize, iris_store: SharedIrisesRef<ArcIris>) ->
 fn worker_thread(ch: Receiver<IrisTask>, iris_store: SharedIrisesRef<ArcIris>) {
     while let Ok(task) = ch.recv() {
         match task {
+            IrisTask::Realloc { iris, rsp } => {
+                // Re-allocate from this thread.
+                // This attempts to use the NUMA-aware first-touch policy of the OS.
+                let new_iris = Arc::new((*iris).clone());
+                let _ = rsp.send(new_iris);
+            }
+
+            IrisTask::Insert {
+                vector_id,
+                iris,
+                rsp,
+            } => {
+                let mut store = iris_store.data.blocking_write();
+                let vector_id = store.insert(vector_id, iris);
+                let _ = rsp.send(vector_id);
+            }
+
             IrisTask::DotProductPairs { pairs, rsp } => {
                 let store = iris_store.data.blocking_read();
 
@@ -167,7 +216,8 @@ fn worker_thread(ch: Receiver<IrisTask>, iris_store: SharedIrisesRef<ArcIris>) {
 const SHARD_COUNT: usize = 2;
 
 pub fn select_core_ids(shard_index: usize) -> Vec<CoreId> {
-    let core_ids = core_affinity::get_core_ids().unwrap();
+    let mut core_ids = core_affinity::get_core_ids().unwrap();
+    core_ids.sort();
     assert!(!core_ids.is_empty());
 
     let shard_count = cmp::min(SHARD_COUNT, core_ids.len());