Fused mirror scan, packed pair kernel, and lane pipelining for the CPU linear scan

.github/workflows/build-and-push-linear-scan-server.yaml

@@ -32,9 +32,11 @@ jobs:
         uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8
 
       - name: Install Docker
-        uses: docker/setup-docker-action@e43656e248c0bd0647d3f5c195d116aacf6fcaf4
-        with:
-          version: v29.7.2
+        run: |
+          curl -fsSL https://get.docker.com | sh
+          sudo usermod -aG docker $USER
+          sudo apt-get install acl
+          sudo setfacl --modify user:$USER:rw /var/run/docker.sock
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435

iris-mpc-cpu/Cargo.toml

@@ -99,6 +99,14 @@ harness = false
 name = "networking"
 harness = false
 
+[[bench]]
+name = "linear_scan_dot_cpu"
+harness = false
+
+[[bench]]
+name = "linear_scan_candidates"
+harness = false
+
 [[bench]]
 name = "set_hash"
 harness = false

iris-mpc-cpu/benches/dot.rs

@@ -632,13 +632,18 @@ pub fn bench_worker_pool(c: &mut Criterion) {
     let num_iris_codes = iris_codes.len();
     let dist = Uniform::new(0, num_iris_codes);
 
-    let points_map: HashMap<VectorId, Arc<GaloisRingSharedIris>> = HashMap::new();
+    let layout = iris_mpc_cpu::protocol::shared_iris::ResidentLayout::U16;
+    let points_map: HashMap<VectorId, iris_mpc_cpu::protocol::shared_iris::ResidentIris> =
+        HashMap::new();
     let shared_irises = SharedIrises::new(
         points_map,
-        Arc::new(GaloisRingSharedIris::default_for_party(0)),
+        iris_mpc_cpu::protocol::shared_iris::ResidentIris::from_arc(
+            Arc::new(GaloisRingSharedIris::default_for_party(0)),
+            layout,
+        ),
     )
     .to_arc();
-    let pool = init_workers(0, shared_irises, true);
+    let pool = init_workers(0, shared_irises, true, layout);
 
     // similar to numa_realloc
     for (idx, iris) in iris_codes.iter().enumerate() {

iris-mpc-cpu/benches/linear_scan_candidates.rs

@@ -0,0 +1,75 @@
+use iris_mpc_common::VectorId;
+use std::{
+    collections::HashSet,
+    env,
+    hint::black_box,
+    mem::size_of,
+    time::{Duration, Instant},
+};
+
+fn env_usize(name: &str, default: usize) -> usize {
+    env::var(name)
+        .ok()
+        .and_then(|value| value.parse().ok())
+        .unwrap_or(default)
+}
+
+fn main() {
+    let records = env_usize("IRIS_MPC_CANDIDATE_BENCH_RECORDS", 2_000_000);
+    let candidate_count = env_usize("IRIS_MPC_CANDIDATE_BENCH_CANDIDATES", 128);
+    let repetitions = env_usize("IRIS_MPC_CANDIDATE_BENCH_REPETITIONS", 5);
+    assert!(records > 0 && records <= u32::MAX as usize);
+    assert!(repetitions > 0);
+
+    let live_ids = (0..records)
+        .map(|index| VectorId::from_0_index(index as u32))
+        .collect::<Vec<_>>();
+    let candidates = (0..candidate_count)
+        .map(|index| {
+            let record = index * records / candidate_count.max(1);
+            VectorId::from_0_index(record.min(records - 1) as u32)
+        })
+        .collect::<Vec<_>>();
+
+    let mut hash_build = Duration::ZERO;
+    let mut hash_lookup = Duration::ZERO;
+    let mut hash_capacity = 0usize;
+    for _ in 0..repetitions {
+        let started = Instant::now();
+        let live_ids_set = live_ids.iter().copied().collect::<HashSet<_>>();
+        hash_build += started.elapsed();
+        hash_capacity = live_ids_set.capacity();
+
+        let started = Instant::now();
+        let found = candidates
+            .iter()
+            .filter(|id| live_ids_set.contains(id))
+            .count();
+        hash_lookup += started.elapsed();
+        black_box(found);
+    }
+
+    let mut binary_lookup = Duration::ZERO;
+    for _ in 0..repetitions {
+        let started = Instant::now();
+        let found = candidates
+            .iter()
+            .filter(|id| live_ids.binary_search(id).is_ok())
+            .count();
+        binary_lookup += started.elapsed();
+        black_box(found);
+    }
+
+    let divisor = repetitions as f64;
+    println!(
+        "LINEAR_SCAN_CANDIDATE_RESULT records={records} candidates={} repetitions={repetitions} \
+         vector_id_bytes={} hash_capacity={hash_capacity} hash_table_lower_bound_bytes={} \
+         hash_build_seconds={:.9} hash_lookup_seconds={:.9} binary_lookup_seconds={:.9}",
+        candidates.len(),
+        size_of::<VectorId>(),
+        hash_capacity * size_of::<VectorId>(),
+        hash_build.as_secs_f64() / divisor,
+        hash_lookup.as_secs_f64() / divisor,
+        binary_lookup.as_secs_f64() / divisor,
+    );
+}

iris-mpc-cpu/benches/linear_scan_dot_cpu.rs

@@ -0,0 +1,219 @@
+use eyre::{ensure, Result};
+use iris_mpc_common::{VectorId, IRIS_CODE_LENGTH, MASK_CODE_LENGTH, ROTATIONS};
+use iris_mpc_cpu::{
+    execution::hawk_main::iris_worker::{
+        init_workers, IrisWorkerPool, LocalIrisWorkerPool, QueryId, QuerySpec,
+    },
+    hawkers::{aby3::aby3_store::DistanceMode, shared_irises::SharedIrises},
+    protocol::shared_iris::{ArcIris, GaloisRingSharedIris},
+};
+use rayon::prelude::*;
+use std::{
+    collections::HashMap,
+    env,
+    hint::black_box,
+    sync::Arc,
+    time::{Duration, Instant},
+};
+
+const DEFAULT_DB_SIZE: usize = 1_572_864;
+const DEFAULT_WARMUP_RUNS: usize = 1;
+const DEFAULT_MEASURED_RUNS: usize = 3;
+const DEFAULT_TOKIO_CORES: usize = 8;
+
+fn env_usize(name: &str, default: usize) -> Result<usize> {
+    match env::var(name) {
+        Ok(value) => Ok(value.parse()?),
+        Err(env::VarError::NotPresent) => Ok(default),
+        Err(err) => Err(err.into()),
+    }
+}
+
+fn build_pool(db_size: usize, numa_shard: usize) -> (LocalIrisWorkerPool, Vec<VectorId>, ArcIris) {
+    let query = Arc::new(GaloisRingSharedIris::default_for_party(0));
+    println!(
+        "BENCH_LOADING db_size={db_size} payload_gib={:.3}",
+        db_size as f64 * 2.0 * (IRIS_CODE_LENGTH + MASK_CODE_LENGTH) as f64
+            / (1024.0 * 1024.0 * 1024.0),
+    );
+    let load_started = Instant::now();
+
+    // Deep-clone every record so the scan reads a production-sized working
+    // set instead of repeatedly hitting one shared cache-resident allocation.
+    let irises = (0..db_size)
+        .into_par_iter()
+        .map(|_| Arc::new((*query).clone()))
+        .collect::<Vec<_>>();
+    let layout = iris_mpc_cpu::protocol::shared_iris::preferred_scan_layout();
+    let mut store = SharedIrises::new(
+        HashMap::new(),
+        iris_mpc_cpu::protocol::shared_iris::ResidentIris::from_arc(
+            Arc::new(GaloisRingSharedIris::default_for_party(0)),
+            layout,
+        ),
+    );
+    store.reserve(db_size);
+    let mut vector_ids = Vec::with_capacity(db_size);
+    for iris in irises {
+        vector_ids.push(
+            store.append(iris_mpc_cpu::protocol::shared_iris::ResidentIris::from_arc(
+                iris, layout,
+            )),
+        );
+    }
+    println!(
+        "BENCH_LOADED seconds={:.3}",
+        load_started.elapsed().as_secs_f64()
+    );
+
+    let store = store.to_arc();
+    let workers = init_workers(numa_shard, store.clone(), true, layout);
+    (
+        LocalIrisWorkerPool::new(workers, store, layout, DistanceMode::MinRotation, 0),
+        vector_ids,
+        query,
+    )
+}
+
+fn median(samples: &[Duration]) -> Duration {
+    let mut sorted = samples.to_vec();
+    sorted.sort_unstable();
+    sorted[sorted.len() / 2]
+}
+
+fn report(backend: &str, db_size: usize, orientations: usize, samples: &[Duration]) {
+    let elapsed = median(samples).as_secs_f64();
+    let comparisons = db_size * orientations;
+    println!(
+        "BENCH_RESULT backend={backend} db_size={db_size} orientations={orientations} \
+         median_seconds={elapsed:.6} logical_comparisons_per_second={:.3} \
+         full_scan_equivalent_seconds={:.6}",
+        comparisons as f64 / elapsed,
+        elapsed,
+    );
+}
+
+fn main() -> Result<()> {
+    let db_size = env_usize("IRIS_MPC_DOT_BENCH_DB_SIZE", DEFAULT_DB_SIZE)?;
+    let numa_shard = env_usize("IRIS_MPC_DOT_BENCH_NUMA_SHARD", 0)?;
+    let warmup_runs = env_usize("IRIS_MPC_DOT_BENCH_WARMUP", DEFAULT_WARMUP_RUNS)?;
+    let measured_runs = env_usize("IRIS_MPC_DOT_BENCH_RUNS", DEFAULT_MEASURED_RUNS)?;
+    let tokio_cores = env_usize("IRIS_MPC_DOT_BENCH_TOKIO_CORES", DEFAULT_TOKIO_CORES)?;
+    let tokio_cores = (tokio_cores > 0).then_some(tokio_cores);
+    ensure!(db_size > 0, "DB size must be positive");
+    ensure!(measured_runs > 0, "measured run count must be positive");
+    iris_mpc_common::helpers::numactl::init(tokio_cores);
+    println!(
+        "BENCH_CONFIG batch_size=1 db_size={db_size} numa_shard={numa_shard} \
+         warmup_runs={warmup_runs} measured_runs={measured_runs} rotations={ROTATIONS} \
+         tokio_cores={} dot_cores={}",
+        tokio_cores.unwrap_or_else(iris_mpc_common::helpers::numactl::get_tokio_worker_threads),
+        iris_mpc_cpu::execution::hawk_main::iris_worker::select_core_ids(numa_shard).len(),
+    );
+
+    let (pool, vector_ids, query) = build_pool(db_size, numa_shard);
+    iris_mpc_common::helpers::numactl::restrict_tokio_runtime();
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(iris_mpc_common::helpers::numactl::get_tokio_worker_threads())
+        .on_thread_start(iris_mpc_common::helpers::numactl::restrict_tokio_runtime)
+        .enable_all()
+        .build()?;
+    let normal_id = QueryId::new();
+    let mirror_id = QueryId::new();
+    runtime.block_on(pool.cache_queries(vec![(normal_id, query.clone()), (mirror_id, query)]))?;
+    let normal = QuerySpec::new(normal_id);
+    let mirror = QuerySpec::new(mirror_id);
+    let expected_len = 2 * ROTATIONS * db_size;
+
+    let run_single = || -> Result<Duration> {
+        let started = Instant::now();
+        let output = runtime
+            .block_on(pool.compute_dot_products_full_rotations(normal, vector_ids.clone()))?;
+        let elapsed = started.elapsed();
+        ensure!(output.len() == expected_len, "unexpected result length");
+        black_box(output);
+        Ok(elapsed)
+    };
+    for _ in 0..warmup_runs {
+        black_box(run_single()?);
+    }
+    let mut single_samples = Vec::with_capacity(measured_runs);
+    for run in 0..measured_runs {
+        let elapsed = run_single()?;
+        println!(
+            "BENCH_SAMPLE backend=cpu_single_orientation run={run} seconds={:.6}",
+            elapsed.as_secs_f64()
+        );
+        single_samples.push(elapsed);
+    }
+
+    let run_mirror = || -> Result<Duration> {
+        let started = Instant::now();
+        let (normal_output, mirror_output) = runtime.block_on(async {
+            tokio::join!(
+                pool.compute_dot_products_full_rotations(normal, vector_ids.clone()),
+                pool.compute_dot_products_full_rotations(mirror, vector_ids.clone()),
+            )
+        });
+        let normal_output = normal_output?;
+        let mirror_output = mirror_output?;
+        let elapsed = started.elapsed();
+        ensure!(
+            normal_output.len() == expected_len && mirror_output.len() == expected_len,
+            "unexpected mirror result length"
+        );
+        black_box((normal_output, mirror_output));
+        Ok(elapsed)
+    };
+    for _ in 0..warmup_runs {
+        black_box(run_mirror()?);
+    }
+    let mut mirror_samples = Vec::with_capacity(measured_runs);
+    for run in 0..measured_runs {
+        let elapsed = run_mirror()?;
+        println!(
+            "BENCH_SAMPLE backend=cpu_normal_and_mirror run={run} seconds={:.6}",
+            elapsed.as_secs_f64()
+        );
+        mirror_samples.push(elapsed);
+    }
+
+    let run_fused = || -> Result<Duration> {
+        let started = Instant::now();
+        let [normal_output, mirror_output] = runtime.block_on(
+            pool.compute_dot_products_full_rotations_pair([normal, mirror], vector_ids.clone()),
+        )?;
+        let elapsed = started.elapsed();
+        ensure!(
+            normal_output.len() == expected_len && mirror_output.len() == expected_len,
+            "unexpected fused result length"
+        );
+        black_box((normal_output, mirror_output));
+        Ok(elapsed)
+    };
+    for _ in 0..warmup_runs {
+        black_box(run_fused()?);
+    }
+    let mut fused_samples = Vec::with_capacity(measured_runs);
+    for run in 0..measured_runs {
+        let elapsed = run_fused()?;
+        println!(
+            "BENCH_SAMPLE backend=cpu_fused_pair run={run} seconds={:.6}",
+            elapsed.as_secs_f64()
+        );
+        fused_samples.push(elapsed);
+    }
+
+    report("cpu_single_orientation", db_size, 1, &single_samples);
+    report("cpu_normal_and_mirror", db_size, 2, &mirror_samples);
+    report("cpu_fused_pair", db_size, 2, &fused_samples);
+    println!(
+        "BENCH_MIRROR_RATIO wall_time_ratio={:.3}",
+        median(&mirror_samples).as_secs_f64() / median(&single_samples).as_secs_f64()
+    );
+    println!(
+        "BENCH_FUSED_RATIO wall_time_ratio={:.3}",
+        median(&fused_samples).as_secs_f64() / median(&single_samples).as_secs_f64()
+    );
+    Ok(())
+}