From b56b88b75a75b625f0fcaa9af444a4de073a7158 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Wed, 19 Aug 2026 13:05:21 +0200 Subject: [PATCH 01/14] perf(linear-scan): add compute and pipeline optimizations --- .../build-and-push-linear-scan-server.yaml | 8 +- iris-mpc-cpu/Cargo.toml | 8 + iris-mpc-cpu/benches/dot.rs | 11 +- .../benches/linear_scan_candidates.rs | 75 + iris-mpc-cpu/benches/linear_scan_dot_cpu.rs | 219 +++ iris-mpc-cpu/src/execution/hawk_main.rs | 132 +- .../src/execution/hawk_main/iris_worker.rs | 430 +++++- iris-mpc-cpu/src/execution/hawk_main/rot.rs | 15 + .../src/execution/hawk_main/search.rs | 1302 ++++++++++++++-- .../hawk_main/worker_pool_initializer.rs | 53 +- iris-mpc-cpu/src/genesis/batch_generator.rs | 15 +- iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs | 326 +++- .../src/hawkers/aby3/aby3_store/tests.rs | 138 ++ iris-mpc-cpu/src/hawkers/aby3/test_utils.rs | 21 +- iris-mpc-cpu/src/hawkers/shared_irises.rs | 17 + iris-mpc-cpu/src/protocol/ops.rs | 1315 ++++++++++++++++- iris-mpc-cpu/src/protocol/shared_iris.rs | 145 ++ iris-mpc-cpu/tests/cold_eye.rs | 156 +- iris-mpc-cpu/tests/full_rotation_dot.rs | 18 +- iris-mpc-gpu/Cargo.toml | 5 + iris-mpc-gpu/benches/linear_scan_dot.rs | 559 +++++++ iris-mpc-utils/src/client/options/types.rs | 199 ++- iris-mpc/src/server/mod.rs | 56 +- iris-mpc/src/services/processors/job.rs | 92 +- .../analyze-linear-scan-server-benchmark.py | 246 +++ scripts/init-moto-linear-scan-benchmark.py | 116 ++ scripts/moto-server-with-sns-sequence.py | 2 +- .../run-distributed-linear-scan-benchmark.sh | 281 ++++ scripts/run-distributed-linear-scan-node.sh | 378 +++++ scripts/run-native-ground-truth.sh | 6 +- scripts/tools/pin-network-irqs.sh | 125 ++ 31 files changed, 5831 insertions(+), 638 deletions(-) create mode 100644 iris-mpc-cpu/benches/linear_scan_candidates.rs create mode 100644 iris-mpc-cpu/benches/linear_scan_dot_cpu.rs create mode 100644 iris-mpc-gpu/benches/linear_scan_dot.rs create mode 100755 scripts/analyze-linear-scan-server-benchmark.py create mode 100755 scripts/init-moto-linear-scan-benchmark.py create mode 100755 scripts/run-distributed-linear-scan-benchmark.sh create mode 100755 scripts/run-distributed-linear-scan-node.sh create mode 100755 scripts/tools/pin-network-irqs.sh diff --git a/.github/workflows/build-and-push-linear-scan-server.yaml b/.github/workflows/build-and-push-linear-scan-server.yaml index d56a83560e..d7a1583ac2 100644 --- a/.github/workflows/build-and-push-linear-scan-server.yaml +++ b/.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 diff --git a/iris-mpc-cpu/Cargo.toml b/iris-mpc-cpu/Cargo.toml index 31fe30a353..37f99a9f1d 100644 --- a/iris-mpc-cpu/Cargo.toml +++ b/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 diff --git a/iris-mpc-cpu/benches/dot.rs b/iris-mpc-cpu/benches/dot.rs index 31c91e3197..d0fdce4ace 100644 --- a/iris-mpc-cpu/benches/dot.rs +++ b/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> = HashMap::new(); + let layout = iris_mpc_cpu::protocol::shared_iris::ResidentLayout::U16; + let points_map: HashMap = + 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() { diff --git a/iris-mpc-cpu/benches/linear_scan_candidates.rs b/iris-mpc-cpu/benches/linear_scan_candidates.rs new file mode 100644 index 0000000000..3f5c2703b3 --- /dev/null +++ b/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::>(); + 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::>(); + + 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::>(); + 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::(), + hash_capacity * size_of::(), + hash_build.as_secs_f64() / divisor, + hash_lookup.as_secs_f64() / divisor, + binary_lookup.as_secs_f64() / divisor, + ); +} diff --git a/iris-mpc-cpu/benches/linear_scan_dot_cpu.rs b/iris-mpc-cpu/benches/linear_scan_dot_cpu.rs new file mode 100644 index 0000000000..1463120458 --- /dev/null +++ b/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 { + 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, 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::>(); + 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 { + 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 { + 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 { + 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(()) +} diff --git a/iris-mpc-cpu/src/execution/hawk_main.rs b/iris-mpc-cpu/src/execution/hawk_main.rs index 4fa437c07f..4074358eeb 100644 --- a/iris-mpc-cpu/src/execution/hawk_main.rs +++ b/iris-mpc-cpu/src/execution/hawk_main.rs @@ -122,7 +122,7 @@ use matching::{ use rand::{thread_rng, Rng, SeedableRng}; use rand_chacha::ChaCha8Rng; use scheduler::parallelize; -use search::{SearchParams, SearchQueries}; +use search::{SearchParams, SearchQueries, SearchResults}; use serde::{Deserialize, Serialize}; use session_groups::SessionGroups; use siphasher::sip::SipHasher13; @@ -667,7 +667,8 @@ impl HawkActor { args.party_index, HAWK_DISTANCE_MODE, args.numa, - ), + ) + .with_resident_layout(crate::protocol::shared_iris::preferred_scan_layout()), ); let networking = build_hawk_network_handle(args, shutdown_ct).await?; let initialized = initializer.initialize().await?; @@ -717,7 +718,8 @@ impl HawkActor { HAWK_DISTANCE_MODE, args.numa, iris_store, - ), + ) + .with_resident_layout(crate::protocol::shared_iris::preferred_scan_layout()), ); let networking = build_hawk_network_handle(args, shutdown_ct).await?; let initialized = initializer.initialize().await?; @@ -2004,6 +2006,95 @@ pub struct HawkHandle { job_queue: mpsc::Sender, } +/// Run the exact linear scan for both orientations with a fused first-eye +/// pass. Gathers each orientation's public candidate inputs (LUC, reauth +/// targets) exactly as the per-orientation search does, then evaluates both +/// orientations while streaming the resident eye once. Results per +/// orientation are identical to independent [`search::linear_scan_cascade`] +/// calls. +async fn linear_scan_search_both_orientations( + hawk_actor: &HawkActor, + sessions: &SessionGroups, + request: &HawkRequest, +) -> Result<[SearchResults; 2]> { + use Orientation::{Mirror, Normal}; + + let (forced_anon_stats_ids, extra_candidate_ids_both) = { + // Choice of LEFT registry is arbitrary — both sides are in sync + // w.r.t. stored vector ids. + let reg = hawk_actor.registry[LEFT].read().await; + let forced_anon_stats_ids = request + .batch + .request_types + .iter() + .enumerate() + .map(|(index, request_type)| { + if request_type.as_str() != REAUTH_MESSAGE_TYPE { + return Vec::new(); + } + let request_id = &request.batch.request_ids[index]; + request + .batch + .reauth_target_indices + .get(request_id) + .map(|&target| reg.from_0_indices(&[target])[0]) + .into_iter() + .collect() + }) + .collect::>(); + let luc_ids = request.luc_ids(®); + let extra_candidate_ids_both = [Normal, Mirror].map(|orient| { + let request_types = request.request_types(®, orient); + izip!(&luc_ids, &request_types) + .map(|(luc, request_type)| { + let mut ids = luc.clone(); + if let RequestType::Reauth { + target: Some((target, _)), + } = request_type + { + ids.push(*target); + } + ids.sort_unstable(); + ids.dedup(); + ids + }) + .collect_vec() + }); + (forced_anon_stats_ids, extra_candidate_ids_both) + }; + + let search_params = [Normal, Mirror].map(|orient| { + #[cfg(not(feature = "phase_trace"))] + let _ = orient; + SearchParams::new( + hawk_actor.searcher(), + hawk_actor.search_mode(), + true, + Some(hawk_actor.args.hnsw_param_ef_supermatch), + hawk_actor.args.hnsw_param_ef_saturation_margin, + hawk_actor.args.return_partial_results, + #[cfg(feature = "phase_trace")] + match orient { + Normal => 'N', + Mirror => 'M', + }, + ) + }); + let queries_normal = request.queries(Normal); + let queries_mirror = request.queries(Mirror); + + search::linear_scan_cascade_paired::( + [sessions.for_search(Normal), sessions.for_search(Mirror)], + [&queries_normal, &queries_mirror], + search_params, + [Normal, Mirror], + hawk_actor.full_scan_side, + [&extra_candidate_ids_both[0], &extra_candidate_ids_both[1]], + &forced_anon_stats_ids, + ) + .await +} + impl JobSubmissionHandle for HawkHandle { type A = HawkMutation; @@ -2106,8 +2197,14 @@ impl HawkHandle { // the other side through the startup configuration. let full_scan_side = hawk_actor.full_scan_side; - // Compute search results for a given orientation and compute matching information - let do_search = async |orient| -> Result<_> { + // Compute search results for a given orientation and compute matching + // information. `precomputed_search` carries this orientation's results + // when both orientations were already evaluated by the fused + // one-pass linear scan. + let do_search = async |orient, + precomputed_search: Option< + SearchResults, + >| -> Result<_> { let search_queries = &request.queries(orient); let (luc_ids, request_types, forced_anon_stats_ids) = { // Choice of LEFT registry is arbitrary — both sides are in sync @@ -2169,7 +2266,9 @@ impl HawkHandle { }, ); - let search_results = if hawk_actor.search_mode() == HawkSearchMode::LinearScan { + let search_results = if let Some(search_results) = precomputed_search { + search_results + } else if hawk_actor.search_mode() == HawkSearchMode::LinearScan { // CUDA unions OR-rule and reauth targets into the first-eye // prefilter result before checking both eyes on the subset. let extra_candidate_ids = izip!(&luc_ids, &request_types) @@ -2230,9 +2329,24 @@ impl HawkHandle { == HawkSearchMode::Hnsw || request.batch.full_face_mirror_attacks_detection_enabled { + // The exact scan evaluates both orientations in one pass over + // the resident eye: every streamed target feeds the normal and + // mirror dot products, while each orientation's threshold + // protocol keeps its own sessions and transcript. HNSW mode + // retains the independent per-orientation searches. + let [precomputed_normal, precomputed_mirror] = + if hawk_actor.search_mode() == HawkSearchMode::LinearScan { + let [normal, mirror] = + linear_scan_search_both_orientations(hawk_actor, sessions, &request) + .instrument(span.clone()) + .await?; + [Some(normal), Some(mirror)] + } else { + [None, None] + }; let ((search_normal, matches_normal), (search_mirror, matches_mirror)) = try_join!( - do_search(Orientation::Normal).instrument(span.clone()), - do_search(Orientation::Mirror).instrument(span.clone()), + do_search(Orientation::Normal, precomputed_normal).instrument(span.clone()), + do_search(Orientation::Mirror, precomputed_mirror).instrument(span.clone()), )?; ( search_normal, @@ -2240,7 +2354,7 @@ impl HawkHandle { matching::ResolvedBatch::decide(matches_normal, matches_mirror), ) } else { - let (search_normal, matches_normal) = do_search(Orientation::Normal) + let (search_normal, matches_normal) = do_search(Orientation::Normal, None) .instrument(span.clone()) .await?; let mirror_request_types = { diff --git a/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs b/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs index e330b9a995..9899c0cd97 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs @@ -7,7 +7,7 @@ use crate::{ galois_ring_pairwise_distance, non_existent_distance, pairwise_distance, rotation_aware_pairwise_distance, rotation_aware_pairwise_distance_rowmajor, }, - shared_iris::{ArcIris, GaloisRingSharedIris}, + shared_iris::{ArcIris, GaloisRingSharedIris, ResidentIris, ResidentLayout}, }, shares::RingElement, }; @@ -43,8 +43,19 @@ use tracing::info; pub const DEFAULT_FULL_ROTATION_TASK_SIZE: usize = 256; fn default_full_rotation_task_size() -> NonZeroUsize { - NonZeroUsize::new(DEFAULT_FULL_ROTATION_TASK_SIZE) - .expect("the default full-rotation task size must be nonzero") + // Overridable for scheduling experiments; the default is the tuned + // production value. + static TASK_SIZE: std::sync::OnceLock = std::sync::OnceLock::new(); + *TASK_SIZE.get_or_init(|| { + std::env::var("IRIS_MPC_FULL_ROTATION_TASK_SIZE") + .ok() + .and_then(|value| value.parse().ok()) + .and_then(NonZeroUsize::new) + .unwrap_or_else(|| { + NonZeroUsize::new(DEFAULT_FULL_ROTATION_TASK_SIZE) + .expect("the default full-rotation task size must be nonzero") + }) + }) } /// Defines the types of tasks that can be offloaded to an `IrisWorker`. @@ -117,6 +128,15 @@ enum IrisTask { range: std::ops::Range, rsp: oneshot::Sender>>, }, + /// Both orientations' 31-rotation dot products in one pass over a resident + /// target range. Targets are looked up and streamed once; each loaded + /// target row feeds both queries' rotation tiles. + FullRotationDotProductPairBatch { + queries: [ArcIris; 2], + vector_ids: Arc<[VectorId]>, + range: std::ops::Range, + rsp: oneshot::Sender<[Vec>; 2]>, + }, /// Computes the pairwise distance for pairs of irises in the Galois Ring. RingPairwiseDistance { input: Vec>, @@ -380,6 +400,47 @@ impl IrisPoolHandle { Ok(results) } + /// Paired-query variant of [`Self::full_rotation_dot_product_batch`]: one + /// target traversal per task computes both queries' full rotation sets. + pub async fn full_rotation_dot_product_pair_batch( + &mut self, + queries: [ArcIris; 2], + vector_ids: &[VectorId], + task_size: NonZeroUsize, + ) -> Result<[Vec>; 2]> { + let start = Instant::now(); + let shared_ids: Arc<[VectorId]> = Arc::from(vector_ids); + let task_size = task_size.get(); + let mut responses = Vec::with_capacity(shared_ids.len().div_ceil(task_size)); + + for (i, _) in shared_ids.chunks(task_size).enumerate() { + let range_start = i * task_size; + let range_end = (range_start + task_size).min(shared_ids.len()); + let (tx, rx) = oneshot::channel(); + self.get_next_worker() + .send(IrisTask::FullRotationDotProductPairBatch { + queries: queries.clone(), + vector_ids: shared_ids.clone(), + range: range_start..range_end, + rsp: tx, + })?; + responses.push(rx); + } + + let mut results = [ + Vec::with_capacity(2 * ROTATIONS * shared_ids.len()), + Vec::with_capacity(2 * ROTATIONS * shared_ids.len()), + ]; + for task_results in futures::future::try_join_all(responses).await? { + let [first, second] = task_results; + results[0].extend(first); + results[1].extend(second); + } + self.metric_rotation_aware_dot_product_latency + .record(start.elapsed().as_secs_f64()); + Ok(results) + } + async fn full_rotation_dot_product_irises_batch( &self, query: ArcIris, @@ -536,15 +597,17 @@ impl IrisPoolHandle { pub fn init_workers( shard_index: usize, - iris_store: SharedIrisesRef, + iris_store: SharedIrisesRef, numa: bool, + layout: ResidentLayout, ) -> IrisPoolHandle { let core_ids = select_core_ids(shard_index); info!( - "Dot product shard {} running on {} cores ({:?})", + "Dot product shard {} running on {} cores ({:?}), resident layout {:?}", shard_index, core_ids.len(), - core_ids + core_ids, + layout, ); let mut channels = vec![]; @@ -554,7 +617,7 @@ pub fn init_workers( let iris_store = iris_store.clone(); std::thread::spawn(move || { let _ = core_affinity::set_for_current(core_id); - worker_thread(rx, iris_store, numa); + worker_thread(rx, iris_store, numa, layout); }); } @@ -570,7 +633,12 @@ pub fn init_workers( } } -fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, numa: bool) { +fn worker_thread( + ch: Receiver, + iris_store: SharedIrisesRef, + numa: bool, + layout: ResidentLayout, +) { while let Ok(task) = ch.recv() { match task { IrisTask::Realloc { iris, rsp } => { @@ -589,14 +657,17 @@ fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, n } IrisTask::Insert { vector_id, iris } => { - let iris = if numa { - Arc::new((*iris).clone()) - } else { - iris + // `from_arc` writes the resident representation from this + // thread, so first-touch places it NUMA-locally. The extra + // u16 clone is only needed when the resident layout keeps + // the incoming allocation. + let resident = match layout { + ResidentLayout::U16 if numa => ResidentIris::U16(Arc::new((*iris).clone())), + _ => ResidentIris::from_arc(iris, layout), }; let mut store = iris_store.data.blocking_write(); - store.insert(vector_id, iris); + store.insert(vector_id, resident); } IrisTask::Reserve { additional } => { @@ -607,9 +678,14 @@ fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, n IrisTask::DotProductPairs { pairs, rsp } => { let store = iris_store.data.blocking_read(); + let targets: Vec> = pairs + .iter() + .map(|(_, vid)| store.get_vector(vid).map(ResidentIris::to_arc)) + .collect(); let iris_pairs = pairs .iter() - .map(|(q, vid)| store.get_vector(vid).map(|iris| (q, iris))); + .zip(&targets) + .map(|((q, _), target)| target.as_ref().map(|iris| (q, iris))); let r = pairwise_distance(iris_pairs); let _ = rsp.send(r); @@ -622,9 +698,13 @@ fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, n } => { let store = iris_store.data.blocking_read(); - let iris_pairs = vector_ids + let targets: Vec> = vector_ids .iter() - .map(|v| store.get_vector(v).map(|iris| (&query, iris))); + .map(|v| store.get_vector(v).map(ResidentIris::to_arc)) + .collect(); + let iris_pairs = targets + .iter() + .map(|target| target.as_ref().map(|iris| (&query, iris))); let r = pairwise_distance(iris_pairs); let _ = rsp.send(r); @@ -647,9 +727,13 @@ fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, n rsp, } => { let store = iris_store.data.blocking_read(); - let targets = vector_ids[range].iter().map(|v| store.get_vector(v)); + let targets: Vec> = vector_ids[range] + .iter() + .map(|v| store.get_vector(v).map(ResidentIris::to_arc)) + .collect(); let result = rotation_aware_pairwise_distance_rowmajor::( - &query, targets, + &query, + targets.iter().map(Option::as_ref), ); let _ = rsp.send(result); } @@ -674,9 +758,10 @@ fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, n rsp, } => { let store = iris_store.data.blocking_read(); - let targets = vector_ids[range].iter().map(|v| store.get_vector(v)); - let result = - rotation_aware_pairwise_distance_rowmajor::(&query, targets); + let result = full_rotation_distance_resident( + &query, + vector_ids[range].iter().map(|v| store.get_vector(v)), + ); let _ = rsp.send(result); } @@ -692,6 +777,20 @@ fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, n let _ = rsp.send(result); } + IrisTask::FullRotationDotProductPairBatch { + queries, + vector_ids, + range, + rsp, + } => { + let store = iris_store.data.blocking_read(); + let result = full_rotation_distance_resident_pair( + &queries, + vector_ids[range].iter().map(|v| store.get_vector(v)), + ); + let _ = rsp.send(result); + } + IrisTask::RingPairwiseDistance { input, rsp } => { let r = galois_ring_pairwise_distance(input); let _ = rsp.send(r); @@ -708,6 +807,87 @@ fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, n } } +/// Full 31-rotation distances against resident targets, dispatching on the +/// resident representation: mixed-plane targets use the UMMLA kernel, +/// u16 targets the MLA kernel. Results are bit-identical between the two. +fn full_rotation_distance_resident<'a, I>(query: &ArcIris, targets: I) -> Vec> +where + I: Iterator> + ExactSizeIterator, +{ + let residents: Vec> = targets.collect(); + + #[cfg(target_arch = "aarch64")] + { + use crate::protocol::ops::rotation_aware_pairwise_distance_mixed; + use crate::protocol::shared_iris::MixedPlaneIris; + + let all_mixed = residents + .iter() + .all(|target| target.is_none_or(|resident| resident.as_mixed().is_some())); + if all_mixed && residents.iter().any(Option::is_some) { + let mixed: Vec> = residents + .iter() + .map(|target| target.and_then(ResidentIris::as_mixed)) + .collect(); + return rotation_aware_pairwise_distance_mixed::(query, &mixed); + } + } + + let owned: Vec> = residents + .iter() + .map(|target| target.map(ResidentIris::to_arc)) + .collect(); + rotation_aware_pairwise_distance_rowmajor::( + query, + owned.iter().map(Option::as_ref), + ) +} + +/// Paired-query variant of [`full_rotation_distance_resident`]: mixed-plane +/// residents use the fused UMMLA pair kernel (targets streamed once for both +/// queries); the u16 fallback evaluates the queries sequentially, which is +/// bit-identical and still benefits from the two-slot prerotation cache. +fn full_rotation_distance_resident_pair<'a, I>( + queries: &[ArcIris; 2], + targets: I, +) -> [Vec>; 2] +where + I: Iterator> + ExactSizeIterator, +{ + let residents: Vec> = targets.collect(); + + #[cfg(target_arch = "aarch64")] + { + use crate::protocol::ops::rotation_aware_pairwise_distance_mixed_pair; + use crate::protocol::shared_iris::MixedPlaneIris; + + let all_mixed = residents + .iter() + .all(|target| target.is_none_or(|resident| resident.as_mixed().is_some())); + if all_mixed && residents.iter().any(Option::is_some) { + let mixed: Vec> = residents + .iter() + .map(|target| target.and_then(ResidentIris::as_mixed)) + .collect(); + return rotation_aware_pairwise_distance_mixed_pair::( + [&queries[0], &queries[1]], + &mixed, + ); + } + } + + let owned: Vec> = residents + .iter() + .map(|target| target.map(ResidentIris::to_arc)) + .collect(); + [&queries[0], &queries[1]].map(|query| { + rotation_aware_pairwise_distance_rowmajor::( + query, + owned.iter().map(Option::as_ref), + ) + }) +} + // --------------------------------------------------------------------------- // IrisWorkerPool trait — abstracts over local/remote worker implementations // --------------------------------------------------------------------------- @@ -820,6 +1000,30 @@ pub trait IrisWorkerPool: Debug + Send + Sync { vector_ids: Vec, ) -> BoxFuture<'a, Result>>>; + /// Compute two queries' full 31-rotation dot products over one shared + /// target traversal. The exact scan uses this to evaluate the normal and + /// mirror orientations while streaming the resident database once. + /// + /// Each returned side is identical to a separate + /// [`IrisWorkerPool::compute_dot_products_full_rotations`] call. The + /// default implementation evaluates the queries sequentially; pools with + /// fused kernels override it. + fn compute_dot_products_full_rotations_pair<'a>( + &'a self, + queries: [QuerySpec; 2], + vector_ids: Vec, + ) -> BoxFuture<'a, Result<[Vec>; 2]>> { + Box::pin(async move { + let first = self + .compute_dot_products_full_rotations(queries[0], vector_ids.clone()) + .await?; + let second = self + .compute_dot_products_full_rotations(queries[1], vector_ids) + .await?; + Ok([first, second]) + }) + } + /// Fetch iris data from the worker's store by vector ID. /// /// Returns one `ArcIris` per input ID in the same order. Database-backed @@ -895,6 +1099,13 @@ impl IrisWorkerPool for Arc { ) -> BoxFuture<'a, Result>>> { (**self).compute_dot_products_full_rotations(query, vector_ids) } + fn compute_dot_products_full_rotations_pair<'a>( + &'a self, + queries: [QuerySpec; 2], + vector_ids: Vec, + ) -> BoxFuture<'a, Result<[Vec>; 2]>> { + (**self).compute_dot_products_full_rotations_pair(queries, vector_ids) + } fn fetch_irises<'a>(&'a self, ids: Vec) -> BoxFuture<'a, Result>> { (**self).fetch_irises(ids) } @@ -992,13 +1203,10 @@ struct CachedQuery { pub struct LocalIrisWorkerPool { inner: IrisPoolHandle, query_cache: Arc>>, - iris_store: SharedIrisesRef, + iris_store: SharedIrisesRef, + layout: ResidentLayout, mode: DistanceMode, party_id: usize, - /// Number of records assigned to one full-rotation dot-product task. - /// Normal constructors retain the production default; benchmarks can opt - /// into another nonzero value through `with_full_rotation_task_size`. - full_rotation_task_size: NonZeroUsize, /// When set, the complete iris column stays in Postgres. RAM holds the /// rolling LUC window, a bounded frequency cache, and mutations awaiting a /// database commit; older explicit candidates are fetched sparsely. @@ -1315,6 +1523,23 @@ async fn run_cold_prefetch_worker( } } +fn cold_db_miss_error(party_id: usize, side: usize, missing: &[VectorId]) -> eyre::Report { + let examples = missing.iter().take(8).copied().collect_vec(); + metrics::counter!("linear_scan_cold_db_misses_total").increment(missing.len() as u64); + tracing::error!( + party_id, + side, + missing_count = missing.len(), + ?examples, + "Cold-eye database did not return the exact requested vector versions" + ); + eyre::eyre!( + "cold-eye database missing {} exact vector ID(s) for party {party_id}, side {side}; \ + examples: {examples:?}", + missing.len() + ) +} + async fn prefetch_cold_irises( store: &Store, side: usize, @@ -1427,23 +1652,6 @@ async fn prefetch_cold_irises( } } -fn cold_db_miss_error(party_id: usize, side: usize, missing: &[VectorId]) -> eyre::Report { - let examples = missing.iter().take(8).copied().collect_vec(); - metrics::counter!("linear_scan_cold_db_misses_total").increment(missing.len() as u64); - tracing::error!( - party_id, - side, - missing_count = missing.len(), - ?examples, - "Cold-eye database did not return the exact requested vector versions" - ); - eyre::eyre!( - "cold-eye database missing {} exact vector ID(s) for party {party_id}, side {side}; \ - examples: {examples:?}", - missing.len() - ) -} - impl Debug for LocalIrisWorkerPool { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("LocalIrisWorkerPool") @@ -1455,7 +1663,8 @@ impl Debug for LocalIrisWorkerPool { impl LocalIrisWorkerPool { pub fn new( inner: IrisPoolHandle, - iris_store: SharedIrisesRef, + iris_store: SharedIrisesRef, + layout: ResidentLayout, mode: DistanceMode, party_id: usize, ) -> Self { @@ -1463,9 +1672,9 @@ impl LocalIrisWorkerPool { inner, query_cache: Arc::new(RwLock::new(HashMap::new())), iris_store, + layout, mode, party_id, - full_rotation_task_size: default_full_rotation_task_size(), cold_storage: None, } } @@ -1477,7 +1686,8 @@ impl LocalIrisWorkerPool { /// persistence commits. pub async fn new_cold( inner: IrisPoolHandle, - iris_store: SharedIrisesRef, + iris_store: SharedIrisesRef, + layout: ResidentLayout, mode: DistanceMode, party_id: usize, init: ColdStorageInit, @@ -1536,9 +1746,9 @@ impl LocalIrisWorkerPool { inner, query_cache: Arc::new(RwLock::new(HashMap::new())), iris_store, + layout, mode, party_id, - full_rotation_task_size: default_full_rotation_task_size(), cold_storage: Some(ColdStorage { store, side, @@ -1551,19 +1761,13 @@ impl LocalIrisWorkerPool { /// Create a local worker pool for shard 0 with NUMA pinning. /// Standard construction for tests, benchmarks, and single-node tools. pub fn new_local( - iris_store: SharedIrisesRef, + iris_store: SharedIrisesRef, + layout: ResidentLayout, mode: DistanceMode, party_id: usize, ) -> Self { - let pool = init_workers(0, iris_store.clone(), true); - Self::new(pool, iris_store, mode, party_id) - } - - /// Override full-rotation task granularity for an explicitly configured - /// benchmark or tool. Production constructors always default to 128. - pub fn with_full_rotation_task_size(mut self, task_size: NonZeroUsize) -> Self { - self.full_rotation_task_size = task_size; - self + let pool = init_workers(0, iris_store.clone(), true, layout); + Self::new(pool, iris_store, layout, mode, party_id) } async fn fetch_irises_resident_or_cold(&self, ids: &[VectorId]) -> Result> { @@ -1571,7 +1775,7 @@ impl LocalIrisWorkerPool { let store = self.iris_store.data.read().await; return Ok(ids .iter() - .map(|id| store.get_vector_or_empty(id).clone()) + .map(|id| store.get_vector_or_empty(id).to_arc()) .collect()); }; @@ -1804,7 +2008,7 @@ impl IrisWorkerPool for LocalIrisWorkerPool { let mut inner = self.inner.clone(); let pool = self.clone(); let is_cold = self.cold_storage.is_some(); - let task_size = self.full_rotation_task_size; + let task_size = default_full_rotation_task_size(); Box::pin(async move { let iris = { let cache = query_cache.read().unwrap(); @@ -1831,6 +2035,55 @@ impl IrisWorkerPool for LocalIrisWorkerPool { }) } + fn compute_dot_products_full_rotations_pair<'a>( + &'a self, + queries: [QuerySpec; 2], + vector_ids: Vec, + ) -> BoxFuture<'a, Result<[Vec>; 2]>> { + let query_cache = self.query_cache.clone(); + let mut inner = self.inner.clone(); + let pool = self.clone(); + let is_cold = self.cold_storage.is_some(); + let task_size = default_full_rotation_task_size(); + Box::pin(async move { + let irises = { + let cache = query_cache.read().unwrap(); + let mut resolved = Vec::with_capacity(2); + for query in queries { + let cached = cache + .get(&query.query_id) + .ok_or_else(|| eyre::eyre!("Query {:?} not cached", query.query_id))?; + let rotations = if query.mirrored { + &cached.mirrored_preprocessed_rotations + } else { + &cached.preprocessed_rotations + }; + resolved.push(rotations[query.rotation].clone()); + } + let second = resolved.pop().expect("two resolved queries"); + let first = resolved.pop().expect("two resolved queries"); + [first, second] + }; + if is_cold { + // Cold targets are fetched once and reused by both queries; + // the dot passes themselves stay sequential on this rare path. + let targets = pool.fetch_irises_resident_or_cold(&vector_ids).await?; + let [query_a, query_b] = irises; + let first = inner + .full_rotation_dot_product_irises_batch(query_a, targets.clone(), task_size) + .await?; + let second = inner + .full_rotation_dot_product_irises_batch(query_b, targets, task_size) + .await?; + Ok([first, second]) + } else { + inner + .full_rotation_dot_product_pair_batch(irises, &vector_ids, task_size) + .await + } + }) + } + fn fetch_irises<'a>(&'a self, ids: Vec) -> BoxFuture<'a, Result>> { Box::pin(async move { self.fetch_irises_resident_or_cold(&ids).await }) } @@ -1881,6 +2134,7 @@ impl IrisWorkerPool for LocalIrisWorkerPool { let query_cache = self.query_cache.clone(); let iris_store = self.iris_store.clone(); let cold_storage = self.cold_storage.clone(); + let layout = self.layout; Box::pin(async move { // Resolve query IDs to irises (release cache lock before await). let resolved: Vec<_> = { @@ -1912,7 +2166,7 @@ impl IrisWorkerPool for LocalIrisWorkerPool { // visible in the store immediately after this returns. let mut store = iris_store.data.write().await; for (vector_id, iris) in resolved { - store.insert(vector_id, iris); + store.insert(vector_id, ResidentIris::from_arc(iris, layout)); } Ok(store.set_hash.checksum()) }) @@ -1982,6 +2236,7 @@ impl IrisWorkerPool for LocalIrisWorkerPool { let iris_store = self.iris_store.clone(); let party_id = self.party_id; let cold_storage = self.cold_storage.clone(); + let layout = self.layout; Box::pin(async move { let dummy = Arc::new(GaloisRingSharedIris::dummy_for_party(party_id)); if let Some(cold) = cold_storage { @@ -1994,9 +2249,10 @@ impl IrisWorkerPool for LocalIrisWorkerPool { } return Ok(()); } + let resident_dummy = ResidentIris::from_arc(dummy, layout); let mut store = iris_store.data.write().await; for id in ids { - store.update(id, dummy.clone()); + store.update(id, resident_dummy.clone()); } Ok(()) }) @@ -2143,26 +2399,40 @@ mod tests { let vector_ids = (0..TEST_TARGETS) .map(|index| VectorId::from_0_index(index as u32)) .collect::>(); - let points = vector_ids - .iter() - .copied() - .map(|id| (id, iris.clone())) - .collect::>(); - let storage = SharedIrises::new(points, iris.clone()).to_arc(); - let workers = init_workers(0, storage, false); - - let mut results = Vec::new(); - for task_size in [64, 128, 256, 512] { - let mut workers = workers.clone(); - results.push(runtime.block_on(workers.full_rotation_dot_product_batch( - iris.clone(), - &vector_ids, - NonZeroUsize::new(task_size).unwrap(), - ))?); - } - assert!(results.iter().all(|result| result == &results[0])); - assert_eq!(results[0].len(), TEST_TARGETS * ROTATIONS * 2); + // Cover every supported resident layout; results must be identical. + let mut per_layout = Vec::new(); + for layout in [ + ResidentLayout::U16, + crate::protocol::shared_iris::preferred_scan_layout(), + ] { + let points = vector_ids + .iter() + .copied() + .map(|id| (id, ResidentIris::from_arc(iris.clone(), layout))) + .collect::>(); + let storage = + SharedIrises::new(points, ResidentIris::from_arc(iris.clone(), layout)).to_arc(); + let workers = init_workers(0, storage, false, layout); + + let mut results = Vec::new(); + for task_size in [64, 128, 256, 512] { + let mut workers = workers.clone(); + results.push(runtime.block_on(workers.full_rotation_dot_product_batch( + iris.clone(), + &vector_ids, + NonZeroUsize::new(task_size).unwrap(), + ))?); + } + + assert!(results.iter().all(|result| result == &results[0])); + assert_eq!(results[0].len(), TEST_TARGETS * ROTATIONS * 2); + per_layout.push(results.remove(0)); + } + assert!( + per_layout.iter().all(|result| result == &per_layout[0]), + "resident layouts must produce identical distances" + ); Ok(()) } } diff --git a/iris-mpc-cpu/src/execution/hawk_main/rot.rs b/iris-mpc-cpu/src/execution/hawk_main/rot.rs index 9fec04e4d5..aeb07c0f06 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/rot.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/rot.rs @@ -53,6 +53,21 @@ impl VecRotationSupport { &self.rotations[self.rotations.len() / 2] } + /// Mutably access the item attached to the center rotation. + pub fn center_mut(&mut self) -> &mut R { + let center = self.rotations.len() / 2; + &mut self.rotations[center] + } + + /// Consume the collection and return its center-rotation item. + pub fn into_center(self) -> R { + let center = self.rotations.len() / 2; + self.rotations + .into_iter() + .nth(center) + .expect("rotation support must contain a center") + } + /// Flatten a batch of something with rotations into a concatenated Vec. /// Attach a copy of the corresponding `B` to each rotation. pub fn flatten_broadcast<'a, B>(batch: impl IntoIterator) -> Vec<(R, B)> diff --git a/iris-mpc-cpu/src/execution/hawk_main/search.rs b/iris-mpc-cpu/src/execution/hawk_main/search.rs index c28be88840..6f27deb5e3 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/search.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/search.rs @@ -10,16 +10,25 @@ use crate::{ scheduler::{collect_results, parallelize}, InsertPlanV, StoreId, }, - hawkers::aby3::aby3_store::{Aby3DistanceRef, Aby3Query, Aby3Store, DistanceOps}, + hawkers::aby3::aby3_store::{ + Aby3DistanceRef, Aby3Query, Aby3Store, DistanceOps, FullRotationThresholdResult, + }, hnsw::{graph::UpdateEntryPoint, GraphMem, HnswSearcher}, + shares::RingElement, }; use ampc_anon_stats::types::Eye; use eyre::{OptionExt, Result}; use iris_mpc_common::iris_db::iris::Threshold; use iris_mpc_common::{VectorId, ROTATIONS}; -use std::sync::Arc; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; use std::time::Instant; -use tokio::sync::mpsc::{unbounded_channel, UnboundedSender}; +use tokio::sync::{ + mpsc::{unbounded_channel, UnboundedSender}, + Notify, +}; use tracing::instrument; /// Keep enough records in each MPC call to amortize its fixed costs without @@ -47,6 +56,57 @@ struct LinearScanChunk { range: std::ops::Range, } +#[derive(Clone)] +struct LinearScanPrefetch { + worker: Arc, +} + +#[derive(Clone, Default)] +struct LinearScanHooks { + prefetch: Option, + progress: Option>, +} + +struct LinearScanProgress { + completed_comparisons: AtomicUsize, + start_candidate_after: usize, + notify: Notify, +} + +impl LinearScanProgress { + fn new(comparisons: usize) -> Self { + Self { + completed_comparisons: AtomicUsize::new(0), + // Start the sparse second-eye work while the final 10% of full-eye + // chunks drain. This leaves enough overlap to hide its latency but + // avoids competing with the peak dot/network fan-out. + start_candidate_after: comparisons.saturating_mul(9).div_ceil(10), + notify: Notify::new(), + } + } + + fn record(&self, comparisons: usize) { + let previous = self + .completed_comparisons + .fetch_add(comparisons, Ordering::AcqRel); + if previous < self.start_candidate_after + && previous.saturating_add(comparisons) >= self.start_candidate_after + { + self.notify.notify_waiters(); + } + } + + async fn wait_for_candidate_start(&self) { + loop { + let notified = self.notify.notified(); + if self.completed_comparisons.load(Ordering::Acquire) >= self.start_candidate_after { + return; + } + notified.await; + } + } +} + #[derive(Clone, Copy, Debug)] enum LinearScanStage { Full, @@ -90,11 +150,6 @@ const fn orientation_label(orientation: Orientation) -> &'static str { } } -#[derive(Clone)] -struct LinearScanPrefetch { - worker: Arc, -} - pub type SearchQueries = Arc>>>; pub type SearchResults = @@ -183,13 +238,6 @@ pub async fn search( let n_requests = search_queries[LEFT].len(); assert_eq!(n_requests, search_queries[RIGHT].len()); - // `search_to_identity_update` invokes this function for every batch, even - // when the batch has no reset/recovery updates. Avoid spawning one empty - // session per eye in that common case. - if n_requests == 0 { - return Ok([Vec::new(), Vec::new()]); - } - let (tx, rx) = unbounded_channel::<(TaskId, HawkInsertPlan)>(); let per_session = |batch: Batch| { @@ -274,6 +322,18 @@ pub async fn linear_scan_cascade( }; let full_scan_ids = Arc::new(vec![live_ids.clone(); n_requests]); let first_eye_comparisons = live_ids.len() * n_requests; + let known_second_stage_ids = Arc::new( + extra_candidate_ids + .iter() + .map(|extras| { + Arc::<[VectorId]>::from(collect_live_second_stage_ids( + &live_ids, + std::iter::empty(), + extras, + )) + }) + .collect::>(), + ); tracing::info!( eye = %full_scan_side, @@ -286,14 +346,13 @@ pub async fn linear_scan_cascade( let store = sessions[second_eye][0].aby3_store.read().await; store.workers.clone() }; - // Known OR/reauth candidates are needed exactly once in stage two. Queue - // them once here; adding them to every full-scan chunk would inflate the - // prefetch reservation use count and retain stale entries after the scan. - for extras in extra_candidate_ids { - let prefetch_ids = collect_live_second_stage_ids(&live_ids, [], extras); - prefetch_worker.prefetch_irises(prefetch_ids).await?; - } - let first_results = linear_scan_eye( + let first_eye_progress = Arc::new(LinearScanProgress::new(first_eye_comparisons)); + // LUC, OR-rule, and reauthentication candidates are public before the + // scan starts. Check them on the cold eye while the full resident-eye scan + // is running; any database I/O is thereby hidden behind the long stage. + // Candidates discovered by the anonymous-statistics threshold are still + // prefetched chunk by chunk and checked below. + let first_eye_scan = linear_scan_eye( sessions, search_queries, &search_params, @@ -304,11 +363,32 @@ pub async fn linear_scan_cascade( }, full_scan_ids, Arc::new(forced_anon_stats_ids.clone()), - Some(LinearScanPrefetch { - worker: prefetch_worker.clone(), - }), - ) - .await?; + LinearScanHooks { + prefetch: Some(LinearScanPrefetch { + worker: prefetch_worker.clone(), + }), + progress: Some(first_eye_progress.clone()), + }, + ); + let known_second_eye_scan = async { + first_eye_progress.wait_for_candidate_start().await; + linear_scan_eye( + sessions, + search_queries, + &search_params, + LinearScanEyeContext { + eye: second_eye_side, + stage: LinearScanStage::Candidate, + orientation, + }, + known_second_stage_ids.clone(), + Arc::new(forced_anon_stats_ids.clone()), + LinearScanHooks::default(), + ) + .await + }; + let (first_results, mut second_results) = + tokio::try_join!(first_eye_scan, known_second_eye_scan)?; let prefetch_wait_start = Instant::now(); prefetch_worker.wait_for_prefetch().await?; let prefetch_wait_seconds = prefetch_wait_start.elapsed().as_secs_f64(); @@ -328,43 +408,54 @@ pub async fn linear_scan_cascade( // serial ID. Avoid materializing a full-database HashSet for this sparse // candidate membership check. debug_assert!(live_ids.windows(2).all(|pair| pair[0] < pair[1])); - let mut second_stage_ids = Vec::with_capacity(n_requests); - let mut candidate_count = 0usize; - for (plans, extras) in first_results.iter().zip(extra_candidate_ids) { - let ids = collect_live_second_stage_ids( + let mut discovered_second_stage_ids = Vec::with_capacity(n_requests); + let mut discovered_candidate_count = 0usize; + for (plans, known_ids) in first_results.iter().zip(known_second_stage_ids.iter()) { + let mut ids = collect_live_second_stage_ids( &live_ids, plans .iter() .flat_map(|plan| &plan.classified.anon_stats_matches.results) .map(|(id, _)| *id), - extras, + &[], ); - candidate_count += ids.len(); - second_stage_ids.push(Arc::<[VectorId]>::from(ids)); + ids.retain(|id| known_ids.binary_search(id).is_err()); + discovered_candidate_count += ids.len(); + discovered_second_stage_ids.push(Arc::<[VectorId]>::from(ids)); } + let known_candidate_count = known_second_stage_ids + .iter() + .map(|ids| ids.len()) + .sum::(); + let candidate_count = known_candidate_count + discovered_candidate_count; tracing::info!( eye = %second_eye_side, orientation = orientation_label(orientation), requests = n_requests, candidates = candidate_count, + known_candidates = known_candidate_count, + discovered_candidates = discovered_candidate_count, "Running candidate-only linear-scan stage" ); metrics::counter!("linear_scan_second_eye_candidates_total").increment(candidate_count as u64); - let second_results = linear_scan_eye( - sessions, - search_queries, - &search_params, - LinearScanEyeContext { - eye: second_eye_side, - stage: LinearScanStage::Candidate, - orientation, - }, - Arc::new(second_stage_ids), - Arc::new(forced_anon_stats_ids.clone()), - None, - ) - .await?; + if discovered_candidate_count > 0 { + let discovered_results = linear_scan_eye( + sessions, + search_queries, + &search_params, + LinearScanEyeContext { + eye: second_eye_side, + stage: LinearScanStage::Candidate, + orientation, + }, + Arc::new(discovered_second_stage_ids), + Arc::new(forced_anon_stats_ids.clone()), + LinearScanHooks::default(), + ) + .await?; + merge_linear_scan_results(&mut second_results, discovered_results); + } let total_comparisons = first_eye_comparisons + candidate_count; let elapsed_seconds = cascade_start.elapsed().as_secs_f64(); @@ -442,7 +533,7 @@ async fn linear_scan_eye( context: LinearScanEyeContext, candidate_ids: Arc>>, forced_anon_stats_ids: Arc>>, - prefetch: Option, + hooks: LinearScanHooks, ) -> Result>> { let stage_start = Instant::now(); let eye_index = eye_index(context.eye); @@ -499,7 +590,7 @@ async fn linear_scan_eye( let search_params = search_params.clone(); let candidate_ids = candidate_ids.clone(); let forced_anon_stats_ids = forced_anon_stats_ids.clone(); - let prefetch = prefetch.clone(); + let hooks = hooks.clone(); async move { let mut vector_store = session.aby3_store.write().await; let graph_store = session.graph_store.clone().read_owned().await; @@ -515,7 +606,7 @@ async fn linear_scan_eye( &forced_anon_stats_ids[chunk.i_request], ) .await?; - if let Some(prefetch) = &prefetch { + if let Some(prefetch) = &hooks.prefetch { let chunk_ids = &candidate_ids[chunk.i_request][chunk.range.clone()]; let prefetch_ids = collect_live_second_stage_ids( chunk_ids, @@ -529,6 +620,9 @@ async fn linear_scan_eye( ); prefetch.worker.prefetch_irises(prefetch_ids).await?; } + if let Some(progress) = &hooks.progress { + progress.record(chunk.range.len()); + } results.push((chunk.i_request, chunk.i_chunk, result)); } Ok(results) @@ -544,7 +638,127 @@ async fn linear_scan_eye( } let graph_store = sessions[eye_index][0].graph_store.read().await; - let results = chunk_results + let results = + assemble_linear_scan_results(chunk_results, &search_queries[eye_index], &graph_store)?; + + let strict_match_records = results + .iter() + .map(|rotations| rotations[central_rotation].classified.matches.results.len()) + .sum::(); + let anon_stats_rotation_matches = results + .iter() + .map(|rotations| { + rotations[central_rotation] + .classified + .anon_stats_matches + .results + .len() + }) + .sum::(); + + let elapsed_seconds = stage_start.elapsed().as_secs_f64(); + let comparisons_per_second = comparisons as f64 / elapsed_seconds.max(f64::EPSILON); + let session_utilization = n_workers as f64 / configured_sessions as f64; + let stage_label = context.stage.as_str(); + let eye_label = eye_label(context.eye); + let orientation_label = orientation_label(context.orientation); + metrics::counter!( + "linear_scan_eye_comparisons_total", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .increment(comparisons as u64); + metrics::histogram!( + "linear_scan_eye_duration", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .record(elapsed_seconds); + metrics::histogram!( + "linear_scan_eye_comparisons", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .record(comparisons as f64); + metrics::histogram!( + "linear_scan_eye_comparisons_per_second", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .record(comparisons_per_second); + metrics::histogram!( + "linear_scan_eye_chunks", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .record(chunk_count as f64); + metrics::histogram!( + "linear_scan_eye_session_utilization", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .record(session_utilization); + metrics::histogram!( + "linear_scan_eye_active_sessions", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .record(n_workers as f64); + metrics::counter!( + "linear_scan_eye_strict_match_records_total", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .increment(strict_match_records as u64); + metrics::counter!( + "linear_scan_eye_anon_stats_rotation_matches_total", + "eye" => eye_label, + "stage" => stage_label, + "orientation" => orientation_label, + ) + .increment(anon_stats_rotation_matches as u64); + tracing::info!( + eye = eye_label, + stage = stage_label, + orientation = orientation_label, + requests = n_requests, + comparisons, + rotations_per_comparison = ROTATIONS, + chunks = chunk_count, + chunk_size = LINEAR_SCAN_CHUNK_SIZE, + configured_sessions, + active_sessions = n_workers, + session_utilization, + min_chunks_per_session, + max_chunks_per_session, + strict_match_records, + anon_stats_rotation_matches, + elapsed_seconds, + comparisons_per_second, + "LINEAR_SCAN_EYE_SUMMARY" + ); + + Ok(results) +} + +/// Merge per-chunk plans into per-request results and package them in the +/// three-slot rotation container (full result in the center slot). +fn assemble_linear_scan_results( + chunk_results: Vec>>, + queries: &VecRequests>, + graph_store: &GraphMem, +) -> Result>> { + let n_rotations = ROTMASK.count_ones() as usize; + let central_rotation = n_rotations / 2; + chunk_results .into_iter() .enumerate() .map(|(i_request, results)| { @@ -569,16 +783,40 @@ async fn linear_scan_eye( .expect("central linear-scan result must be consumed once")) } else { Ok(empty_linear_scan_plan( - search_queries[eye_index][i_request][i_rotation], - &graph_store, + queries[i_request][i_rotation], + graph_store, )) } }) .collect::>>() .map(VecRotationSupport::from) }) - .collect::>>()?; + .collect() +} +/// Scheduling shape of one eye-stage scan, shared by the metric emitters. +#[derive(Clone, Copy)] +struct LinearScanStageShape { + n_requests: usize, + comparisons: usize, + chunk_count: usize, + configured_sessions: usize, + n_workers: usize, + min_chunks_per_session: usize, + max_chunks_per_session: usize, +} + +/// Emit the per-stage metrics and `LINEAR_SCAN_EYE_SUMMARY` line for one +/// orientation. Shared by the single and paired stage implementations so both +/// produce identical observability output. +fn emit_linear_scan_eye_summary( + context: LinearScanEyeContext, + shape: LinearScanStageShape, + results: &VecRequests>, + elapsed_seconds: f64, +) { + let n_rotations = ROTMASK.count_ones() as usize; + let central_rotation = n_rotations / 2; let strict_match_records = results .iter() .map(|rotations| rotations[central_rotation].classified.matches.results.len()) @@ -594,9 +832,8 @@ async fn linear_scan_eye( }) .sum::(); - let elapsed_seconds = stage_start.elapsed().as_secs_f64(); - let comparisons_per_second = comparisons as f64 / elapsed_seconds.max(f64::EPSILON); - let session_utilization = n_workers as f64 / configured_sessions as f64; + let comparisons_per_second = shape.comparisons as f64 / elapsed_seconds.max(f64::EPSILON); + let session_utilization = shape.n_workers as f64 / shape.configured_sessions as f64; let stage_label = context.stage.as_str(); let eye_label = eye_label(context.eye); let orientation_label = orientation_label(context.orientation); @@ -606,7 +843,7 @@ async fn linear_scan_eye( "stage" => stage_label, "orientation" => orientation_label, ) - .increment(comparisons as u64); + .increment(shape.comparisons as u64); metrics::histogram!( "linear_scan_eye_duration", "eye" => eye_label, @@ -620,7 +857,7 @@ async fn linear_scan_eye( "stage" => stage_label, "orientation" => orientation_label, ) - .record(comparisons as f64); + .record(shape.comparisons as f64); metrics::histogram!( "linear_scan_eye_comparisons_per_second", "eye" => eye_label, @@ -634,7 +871,7 @@ async fn linear_scan_eye( "stage" => stage_label, "orientation" => orientation_label, ) - .record(chunk_count as f64); + .record(shape.chunk_count as f64); metrics::histogram!( "linear_scan_eye_session_utilization", "eye" => eye_label, @@ -648,7 +885,7 @@ async fn linear_scan_eye( "stage" => stage_label, "orientation" => orientation_label, ) - .record(n_workers as f64); + .record(shape.n_workers as f64); metrics::counter!( "linear_scan_eye_strict_match_records_total", "eye" => eye_label, @@ -667,26 +904,535 @@ async fn linear_scan_eye( eye = eye_label, stage = stage_label, orientation = orientation_label, - requests = n_requests, - comparisons, + requests = shape.n_requests, + comparisons = shape.comparisons, rotations_per_comparison = ROTATIONS, - chunks = chunk_count, + chunks = shape.chunk_count, chunk_size = LINEAR_SCAN_CHUNK_SIZE, - configured_sessions, - active_sessions = n_workers, + configured_sessions = shape.configured_sessions, + active_sessions = shape.n_workers, session_utilization, - min_chunks_per_session, - max_chunks_per_session, + min_chunks_per_session = shape.min_chunks_per_session, + max_chunks_per_session = shape.max_chunks_per_session, strict_match_records, anon_stats_rotation_matches, elapsed_seconds, comparisons_per_second, "LINEAR_SCAN_EYE_SUMMARY" ); +} + +/// Fused full-eye stage for both orientations: one chunk grid over the shared +/// live-ID list, with paired sessions. Each chunk streams its targets once for +/// both orientations' dot products; each orientation's threshold rounds then +/// run on that orientation's own session. Chunk-to-session assignment matches +/// [`linear_scan_eye`], so every per-orientation session sees the same chunk +/// sequence (and therefore the same network transcript) as two independent +/// stages. +#[allow(clippy::too_many_arguments)] +async fn linear_scan_full_stage_paired( + sessions_both: [&BothEyes>; 2], + search_queries_both: [&SearchQueries; 2], + search_params_both: [&SearchParams; 2], + contexts: [LinearScanEyeContext; 2], + full_scan_ids: Arc>>, + forced_anon_stats_ids: Arc>>, + hooks: LinearScanHooks, +) -> Result<[VecRequests>; 2]> { + let stage_start = Instant::now(); + let eye_index = eye_index(contexts[0].eye); + debug_assert_eq!(eye_index, self::eye_index(contexts[1].eye)); + let n_requests = search_queries_both[0][eye_index].len(); + assert_eq!(n_requests, search_queries_both[1][eye_index].len()); + assert_eq!(n_requests, full_scan_ids.len()); + assert_eq!(n_requests, forced_anon_stats_ids.len()); + let comparisons = full_scan_ids.iter().map(|ids| ids.len()).sum::(); + let central_rotation = ROTMASK.count_ones() as usize / 2; + + let mut chunks_per_request = Vec::with_capacity(n_requests); + let mut chunks = Vec::new(); + for (i_request, ids) in full_scan_ids.iter().enumerate() { + let n_chunks = ids.len().div_ceil(LINEAR_SCAN_CHUNK_SIZE).max(1); + chunks_per_request.push(n_chunks); + for i_chunk in 0..n_chunks { + let start = (i_chunk * LINEAR_SCAN_CHUNK_SIZE).min(ids.len()); + let end = (start + LINEAR_SCAN_CHUNK_SIZE).min(ids.len()); + chunks.push(LinearScanChunk { + i_request, + i_chunk, + range: start..end, + }); + } + } + + let chunk_count = chunks.len(); + let configured_sessions = sessions_both[0][eye_index] + .len() + .min(sessions_both[1][eye_index].len()); + let n_workers = configured_sessions + .min(LINEAR_SCAN_MAX_IN_FLIGHT_CHUNKS) + .min(chunks.len()) + .max(1); + let mut batches = vec![Vec::new(); n_workers]; + for (index, chunk) in chunks.into_iter().enumerate() { + batches[index % n_workers].push(chunk); + } + let min_chunks_per_session = batches.iter().map(Vec::len).min().unwrap_or(0); + let max_chunks_per_session = batches.iter().map(Vec::len).max().unwrap_or(0); + + let jobs = batches + .into_iter() + .enumerate() + .filter(|(_, batch)| !batch.is_empty()) + .map(|(i_session, batch)| { + let session_a = sessions_both[0][eye_index][i_session].clone(); + let session_b = sessions_both[1][eye_index][i_session].clone(); + let search_queries_a = search_queries_both[0].clone(); + let search_queries_b = search_queries_both[1].clone(); + let search_params_a = search_params_both[0].clone(); + let search_params_b = search_params_both[1].clone(); + let full_scan_ids = full_scan_ids.clone(); + let forced_anon_stats_ids = forced_anon_stats_ids.clone(); + let hooks = hooks.clone(); + async move { + let mut store_a = session_a.aby3_store.write().await; + let mut store_b = session_b.aby3_store.write().await; + let graph_a = session_a.graph_store.clone().read_owned().await; + let graph_b = session_b.graph_store.clone().read_owned().await; + let mut results = Vec::with_capacity(batch.len()); + // Software-pipeline this lane: the next chunk's fused dot + // products run on the worker pool while the current chunk's + // threshold rounds are in flight, so the dot workers stay fed + // instead of idling for a round trip per chunk. + let queries_for = |chunk: &LinearScanChunk| { + ( + search_queries_a[eye_index][chunk.i_request][central_rotation], + search_queries_b[eye_index][chunk.i_request][central_rotation], + ) + }; + let do_match = search_params_a.do_match; + let dispatch = |store: &Aby3Store, chunk: &LinearScanChunk| { + let (query_a, query_b) = queries_for(chunk); + store.spawn_full_rotation_dot_contributions_pair( + [&query_a, &query_b], + &full_scan_ids[chunk.i_request][chunk.range.clone()], + ) + }; + // Keep one chunk of dot work buffered per lane: the next + // chunk's dot products run while this chunk's threshold + // rounds are in flight. Deeper buffering measures worse — all + // dot work then completes early and the stage drains on + // thresholds alone with idle dot workers. + const DOT_PIPELINE_DEPTH: usize = 1; + let mut pending_dots = std::collections::VecDeque::new(); + if do_match { + for chunk in batch.iter().take(DOT_PIPELINE_DEPTH) { + pending_dots.push_back(dispatch(&store_a, chunk)); + } + } + for (index, chunk) in batch.iter().enumerate() { + let contributions = match pending_dots.pop_front() { + Some(handle) => handle + .await + .map_err(|error| eyre::eyre!("fused dot task failed: {error}"))??, + None => [Vec::new(), Vec::new()], + }; + if do_match { + if let Some(next) = batch.get(index + DOT_PIPELINE_DEPTH) { + pending_dots.push_back(dispatch(&store_a, next)); + } + } + let (query_a, query_b) = queries_for(chunk); + let chunk_ids = &full_scan_ids[chunk.i_request][chunk.range.clone()]; + let plans = per_linear_scan_chunk_pair( + [query_a, query_b], + [&search_params_a, &search_params_b], + (&mut store_a, &mut store_b), + (&graph_a, &graph_b), + contributions, + chunk_ids, + &forced_anon_stats_ids[chunk.i_request], + ) + .await?; + if let Some(prefetch) = &hooks.prefetch { + // One union prefetch warms the cold eye for both + // orientations' second-stage candidates. + let prefetch_ids = collect_live_second_stage_ids( + chunk_ids, + plans.iter().flat_map(|plan| { + plan.classified + .anon_stats_matches + .results + .iter() + .map(|(id, _)| *id) + }), + &[], + ); + prefetch.worker.prefetch_irises(prefetch_ids).await?; + } + if let Some(progress) = &hooks.progress { + progress.record(chunk.range.len()); + } + results.push((chunk.i_request, chunk.i_chunk, plans)); + } + Ok(results) + } + }); + + let mut chunk_results: [Vec>>; 2] = [ + chunks_per_request + .iter() + .map(|&len| vec![None; len]) + .collect(), + chunks_per_request + .iter() + .map(|&len| vec![None; len]) + .collect(), + ]; + for (i_request, i_chunk, plans) in parallelize(jobs).await?.into_iter().flatten() { + let [plan_a, plan_b] = plans; + chunk_results[0][i_request][i_chunk] = Some(plan_a); + chunk_results[1][i_request][i_chunk] = Some(plan_b); + } + let [chunk_results_a, chunk_results_b] = chunk_results; + + let graph_a = sessions_both[0][eye_index][0].graph_store.read().await; + let graph_b = sessions_both[1][eye_index][0].graph_store.read().await; + let results = [ + assemble_linear_scan_results( + chunk_results_a, + &search_queries_both[0][eye_index], + &graph_a, + )?, + assemble_linear_scan_results( + chunk_results_b, + &search_queries_both[1][eye_index], + &graph_b, + )?, + ]; + + let elapsed_seconds = stage_start.elapsed().as_secs_f64(); + let shape = LinearScanStageShape { + n_requests, + comparisons, + chunk_count, + configured_sessions, + n_workers, + min_chunks_per_session, + max_chunks_per_session, + }; + for (context, results) in contexts.iter().zip(&results) { + emit_linear_scan_eye_summary(*context, shape, results, elapsed_seconds); + } Ok(results) } +/// Run both orientations' two-eye linear-scan cascades with a fused first-eye +/// stage: the resident full-scan eye is streamed once and every loaded target +/// feeds both orientations' 31-rotation dot products. All MPC threshold work +/// stays on each orientation's own sessions, so results and per-orientation +/// transcripts are identical to two concurrent [`linear_scan_cascade`] calls. +/// The sparse second-eye candidate stages remain per-orientation. +#[instrument(level = "trace", target = "searcher::network", skip_all)] +pub async fn linear_scan_cascade_paired( + sessions_both: [&BothEyes>; 2], + search_queries_both: [&SearchQueries; 2], + search_params_both: [SearchParams; 2], + orientations: [Orientation; 2], + full_scan_side: Eye, + extra_candidate_ids_both: [&VecRequests>; 2], + forced_anon_stats_ids: &VecRequests>, +) -> Result<[SearchResults; 2]> { + let cascade_start = Instant::now(); + for params in &search_params_both { + debug_assert_eq!(params.mode, HawkSearchMode::LinearScan); + } + + let first_eye = eye_index(full_scan_side); + let second_eye_side = full_scan_side.other(); + let second_eye = eye_index(second_eye_side); + + let n_requests = search_queries_both[0][first_eye].len(); + for sessions in sessions_both { + let n_sessions = sessions[LEFT].len(); + assert!(n_sessions > 0, "linear scan requires at least one session"); + assert_eq!(n_sessions, sessions[RIGHT].len()); + } + for search_queries in search_queries_both { + assert_eq!(n_requests, search_queries[LEFT].len()); + assert_eq!(n_requests, search_queries[RIGHT].len()); + } + for extra_candidate_ids in extra_candidate_ids_both { + assert_eq!(n_requests, extra_candidate_ids.len()); + } + assert_eq!(n_requests, forced_anon_stats_ids.len()); + + // Both eye registries and both orientation groups observe the same live + // VectorIds. Build the list once and share it across the fused stage. + let live_ids = { + let vector_store = sessions_both[0][first_eye][0].aby3_store.read().await; + let registry = vector_store.registry.read().await; + Arc::<[VectorId]>::from( + registry + .get_points() + .iter() + .enumerate() + .filter_map(|(serial_id, entry)| { + entry + .as_ref() + .map(|(version, ())| VectorId::new(serial_id as u32, *version)) + }) + .collect::>(), + ) + }; + let full_scan_ids = Arc::new(vec![live_ids.clone(); n_requests]); + let first_eye_comparisons = live_ids.len() * n_requests; + debug_assert!(live_ids.windows(2).all(|pair| pair[0] < pair[1])); + + let known_second_stage_ids_both = extra_candidate_ids_both.map(|extra_candidate_ids| { + Arc::new( + extra_candidate_ids + .iter() + .map(|extras| { + Arc::<[VectorId]>::from(collect_live_second_stage_ids( + &live_ids, + std::iter::empty(), + extras, + )) + }) + .collect::>(), + ) + }); + + for orientation in orientations { + tracing::info!( + eye = %full_scan_side, + orientation = orientation_label(orientation), + requests = n_requests, + vectors = live_ids.len(), + "Running full linear-scan stage" + ); + } + let prefetch_worker = { + let store = sessions_both[0][second_eye][0].aby3_store.read().await; + store.workers.clone() + }; + let first_eye_progress = Arc::new(LinearScanProgress::new(first_eye_comparisons)); + let contexts = orientations.map(|orientation| LinearScanEyeContext { + eye: full_scan_side, + stage: LinearScanStage::Full, + orientation, + }); + let forced_anon_stats_ids_shared = Arc::new(forced_anon_stats_ids.clone()); + let first_eye_scan = linear_scan_full_stage_paired( + sessions_both, + search_queries_both, + [&search_params_both[0], &search_params_both[1]], + contexts, + full_scan_ids, + forced_anon_stats_ids_shared.clone(), + LinearScanHooks { + prefetch: Some(LinearScanPrefetch { + worker: prefetch_worker.clone(), + }), + progress: Some(first_eye_progress.clone()), + }, + ); + // LUC and reauthentication candidates are public before the scan starts. + // Check them on the cold eye per orientation while the fused resident-eye + // scan is running. + let known_second_eye_scan = |index: usize| { + let progress = first_eye_progress.clone(); + let known_ids = known_second_stage_ids_both[index].clone(); + let forced = forced_anon_stats_ids_shared.clone(); + let search_params = search_params_both[index].clone(); + async move { + progress.wait_for_candidate_start().await; + linear_scan_eye( + sessions_both[index], + search_queries_both[index], + &search_params, + LinearScanEyeContext { + eye: second_eye_side, + stage: LinearScanStage::Candidate, + orientation: orientations[index], + }, + known_ids, + forced, + LinearScanHooks::default(), + ) + .await + } + }; + let (first_results_both, known_results_a, known_results_b) = tokio::try_join!( + first_eye_scan, + known_second_eye_scan(0), + known_second_eye_scan(1), + )?; + let mut second_results_both = [known_results_a, known_results_b]; + let prefetch_wait_start = Instant::now(); + prefetch_worker.wait_for_prefetch().await?; + let prefetch_wait_seconds = prefetch_wait_start.elapsed().as_secs_f64(); + metrics::histogram!("linear_scan_cold_prefetch_wait_duration").record(prefetch_wait_seconds); + for orientation in orientations { + metrics::histogram!( + "linear_scan_cascade_prefetch_wait_duration", + "eye" => eye_label(second_eye_side), + "orientation" => orientation_label(orientation), + ) + .record(prefetch_wait_seconds); + } + + // Discovered candidates per orientation, exactly as in the single-cascade + // path: retain live IDs, drop already-checked known candidates. + let mut discovered_ids_both = Vec::with_capacity(2); + let mut candidate_counts = [0usize; 2]; + for index in 0..2 { + let mut discovered_second_stage_ids = Vec::with_capacity(n_requests); + let mut discovered_candidate_count = 0usize; + for (plans, known_ids) in first_results_both[index] + .iter() + .zip(known_second_stage_ids_both[index].iter()) + { + let mut ids = collect_live_second_stage_ids( + &live_ids, + plans + .iter() + .flat_map(|plan| &plan.classified.anon_stats_matches.results) + .map(|(id, _)| *id), + &[], + ); + ids.retain(|id| known_ids.binary_search(id).is_err()); + discovered_candidate_count += ids.len(); + discovered_second_stage_ids.push(Arc::<[VectorId]>::from(ids)); + } + let known_candidate_count = known_second_stage_ids_both[index] + .iter() + .map(|ids| ids.len()) + .sum::(); + let candidate_count = known_candidate_count + discovered_candidate_count; + candidate_counts[index] = candidate_count; + + tracing::info!( + eye = %second_eye_side, + orientation = orientation_label(orientations[index]), + requests = n_requests, + candidates = candidate_count, + known_candidates = known_candidate_count, + discovered_candidates = discovered_candidate_count, + "Running candidate-only linear-scan stage" + ); + metrics::counter!("linear_scan_second_eye_candidates_total") + .increment(candidate_count as u64); + discovered_ids_both.push((discovered_candidate_count, discovered_second_stage_ids)); + } + + let discovered_scan = |index: usize, discovered: Vec>| { + let search_params = search_params_both[index].clone(); + let forced = forced_anon_stats_ids_shared.clone(); + async move { + linear_scan_eye( + sessions_both[index], + search_queries_both[index], + &search_params, + LinearScanEyeContext { + eye: second_eye_side, + stage: LinearScanStage::Candidate, + orientation: orientations[index], + }, + Arc::new(discovered), + forced, + LinearScanHooks::default(), + ) + .await + } + }; + let mut discovered_iter = discovered_ids_both.into_iter(); + let (discovered_count_a, discovered_ids_a) = discovered_iter.next().expect("two orientations"); + let (discovered_count_b, discovered_ids_b) = discovered_iter.next().expect("two orientations"); + let (discovered_results_a, discovered_results_b) = tokio::try_join!( + async { + if discovered_count_a > 0 { + discovered_scan(0, discovered_ids_a).await.map(Some) + } else { + Ok(None) + } + }, + async { + if discovered_count_b > 0 { + discovered_scan(1, discovered_ids_b).await.map(Some) + } else { + Ok(None) + } + }, + )?; + if let Some(results) = discovered_results_a { + merge_linear_scan_results(&mut second_results_both[0], results); + } + if let Some(results) = discovered_results_b { + merge_linear_scan_results(&mut second_results_both[1], results); + } + + let elapsed_seconds = cascade_start.elapsed().as_secs_f64(); + for index in 0..2 { + let orientation = orientation_label(orientations[index]); + let candidate_count = candidate_counts[index]; + let total_comparisons = first_eye_comparisons + candidate_count; + let comparisons_per_second = total_comparisons as f64 / elapsed_seconds.max(f64::EPSILON); + let second_eye_candidate_fraction = if first_eye_comparisons == 0 { + 0.0 + } else { + candidate_count as f64 / first_eye_comparisons as f64 + }; + metrics::counter!( + "linear_scan_cascade_comparisons_total", + "orientation" => orientation, + ) + .increment(total_comparisons as u64); + metrics::histogram!( + "linear_scan_cascade_duration", + "orientation" => orientation, + ) + .record(elapsed_seconds); + metrics::histogram!( + "linear_scan_cascade_comparisons_per_second", + "orientation" => orientation, + ) + .record(comparisons_per_second); + metrics::histogram!( + "linear_scan_second_eye_candidate_fraction", + "eye" => eye_label(second_eye_side), + "orientation" => orientation, + ) + .record(second_eye_candidate_fraction); + tracing::info!( + orientation = orientation, + full_scan_eye = eye_label(full_scan_side), + candidate_eye = eye_label(second_eye_side), + requests = n_requests, + database_records = live_ids.len(), + rotations_per_comparison = ROTATIONS, + first_eye_comparisons, + second_eye_comparisons = candidate_count, + total_comparisons, + second_eye_candidate_fraction, + prefetch_wait_seconds, + elapsed_seconds, + comparisons_per_second, + "LINEAR_SCAN_CASCADE_SUMMARY" + ); + } + + let [first_a, first_b] = first_results_both; + let [second_a, second_b] = second_results_both; + let pack = |first_results, second_results| match full_scan_side { + Eye::Left => [first_results, second_results], + Eye::Right => [second_results, first_results], + }; + Ok([pack(first_a, second_a), pack(first_b, second_b)]) +} + fn merge_linear_scan_plan(target: &mut HawkInsertPlan, mut source: HawkInsertPlan) { fn merge_matches(target: &mut SaturableMatches, mut source: SaturableMatches) { target.results.append(&mut source.results); @@ -716,6 +1462,34 @@ fn merge_linear_scan_plan(target: &mut HawkInsertPlan, mut source: HawkInsertPla .append(&mut source.classified.partial_match_rotations); } +fn merge_linear_scan_results( + target: &mut VecRequests>, + source: VecRequests>, +) { + for (target, source) in target.iter_mut().zip(source) { + let target = target.center_mut(); + merge_linear_scan_plan(target, source.into_center()); + + // The early known set and the later threshold-discovered set are each + // serial-ID ordered but may interleave. Restore the single sorted order + // produced by the non-overlapped scan; stable sorting also preserves + // rotation order for repeated anonymous-statistics IDs. + target.classified.matches.results.sort_by_key(|(id, _)| *id); + target + .classified + .anon_stats_matches + .results + .sort_by_key(|(id, _)| *id); + if let Some(pre_extension) = &mut target.classified.pre_extension { + pre_extension.results.sort_by_key(|(id, _)| *id); + } + target + .classified + .partial_match_rotations + .sort_by_key(|(id, _)| *id); + } +} + #[instrument(level = "trace", target = "searcher::network", skip_all)] async fn per_session( session: &HawkSession, @@ -726,18 +1500,6 @@ async fn per_session( batch: Batch, ) -> Result<()> { let inner = async { - // Linear scan does not build graph links for identity updates. The - // shared HNSW path expresses that operation as a no-match search, but - // materializing every live VectorId would be pure overhead here. - if search_params.mode == HawkSearchMode::LinearScan && !search_params.do_match { - let graph_store = session.graph_store.clone().read_owned().await; - for task in batch.tasks { - let query = search_queries[batch.i_eye][task.i_request][task.i_rotation]; - tx.send((task.id(), empty_linear_scan_plan(query, &graph_store)))?; - } - return Ok(()); - } - let mut vector_store = session.aby3_store.write().await; let graph_store = session.graph_store.clone().read_owned().await; @@ -841,47 +1603,33 @@ async fn per_linear_scan_query( let mut classified = ClassifiedMatches::default(); if search_params.do_match { + classified + .matches + .results + .reserve(vector_ids.len().min(4096)); + classified + .anon_stats_matches + .results + .reserve(vector_ids.len().min(4096)); + for ids in vector_ids.chunks(LINEAR_SCAN_CHUNK_SIZE) { let forced_anon_stats_vectors = forced_anon_stats_ids .iter() .filter_map(|id| ids.binary_search(id).ok()) .collect::>(); let thresholds = aby3_store - .eval_distance_batch_full_rotation_thresholds_with_forced_anon_stats( + .eval_distance_batch_full_rotation_thresholds_fused_with_forced_anon_stats( &query, ids, &forced_anon_stats_vectors, ) .await?; - classified.matches.results.extend( - ids.iter() - .copied() - .zip(&thresholds.matches) - .filter_map(|(id, &distance)| distance.map(|distance| (id, distance))), - ); - classified.anon_stats_matches.results.extend( - thresholds - .anon_stats_matches - .into_iter() - .map(|(vector, _rotation, distance)| (ids[vector], distance)), + extend_classified_from_thresholds( + &mut classified, + ids, + thresholds, + search_params.return_partial_results, ); - - if search_params.return_partial_results { - classified.partial_match_rotations.extend( - ids.iter() - .copied() - .zip(thresholds.match_rotations) - .filter_map(|(id, rotations)| { - (!rotations.is_empty()).then(|| { - let rotations = rotations - .into_iter() - .map(|rotation| rotation as i8 - 15) - .collect(); - (id, rotations) - }) - }), - ); - } } classified.linear_scan_supermatch_threshold = search_params @@ -907,6 +1655,133 @@ async fn per_linear_scan_query( }) } +/// Classify one chunk's opened threshold results into the accumulating +/// [`ClassifiedMatches`]. Shared by the single-orientation and paired scans. +fn extend_classified_from_thresholds( + classified: &mut ClassifiedMatches, + ids: &[VectorId], + thresholds: FullRotationThresholdResult, + return_partial_results: bool, +) { + classified.matches.results.extend( + ids.iter() + .copied() + .zip(&thresholds.matches) + .filter_map(|(id, &distance)| distance.map(|distance| (id, distance))), + ); + classified.anon_stats_matches.results.extend( + thresholds + .anon_stats_matches + .into_iter() + .map(|(vector, _rotation, distance)| (ids[vector], distance)), + ); + + if return_partial_results { + classified.partial_match_rotations.extend( + ids.iter() + .copied() + .zip(thresholds.match_rotations) + .filter_map(|(id, rotations)| { + (!rotations.is_empty()).then(|| { + let rotations = rotations + .into_iter() + .map(|rotation| rotation as i8 - 15) + .collect(); + (id, rotations) + }) + }), + ); + } +} + +/// Threshold rounds and classification for one fused chunk whose local dot +/// contributions were already computed (typically pipelined on the worker +/// pool while the previous chunk's thresholds ran). Each orientation's +/// threshold protocol runs on its own session, so the per-orientation network +/// transcript is identical to the unfused path. +#[allow(clippy::too_many_arguments)] +async fn per_linear_scan_chunk_pair( + queries: [Aby3Query; 2], + search_params: [&SearchParams; 2], + stores: (&mut Aby3Store, &mut Aby3Store), + graph_stores: (&GraphMem, &GraphMem), + contributions: [Vec>; 2], + vector_ids: &[VectorId], + forced_anon_stats_ids: &[VectorId], +) -> Result<[HawkInsertPlan; 2]> { + let start = Instant::now(); + let (store_a, store_b) = stores; + let mut classified = [ClassifiedMatches::default(), ClassifiedMatches::default()]; + debug_assert_eq!(search_params[0].do_match, search_params[1].do_match); + debug_assert!(vector_ids.len() <= LINEAR_SCAN_CHUNK_SIZE); + + if search_params[0].do_match { + let ids = vector_ids; + let forced_anon_stats_vectors = forced_anon_stats_ids + .iter() + .filter_map(|id| ids.binary_search(id).ok()) + .collect::>(); + let [contributions_a, contributions_b] = contributions; + let (thresholds_a, thresholds_b) = tokio::try_join!( + store_a.eval_full_rotation_thresholds_fused_from_contributions_with_forced_anon_stats( + contributions_a, + ids.len(), + &forced_anon_stats_vectors, + ), + store_b.eval_full_rotation_thresholds_fused_from_contributions_with_forced_anon_stats( + contributions_b, + ids.len(), + &forced_anon_stats_vectors, + ), + )?; + extend_classified_from_thresholds( + &mut classified[0], + ids, + thresholds_a, + search_params[0].return_partial_results, + ); + extend_classified_from_thresholds( + &mut classified[1], + ids, + thresholds_b, + search_params[1].return_partial_results, + ); + + for (side, params) in classified.iter_mut().zip(&search_params) { + side.linear_scan_supermatch_threshold = params + .hnsw_supermatch + .as_ref() + .map(|searcher| searcher.params.get_ef_search(0)); + } + } + + metrics::histogram!("linear_scan_query_duration").record(start.elapsed().as_secs_f64()); + metrics::counter!("linear_scan_vectors_total").increment(2 * vector_ids.len() as u64); + + let [classified_a, classified_b] = classified; + let as_plan = |query: Aby3Query, as_of, classified| HawkInsertPlan { + plan: InsertPlanV { + query, + links: Vec::new(), + update_ep: UpdateEntryPoint::False, + as_of, + }, + classified, + }; + Ok([ + as_plan( + queries[0], + graph_stores.0.last_update_seq_no, + classified_a, + ), + as_plan( + queries[1], + graph_stores.1.last_update_seq_no, + classified_b, + ), + ]) +} + /// Preserve the three-slot HNSW-shaped result container without rescanning the /// database for the two non-central base rotations. Matching merges all slots, /// so placing the full 31-rotation result in the center is behaviorally @@ -1201,22 +2076,6 @@ mod tests { ); } - #[tokio::test] - async fn empty_search_does_not_require_sessions() -> Result<()> { - let sessions: BothEyes> = [Vec::new(), Vec::new()]; - let queries: SearchQueries<0> = Arc::new([Vec::new(), Vec::new()]); - let request_ids: SearchIds = Arc::new(Vec::new()); - let params = SearchParams::new_no_match( - Arc::new(HnswSearcher::new_with_test_parameters()), - HawkSearchMode::LinearScan, - ); - - let results = search(&sessions, &queries, &request_ids, params).await?; - - assert!(results.iter().all(Vec::is_empty)); - Ok(()) - } - #[test] fn second_stage_ids_retain_only_current_live_versions() { let live_v1 = VectorId::new(2, 1); @@ -1304,7 +2163,20 @@ mod tests { 'L', ); - let extra_candidate_ids = vec![Vec::new(); batch_size]; + // Exercise both overlapped known candidates and candidates discovered + // only by the full-eye threshold. Even-numbered requests include their + // exact match up front; odd-numbered requests include a different + // live record so their exact match must be merged in afterward. + let extra_candidate_ids = (0..batch_size) + .map(|index| { + let extra = if index.is_multiple_of(2) { + index + } else { + (index + 1) % batch_size + }; + vec![VectorId::from_0_index(extra as u32)] + }) + .collect::>(); let forced_anon_stats_ids = vec![Vec::new(); batch_size]; let result = linear_scan_cascade( &sessions, @@ -1319,6 +2191,20 @@ mod tests { for side in result { for (query_index, rotations) in side.iter().enumerate() { + for rotation in rotations.iter() { + assert!(rotation + .classified + .matches + .results + .windows(2) + .all(|pair| pair[0].0 <= pair[1].0)); + assert!(rotation + .classified + .anon_stats_matches + .results + .windows(2) + .all(|pair| pair[0].0 <= pair[1].0)); + } assert!(rotations.iter().any(|rotation| { rotation .classified @@ -1337,6 +2223,166 @@ mod tests { Ok(actor) } + #[tokio::test] + async fn paired_linear_scan_matches_single_cascades() -> Result<()> { + let actors = setup_linear_scan_actors().await?; + + parallelize(actors.into_iter().map(go_paired_linear_scan)).await?; + + Ok(()) + } + + /// Opened-result projection of one plan: everything that is deterministic + /// across protocol runs. Distance *shares* differ between session sets by + /// construction, so they are excluded. + #[allow(clippy::type_complexity)] + fn opened_projection( + results: &SearchResults, + ) -> Vec< + Vec< + Vec<( + Aby3Query, + Vec, + bool, + Vec, + bool, + Vec<(VectorId, Vec)>, + Option, + )>, + >, + > { + results + .iter() + .map(|side| { + side.iter() + .map(|rotations| { + rotations + .iter() + .map(|plan| { + ( + plan.plan.query, + plan.classified + .matches + .results + .iter() + .map(|(id, _)| *id) + .collect(), + plan.classified.matches.saturated, + plan.classified + .anon_stats_matches + .results + .iter() + .map(|(id, _)| *id) + .collect(), + plan.classified.anon_stats_matches.saturated, + plan.classified.partial_match_rotations.clone(), + plan.classified.linear_scan_supermatch_threshold, + ) + }) + .collect() + }) + .collect() + }) + .collect() + } + + async fn go_paired_linear_scan(mut actor: HawkActor) -> Result { + init_iris_db(&mut actor).await?; + + let sessions_normal = actor.new_sessions().await?; + let sessions_mirror = actor.new_sessions().await?; + let sessions_reference = actor.new_sessions().await?; + let batch_size = 3; + let request = make_request(batch_size, actor.party_id); + request.cache_into(&actor.worker_pools).await?; + let search_params = SearchParams::new( + actor.searcher(), + HawkSearchMode::LinearScan, + true, + None, + 0, + true, + #[cfg(feature = "phase_trace")] + 'L', + ); + + // Same candidate shapes as `go_linear_scan`: overlapped known + // candidates plus threshold-discovered ones. + let extra_candidate_ids = (0..batch_size) + .map(|index| { + let extra = if index.is_multiple_of(2) { + index + } else { + (index + 1) % batch_size + }; + vec![VectorId::from_0_index(extra as u32)] + }) + .collect::>(); + let forced_anon_stats_ids = vec![Vec::new(); batch_size]; + + let queries_normal = request.queries(Orientation::Normal); + let queries_mirror = request.queries(Orientation::Mirror); + let [paired_normal, paired_mirror] = linear_scan_cascade_paired( + [&sessions_normal, &sessions_mirror], + [&queries_normal, &queries_mirror], + [search_params.clone(), search_params.clone()], + [Orientation::Normal, Orientation::Mirror], + Eye::Left, + [&extra_candidate_ids, &extra_candidate_ids], + &forced_anon_stats_ids, + ) + .await?; + + let reference_normal = linear_scan_cascade( + &sessions_reference, + &queries_normal, + search_params.clone(), + Orientation::Normal, + Eye::Left, + &extra_candidate_ids, + &forced_anon_stats_ids, + ) + .await?; + let reference_mirror = linear_scan_cascade( + &sessions_reference, + &queries_mirror, + search_params, + Orientation::Mirror, + Eye::Left, + &extra_candidate_ids, + &forced_anon_stats_ids, + ) + .await?; + + assert_eq!( + opened_projection(&paired_normal), + opened_projection(&reference_normal), + "normal orientation" + ); + assert_eq!( + opened_projection(&paired_mirror), + opened_projection(&reference_mirror), + "mirror orientation" + ); + // The fused pass must find the planted matches, not merely agree with + // an equally-empty reference. + for side in &paired_normal { + for (query_index, rotations) in side.iter().enumerate() { + assert!(rotations.iter().any(|rotation| { + rotation + .classified + .matches + .results + .iter() + .any(|(id, _)| *id == VectorId::from_0_index(query_index as u32)) + })); + } + } + + actor.sync_peers().await?; + Ok(actor) + } + async fn go_search(mut actor: HawkActor) -> Result { init_iris_db(&mut actor).await?; init_graph(&mut actor).await?; diff --git a/iris-mpc-cpu/src/execution/hawk_main/worker_pool_initializer.rs b/iris-mpc-cpu/src/execution/hawk_main/worker_pool_initializer.rs index 9064150ebb..d85137de45 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/worker_pool_initializer.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/worker_pool_initializer.rs @@ -7,10 +7,10 @@ use crate::execution::hawk_main::iris_worker::{ }; use crate::execution::hawk_main::{BothEyes, HawkOps, LEFT, RIGHT}; use crate::hawkers::aby3::aby3_store::{ - Aby3SharedIrises, Aby3SharedIrisesRef, Aby3Store, DistanceMode, VectorIdRegistryRef, + Aby3SharedIrises, Aby3Store, DistanceMode, VectorIdRegistryRef, }; -use crate::hawkers::shared_irises::SharedIrises; -use crate::protocol::shared_iris::GaloisRingSharedIris; +use crate::hawkers::shared_irises::{SharedIrises, SharedIrisesRef}; +use crate::protocol::shared_iris::{GaloisRingSharedIris, ResidentIris, ResidentLayout}; use ampc_server_utils::shutdown_handler::ShutdownHandler; use async_trait::async_trait; use eyre::Result; @@ -64,6 +64,10 @@ pub struct LocalWorkerPoolInitializer { pub distance_mode: DistanceMode, pub numa: bool, pub mode: LocalInitMode, + /// Resident representation of the pools' iris stores. `U16` (default) + /// keeps plain `ArcIris` values as required by the HNSW hot paths; + /// exact-scan actors opt into `preferred_scan_layout()`. + pub layout: ResidentLayout, } impl LocalWorkerPoolInitializer { @@ -73,9 +77,16 @@ impl LocalWorkerPoolInitializer { distance_mode, numa, mode: LocalInitMode::Empty, + layout: ResidentLayout::U16, } } + /// Choose the resident representation of the pools' iris stores. + pub fn with_resident_layout(mut self, layout: ResidentLayout) -> Self { + self.layout = layout; + self + } + pub fn new_seeded( party_id: usize, distance_mode: DistanceMode, @@ -87,6 +98,7 @@ impl LocalWorkerPoolInitializer { distance_mode, numa, mode: LocalInitMode::Seeded(seed_stores), + layout: ResidentLayout::U16, } } @@ -101,6 +113,7 @@ impl LocalWorkerPoolInitializer { distance_mode, numa, mode: LocalInitMode::LoadFromDb(params), + layout: ResidentLayout::U16, } } } @@ -113,23 +126,25 @@ impl WorkerPoolInitializer for LocalWorkerPoolInitializer { distance_mode, numa, mode, + layout, } = *self; // Materialize the iris stores. `Seeded` installs caller-provided // stores; the rest start blank. - let iris_stores: BothEyes = match &mode { - LocalInitMode::Seeded(seeds) => { - let [left, right] = seeds.clone(); - [SharedIrises::to_arc(left), SharedIrises::to_arc(right)] - } - _ => [ - Aby3Store::::new_storage(None).to_arc(), - Aby3Store::::new_storage(None).to_arc(), - ], + let iris_stores: BothEyes> = match &mode { + LocalInitMode::Seeded(seeds) => seeds.clone().map(|seed| { + seed.map_values(|iris| ResidentIris::from_arc(iris, layout)) + .to_arc() + }), + _ => [LEFT, RIGHT].map(|_| { + Aby3Store::::new_storage(None) + .map_values(|iris| ResidentIris::from_arc(iris, layout)) + .to_arc() + }), }; let workers_handle: BothEyes = - [LEFT, RIGHT].map(|side| init_workers(side, iris_stores[side].clone(), numa)); + [LEFT, RIGHT].map(|side| init_workers(side, iris_stores[side].clone(), numa, layout)); let mut db_size: usize = 0; let mut cold_storage: Option<(Store, usize, usize, usize)> = None; @@ -219,6 +234,7 @@ impl WorkerPoolInitializer for LocalWorkerPoolInitializer { LocalIrisWorkerPool::new_cold( workers_handle[cold_side].clone(), iris_stores[cold_side].clone(), + layout, distance_mode, party_id, ColdStorageInit { @@ -246,6 +262,7 @@ impl WorkerPoolInitializer for LocalWorkerPoolInitializer { LocalIrisWorkerPool::new( workers_handle[side].clone(), iris_stores[side].clone(), + layout, distance_mode, party_id, ) @@ -336,9 +353,13 @@ mod tests { #[tokio::test] async fn single_eye_loader_does_not_materialize_cold_eye() -> Result<()> { - let stores: BothEyes = - [LEFT, RIGHT].map(|_| Aby3Store::::new_storage(None).to_arc()); - let handles = [LEFT, RIGHT].map(|side| init_workers(side, stores[side].clone(), false)); + let stores: BothEyes> = [LEFT, RIGHT].map(|_| { + Aby3Store::::new_storage(None) + .map_values(|iris| ResidentIris::from_arc(iris, ResidentLayout::U16)) + .to_arc() + }); + let handles = [LEFT, RIGHT] + .map(|side| init_workers(side, stores[side].clone(), false, ResidentLayout::U16)); let iris = GaloisRingSharedIris::default_for_party(0); let id = VectorId::from_0_index(7); let mut loader = FanoutLoader { diff --git a/iris-mpc-cpu/src/genesis/batch_generator.rs b/iris-mpc-cpu/src/genesis/batch_generator.rs index 7096338d73..ae506dc050 100644 --- a/iris-mpc-cpu/src/genesis/batch_generator.rs +++ b/iris-mpc-cpu/src/genesis/batch_generator.rs @@ -499,7 +499,20 @@ mod tests { .to_arc() }); let worker_pools = [LEFT, RIGHT].map(|side| { - LocalIrisWorkerPool::new_local(iris_stores[side].clone(), HAWK_DISTANCE_MODE, PARTY_ID) + use crate::protocol::shared_iris::{ResidentIris, ResidentLayout}; + let resident_store = iris_stores[side] + .data + .try_read() + .unwrap() + .clone() + .map_values(|iris| ResidentIris::from_arc(iris, ResidentLayout::U16)) + .to_arc(); + LocalIrisWorkerPool::new_local( + resident_store, + ResidentLayout::U16, + HAWK_DISTANCE_MODE, + PARTY_ID, + ) }); (registries, worker_pools, SIZE_OF_IRIS_DB) } diff --git a/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs b/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs index d15e6f3131..abb9a5ead8 100644 --- a/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs +++ b/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs @@ -27,10 +27,12 @@ use crate::{ RingElement, }, }; +#[cfg(test)] +use ampc_actor_utils::protocol::fhd_ops::fhd_greater_than_anon_stats_threshold; use ampc_actor_utils::protocol::{ binary::open_bin, fhd_ops::{ - fhd_greater_than_anon_stats_threshold, fhd_greater_than_threshold_pre_lifted_masks, + fhd_greater_than_anon_stats_from_galois, fhd_greater_than_threshold_pre_lifted_masks, lift_fhd_mask_dots, }, ops::batch_signed_lift_vec, @@ -639,29 +641,9 @@ where Ok((distances, rotation_matches)) } - /// Evaluate all 31 rotations in one worker pass over `vectors`, then run - /// the same AMPC lift and oblivious minimum as the HNSW distance path. - /// This is the exact-scan counterpart of three overlapping 11-rotation - /// HNSW windows. - #[instrument(level = "trace", target = "searcher::network", skip_all)] - pub async fn eval_distance_batch_full_rotations( - &mut self, - query: &Aby3Query, - vectors: &[VectorId], - ) -> Result>> { - if vectors.is_empty() { - return Ok(Vec::new()); - } - let rotation_distances = self.full_rotation_distances(query, vectors).await?; - self.oblivious_min_distance_batch(distance_fn::transpose_from_flat_with_rotations( - &rotation_distances, - ROTATIONS, - )) - .await - } - /// Full-rotation exact scan with the opened per-rotation match metadata - /// required by the GPU-compatible partial-results response. + /// retained as a correctness oracle for the fused production path. + #[cfg(test)] #[instrument(level = "trace", target = "searcher::network", skip_all)] pub async fn eval_distance_batch_full_rotations_with_rotation_matches( &mut self, @@ -692,6 +674,7 @@ where Ok((distances, rotation_matches)) } + #[cfg(test)] async fn full_rotation_distances( &mut self, query: &Aby3Query, @@ -713,12 +696,27 @@ where self.lift_distances(dot_shares).await } + #[cfg(test)] #[instrument(level = "trace", target = "searcher::network", skip_all)] async fn full_rotation_dot_shares( &mut self, query: &Aby3Query, vectors: &[VectorId], ) -> Result>> { + let ds_and_ts = self.full_rotation_dot_contributions(query, vectors).await?; + galois_ring_to_rep3(&mut self.session, ds_and_ts).await + } + + /// Compute the local additive dot-product contributions before refreshing + /// them into replicated shares. The fused exact-scan path consumes this + /// representation directly and materializes scalar shares only for public + /// candidates. + #[instrument(level = "trace", target = "searcher::network", skip_all)] + async fn full_rotation_dot_contributions( + &mut self, + query: &Aby3Query, + vectors: &[VectorId], + ) -> Result>> { eyre::ensure!( self.distance_fn.mode == DistanceMode::MinRotation, "full-rotation scan requires min-rotation distance mode" @@ -729,11 +727,45 @@ where ); metrics::counter!("distance_evaluations_total").increment(vectors.len() as u64); metrics::histogram!("distance_evaluations_batch_size").record(vectors.len() as f64); - let ds_and_ts = self - .workers + self.workers .compute_dot_products_full_rotations(*query, vectors.to_vec()) - .await?; - galois_ring_to_rep3(&mut self.session, ds_and_ts).await + .await + } + + /// Dispatch both orientations' local dot contributions for one chunk as a + /// spawned task on the worker pool. The caller can drive the previous + /// chunk's threshold rounds while this chunk's dot products compute, + /// keeping the dot workers fed. Each returned side is identical to a + /// separate [`Self::full_rotation_dot_contributions`] call; only the + /// worker-level target streaming is shared. This performs no network + /// communication, so fusing and pipelining the dot passes is invisible to + /// the MPC transcript. + pub fn spawn_full_rotation_dot_contributions_pair( + &self, + queries: [&Aby3Query; 2], + vectors: &[VectorId], + ) -> tokio::task::JoinHandle>; 2]>> { + let mode = self.distance_fn.mode; + let specs = [*queries[0], *queries[1]]; + let workers = self.workers.clone(); + let vectors = vectors.to_vec(); + tokio::spawn(async move { + eyre::ensure!( + mode == DistanceMode::MinRotation, + "full-rotation scan requires min-rotation distance mode" + ); + for query in specs { + eyre::ensure!( + query.rotation == crate::execution::hawk_main::iris_worker::CENTER_ROTATION, + "full-rotation scan must start from the center query rotation" + ); + } + metrics::counter!("distance_evaluations_total").increment(2 * vectors.len() as u64); + metrics::histogram!("distance_evaluations_batch_size").record(vectors.len() as f64); + workers + .compute_dot_products_full_rotations_pair(specs, vectors) + .await + }) } /// Check whether a batch of distances are matches at the given threshold. @@ -751,32 +783,14 @@ where } impl Aby3Store { - /// Evaluate the exact-scan thresholds exactly like the GPU actor: compare - /// every rotation in a batch, OR the opened rotation bits per vector, and - /// never materialize an oblivious minimum distance. + /// Unfused exact-scan threshold implementation retained as a correctness + /// oracle for the production path. + #[cfg(test)] #[instrument(level = "trace", target = "searcher::network", skip_all)] pub async fn eval_distance_batch_full_rotation_thresholds( &mut self, query: &Aby3Query, vectors: &[VectorId], - ) -> Result { - self.eval_distance_batch_full_rotation_thresholds_with_forced_anon_stats( - query, - vectors, - &[], - ) - .await - } - - /// GPU-compatible exact scan with selected records retained for anonymous - /// statistics regardless of the anonymous-statistics threshold. The CUDA - /// actor uses this for reauthentication targets. - #[instrument(level = "trace", target = "searcher::network", skip_all)] - pub async fn eval_distance_batch_full_rotation_thresholds_with_forced_anon_stats( - &mut self, - query: &Aby3Query, - vectors: &[VectorId], - forced_anon_stats_vectors: &[usize], ) -> Result { if vectors.is_empty() { return Ok(FullRotationThresholdResult { @@ -806,7 +820,7 @@ impl Aby3Store { let anon_gt = fhd_greater_than_anon_stats_threshold(&mut self.session, &code_dots, &mask_dots) .await?; - let mut anon_rotation_bits = open_bin(&mut self.session, &anon_gt) + let anon_rotation_bits = open_bin(&mut self.session, &anon_gt) .await? .into_iter() .map(|bit| !bool::from(bit)) @@ -817,17 +831,6 @@ impl Aby3Store { "anonymous threshold result has unexpected length" ); - // CUDA unions the reauthentication target into the public candidate - // bitmap and stores all of its rotations, even those outside the - // anonymous-statistics threshold. - for &vector in forced_anon_stats_vectors { - eyre::ensure!( - vector < vectors.len(), - "forced anonymous-statistics vector index is out of bounds" - ); - anon_rotation_bits[vector * ROTATIONS..(vector + 1) * ROTATIONS].fill(true); - } - // Exactly like the GPU actor, collapse the wider public prefilter to a // record bitmap, then run the strict threshold over all 31 rotations // of every surviving record. Although the strict threshold is a subset @@ -922,6 +925,209 @@ impl Aby3Store { match_rotations, }) } + + /// Allocation-fused threshold implementation used by the production CPU + /// exact scan. + /// + /// It preserves the Galois-to-Rep3 refresh and the threshold circuit's + /// network transcript, but bit-transposes the two refreshed components + /// directly instead of first allocating a dense scalar `Share` batch + /// and three mostly-zero packed component vectors. + #[instrument(level = "trace", target = "searcher::network", skip_all)] + pub async fn eval_distance_batch_full_rotation_thresholds_fused( + &mut self, + query: &Aby3Query, + vectors: &[VectorId], + ) -> Result { + self.eval_distance_batch_full_rotation_thresholds_fused_with_forced_anon_stats( + query, + vectors, + &[], + ) + .await + } + + /// GPU-compatible fused scan with selected records retained for anonymous + /// statistics regardless of the anonymous-statistics threshold. The CUDA + /// actor uses this for reauthentication targets. + #[instrument(level = "trace", target = "searcher::network", skip_all)] + pub async fn eval_distance_batch_full_rotation_thresholds_fused_with_forced_anon_stats( + &mut self, + query: &Aby3Query, + vectors: &[VectorId], + forced_anon_stats_vectors: &[usize], + ) -> Result { + if vectors.is_empty() { + return Ok(FullRotationThresholdResult { + matches: Vec::new(), + anon_stats_matches: Vec::new(), + match_rotations: Vec::new(), + }); + } + + let dot_contributions = self.full_rotation_dot_contributions(query, vectors).await?; + self.eval_full_rotation_thresholds_fused_from_contributions_with_forced_anon_stats( + dot_contributions, + vectors.len(), + forced_anon_stats_vectors, + ) + .await + } + + /// Threshold half of [`Self::eval_distance_batch_full_rotation_thresholds_fused`], + /// taking precomputed full-rotation dot contributions for `n_vectors` + /// records. The dot phase is pure local worker compute, so callers can + /// overlap the next chunk's dot products with this chunk's threshold + /// network rounds without changing the per-session wire transcript. + #[instrument(level = "trace", target = "searcher::network", skip_all)] + pub async fn eval_full_rotation_thresholds_fused_from_contributions( + &mut self, + dot_contributions: Vec>, + n_vectors: usize, + ) -> Result { + self.eval_full_rotation_thresholds_fused_from_contributions_with_forced_anon_stats( + dot_contributions, + n_vectors, + &[], + ) + .await + } + + #[instrument(level = "trace", target = "searcher::network", skip_all)] + pub async fn eval_full_rotation_thresholds_fused_from_contributions_with_forced_anon_stats( + &mut self, + dot_contributions: Vec>, + n_vectors: usize, + forced_anon_stats_vectors: &[usize], + ) -> Result { + if n_vectors == 0 { + return Ok(FullRotationThresholdResult { + matches: Vec::new(), + anon_stats_matches: Vec::new(), + match_rotations: Vec::new(), + }); + } + let vectors_len = n_vectors; + let expected_dots = vectors_len * ROTATIONS * 2; + eyre::ensure!( + dot_contributions.len() == expected_dots, + "full-rotation dot result has unexpected length" + ); + let (anon_gt, dot_shares) = + fhd_greater_than_anon_stats_from_galois(&mut self.session, dot_contributions).await?; + eyre::ensure!( + dot_shares.len() == vectors_len * ROTATIONS, + "fused full-rotation dot result has unexpected length" + ); + let mut anon_rotation_bits = open_bin(&mut self.session, &anon_gt) + .await? + .into_iter() + .map(|bit| !bool::from(bit)) + .collect::>(); + + eyre::ensure!( + anon_rotation_bits.len() == dot_shares.len(), + "anonymous threshold result has unexpected length" + ); + + // CUDA unions the reauthentication target into the public candidate + // bitmap and stores all of its rotations, even those outside the + // anonymous-statistics threshold. + for &vector in forced_anon_stats_vectors { + eyre::ensure!( + vector < vectors_len, + "forced anonymous-statistics vector index is out of bounds" + ); + anon_rotation_bits[vector * ROTATIONS..(vector + 1) * ROTATIONS].fill(true); + } + + let dot_count = dot_shares.len(); + let candidate_rotation_indices = gpu_candidate_rotation_indices(&anon_rotation_bits); + let (candidate_codes, candidate_raw_masks) = + dot_shares.select(&candidate_rotation_indices)?; + drop(dot_shares); + let candidate_lifted_masks = if candidate_raw_masks.is_empty() { + Vec::new() + } else { + lift_fhd_mask_dots(&mut self.session, &candidate_raw_masks).await? + }; + let mut match_rotation_bits = vec![false; dot_count]; + if !candidate_rotation_indices.is_empty() { + let match_gt = fhd_greater_than_threshold_pre_lifted_masks( + &mut self.session, + &candidate_codes, + &candidate_lifted_masks, + Threshold::Match.ratio(), + ) + .await?; + let candidate_match_bits = open_bin(&mut self.session, &match_gt) + .await? + .into_iter() + .map(|bit| !bool::from(bit)); + for (&index, is_match) in candidate_rotation_indices.iter().zip(candidate_match_bits) { + match_rotation_bits[index] = is_match; + } + } + + eyre::ensure!( + match_rotation_bits + .iter() + .zip(&anon_rotation_bits) + .all(|(&is_match, &is_anon_match)| !is_match || is_anon_match), + "strict match threshold produced a result outside the anonymous prefilter" + ); + + let anon_rotation_indices = anon_rotation_bits + .iter() + .enumerate() + .filter_map(|(index, &is_match)| is_match.then_some(index)) + .collect::>(); + let anon_codes = anon_rotation_indices + .iter() + .map(|index| { + let candidate_index = candidate_rotation_indices + .binary_search(index) + .expect("anonymous rotation must belong to a candidate record"); + candidate_codes[candidate_index] + }) + .collect::>(); + let lifted_anon_codes = if anon_codes.is_empty() { + Vec::new() + } else { + batch_signed_lift_vec(&mut self.session, anon_codes).await? + }; + + let mut matches = vec![None; vectors_len]; + let mut anon_stats_matches = Vec::with_capacity(anon_rotation_indices.len()); + for (&index, code_dot) in anon_rotation_indices.iter().zip(lifted_anon_codes) { + let vector = index / ROTATIONS; + let rotation = index % ROTATIONS; + let candidate_index = candidate_rotation_indices + .binary_search(&index) + .expect("anonymous rotation must belong to a candidate record"); + let distance = DistanceShare::new(code_dot, candidate_lifted_masks[candidate_index]); + anon_stats_matches.push((vector, rotation, distance)); + if match_rotation_bits[index] && matches[vector].is_none() { + matches[vector] = Some(distance); + } + } + + let match_rotations = match_rotation_bits + .chunks_exact(ROTATIONS) + .map(|bits| { + bits.iter() + .enumerate() + .filter_map(|(rotation, &is_match)| is_match.then_some(rotation)) + .collect() + }) + .collect(); + + Ok(FullRotationThresholdResult { + matches, + anon_stats_matches, + match_rotations, + }) + } } impl VectorStore for Aby3Store diff --git a/iris-mpc-cpu/src/hawkers/aby3/aby3_store/tests.rs b/iris-mpc-cpu/src/hawkers/aby3/aby3_store/tests.rs index ddd901dead..21c358494f 100644 --- a/iris-mpc-cpu/src/hawkers/aby3/aby3_store/tests.rs +++ b/iris-mpc-cpu/src/hawkers/aby3/aby3_store/tests.rs @@ -93,6 +93,144 @@ async fn full_rotation_threshold_scan_matches_min_distance_reference() -> Result Ok(()) } +#[tokio::test(flavor = "multi_thread")] +async fn fused_full_rotation_threshold_scan_matches_legacy() -> Result<()> { + let mut rng = AesRng::seed_from_u64(0x6675_7365_645f_6668); + let vectors_and_graphs = shared_random_setup(&mut rng, 3, NetworkType::Local).await?; + + let tasks = vectors_and_graphs + .into_iter() + .map(|(store, _graph)| async move { + let mut store = store.lock_owned().await; + let ids = [ + VectorId::from_0_index(0), + VectorId::from_0_index(1), + VectorId::from_0_index(2), + ]; + // The self-comparison exercises candidate reconstruction, strict + // matching, and anonymous-statistics distance persistence; the + // random records exercise the dense no-match path. + let query = store.cache_query_from_store(&ids[0]).await?; + let legacy = store + .eval_distance_batch_full_rotation_thresholds(&query, &ids) + .await?; + let fused = store + .eval_distance_batch_full_rotation_thresholds_fused(&query, &ids) + .await?; + Ok((legacy, fused)) + }); + let results = parallelize(tasks).await?; + + let open = |shares: Vec>| { + shares + .into_iter() + .reduce(|sum, share| sum + share) + .expect("one share per party") + .get_a() + .convert() + }; + + for (legacy, fused) in &results { + assert_eq!(legacy.match_rotations, fused.match_rotations); + assert_eq!(legacy.matches.len(), fused.matches.len()); + assert_eq!( + legacy + .matches + .iter() + .map(Option::is_some) + .collect::>(), + fused + .matches + .iter() + .map(Option::is_some) + .collect::>() + ); + assert_eq!( + legacy + .anon_stats_matches + .iter() + .map(|(vector, rotation, _)| (*vector, *rotation)) + .collect::>(), + fused + .anon_stats_matches + .iter() + .map(|(vector, rotation, _)| (*vector, *rotation)) + .collect::>() + ); + } + + for vector in 0..3 { + let legacy_present = results[0].0.matches[vector].is_some(); + assert!(results + .iter() + .all(|(legacy, _)| legacy.matches[vector].is_some() == legacy_present)); + assert!(results + .iter() + .all(|(_, fused)| fused.matches[vector].is_some() == legacy_present)); + if legacy_present { + let legacy_code = open( + results + .iter() + .map(|(legacy, _)| legacy.matches[vector].unwrap().code_dot) + .collect(), + ); + let fused_code = open( + results + .iter() + .map(|(_, fused)| fused.matches[vector].unwrap().code_dot) + .collect(), + ); + let legacy_mask = open( + results + .iter() + .map(|(legacy, _)| legacy.matches[vector].unwrap().mask_dot) + .collect(), + ); + let fused_mask = open( + results + .iter() + .map(|(_, fused)| fused.matches[vector].unwrap().mask_dot) + .collect(), + ); + assert_eq!((legacy_code, legacy_mask), (fused_code, fused_mask)); + } + } + + let anon_len = results[0].0.anon_stats_matches.len(); + assert!(results.iter().all(|(legacy, fused)| { + legacy.anon_stats_matches.len() == anon_len && fused.anon_stats_matches.len() == anon_len + })); + for index in 0..anon_len { + let legacy_code = open( + results + .iter() + .map(|(legacy, _)| legacy.anon_stats_matches[index].2.code_dot) + .collect(), + ); + let fused_code = open( + results + .iter() + .map(|(_, fused)| fused.anon_stats_matches[index].2.code_dot) + .collect(), + ); + let legacy_mask = open( + results + .iter() + .map(|(legacy, _)| legacy.anon_stats_matches[index].2.mask_dot) + .collect(), + ); + let fused_mask = open( + results + .iter() + .map(|(_, fused)| fused.anon_stats_matches[index].2.mask_dot) + .collect(), + ); + assert_eq!((legacy_code, legacy_mask), (fused_code, fused_mask)); + } + + Ok(()) +} + #[tokio::test(flavor = "multi_thread")] async fn test_gr_hnsw() -> Result<()> { let mut rng = AesRng::seed_from_u64(0_u64); diff --git a/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs b/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs index f8a45bc89d..199b169660 100644 --- a/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs +++ b/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs @@ -76,8 +76,19 @@ pub async fn setup_local_aby3_players_with_preloaded_db( .zip(storages) .map(|(session, storage)| { let party_id = session.network_session.own_role.index(); + let layout = crate::protocol::shared_iris::preferred_scan_layout(); + let resident_storage = storage + .data + .try_read() + .unwrap() + .clone() + .map_values(|iris| { + crate::protocol::shared_iris::ResidentIris::from_arc(iris, layout) + }) + .to_arc(); let workers: Arc = Arc::new(LocalIrisWorkerPool::new_local( - storage.clone(), + resident_storage, + layout, plain_store.distance_mode, party_id, )); @@ -99,9 +110,15 @@ pub async fn setup_local_store_aby3_players(network_t: NetworkType) -> Result::new_storage(None).to_arc(); + let layout = crate::protocol::shared_iris::preferred_scan_layout(); + let storage = Aby3Store::::new_storage(None) + .map_values(|iris| { + crate::protocol::shared_iris::ResidentIris::from_arc(iris, layout) + }) + .to_arc(); let workers: Arc = Arc::new(LocalIrisWorkerPool::new_local( storage.clone(), + layout, TEST_DISTANCE_MODE, party_id, )); diff --git a/iris-mpc-cpu/src/hawkers/shared_irises.rs b/iris-mpc-cpu/src/hawkers/shared_irises.rs index d6e9a3f869..7cd55d14c3 100644 --- a/iris-mpc-cpu/src/hawkers/shared_irises.rs +++ b/iris-mpc-cpu/src/hawkers/shared_irises.rs @@ -173,6 +173,23 @@ impl SharedIrises { } } + /// Convert every stored value (including the empty-iris template) while + /// preserving ids, versions, and checksums. Used to materialize a worker + /// pool store in its resident layout from a seed store. + pub fn map_values(self, f: impl Fn(I) -> J) -> SharedIrises { + SharedIrises { + points: self + .points + .into_iter() + .map(|opt| opt.map(|(v, iris)| (v, f(iris)))) + .collect(), + size: self.size, + next_id: self.next_id, + empty_iris: f(self.empty_iris), + set_hash: self.set_hash, + } + } + /// Create a metadata-only registry from this store. /// /// Preserves all VectorId presence, version, and checksum data but diff --git a/iris-mpc-cpu/src/protocol/ops.rs b/iris-mpc-cpu/src/protocol/ops.rs index 6e76c2bbfd..fb0dd9580f 100644 --- a/iris-mpc-cpu/src/protocol/ops.rs +++ b/iris-mpc-cpu/src/protocol/ops.rs @@ -458,23 +458,64 @@ fn accumulate_component_tiled_6x4( result_lane: usize, additive_shares: &mut [RingElement], ) { - const ROW_SIZE: usize = PrerotatedQueryRowMajor::ROW_SIZE; + // Iterate target groups in the outer loop so each target is streamed as + // one contiguous `rows * ROW_SIZE` run. Row-outer ordering instead visits + // every target once per row, which restarts the hardware prefetcher on a + // short 1600-byte burst for each (row, target) pair. The query rotation + // matrix is re-streamed per group, but it stays L2-resident and its + // (row-major) traversal is itself perfectly sequential. + let mut target_idx = 0; + while target_idx + 4 <= targets.len() { + accumulate_group_rows::( + query, + &targets[target_idx..target_idx + 4], + rows, + target_idx, + result_lane, + additive_shares, + ); + target_idx += 4; + } - for row_idx in 0..rows { - let query_rows_start = row_idx * ROTATIONS * ROW_SIZE; - let query_rows = &query[query_rows_start..query_rows_start + ROTATIONS * ROW_SIZE]; - let target_row_start = row_idx * ROW_SIZE; + for (target_offset, target) in targets[target_idx..].iter().enumerate() { + let Some(target) = target else { + continue; + }; + accumulate_scalar_target_all_rows::( + query, + target, + rows, + target_idx + target_offset, + result_lane, + additive_shares, + ); + } +} - let mut target_idx = 0; - while target_idx + 4 <= targets.len() { - let group = &targets[target_idx..target_idx + 4]; - if let (Some(target0), Some(target1), Some(target2), Some(target3)) = - (group[0], group[1], group[2], group[3]) - { - let target0 = &target0[target_row_start..target_row_start + ROW_SIZE]; - let target1 = &target1[target_row_start..target_row_start + ROW_SIZE]; - let target2 = &target2[target_row_start..target_row_start + ROW_SIZE]; - let target3 = &target3[target_row_start..target_row_start + ROW_SIZE]; +/// One group of four targets against every rotation tile, row by row. +#[cfg(target_arch = "aarch64")] +fn accumulate_group_rows( + query: &[u16], + group: &[Option<&[u16]>], + rows: usize, + target_idx: usize, + result_lane: usize, + additive_shares: &mut [RingElement], +) { + const ROW_SIZE: usize = PrerotatedQueryRowMajor::ROW_SIZE; + + { + if let (Some(full_target0), Some(full_target1), Some(full_target2), Some(full_target3)) = + (group[0], group[1], group[2], group[3]) + { + for row_idx in 0..rows { + let query_rows_start = row_idx * ROTATIONS * ROW_SIZE; + let query_rows = &query[query_rows_start..query_rows_start + ROTATIONS * ROW_SIZE]; + let target_row_start = row_idx * ROW_SIZE; + let target0 = &full_target0[target_row_start..target_row_start + ROW_SIZE]; + let target1 = &full_target1[target_row_start..target_row_start + ROW_SIZE]; + let target2 = &full_target2[target_row_start..target_row_start + ROW_SIZE]; + let target3 = &full_target3[target_row_start..target_row_start + ROW_SIZE]; let mut rotation_idx = 0; while rotation_idx + 6 <= ROTATIONS { @@ -541,10 +582,9 @@ fn accumulate_component_tiled_6x4( while rotation_idx < ROTATIONS { let query_row = &query_rows[rotation_idx * ROW_SIZE..(rotation_idx + 1) * ROW_SIZE]; - for (target_offset, target_row) in - [target0, target1, target2, target3].into_iter().enumerate() - { - let partial = simple_dot_product(query_row, target_row); + let partials = + dot_product_1x4_u16(query_row, [target0, target1, target2, target3]); + for (target_offset, partial) in partials.into_iter().enumerate() { let result_idx = (target_idx + target_offset) * ROTATIONS * 2 + rotation_idx * 2 + result_lane; @@ -553,40 +593,51 @@ fn accumulate_component_tiled_6x4( } rotation_idx += 1; } - } else { - // Missing vectors are uncommon in a full scan. Preserve their - // sentinel handling below while still evaluating live members - // of a mixed group exactly. - for (target_offset, target) in group.iter().enumerate() { - let Some(target) = target else { - continue; - }; - let target_row = &target[target_row_start..target_row_start + ROW_SIZE]; - accumulate_scalar_target::( - query_rows, - target_row, - target_idx + target_offset, - result_lane, - additive_shares, - ); - } } - target_idx += 4; + } else { + // Missing vectors are uncommon in a full scan. Preserve their + // sentinel handling below while still evaluating live members + // of a mixed group exactly. + for (target_offset, target) in group.iter().enumerate() { + let Some(target) = target else { + continue; + }; + accumulate_scalar_target_all_rows::( + query, + target, + rows, + target_idx + target_offset, + result_lane, + additive_shares, + ); + } } + } +} - for (target_offset, target) in targets[target_idx..].iter().enumerate() { - let Some(target) = target else { - continue; - }; - let target_row = &target[target_row_start..target_row_start + ROW_SIZE]; - accumulate_scalar_target::( - query_rows, - target_row, - target_idx + target_offset, - result_lane, - additive_shares, - ); - } +#[cfg(target_arch = "aarch64")] +#[inline] +fn accumulate_scalar_target_all_rows( + query: &[u16], + target: &[u16], + rows: usize, + target_idx: usize, + result_lane: usize, + additive_shares: &mut [RingElement], +) { + const ROW_SIZE: usize = PrerotatedQueryRowMajor::ROW_SIZE; + for row_idx in 0..rows { + let query_rows_start = row_idx * ROTATIONS * ROW_SIZE; + let query_rows = &query[query_rows_start..query_rows_start + ROTATIONS * ROW_SIZE]; + let target_row_start = row_idx * ROW_SIZE; + let target_row = &target[target_row_start..target_row_start + ROW_SIZE]; + accumulate_scalar_target::( + query_rows, + target_row, + target_idx, + result_lane, + additive_shares, + ); } } @@ -608,6 +659,51 @@ fn accumulate_scalar_target( } } +/// One rotation against four targets: the remainder tile after the 6-wide +/// (and optional 4-wide) passes. Keeping the four targets in one pass reuses +/// each query vector load for four MLAs instead of evaluating four +/// independent dot products. +#[cfg(target_arch = "aarch64")] +#[inline(always)] +fn dot_product_1x4_u16(query: &[u16], targets: [&[u16]; 4]) -> [u16; 4] { + use std::arch::aarch64::{uint16x8_t, vaddvq_u16, vdupq_n_u16, vld1q_u16, vmlaq_u16}; + + debug_assert_eq!(query.len(), 800); + debug_assert!(targets.iter().all(|target| target.len() == 800)); + + // SAFETY: AArch64 guarantees Advanced SIMD. Each pointer is derived from + // an 800-element slice and the loop only issues unaligned-safe 8-lane + // loads at offsets 0..792. All arithmetic deliberately wraps in u16. + unsafe { + let zero = vdupq_n_u16(0); + let mut acc0: uint16x8_t = zero; + let mut acc1: uint16x8_t = zero; + let mut acc2: uint16x8_t = zero; + let mut acc3: uint16x8_t = zero; + + let mut idx = 0; + while idx < 800 { + let query0 = vld1q_u16(query.as_ptr().add(idx)); + let target0 = vld1q_u16(targets[0].as_ptr().add(idx)); + let target1 = vld1q_u16(targets[1].as_ptr().add(idx)); + let target2 = vld1q_u16(targets[2].as_ptr().add(idx)); + let target3 = vld1q_u16(targets[3].as_ptr().add(idx)); + acc0 = vmlaq_u16(acc0, query0, target0); + acc1 = vmlaq_u16(acc1, query0, target1); + acc2 = vmlaq_u16(acc2, query0, target2); + acc3 = vmlaq_u16(acc3, query0, target3); + idx += 8; + } + + [ + vaddvq_u16(acc0), + vaddvq_u16(acc1), + vaddvq_u16(acc2), + vaddvq_u16(acc3), + ] + } +} + #[cfg(target_arch = "aarch64")] #[inline(always)] fn dot_product_4x4_u16(queries: [&[u16]; 4], targets: [&[u16]; 4]) -> [[u16; 4]; 4] { @@ -839,6 +935,993 @@ pub fn non_existent_distance() -> Vec> { ] } +#[cfg(target_arch = "aarch64")] +pub use mixed_scan::{ + rotation_aware_pairwise_distance_mixed, rotation_aware_pairwise_distance_mixed_pair, +}; + +/// UMMLA-based full-rotation scan over mixed lo/hi plane residents. +/// +/// Every u16 product is decomposed as +/// `x*y mod 2^16 = xl*yl + 2^8*(xl*yh + xh*yl)` — the `xh*yh` term vanishes +/// modulo 2^16. With rows stored as 8-byte-interleaved planes +/// `[lo0..7 | hi0..7]` on both the query and target side, a single UMMLA per +/// (rotation, target) per 8 coefficients accumulates all three needed +/// partial products into separate u32 lanes (`[ll, lh, hl, discard]`), and +/// the discarded lane is exactly the vanishing high-high term. Results are +/// bit-identical to the u16 MLA kernel. +#[cfg(target_arch = "aarch64")] +mod mixed_scan { + use super::{ + PrerotatedQueryRowMajor, PrerotatedQueryRowMajorView, RingElement, SHARE_OF_MAX_DISTANCE, + }; + use crate::protocol::shared_iris::{ArcIris, MixedPlaneIris}; + use std::arch::asm; + use std::cell::RefCell; + + const ROW_SIZE: usize = PrerotatedQueryRowMajor::ROW_SIZE; + /// A mixed-plane row occupies the same bytes as the u16 row. + const MIXED_ROW_BYTES: usize = 2 * ROW_SIZE; + /// Two copies of a row, used to expose every circular 800-element window. + const DOUBLED_MIXED_ROW_BYTES: usize = 2 * MIXED_ROW_BYTES; + const QUERY_PHASES: usize = 2; + const CODE_ROWS: usize = PrerotatedQueryRowMajor::CODE_ROWS; + const MASK_ROWS: usize = PrerotatedQueryRowMajor::MASK_ROWS; + const GROUP_TARGETS: usize = 4; + const TILE_ROTATIONS: usize = 6; + + /// Compact mixed-plane query. Each row has two doubled copies: one starts + /// at coefficient 0 and one at coefficient 4. Since every supported + /// rotation is a multiple of four, selecting a phase makes its start + /// 8-element aligned and therefore directly loadable by UMMLA. + struct DoubledQueryMixed { + code: Vec, + mask: Vec, + cached_query: Option, + cached_rotations: usize, + } + + impl DoubledQueryMixed { + fn new_buffer() -> Self { + Self { + code: vec![0u8; CODE_ROWS * QUERY_PHASES * DOUBLED_MIXED_ROW_BYTES], + mask: vec![0u8; MASK_ROWS * QUERY_PHASES * DOUBLED_MIXED_ROW_BYTES], + cached_query: None, + cached_rotations: 0, + } + } + + fn matches(&self, query: &ArcIris) -> bool { + self.cached_rotations == ROTATIONS + && self + .cached_query + .as_ref() + .is_some_and(|cached| std::sync::Arc::ptr_eq(cached, query)) + } + + fn fill_if_changed(&mut self, query: &ArcIris) { + if self.matches::(query) { + return; + } + let fill_component = |component: &mut [u8], rows: usize, coefs: &[u16]| { + for row_idx in 0..rows { + let src_row = &coefs[row_idx * ROW_SIZE..(row_idx + 1) * ROW_SIZE]; + for phase in 0..QUERY_PHASES { + let dst_start = (row_idx * QUERY_PHASES + phase) * DOUBLED_MIXED_ROW_BYTES; + let dst = &mut component[dst_start..dst_start + DOUBLED_MIXED_ROW_BYTES]; + for group in 0..(2 * ROW_SIZE / 8) { + for lane in 0..8 { + let src_idx = (phase * 4 + group * 8 + lane) % ROW_SIZE; + let value = src_row[src_idx]; + dst[group * 16 + lane] = value as u8; + dst[group * 16 + 8 + lane] = (value >> 8) as u8; + } + } + } + } + }; + fill_component(&mut self.code, CODE_ROWS, &query.code.coefs); + fill_component(&mut self.mask, MASK_ROWS, &query.mask.coefs); + self.cached_query = Some(query.clone()); + self.cached_rotations = ROTATIONS; + } + } + + thread_local! { + static DOUBLED_MIXED: RefCell<[Option; 2]> = + const { RefCell::new([None, None]) }; + static MIXED_LRU: RefCell = const { RefCell::new(0) }; + static PAIR_PACKED: RefCell> = const { RefCell::new(None) }; + } + + /// Cross-query packed doubled rows for the fused two-query scan. For each + /// 8-coefficient group, one 16-byte block holds a single byte plane of + /// BOTH queries: `[qa_plane(8) | qb_plane(8)]`. Loaded as the UMMLA "B" + /// operand against a target block `[t_lo(8) | t_hi(8)]` (or a + /// `trn1`-combined `[t1_lo | t2_lo]`), one instruction then produces + /// partial products for both queries at once, eliminating the discarded + /// high-high lane of the single-query scheme: three UMMLA cover what + /// four cover in the unpacked layout. + /// + /// Both orientations rotate by the same amounts, so the two-phase doubled + /// window addressing is identical to [`DoubledQueryMixed`]. + struct PairPackedQueryMixed { + code_lo: Vec, + code_hi: Vec, + mask_lo: Vec, + mask_hi: Vec, + cached_queries: Option<(ArcIris, ArcIris)>, + cached_rotations: usize, + } + + impl PairPackedQueryMixed { + fn new_buffer() -> Self { + Self { + code_lo: vec![0u8; CODE_ROWS * QUERY_PHASES * DOUBLED_MIXED_ROW_BYTES], + code_hi: vec![0u8; CODE_ROWS * QUERY_PHASES * DOUBLED_MIXED_ROW_BYTES], + mask_lo: vec![0u8; MASK_ROWS * QUERY_PHASES * DOUBLED_MIXED_ROW_BYTES], + mask_hi: vec![0u8; MASK_ROWS * QUERY_PHASES * DOUBLED_MIXED_ROW_BYTES], + cached_queries: None, + cached_rotations: 0, + } + } + + fn matches(&self, queries: [&ArcIris; 2]) -> bool { + self.cached_rotations == ROTATIONS + && self.cached_queries.as_ref().is_some_and(|(a, b)| { + std::sync::Arc::ptr_eq(a, queries[0]) && std::sync::Arc::ptr_eq(b, queries[1]) + }) + } + + fn fill_if_changed(&mut self, queries: [&ArcIris; 2]) { + if self.matches::(queries) { + return; + } + let fill_component = + |lo: &mut [u8], hi: &mut [u8], rows: usize, coefs_a: &[u16], coefs_b: &[u16]| { + for row_idx in 0..rows { + let src_a = &coefs_a[row_idx * ROW_SIZE..(row_idx + 1) * ROW_SIZE]; + let src_b = &coefs_b[row_idx * ROW_SIZE..(row_idx + 1) * ROW_SIZE]; + for phase in 0..QUERY_PHASES { + let dst_start = + (row_idx * QUERY_PHASES + phase) * DOUBLED_MIXED_ROW_BYTES; + let lo = &mut lo[dst_start..dst_start + DOUBLED_MIXED_ROW_BYTES]; + let hi = &mut hi[dst_start..dst_start + DOUBLED_MIXED_ROW_BYTES]; + for group in 0..(2 * ROW_SIZE / 8) { + for lane in 0..8 { + let src_idx = (phase * 4 + group * 8 + lane) % ROW_SIZE; + let value_a = src_a[src_idx]; + let value_b = src_b[src_idx]; + lo[group * 16 + lane] = value_a as u8; + lo[group * 16 + 8 + lane] = value_b as u8; + hi[group * 16 + lane] = (value_a >> 8) as u8; + hi[group * 16 + 8 + lane] = (value_b >> 8) as u8; + } + } + } + } + }; + fill_component( + &mut self.code_lo, + &mut self.code_hi, + CODE_ROWS, + &queries[0].code.coefs, + &queries[1].code.coefs, + ); + fill_component( + &mut self.mask_lo, + &mut self.mask_hi, + MASK_ROWS, + &queries[0].mask.coefs, + &queries[1].mask.coefs, + ); + self.cached_queries = Some((queries[0].clone(), queries[1].clone())); + self.cached_rotations = ROTATIONS; + } + } + + #[inline(always)] + fn reduce_ummla(lanes: [u32; 4]) -> u16 { + lanes[0].wrapping_add(lanes[1].wrapping_add(lanes[2]) << 8) as u16 + } + + /// Six rotations against four targets. The explicit assembly keeps all + /// 24 accumulators in registers and lets Neoverse V2 issue UMMLA on all + /// four SIMD pipes. The stable Rust i8mm intrinsic is not available yet. + #[target_feature(enable = "i8mm")] + unsafe fn dot_product_6x4_ummla( + queries: [*const u8; TILE_ROTATIONS], + targets: [*const u8; GROUP_TARGETS], + ) -> [[u16; GROUP_TARGETS]; TILE_ROTATIONS] { + let mut raw = [[0u32; 4]; TILE_ROTATIONS * GROUP_TARGETS]; + let q0 = queries[0]; + let q1 = queries[1]; + let q2 = queries[2]; + let q3 = queries[3]; + let q4 = queries[4]; + let q5 = queries[5]; + let t0 = targets[0]; + let t1 = targets[1]; + let t2 = targets[2]; + let t3 = targets[3]; + asm!( + "movi v0.4s, #0", "movi v1.4s, #0", "movi v2.4s, #0", "movi v3.4s, #0", + "movi v4.4s, #0", "movi v5.4s, #0", "movi v6.4s, #0", "movi v7.4s, #0", + "movi v8.4s, #0", "movi v9.4s, #0", "movi v10.4s, #0", "movi v11.4s, #0", + "movi v12.4s, #0", "movi v13.4s, #0", "movi v14.4s, #0", "movi v15.4s, #0", + "movi v16.4s, #0", "movi v17.4s, #0", "movi v18.4s, #0", "movi v19.4s, #0", + "movi v20.4s, #0", "movi v21.4s, #0", "movi v22.4s, #0", "movi v23.4s, #0", + "mov {groups}, #100", + "2:", + "ldr q24, [{q0}], #16", "ldr q25, [{q1}], #16", "ldr q26, [{q2}], #16", + "ldr q27, [{q3}], #16", "ldr q28, [{q4}], #16", "ldr q29, [{q5}], #16", + "ldr q30, [{t0}], #16", + "ummla v0.4s, v24.16b, v30.16b", "ummla v4.4s, v25.16b, v30.16b", + "ummla v8.4s, v26.16b, v30.16b", "ummla v12.4s, v27.16b, v30.16b", + "ummla v16.4s, v28.16b, v30.16b", "ummla v20.4s, v29.16b, v30.16b", + "ldr q30, [{t1}], #16", + "ummla v1.4s, v24.16b, v30.16b", "ummla v5.4s, v25.16b, v30.16b", + "ummla v9.4s, v26.16b, v30.16b", "ummla v13.4s, v27.16b, v30.16b", + "ummla v17.4s, v28.16b, v30.16b", "ummla v21.4s, v29.16b, v30.16b", + "ldr q30, [{t2}], #16", + "ummla v2.4s, v24.16b, v30.16b", "ummla v6.4s, v25.16b, v30.16b", + "ummla v10.4s, v26.16b, v30.16b", "ummla v14.4s, v27.16b, v30.16b", + "ummla v18.4s, v28.16b, v30.16b", "ummla v22.4s, v29.16b, v30.16b", + "ldr q30, [{t3}], #16", + "ummla v3.4s, v24.16b, v30.16b", "ummla v7.4s, v25.16b, v30.16b", + "ummla v11.4s, v26.16b, v30.16b", "ummla v15.4s, v27.16b, v30.16b", + "ummla v19.4s, v28.16b, v30.16b", "ummla v23.4s, v29.16b, v30.16b", + "subs {groups}, {groups}, #1", "b.ne 2b", + "stp q0, q1, [{out}, #0]", "stp q2, q3, [{out}, #32]", + "stp q4, q5, [{out}, #64]", "stp q6, q7, [{out}, #96]", + "stp q8, q9, [{out}, #128]", "stp q10, q11, [{out}, #160]", + "stp q12, q13, [{out}, #192]", "stp q14, q15, [{out}, #224]", + "stp q16, q17, [{out}, #256]", "stp q18, q19, [{out}, #288]", + "stp q20, q21, [{out}, #320]", "stp q22, q23, [{out}, #352]", + q0 = inout(reg) q0 => _, q1 = inout(reg) q1 => _, q2 = inout(reg) q2 => _, + q3 = inout(reg) q3 => _, q4 = inout(reg) q4 => _, q5 = inout(reg) q5 => _, + t0 = inout(reg) t0 => _, t1 = inout(reg) t1 => _, t2 = inout(reg) t2 => _, + t3 = inout(reg) t3 => _, out = in(reg) raw.as_mut_ptr(), groups = out(reg) _, + out("v0") _, out("v1") _, out("v2") _, out("v3") _, out("v4") _, out("v5") _, + out("v6") _, out("v7") _, out("v8") _, out("v9") _, out("v10") _, out("v11") _, + out("v12") _, out("v13") _, out("v14") _, out("v15") _, out("v16") _, out("v17") _, + out("v18") _, out("v19") _, out("v20") _, out("v21") _, out("v22") _, out("v23") _, + out("v24") _, out("v25") _, out("v26") _, out("v27") _, out("v28") _, out("v29") _, + out("v30") _, options(nostack), + ); + + std::array::from_fn(|rotation| { + std::array::from_fn(|target| reduce_ummla(raw[rotation * GROUP_TARGETS + target])) + }) + } + + #[target_feature(enable = "i8mm")] + unsafe fn dot_product_1x4_ummla( + query: *const u8, + targets: [*const u8; GROUP_TARGETS], + ) -> [u16; GROUP_TARGETS] { + let mut raw = [[0u32; 4]; GROUP_TARGETS]; + let t0 = targets[0]; + let t1 = targets[1]; + let t2 = targets[2]; + let t3 = targets[3]; + asm!( + "movi v0.4s, #0", "movi v1.4s, #0", "movi v2.4s, #0", "movi v3.4s, #0", + "mov {groups}, #100", + "2:", + "ldr q24, [{query}], #16", + "ldr q30, [{t0}], #16", "ummla v0.4s, v24.16b, v30.16b", + "ldr q30, [{t1}], #16", "ummla v1.4s, v24.16b, v30.16b", + "ldr q30, [{t2}], #16", "ummla v2.4s, v24.16b, v30.16b", + "ldr q30, [{t3}], #16", "ummla v3.4s, v24.16b, v30.16b", + "subs {groups}, {groups}, #1", "b.ne 2b", + "stp q0, q1, [{out}, #0]", "stp q2, q3, [{out}, #32]", + query = inout(reg) query => _, + t0 = inout(reg) t0 => _, t1 = inout(reg) t1 => _, + t2 = inout(reg) t2 => _, t3 = inout(reg) t3 => _, + out = in(reg) raw.as_mut_ptr(), groups = out(reg) _, + out("v0") _, out("v1") _, out("v2") _, out("v3") _, + out("v24") _, out("v30") _, options(nostack), + ); + std::array::from_fn(|target| reduce_ummla(raw[target])) + } + + /// Four rotations of a packed query pair against four targets (two + /// target pairs). Per (rotation, target pair), three UMMLA produce all + /// twelve needed byte-plane products for both queries and both targets — + /// none of the multiplier work lands in a discarded lane: + /// - `acc_a = [t_lo|t_hi] x [qa_lo|qb_lo]` -> ll and lo*hi for target 2p + /// - `acc_b = same for target 2p+1` + /// - `acc_c = trn1(t_2p, t_2p+1) x [qa_hi|qb_hi]` -> hi*lo for both + /// + /// The trn1-combined lo rows are built from the already-loaded target + /// registers, so targets are still loaded once per group step. Exactly + /// fills the 32-register file: 24 accumulators + 4 targets + 2 trn + + /// 2 query operands. Adds this row's products into `raw` so callers + /// accumulate a whole component across rows before one reduction. + #[target_feature(enable = "i8mm")] + unsafe fn dot_product_pair_4x2p_acc( + query_lo: [*const u8; 4], + query_hi: [*const u8; 4], + targets: [*const u8; GROUP_TARGETS], + raw: &mut [[u32; 4]; 24], + ) { + let ql0 = query_lo[0]; + let ql1 = query_lo[1]; + let ql2 = query_lo[2]; + let ql3 = query_lo[3]; + let qh0 = query_hi[0]; + let qh1 = query_hi[1]; + let qh2 = query_hi[2]; + let qh3 = query_hi[3]; + let t0 = targets[0]; + let t1 = targets[1]; + let t2 = targets[2]; + let t3 = targets[3]; + asm!( + // Accumulators are loaded and stored rather than zero-initialized: + // callers accumulate a whole component's rows in u32 and reduce + // once per tile. ldp/stp run on the load/store pipes, which have + // headroom, instead of movi on the contended SIMD pipes. + "ldp q0, q1, [{out}, #0]", "ldp q2, q3, [{out}, #32]", + "ldp q4, q5, [{out}, #64]", "ldp q6, q7, [{out}, #96]", + "ldp q8, q9, [{out}, #128]", "ldp q10, q11, [{out}, #160]", + "ldp q12, q13, [{out}, #192]", "ldp q14, q15, [{out}, #224]", + "ldp q16, q17, [{out}, #256]", "ldp q18, q19, [{out}, #288]", + "ldp q20, q21, [{out}, #320]", "ldp q22, q23, [{out}, #352]", + "mov {groups}, #100", + "2:", + "ldr q24, [{t0}], #16", "ldr q25, [{t1}], #16", + "ldr q26, [{t2}], #16", "ldr q27, [{t3}], #16", + "trn1 v28.2d, v24.2d, v25.2d", + "trn1 v29.2d, v26.2d, v27.2d", + "ldr q30, [{ql0}], #16", "ldr q31, [{qh0}], #16", + "ummla v0.4s, v24.16b, v30.16b", "ummla v1.4s, v25.16b, v30.16b", + "ummla v2.4s, v28.16b, v31.16b", "ummla v3.4s, v26.16b, v30.16b", + "ummla v4.4s, v27.16b, v30.16b", "ummla v5.4s, v29.16b, v31.16b", + "ldr q30, [{ql1}], #16", "ldr q31, [{qh1}], #16", + "ummla v6.4s, v24.16b, v30.16b", "ummla v7.4s, v25.16b, v30.16b", + "ummla v8.4s, v28.16b, v31.16b", "ummla v9.4s, v26.16b, v30.16b", + "ummla v10.4s, v27.16b, v30.16b", "ummla v11.4s, v29.16b, v31.16b", + "ldr q30, [{ql2}], #16", "ldr q31, [{qh2}], #16", + "ummla v12.4s, v24.16b, v30.16b", "ummla v13.4s, v25.16b, v30.16b", + "ummla v14.4s, v28.16b, v31.16b", "ummla v15.4s, v26.16b, v30.16b", + "ummla v16.4s, v27.16b, v30.16b", "ummla v17.4s, v29.16b, v31.16b", + "ldr q30, [{ql3}], #16", "ldr q31, [{qh3}], #16", + "ummla v18.4s, v24.16b, v30.16b", "ummla v19.4s, v25.16b, v30.16b", + "ummla v20.4s, v28.16b, v31.16b", "ummla v21.4s, v26.16b, v30.16b", + "ummla v22.4s, v27.16b, v30.16b", "ummla v23.4s, v29.16b, v31.16b", + "subs {groups}, {groups}, #1", "b.ne 2b", + "stp q0, q1, [{out}, #0]", "stp q2, q3, [{out}, #32]", + "stp q4, q5, [{out}, #64]", "stp q6, q7, [{out}, #96]", + "stp q8, q9, [{out}, #128]", "stp q10, q11, [{out}, #160]", + "stp q12, q13, [{out}, #192]", "stp q14, q15, [{out}, #224]", + "stp q16, q17, [{out}, #256]", "stp q18, q19, [{out}, #288]", + "stp q20, q21, [{out}, #320]", "stp q22, q23, [{out}, #352]", + ql0 = inout(reg) ql0 => _, ql1 = inout(reg) ql1 => _, + ql2 = inout(reg) ql2 => _, ql3 = inout(reg) ql3 => _, + qh0 = inout(reg) qh0 => _, qh1 = inout(reg) qh1 => _, + qh2 = inout(reg) qh2 => _, qh3 = inout(reg) qh3 => _, + t0 = inout(reg) t0 => _, t1 = inout(reg) t1 => _, t2 = inout(reg) t2 => _, + t3 = inout(reg) t3 => _, out = in(reg) raw.as_mut_ptr(), groups = out(reg) _, + out("v0") _, out("v1") _, out("v2") _, out("v3") _, out("v4") _, out("v5") _, + out("v6") _, out("v7") _, out("v8") _, out("v9") _, out("v10") _, out("v11") _, + out("v12") _, out("v13") _, out("v14") _, out("v15") _, out("v16") _, out("v17") _, + out("v18") _, out("v19") _, out("v20") _, out("v21") _, out("v22") _, out("v23") _, + out("v24") _, out("v25") _, out("v26") _, out("v27") _, out("v28") _, out("v29") _, + out("v30") _, out("v31") _, options(nostack), + ); + } + + /// Scatter one rotation's packed-pair accumulator block (`[a, b, c]` for + /// two targets x two queries) into both queries' share vectors. The + /// packed path writes each result exactly once, so the mask lane's + /// doubling (the epilogue's job on the other paths) is folded in here. + #[inline(always)] + fn scatter_pair_block( + block: &[[u32; 4]], + base_target_idx: usize, + pair: usize, + rotation: usize, + result_lane: usize, + additive_shares: &mut [Vec>; 2], + ) { + let acc_a = &block[0]; + let acc_b = &block[1]; + let acc_c = &block[2]; + let lane_scale = 1 + result_lane as u16; + for (query, shares) in additive_shares.iter_mut().enumerate() { + let first = (acc_a[query].wrapping_add(acc_a[2 + query].wrapping_add(acc_c[query]) << 8) + as u16) + .wrapping_mul(lane_scale); + let second = (acc_b[query] + .wrapping_add(acc_b[2 + query].wrapping_add(acc_c[2 + query]) << 8) + as u16) + .wrapping_mul(lane_scale); + let first_idx = (base_target_idx + pair * 2) * ROTATIONS * 2 + rotation * 2 + result_lane; + let second_idx = + (base_target_idx + pair * 2 + 1) * ROTATIONS * 2 + rotation * 2 + result_lane; + shares[first_idx].0 = shares[first_idx].0.wrapping_add(first); + shares[second_idx].0 = shares[second_idx].0.wrapping_add(second); + } + } + + /// Packed-pair scan of one component of four present targets. Rows + /// accumulate into persistent per-tile u32 buffers (u32 lanes cannot + /// overflow: 16 rows x 800 coefficients x 255^2 < 2^32) and each tile is + /// reduced and scattered once at the end, instead of once per row. + #[target_feature(enable = "i8mm")] + unsafe fn scan_four_targets_pair_packed( + query_lo: &[u8], + query_hi: &[u8], + rows: usize, + targets: [&[u8]; GROUP_TARGETS], + base_target_idx: usize, + result_lane: usize, + additive_shares: &mut [Vec>; 2], + ) { + const PAIR_TILE_ROTATIONS: usize = 4; + const MAX_TILES: usize = 8; + let n_tiles = ROTATIONS.div_ceil(PAIR_TILE_ROTATIONS); + assert!(n_tiles <= MAX_TILES, "unsupported rotation count"); + let mut raws = [[[0u32; 4]; 24]; MAX_TILES]; + + // Row `r`'s window for a rotation is the row-0 window plus a constant + // stride, so the 2 x 4 window pointers per tile are derived once. + // The final tile is padded: rotation slots beyond `ROTATIONS` reuse + // the last valid window and their results are simply not scattered. + // For the 31-rotation scan this wastes 1/32 of the tile work, far + // less than single-rotation remainder passes. + let query_row_stride = QUERY_PHASES * DOUBLED_MIXED_ROW_BYTES; + let window_bases: [([*const u8; 4], [*const u8; 4]); MAX_TILES] = + std::array::from_fn(|tile| { + let rotation = (tile * PAIR_TILE_ROTATIONS).min(ROTATIONS - 1); + let live_rotations = PAIR_TILE_ROTATIONS.min(ROTATIONS - rotation); + ( + std::array::from_fn(|offset| { + query_window::( + query_lo, + 0, + rotation + offset.min(live_rotations - 1), + ) + }), + std::array::from_fn(|offset| { + query_window::( + query_hi, + 0, + rotation + offset.min(live_rotations - 1), + ) + }), + ) + }); + + for row in 0..rows { + let target_row_start = row * MIXED_ROW_BYTES; + let target_ptrs = + std::array::from_fn(|target| targets[target][target_row_start..].as_ptr()); + let row_offset = row * query_row_stride; + for (tile, raw) in raws.iter_mut().enumerate().take(n_tiles) { + let (lo_bases, hi_bases) = &window_bases[tile]; + let lo_ptrs = std::array::from_fn(|offset| lo_bases[offset].add(row_offset)); + let hi_ptrs = std::array::from_fn(|offset| hi_bases[offset].add(row_offset)); + dot_product_pair_4x2p_acc(lo_ptrs, hi_ptrs, target_ptrs, raw); + } + } + + for (tile, raw) in raws.iter().enumerate().take(n_tiles) { + let rotation = tile * PAIR_TILE_ROTATIONS; + let live_rotations = PAIR_TILE_ROTATIONS.min(ROTATIONS - rotation); + for rotation_offset in 0..live_rotations { + for pair in 0..2 { + scatter_pair_block::( + &raw[rotation_offset * 6 + pair * 3..rotation_offset * 6 + pair * 3 + 3], + base_target_idx, + pair, + rotation + rotation_offset, + result_lane, + additive_shares, + ); + } + } + } + } + + /// Packed-pair scan of a whole component: all groups of four targets. + /// Callers guarantee `targets.len()` is a multiple of four with every + /// target present. + #[target_feature(enable = "i8mm")] + unsafe fn accumulate_component_pair_packed( + query_lo: &[u8], + query_hi: &[u8], + targets: &[&[u8]], + rows: usize, + result_lane: usize, + additive_shares: &mut [Vec>; 2], + ) { + debug_assert_eq!(targets.len() % GROUP_TARGETS, 0); + let mut target_idx = 0; + while target_idx + GROUP_TARGETS <= targets.len() { + let group = std::array::from_fn(|offset| targets[target_idx + offset]); + scan_four_targets_pair_packed::( + query_lo, + query_hi, + rows, + group, + target_idx, + result_lane, + additive_shares, + ); + target_idx += GROUP_TARGETS; + } + } + + #[inline(always)] + unsafe fn query_window( + query: &[u8], + row: usize, + rotation: usize, + ) -> *const u8 { + let amount = PrerotatedQueryRowMajorView::::ROTATION_AMOUNTS[rotation]; + let phase = (amount / 4) & 1; + let aligned_amount = amount - phase * 4; + query + .as_ptr() + .add((row * QUERY_PHASES + phase) * DOUBLED_MIXED_ROW_BYTES + aligned_amount * 2) + } + + /// All rotation tiles of one query row against four loaded target rows. + /// Factored out so the single-query and paired-query scans accumulate in + /// exactly the same instruction order (bit-identical results). + #[target_feature(enable = "i8mm")] + unsafe fn scan_row_rotations( + query: &[u8], + row: usize, + target_ptrs: [*const u8; GROUP_TARGETS], + base_target_idx: usize, + result_lane: usize, + additive_shares: &mut [RingElement], + ) { + let mut rotation = 0; + while rotation + TILE_ROTATIONS <= ROTATIONS { + let query_ptrs = std::array::from_fn(|offset| { + query_window::(query, row, rotation + offset) + }); + let partials = dot_product_6x4_ummla(query_ptrs, target_ptrs); + for (rotation_offset, values) in partials.into_iter().enumerate() { + for (target_offset, partial) in values.into_iter().enumerate() { + let result_idx = (base_target_idx + target_offset) * ROTATIONS * 2 + + (rotation + rotation_offset) * 2 + + result_lane; + additive_shares[result_idx].0 = + additive_shares[result_idx].0.wrapping_add(partial); + } + } + rotation += TILE_ROTATIONS; + } + while rotation < ROTATIONS { + let partials = + dot_product_1x4_ummla(query_window::(query, row, rotation), target_ptrs); + for (target_offset, partial) in partials.into_iter().enumerate() { + let result_idx = + (base_target_idx + target_offset) * ROTATIONS * 2 + rotation * 2 + result_lane; + additive_shares[result_idx].0 = additive_shares[result_idx].0.wrapping_add(partial); + } + rotation += 1; + } + } + + #[target_feature(enable = "i8mm")] + unsafe fn scan_four_targets( + query: &[u8], + rows: usize, + targets: [&[u8]; GROUP_TARGETS], + base_target_idx: usize, + result_lane: usize, + additive_shares: &mut [RingElement], + ) { + for row in 0..rows { + let target_row_start = row * MIXED_ROW_BYTES; + let target_ptrs = + std::array::from_fn(|target| targets[target][target_row_start..].as_ptr()); + scan_row_rotations::( + query, + row, + target_ptrs, + base_target_idx, + result_lane, + additive_shares, + ); + } + } + + /// Two queries against the same four targets in one target traversal. The + /// four target rows (6.4 KB) stay L1-resident across both queries' tiles, + /// so the second query's rotations cost no additional target streaming. + #[target_feature(enable = "i8mm")] + unsafe fn scan_four_targets_pair( + queries: [&[u8]; 2], + rows: usize, + targets: [&[u8]; GROUP_TARGETS], + base_target_idx: usize, + result_lane: usize, + additive_shares: &mut [Vec>; 2], + ) { + for row in 0..rows { + let target_row_start = row * MIXED_ROW_BYTES; + let target_ptrs = + std::array::from_fn(|target| targets[target][target_row_start..].as_ptr()); + for (query, shares) in queries.iter().zip(additive_shares.iter_mut()) { + scan_row_rotations::( + query, + row, + target_ptrs, + base_target_idx, + result_lane, + shares, + ); + } + } + } + + #[target_feature(enable = "i8mm")] + unsafe fn scan_one_target( + query: &[u8], + rows: usize, + target: &[u8], + target_idx: usize, + result_lane: usize, + additive_shares: &mut [RingElement], + ) { + for row in 0..rows { + let target_ptr = target[row * MIXED_ROW_BYTES..].as_ptr(); + let target_ptrs = [target_ptr; GROUP_TARGETS]; + for rotation in 0..ROTATIONS { + let partial = dot_product_1x4_ummla( + query_window::(query, row, rotation), + target_ptrs, + )[0]; + let result_idx = target_idx * ROTATIONS * 2 + rotation * 2 + result_lane; + additive_shares[result_idx].0 = additive_shares[result_idx].0.wrapping_add(partial); + } + } + } + + #[target_feature(enable = "i8mm")] + unsafe fn accumulate_component_mixed( + query: &[u8], + targets: &[Option<&[u8]>], + rows: usize, + result_lane: usize, + additive_shares: &mut [RingElement], + ) { + let mut target_idx = 0; + while target_idx < targets.len() { + if target_idx + GROUP_TARGETS <= targets.len() + && targets[target_idx..target_idx + GROUP_TARGETS] + .iter() + .all(Option::is_some) + { + let group = std::array::from_fn(|offset| { + targets[target_idx + offset].expect("checked present target") + }); + scan_four_targets::( + query, + rows, + group, + target_idx, + result_lane, + additive_shares, + ); + target_idx += GROUP_TARGETS; + } else { + if let Some(target) = targets[target_idx] { + scan_one_target::( + query, + rows, + target, + target_idx, + result_lane, + additive_shares, + ); + } + target_idx += 1; + } + } + } + + /// Paired-query counterpart of [`accumulate_component_mixed`]: every + /// target group is streamed once and evaluated by both queries. + #[target_feature(enable = "i8mm")] + unsafe fn accumulate_component_mixed_pair( + queries: [&[u8]; 2], + targets: &[Option<&[u8]>], + rows: usize, + result_lane: usize, + additive_shares: &mut [Vec>; 2], + ) { + let mut target_idx = 0; + while target_idx < targets.len() { + if target_idx + GROUP_TARGETS <= targets.len() + && targets[target_idx..target_idx + GROUP_TARGETS] + .iter() + .all(Option::is_some) + { + let group = std::array::from_fn(|offset| { + targets[target_idx + offset].expect("checked present target") + }); + scan_four_targets_pair::( + queries, + rows, + group, + target_idx, + result_lane, + additive_shares, + ); + target_idx += GROUP_TARGETS; + } else { + if let Some(target) = targets[target_idx] { + for (query, shares) in queries.iter().zip(additive_shares.iter_mut()) { + scan_one_target::( + query, + rows, + target, + target_idx, + result_lane, + shares, + ); + } + } + target_idx += 1; + } + } + } + + /// Mixed-plane counterpart of + /// [`super::rotation_aware_pairwise_distance_rowmajor`]: identical inputs, + /// outputs, and sentinel semantics, operating on plane residents. + /// + /// # Panics + /// The caller must only invoke this when the `i8mm` CPU feature is + /// present (pools only adopt the mixed layout in that case). + pub fn rotation_aware_pairwise_distance_mixed( + query: &ArcIris, + targets: &[Option<&MixedPlaneIris>], + ) -> Vec> { + assert!( + std::arch::is_aarch64_feature_detected!("i8mm"), + "mixed-plane scan kernel requires the i8mm CPU feature" + ); + let mut additive_shares = vec![RingElement(0u16); 2 * ROTATIONS * targets.len()]; + + DOUBLED_MIXED.with(|cell| { + let mut entries = cell.borrow_mut(); + let hit = entries.iter().position(|entry| { + entry + .as_ref() + .is_some_and(|doubled| doubled.matches::(query)) + }); + let index = hit.unwrap_or_else(|| { + entries + .iter() + .position(Option::is_none) + .unwrap_or_else(|| MIXED_LRU.with(|lru| *lru.borrow())) + }); + MIXED_LRU.with(|lru| *lru.borrow_mut() = 1 - index); + let doubled = entries[index].get_or_insert_with(DoubledQueryMixed::new_buffer); + doubled.fill_if_changed::(query); + + let code_targets: Vec> = targets + .iter() + .map(|target| target.map(MixedPlaneIris::code_planes)) + .collect(); + let mask_targets: Vec> = targets + .iter() + .map(|target| target.map(MixedPlaneIris::mask_planes)) + .collect(); + // SAFETY: i8mm presence asserted above. + unsafe { + accumulate_component_mixed::( + &doubled.code, + &code_targets, + CODE_ROWS, + 0, + &mut additive_shares, + ); + accumulate_component_mixed::( + &doubled.mask, + &mask_targets, + MASK_ROWS, + 1, + &mut additive_shares, + ); + } + }); + + apply_scan_epilogue::(targets, &mut additive_shares); + additive_shares + } + + /// Same epilogue as the u16 kernel: double the mask lanes of present + /// targets, fill sentinel distances for missing ones. + fn apply_scan_epilogue( + targets: &[Option<&MixedPlaneIris>], + additive_shares: &mut [RingElement], + ) { + for (target_idx, target) in targets.iter().enumerate() { + let base_idx = target_idx * ROTATIONS * 2; + if target.is_some() { + for rot_idx in 0..ROTATIONS { + let mask_idx = base_idx + rot_idx * 2 + 1; + additive_shares[mask_idx] = RingElement(2) * additive_shares[mask_idx]; + } + } else { + let (a, b) = SHARE_OF_MAX_DISTANCE; + for rot_idx in 0..ROTATIONS { + let code_idx = base_idx + rot_idx * 2; + additive_shares[code_idx] = RingElement(a); + additive_shares[code_idx + 1] = RingElement(b); + } + } + } + } + + /// Fused two-query scan: identical outputs to two independent + /// [`rotation_aware_pairwise_distance_mixed`] calls, but each target row + /// is streamed once and feeds both queries' rotation tiles. Used by the + /// exact scan to evaluate the normal and mirror orientations in one pass + /// over the resident database. + /// + /// Full four-target groups with every target present take the packed-pair + /// kernel (three UMMLA per four query/target/rotation results); calls + /// containing a missing target or a tail group fall back to the unpacked + /// pair scan. Both produce bit-identical results. + /// + /// # Panics + /// The caller must only invoke this when the `i8mm` CPU feature is + /// present (pools only adopt the mixed layout in that case). + pub fn rotation_aware_pairwise_distance_mixed_pair( + queries: [&ArcIris; 2], + targets: &[Option<&MixedPlaneIris>], + ) -> [Vec>; 2] { + assert!( + std::arch::is_aarch64_feature_detected!("i8mm"), + "mixed-plane scan kernel requires the i8mm CPU feature" + ); + let packed_eligible = targets.len().is_multiple_of(GROUP_TARGETS) + && !targets.is_empty() + && targets.iter().all(Option::is_some); + if packed_eligible { + return rotation_aware_pairwise_distance_mixed_pair_packed::( + queries, targets, + ); + } + rotation_aware_pairwise_distance_mixed_pair_unpacked::(queries, targets) + } + + fn rotation_aware_pairwise_distance_mixed_pair_packed( + queries: [&ArcIris; 2], + targets: &[Option<&MixedPlaneIris>], + ) -> [Vec>; 2] { + let mut additive_shares = [ + vec![RingElement(0u16); 2 * ROTATIONS * targets.len()], + vec![RingElement(0u16); 2 * ROTATIONS * targets.len()], + ]; + + PAIR_PACKED.with(|cell| { + let mut entry = cell.borrow_mut(); + let packed = entry.get_or_insert_with(PairPackedQueryMixed::new_buffer); + packed.fill_if_changed::(queries); + + let code_targets: Vec<&[u8]> = targets + .iter() + .map(|target| target.expect("packed pair scan requires present targets")) + .map(MixedPlaneIris::code_planes) + .collect(); + let mask_targets: Vec<&[u8]> = targets + .iter() + .map(|target| target.expect("packed pair scan requires present targets")) + .map(MixedPlaneIris::mask_planes) + .collect(); + // SAFETY: i8mm presence asserted by the public entry point. + unsafe { + accumulate_component_pair_packed::( + &packed.code_lo, + &packed.code_hi, + &code_targets, + CODE_ROWS, + 0, + &mut additive_shares, + ); + accumulate_component_pair_packed::( + &packed.mask_lo, + &packed.mask_hi, + &mask_targets, + MASK_ROWS, + 1, + &mut additive_shares, + ); + } + }); + + // No epilogue: every target is present by construction and the mask + // doubling is applied inside the single per-result scatter. + additive_shares + } + + fn rotation_aware_pairwise_distance_mixed_pair_unpacked( + queries: [&ArcIris; 2], + targets: &[Option<&MixedPlaneIris>], + ) -> [Vec>; 2] { + let mut additive_shares = [ + vec![RingElement(0u16); 2 * ROTATIONS * targets.len()], + vec![RingElement(0u16); 2 * ROTATIONS * targets.len()], + ]; + + DOUBLED_MIXED.with(|cell| { + let mut entries = cell.borrow_mut(); + // Materialize both queries in the two cache slots. If a query is + // already cached, keep its slot; otherwise fill the slot that the + // other query does not occupy. + let slot_of = |entries: &[Option; 2], query: &ArcIris| { + entries.iter().position(|entry| { + entry + .as_ref() + .is_some_and(|doubled| doubled.matches::(query)) + }) + }; + let index_a = slot_of(&entries, queries[0]).unwrap_or_else(|| { + let index = match slot_of(&entries, queries[1]) { + Some(index_b) => 1 - index_b, + None => 0, + }; + let doubled = entries[index].get_or_insert_with(DoubledQueryMixed::new_buffer); + doubled.fill_if_changed::(queries[0]); + index + }); + let index_b = slot_of(&entries, queries[1]).unwrap_or_else(|| { + let index = 1 - index_a; + let doubled = entries[index].get_or_insert_with(DoubledQueryMixed::new_buffer); + doubled.fill_if_changed::(queries[1]); + index + }); + + let doubled_a = entries[index_a].as_ref().expect("slot filled above"); + let doubled_b = entries[index_b].as_ref().expect("slot filled above"); + + let code_targets: Vec> = targets + .iter() + .map(|target| target.map(MixedPlaneIris::code_planes)) + .collect(); + let mask_targets: Vec> = targets + .iter() + .map(|target| target.map(MixedPlaneIris::mask_planes)) + .collect(); + // SAFETY: i8mm presence asserted above. + unsafe { + accumulate_component_mixed_pair::( + [&doubled_a.code, &doubled_b.code], + &code_targets, + CODE_ROWS, + 0, + &mut additive_shares, + ); + accumulate_component_mixed_pair::( + [&doubled_a.mask, &doubled_b.mask], + &mask_targets, + MASK_ROWS, + 1, + &mut additive_shares, + ); + } + }); + + for shares in &mut additive_shares { + apply_scan_epilogue::(targets, shares); + } + additive_shares + } +} + #[cfg(test)] mod tests { use super::*; @@ -1099,6 +2182,138 @@ mod tests { } } + #[cfg(target_arch = "aarch64")] + #[test] + fn mixed_ummla_scan_matches_u16_scan_with_tail_and_missing_targets() { + use crate::protocol::shared_iris::{GaloisRingSharedIris, MixedPlaneIris}; + + if !std::arch::is_aarch64_feature_detected!("i8mm") { + return; + } + + let mut rng = AesRng::seed_from_u64(0x1_8_8_4); + let iris_db = IrisDB::new_random_rng(10, &mut rng).db; + let query_shares = + GaloisRingSharedIris::generate_shares_locally(&mut rng, iris_db[0].clone()); + let target_shares: Vec<_> = iris_db[1..] + .iter() + .map(|iris| GaloisRingSharedIris::generate_shares_locally(&mut rng, iris.clone())) + .collect(); + + for party in 0..3 { + let mut query = query_shares[party].clone(); + query.code.preprocess_iris_code_query_share(); + query.mask.preprocess_mask_code_query_share(); + let query = Arc::new(query); + let targets: Vec = target_shares + .iter() + .map(|shares| Arc::new(shares[party].clone())) + .collect(); + let mixed: Vec = targets + .iter() + .map(|target| MixedPlaneIris::from_iris(target)) + .collect(); + let present = [true, false, true, true, true, true, false, true, true]; + + let u16_targets = targets + .iter() + .zip(present) + .map(|(target, present)| present.then_some(target)); + let mixed_targets: Vec> = mixed + .iter() + .zip(present) + .map(|(target, present)| present.then_some(target)) + .collect(); + + let expected = rotation_aware_pairwise_distance_rowmajor::<31, _>(&query, u16_targets); + let actual = rotation_aware_pairwise_distance_mixed::<31>(&query, &mixed_targets); + assert_eq!(actual, expected, "party {party}"); + } + } + + /// The fused two-query pass must be byte-identical to two independent + /// single-query passes, across group/tail/missing-target shapes. The + /// query pair mirrors production: one normal-preprocessed query and one + /// mirrored-preprocessed query of a different iris. + #[cfg(target_arch = "aarch64")] + #[test] + fn mixed_ummla_pair_scan_matches_two_single_scans() { + use crate::protocol::shared_iris::{GaloisRingSharedIris, MixedPlaneIris}; + + if !std::arch::is_aarch64_feature_detected!("i8mm") { + return; + } + + let mut rng = AesRng::seed_from_u64(0x1_8_8_5); + let iris_db = IrisDB::new_random_rng(12, &mut rng).db; + let query_shares = + GaloisRingSharedIris::generate_shares_locally(&mut rng, iris_db[0].clone()); + let mirrored_query_shares = + GaloisRingSharedIris::generate_mirrored_shares_locally(&mut rng, iris_db[1].clone()); + let target_shares: Vec<_> = iris_db[2..] + .iter() + .map(|iris| GaloisRingSharedIris::generate_shares_locally(&mut rng, iris.clone())) + .collect(); + + for party in 0..3 { + let mut query_a = query_shares[party].clone(); + query_a.code.preprocess_iris_code_query_share(); + query_a.mask.preprocess_mask_code_query_share(); + let query_a = Arc::new(query_a); + let mut query_b = mirrored_query_shares[party].clone(); + query_b.code.preprocess_iris_code_query_share(); + query_b.mask.preprocess_mask_code_query_share(); + let query_b = Arc::new(query_b); + + let targets: Vec = target_shares + .iter() + .map(|shares| Arc::new(shares[party].clone())) + .collect(); + let mixed: Vec = targets + .iter() + .map(|target| MixedPlaneIris::from_iris(target)) + .collect(); + + // Cover: the packed fast path (all present, multiple of four — + // one and two groups), a missing target inside a group, a tail + // shorter than a group, and an all-missing prefix. + let present_patterns: [&[bool]; 6] = [ + &[true; 4], + &[true; 8], + &[true; 10], + &[true, false, true, true, true, true, false, true, true, true], + &[true, true, true], + &[false, false, true, true, true, true, true], + ]; + for present in present_patterns { + let mixed_targets: Vec> = mixed + .iter() + .zip(present) + .map(|(target, &present)| present.then_some(target)) + .collect(); + + let expected_a = + rotation_aware_pairwise_distance_mixed::<31>(&query_a, &mixed_targets); + let expected_b = + rotation_aware_pairwise_distance_mixed::<31>(&query_b, &mixed_targets); + let [actual_a, actual_b] = rotation_aware_pairwise_distance_mixed_pair::<31>( + [&query_a, &query_b], + &mixed_targets, + ); + assert_eq!(actual_a, expected_a, "party {party} query A"); + assert_eq!(actual_b, expected_b, "party {party} query B"); + + // Same-query pairing (both slots resolve to one cache entry). + let [same_a, same_b] = rotation_aware_pairwise_distance_mixed_pair::<31>( + [&query_a, &query_a], + &mixed_targets, + ); + assert_eq!(same_a, expected_a, "party {party} same-query A"); + assert_eq!(same_b, expected_a, "party {party} same-query B"); + } + } + } + #[rstest] #[case(0, 1)] #[case(1, 5)] diff --git a/iris-mpc-cpu/src/protocol/shared_iris.rs b/iris-mpc-cpu/src/protocol/shared_iris.rs index 65792e3e85..3e4ae7bf14 100644 --- a/iris-mpc-cpu/src/protocol/shared_iris.rs +++ b/iris-mpc-cpu/src/protocol/shared_iris.rs @@ -126,3 +126,148 @@ impl GaloisRingSharedIris { ] } } + +/// 8-byte-interleaved lo/hi plane representation of an iris share for the +/// UMMLA-based exact-scan kernel: every 8 consecutive u16 coefficients are +/// stored as `[lo0..lo7 | hi0..hi7]`. Same byte count as the u16 form; the +/// original share is reconstructed exactly by [`MixedPlaneIris::to_iris`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MixedPlaneIris { + /// `id` of the code share (party_id + 1), preserved for reconstruction. + code_id: usize, + /// `id` of the mask share (party_id + 1), preserved for reconstruction. + mask_id: usize, + /// 16 rows x 1600 bytes. + code: Box<[u8]>, + /// 8 rows x 1600 bytes. + mask: Box<[u8]>, +} + +fn mix_planes(src: &[u16], dst: &mut [u8]) { + debug_assert_eq!(src.len() * 2, dst.len()); + for (src8, dst16) in src.chunks_exact(8).zip(dst.chunks_exact_mut(16)) { + for k in 0..8 { + dst16[k] = src8[k] as u8; + dst16[8 + k] = (src8[k] >> 8) as u8; + } + } +} + +fn unmix_planes(src: &[u8], dst: &mut [u16]) { + debug_assert_eq!(src.len(), dst.len() * 2); + for (dst8, src16) in dst.chunks_exact_mut(8).zip(src.chunks_exact(16)) { + for k in 0..8 { + dst8[k] = src16[k] as u16 | ((src16[8 + k] as u16) << 8); + } + } +} + +impl MixedPlaneIris { + pub fn from_iris(iris: &GaloisRingSharedIris) -> Self { + let mut code = vec![0u8; iris.code.coefs.len() * 2].into_boxed_slice(); + let mut mask = vec![0u8; iris.mask.coefs.len() * 2].into_boxed_slice(); + mix_planes(&iris.code.coefs, &mut code); + mix_planes(&iris.mask.coefs, &mut mask); + Self { + code_id: iris.code.id, + mask_id: iris.mask.id, + code, + mask, + } + } + + pub fn to_iris(&self) -> GaloisRingSharedIris { + let mut iris = GaloisRingSharedIris::default_for_party(0); + iris.code.id = self.code_id; + iris.mask.id = self.mask_id; + unmix_planes(&self.code, &mut iris.code.coefs); + unmix_planes(&self.mask, &mut iris.mask.coefs); + iris + } + + #[inline(always)] + pub fn code_planes(&self) -> &[u8] { + &self.code + } + + #[inline(always)] + pub fn mask_planes(&self) -> &[u8] { + &self.mask + } +} + +/// Resident layout of a worker pool's iris store. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResidentLayout { + /// Plain `ArcIris` values; all access paths borrow them at zero cost. + U16, + /// Mixed lo/hi plane values for the UMMLA exact-scan kernel. Non-scan + /// access paths reconstruct the u16 share on demand. + MixedPlane, +} + +/// Layout to use for exact-scan worker pools on this machine: mixed planes +/// when the UMMLA kernel is available (aarch64 with i8mm), unless disabled +/// via `IRIS_MPC_DISABLE_MIXED_SCAN=1`. +pub fn preferred_scan_layout() -> ResidentLayout { + #[cfg(target_arch = "aarch64")] + { + let disabled = std::env::var("IRIS_MPC_DISABLE_MIXED_SCAN") + .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) + .unwrap_or(false); + if !disabled && std::arch::is_aarch64_feature_detected!("i8mm") { + return ResidentLayout::MixedPlane; + } + } + ResidentLayout::U16 +} + +/// An iris share as resident in a worker pool store, in the pool's layout. +#[derive(Debug, Clone)] +pub enum ResidentIris { + U16(ArcIris), + Mixed(Arc), +} + +impl ResidentIris { + pub fn from_arc(iris: ArcIris, layout: ResidentLayout) -> Self { + match layout { + ResidentLayout::U16 => Self::U16(iris), + ResidentLayout::MixedPlane => Self::Mixed(Arc::new(MixedPlaneIris::from_iris(&iris))), + } + } + + /// The u16 form: a cheap handle clone for `U16`, an exact reconstruction + /// for `Mixed`. + pub fn to_arc(&self) -> ArcIris { + match self { + Self::U16(iris) => iris.clone(), + Self::Mixed(planes) => Arc::new(planes.to_iris()), + } + } + + #[inline(always)] + pub fn as_mixed(&self) -> Option<&MixedPlaneIris> { + match self { + Self::U16(_) => None, + Self::Mixed(planes) => Some(planes), + } + } +} + +#[cfg(test)] +mod mixed_plane_tests { + use super::*; + use iris_mpc_common::iris_db::iris::IrisCode; + use rand::{rngs::StdRng, SeedableRng}; + + #[test] + fn mixed_plane_round_trip() { + let mut rng = StdRng::seed_from_u64(7); + let iris = IrisCode::random_rng(&mut rng); + for share in GaloisRingSharedIris::generate_shares_locally(&mut rng, iris) { + let planes = MixedPlaneIris::from_iris(&share); + assert_eq!(planes.to_iris(), share); + } + } +} diff --git a/iris-mpc-cpu/tests/cold_eye.rs b/iris-mpc-cpu/tests/cold_eye.rs index 8b271ca7f7..5f656c4aee 100644 --- a/iris-mpc-cpu/tests/cold_eye.rs +++ b/iris-mpc-cpu/tests/cold_eye.rs @@ -16,7 +16,7 @@ use iris_mpc_cpu::{ aby3::aby3_store::{Aby3Store, DistanceMode, FhdOps}, shared_irises::SharedIrises, }, - protocol::shared_iris::GaloisRingSharedIris, + protocol::shared_iris::{GaloisRingSharedIris, ResidentIris, ResidentLayout}, }; use iris_mpc_store::{ test_utils::{cleanup, temporary_name, test_db_url}, @@ -50,18 +50,28 @@ async fn cold_eye_prefetched_dot_product_matches_resident_and_populates_lfu() -> let resident_store = Aby3Store::::new_storage(Some(HashMap::from([(vector_id, target.clone())]))) + .map_values(|iris| ResidentIris::from_arc(iris, ResidentLayout::U16)) .to_arc(); let resident: Arc = Arc::new(LocalIrisWorkerPool::new_local( resident_store, + ResidentLayout::U16, DistanceMode::MinRotation, 0, )); - let cold_store = SharedIrises::to_arc(Aby3Store::::new_storage(None)); + let cold_store = SharedIrises::new( + HashMap::new(), + ResidentIris::from_arc( + Arc::new(GaloisRingSharedIris::default_for_party(0)), + ResidentLayout::U16, + ), + ) + .to_arc(); let cold: Arc = Arc::new( LocalIrisWorkerPool::new_cold( - init_workers(RIGHT, cold_store.clone(), false), + init_workers(RIGHT, cold_store.clone(), false, ResidentLayout::U16), cold_store.clone(), + ResidentLayout::U16, DistanceMode::MinRotation, 0, ColdStorageInit { @@ -133,137 +143,153 @@ async fn cold_eye_prefetched_dot_product_matches_resident_and_populates_lfu() -> } #[tokio::test] -async fn cold_eye_luc_window_rolls_forward_and_survives_persistence_ack() -> Result<()> { +async fn cold_eye_version_miss_fails_prefetch_and_foreground_fetch() -> Result<()> { let schema_name = temporary_name(); let postgres = PostgresClient::new(&test_db_url()?, &schema_name, AccessMode::ReadWrite).await?; run_migrations(&postgres.pool, false).await?; let db = Store::new(&postgres).await?; - let original = Arc::new(GaloisRingSharedIris::default_for_party(0)); - let vector_id = VectorId::new(1, 0); + let stored = Arc::new(GaloisRingSharedIris::dummy_for_party(0)); + let current_id = VectorId::new(1, 0); let mut tx = db.tx().await?; db.insert_irises( &mut tx, &[StoredIrisRef { id: 1, - left_code: &original.code.coefs, - left_mask: &original.mask.coefs, - right_code: &original.code.coefs, - right_mask: &original.mask.coefs, + left_code: &stored.code.coefs, + left_mask: &stored.mask.coefs, + right_code: &stored.code.coefs, + right_mask: &stored.mask.coefs, }], ) .await?; tx.commit().await?; - let cold_store = SharedIrises::to_arc(Aby3Store::::new_storage(None)); + let cold_store = SharedIrises::new( + HashMap::new(), + ResidentIris::from_arc( + Arc::new(GaloisRingSharedIris::default_for_party(0)), + ResidentLayout::U16, + ), + ) + .to_arc(); let cold: Arc = Arc::new( LocalIrisWorkerPool::new_cold( - init_workers(RIGHT, cold_store.clone(), false), + init_workers(RIGHT, cold_store.clone(), false, ResidentLayout::U16), cold_store, + ResidentLayout::U16, DistanceMode::MinRotation, 0, ColdStorageInit { store: db.clone(), side: RIGHT, - luc_window_ids: vec![vector_id], + luc_window_ids: Vec::new(), latest_serial_id: 1, - luc_window_capacity: 1, + luc_window_capacity: 0, lfu_cache_capacity: 8, }, ) .await?, ); - // Change Postgres after startup. The old current LUC version must still be - // served from RAM rather than becoming a miss against the newer DB row. - let replacement = Arc::new(GaloisRingSharedIris::dummy_for_party(0)); - db.update_iris( - None, - 1, - &replacement.code, - &replacement.mask, - &replacement.code, - &replacement.mask, - ) - .await?; - assert_eq!(cold.fetch_irises(vec![vector_id]).await?[0], original); + // Model the registry being one version ahead of Postgres. Neither the + // asynchronous prefetch nor its foreground fallback may fabricate an + // all-zero iris for this exact-version miss. + let registry_id = current_id.next_version(); + cold.prefetch_irises(vec![registry_id]).await?; + let prefetch_error = cold.wait_for_prefetch().await.unwrap_err(); + assert!( + format!("{prefetch_error:#}").contains("cold-eye database missing"), + "unexpected prefetch error: {prefetch_error:#}" + ); - // The actor-side update advances the rolling cache immediately. Releasing - // its pending-write entry after commit must not evict the LUC copy. - let query_id = QueryId::new(); - cold.cache_queries(vec![(query_id, replacement.clone())]) - .await?; - let next_id = vector_id.next_version(); - cold.insert_irises(vec![(query_id, next_id)]).await?; - assert_eq!(cold.fetch_irises(vec![next_id]).await?[0], replacement); - cold.acknowledge_persisted(vec![next_id.serial_id()]) - .await?; - assert_eq!(cold.fetch_irises(vec![next_id]).await?[0], replacement); + let foreground_error = cold.fetch_irises(vec![registry_id]).await.unwrap_err(); + assert!( + format!("{foreground_error:#}").contains("cold-eye database missing"), + "unexpected foreground error: {foreground_error:#}" + ); cleanup(&postgres, &schema_name).await?; Ok(()) } #[tokio::test] -async fn cold_eye_version_miss_fails_prefetch_and_foreground_fetch() -> Result<()> { +async fn cold_eye_luc_window_rolls_forward_and_survives_persistence_ack() -> Result<()> { let schema_name = temporary_name(); let postgres = PostgresClient::new(&test_db_url()?, &schema_name, AccessMode::ReadWrite).await?; run_migrations(&postgres.pool, false).await?; let db = Store::new(&postgres).await?; - let stored = GaloisRingSharedIris::dummy_for_party(0); - let current_id = VectorId::new(1, 0); + let original = Arc::new(GaloisRingSharedIris::default_for_party(0)); + let vector_id = VectorId::new(1, 0); let mut tx = db.tx().await?; db.insert_irises( &mut tx, &[StoredIrisRef { id: 1, - left_code: &stored.code.coefs, - left_mask: &stored.mask.coefs, - right_code: &stored.code.coefs, - right_mask: &stored.mask.coefs, + left_code: &original.code.coefs, + left_mask: &original.mask.coefs, + right_code: &original.code.coefs, + right_mask: &original.mask.coefs, }], ) .await?; tx.commit().await?; - let cold_store = SharedIrises::to_arc(Aby3Store::::new_storage(None)); + let cold_store = SharedIrises::new( + HashMap::new(), + ResidentIris::from_arc( + Arc::new(GaloisRingSharedIris::default_for_party(0)), + ResidentLayout::U16, + ), + ) + .to_arc(); let cold: Arc = Arc::new( LocalIrisWorkerPool::new_cold( - init_workers(RIGHT, cold_store.clone(), false), + init_workers(RIGHT, cold_store.clone(), false, ResidentLayout::U16), cold_store, + ResidentLayout::U16, DistanceMode::MinRotation, 0, ColdStorageInit { - store: db, + store: db.clone(), side: RIGHT, - luc_window_ids: Vec::new(), + luc_window_ids: vec![vector_id], latest_serial_id: 1, - luc_window_capacity: 0, + luc_window_capacity: 1, lfu_cache_capacity: 8, }, ) .await?, ); - // Model the registry being one version ahead of Postgres. Neither the - // asynchronous prefetch nor its foreground fallback may fabricate an - // all-zero iris for this exact-version miss. - let registry_id = current_id.next_version(); - cold.prefetch_irises(vec![registry_id]).await?; - let prefetch_error = cold.wait_for_prefetch().await.unwrap_err(); - assert!( - format!("{prefetch_error:#}").contains("cold-eye database missing"), - "unexpected prefetch error: {prefetch_error:#}" - ); + // Change Postgres after startup. The old current LUC version must still be + // served from RAM rather than becoming a miss against the newer DB row. + let replacement = Arc::new(GaloisRingSharedIris::dummy_for_party(0)); + db.update_iris( + None, + 1, + &replacement.code, + &replacement.mask, + &replacement.code, + &replacement.mask, + ) + .await?; + assert_eq!(cold.fetch_irises(vec![vector_id]).await?[0], original); - let foreground_error = cold.fetch_irises(vec![registry_id]).await.unwrap_err(); - assert!( - format!("{foreground_error:#}").contains("cold-eye database missing"), - "unexpected foreground error: {foreground_error:#}" - ); + // The actor-side update advances the rolling cache immediately. Releasing + // its pending-write entry after commit must not evict the LUC copy. + let query_id = QueryId::new(); + cold.cache_queries(vec![(query_id, replacement.clone())]) + .await?; + let next_id = vector_id.next_version(); + cold.insert_irises(vec![(query_id, next_id)]).await?; + assert_eq!(cold.fetch_irises(vec![next_id]).await?[0], replacement); + cold.acknowledge_persisted(vec![next_id.serial_id()]) + .await?; + assert_eq!(cold.fetch_irises(vec![next_id]).await?[0], replacement); cleanup(&postgres, &schema_name).await?; Ok(()) diff --git a/iris-mpc-cpu/tests/full_rotation_dot.rs b/iris-mpc-cpu/tests/full_rotation_dot.rs index b59cb861bd..eb8436a613 100644 --- a/iris-mpc-cpu/tests/full_rotation_dot.rs +++ b/iris-mpc-cpu/tests/full_rotation_dot.rs @@ -46,14 +46,26 @@ fn fused_full_rotation_dot_matches_three_hnsw_windows() -> Result<()> { fn run_test() -> Result<()> { let runtime = tokio::runtime::Runtime::new()?; + // `preferred_scan_layout()` stores mixed planes where the UMMLA kernel is + // available, so the fused scan below runs the mixed kernel while the + // windowed comparison paths run the u16 kernel — a cross-kernel check. + let layout = iris_mpc_cpu::protocol::shared_iris::preferred_scan_layout(); let mut store = SharedIrises::new( HashMap::new(), - Arc::new(GaloisRingSharedIris::default_for_party(0)), + iris_mpc_cpu::protocol::shared_iris::ResidentIris::from_arc( + Arc::new(GaloisRingSharedIris::default_for_party(0)), + layout, + ), ); let vector_ids = (0..TARGETS) - .map(|idx| store.append(Arc::new(deterministic_iris(idx as u16 + 1)))) + .map(|idx| { + store.append(iris_mpc_cpu::protocol::shared_iris::ResidentIris::from_arc( + Arc::new(deterministic_iris(idx as u16 + 1)), + layout, + )) + }) .collect::>(); - let pool = LocalIrisWorkerPool::new_local(store.to_arc(), DistanceMode::MinRotation, 0); + let pool = LocalIrisWorkerPool::new_local(store.to_arc(), layout, DistanceMode::MinRotation, 0); let query_id = QueryId::new(); runtime.block_on(pool.cache_queries(vec![(query_id, Arc::new(deterministic_iris(0x5a5a)))]))?; diff --git a/iris-mpc-gpu/Cargo.toml b/iris-mpc-gpu/Cargo.toml index ab9aa03586..078d3200f5 100644 --- a/iris-mpc-gpu/Cargo.toml +++ b/iris-mpc-gpu/Cargo.toml @@ -52,3 +52,8 @@ harness = false [[bench]] name = "transpose" harness = false + +[[bench]] +name = "linear_scan_dot" +harness = false +required-features = ["gpu_dependent"] diff --git a/iris-mpc-gpu/benches/linear_scan_dot.rs b/iris-mpc-gpu/benches/linear_scan_dot.rs new file mode 100644 index 0000000000..50314b3c10 --- /dev/null +++ b/iris-mpc-gpu/benches/linear_scan_dot.rs @@ -0,0 +1,559 @@ +use eyre::{ensure, Result}; +use iris_mpc_common::{ + galois_engine::degree4::{GaloisRingIrisCodeShare, GaloisRingTrimmedMaskCodeShare}, + VectorId, IRIS_CODE_LENGTH, MASK_CODE_LENGTH, ROTATIONS, +}; +use iris_mpc_cpu::{ + execution::hawk_main::iris_worker::{IrisWorkerPool, LocalIrisWorkerPool, QueryId, QuerySpec}, + hawkers::{aby3::aby3_store::DistanceMode, shared_irises::SharedIrises}, + protocol::shared_iris::{ArcIris, GaloisRingSharedIris}, +}; +use iris_mpc_gpu::{ + dot::share_db::{preprocess_query, ProcessedDatabase, ShareDB, SlicedProcessedDatabase}, + helpers::{ + device_manager::DeviceManager, + query_processor::{ + CudaVec2DSlicerRawPointer, CudaVec2DSlicerU32, CudaVec2DSlicerU8, StreamAwareCudaSlice, + }, + }, +}; +use rayon::prelude::*; +use std::{ + collections::HashMap, + env, + hint::black_box, + sync::Arc, + time::{Duration, Instant}, +}; + +/// This is the chunk size used by `ServerActor` in production. Keeping it here +/// (instead of exporting an actor implementation detail) makes the benchmark +/// fail visibly if its production-parity assumption is changed. +const PRODUCTION_DB_CHUNK_SIZE: usize = 1 << 15; +const DEFAULT_DB_SIZE: usize = 3 * 4 * PRODUCTION_DB_CHUNK_SIZE; +const DEFAULT_WARMUP_RUNS: usize = 1; +const DEFAULT_MEASURED_RUNS: usize = 5; + +struct GpuScan<'a> { + device_manager: Arc, + code_engine: ShareDB, + mask_engine: ShareDB, + code_db: SlicedProcessedDatabase, + mask_db: SlicedProcessedDatabase, + db_sizes: Vec, + code_queries: CudaVec2DSlicerU8, + mask_queries: CudaVec2DSlicerU8, + code_query_sums: CudaVec2DSlicerU32, + mask_query_sums: CudaVec2DSlicerU32, + streams: [Vec; 2], + blass: [Vec; 2], + code_buffers: [iris_mpc_gpu::dot::share_db::DBChunkBuffers; 2], + mask_buffers: [iris_mpc_gpu::dot::share_db::DBChunkBuffers; 2], + _lifetime: std::marker::PhantomData<&'a ()>, +} + +impl GpuScan<'_> { + fn chunk_sizes(&self, chunk_idx: usize) -> Vec { + self.db_sizes + .iter() + .map(|&size| { + size.saturating_sub(PRODUCTION_DB_CHUNK_SIZE * chunk_idx) + .min(PRODUCTION_DB_CHUNK_SIZE) + }) + .collect() + } + + fn n_chunks(&self) -> usize { + self.db_sizes + .iter() + .copied() + .max() + .unwrap_or(0) + .div_ceil(PRODUCTION_DB_CHUNK_SIZE) + } + + fn prefetch(&self, chunk_idx: usize, buffer: usize, stream: usize) { + let sizes = self.chunk_sizes(chunk_idx); + let offsets = self + .db_sizes + .iter() + .map(|_| chunk_idx * PRODUCTION_DB_CHUNK_SIZE) + .collect::>(); + self.code_db.prefetch_chunk( + &self.code_engine, + &self.code_buffers[buffer], + &sizes, + &offsets, + &self.db_sizes, + &self.streams[stream], + ); + self.mask_db.prefetch_chunk( + &self.mask_engine, + &self.mask_buffers[buffer], + &sizes, + &offsets, + &self.db_sizes, + &self.streams[stream], + ); + } + + fn dot_chunk(&mut self, logical_chunk_idx: usize, buffer: usize, stream: usize) { + let sizes = self.chunk_sizes(logical_chunk_idx); + let offset = logical_chunk_idx * PRODUCTION_DB_CHUNK_SIZE; + self.code_engine.dot( + &self.code_queries, + &CudaVec2DSlicerRawPointer::from(&self.code_buffers[buffer]), + &sizes, + 0, + &self.streams[stream], + &self.blass[stream], + ); + self.mask_engine.dot( + &self.mask_queries, + &CudaVec2DSlicerRawPointer::from(&self.mask_buffers[buffer]), + &sizes, + 0, + &self.streams[stream], + &self.blass[stream], + ); + self.code_engine.dot_reduce( + &self.code_query_sums, + &self.code_db.code_sums_gr, + &sizes, + offset, + &self.streams[stream], + ); + self.mask_engine.dot_reduce_and_multiply( + &self.mask_query_sums, + &self.mask_db.code_sums_gr, + &sizes, + offset, + &self.streams[stream], + 2, + ); + } + + /// Production-shaped copy-only scan: page-locked host DB, async HtoD, and + /// alternating buffers/streams, without any dot kernels. + fn run_copy_only_once(&mut self) -> Duration { + let started = Instant::now(); + self.prefetch(0, 0, 0); + for chunk_idx in 1..self.n_chunks() { + let slot = chunk_idx % 2; + self.prefetch(chunk_idx, slot, slot); + } + self.device_manager.await_streams(&self.streams[0]); + self.device_manager.await_streams(&self.streams[1]); + started.elapsed() + } + + /// Put two full chunks in the production buffers before the timed + /// resident-dot measurement. Each buffer is much larger than GPU L2, so + /// alternating them cannot turn the scan into an L2-cache benchmark. + fn prepare_resident_dot(&self) { + self.prefetch(0, 0, 0); + self.prefetch(1.min(self.n_chunks() - 1), 1, 1); + self.device_manager.await_streams(&self.streams[0]); + self.device_manager.await_streams(&self.streams[1]); + } + + fn run_resident_dot_once(&mut self) -> Duration { + let started = Instant::now(); + for chunk_idx in 0..self.n_chunks() { + let slot = chunk_idx % 2; + self.dot_chunk(chunk_idx, slot, slot); + // ShareDB reuses one result/intermediate allocation per device. + // Production events serialize that reuse while the other stream + // prefetches; a stream sync provides the same dependency here. + self.device_manager.await_streams(&self.streams[slot]); + } + started.elapsed() + } + + fn run_combined_once(&mut self) -> Duration { + let n_chunks = self.n_chunks(); + + let started = Instant::now(); + + // Production starts chunk zero on stream set zero. Every later chunk + // is loaded into the other buffer/stream while the current chunk's + // code and mask dots and reductions execute. + self.prefetch(0, 0, 0); + + for chunk_idx in 0..n_chunks { + let current = chunk_idx % 2; + let next = (chunk_idx + 1) % 2; + if chunk_idx + 1 < n_chunks { + self.prefetch(chunk_idx + 1, next, next); + } + self.dot_chunk(chunk_idx, current, current); + self.device_manager.await_streams(&self.streams[current]); + } + + self.device_manager.await_streams(&self.streams[0]); + self.device_manager.await_streams(&self.streams[1]); + started.elapsed() + } +} + +fn env_usize(name: &str, default: usize) -> Result { + match env::var(name) { + Ok(value) => Ok(value.parse()?), + Err(env::VarError::NotPresent) => Ok(default), + Err(err) => Err(err.into()), + } +} + +fn deterministic_iris(party_id: usize) -> GaloisRingSharedIris { + let mut code = [0u16; IRIS_CODE_LENGTH]; + let mut mask = [0u16; MASK_CODE_LENGTH]; + for (idx, value) in code.iter_mut().enumerate() { + *value = (idx as u16).wrapping_mul(17).wrapping_add(23); + } + for (idx, value) in mask.iter_mut().enumerate() { + *value = (idx as u16).wrapping_mul(29).wrapping_add(11); + } + GaloisRingSharedIris { + code: GaloisRingIrisCodeShare::new(code, party_id), + mask: GaloisRingTrimmedMaskCodeShare::new(mask, party_id), + } +} + +/// Upload the signed-limb row sums consumed by the production reduction +/// kernel. `ShareDB::query_sums` obtains the same values with GEMM, but that +/// helper requires its row count to be divisible by four. A literal B=1 query +/// has 31 rows, so doing this once on the host preserves the exact 31-row dot +/// workload instead of padding the measured GPU scan with phantom queries. +fn upload_query_sums( + device_manager: &DeviceManager, + preprocessed_query: &[Vec], + query_rows: usize, + row_width: usize, +) -> Result { + ensure!(preprocessed_query.len() == 2, "a ring share has two limbs"); + ensure!( + preprocessed_query + .iter() + .all(|limb| limb.len() == query_rows * row_width), + "unexpected preprocessed query shape" + ); + let sums = preprocessed_query + .iter() + .map(|limb| { + limb.chunks_exact(row_width) + .map(|row| row.iter().map(|&value| value as i8 as i32).sum::() as u32) + .collect::>() + }) + .collect::>(); + + let limb_0 = device_manager + .devices() + .iter() + .map(|device| { + device + .htod_sync_copy(&sums[0]) + .map(StreamAwareCudaSlice::from) + .map_err(Into::into) + }) + .collect::>>()?; + let limb_1 = device_manager + .devices() + .iter() + .map(|device| { + device + .htod_sync_copy(&sums[1]) + .map(StreamAwareCudaSlice::from) + .map_err(Into::into) + }) + .collect::>>()?; + Ok(CudaVec2DSlicerU32 { limb_0, limb_1 }) +} + +fn build_gpu_scan(db_size: usize, query: &GaloisRingSharedIris) -> Result> { + let device_manager = Arc::new(DeviceManager::init()); + let n_devices = device_manager.device_count(); + ensure!(n_devices > 0, "benchmark requires at least one GPU"); + ensure!( + db_size.is_multiple_of(n_devices * PRODUCTION_DB_CHUNK_SIZE), + "DB size must be a multiple of {} so every GPU chunk is full", + n_devices * PRODUCTION_DB_CHUNK_SIZE + ); + + let code_engine = ShareDB::init( + 0, + device_manager.clone(), + PRODUCTION_DB_CHUNK_SIZE, + ROTATIONS, + IRIS_CODE_LENGTH, + ([0u32; 8], [0u32; 8]), + vec![], + ); + let mask_engine = ShareDB::init( + 0, + device_manager.clone(), + PRODUCTION_DB_CHUNK_SIZE, + ROTATIONS, + MASK_CODE_LENGTH, + ([0u32; 8], [0u32; 8]), + vec![], + ); + + let streams = [device_manager.fork_streams(), device_manager.fork_streams()]; + let blass = [ + device_manager.create_cublas(&streams[0]), + device_manager.create_cublas(&streams[1]), + ]; + let code_buffers = [ + code_engine.alloc_db_chunk_buffer(PRODUCTION_DB_CHUNK_SIZE), + code_engine.alloc_db_chunk_buffer(PRODUCTION_DB_CHUNK_SIZE), + ]; + let mask_buffers = [ + mask_engine.alloc_db_chunk_buffer(PRODUCTION_DB_CHUNK_SIZE), + mask_engine.alloc_db_chunk_buffer(PRODUCTION_DB_CHUNK_SIZE), + ]; + + let mut code_db = code_engine.alloc_db(db_size); + let mut mask_db = mask_engine.alloc_db(db_size); + let code_limb_0 = query + .code + .coefs + .iter() + .map(|x| *x as u8) + .collect::>(); + let code_limb_1 = query + .code + .coefs + .iter() + .map(|x| (*x >> 8) as u8) + .collect::>(); + let mask_limb_0 = query + .mask + .coefs + .iter() + .map(|x| *x as u8) + .collect::>(); + let mask_limb_1 = query + .mask + .coefs + .iter() + .map(|x| (*x >> 8) as u8) + .collect::>(); + + println!("loading GPU host database ({db_size} records)"); + (0..db_size).into_par_iter().for_each(|idx| { + ShareDB::load_single_record_from_s3( + idx, + &code_db.code_gr, + &code_limb_0, + &code_limb_1, + n_devices, + IRIS_CODE_LENGTH, + ); + ShareDB::load_single_record_from_s3( + idx, + &mask_db.code_gr, + &mask_limb_0, + &mask_limb_1, + n_devices, + MASK_CODE_LENGTH, + ); + }); + let db_sizes = vec![db_size / n_devices; n_devices]; + code_db.preprocess(&code_engine, &db_sizes); + mask_db.preprocess(&mask_engine, &db_sizes); + device_manager.register_host_memory(&code_db, db_size, IRIS_CODE_LENGTH); + device_manager.register_host_memory(&mask_db, db_size, MASK_CODE_LENGTH); + + let code_query = (0..ROTATIONS) + .flat_map(|rotation| { + query + .code + .coefs + .iter() + .map(move |value| value.wrapping_add(rotation as u16)) + }) + .collect::>(); + let mask_query = (0..ROTATIONS) + .flat_map(|rotation| { + query + .mask + .coefs + .iter() + .map(move |value| value.wrapping_add(rotation as u16)) + }) + .collect::>(); + let code_query = preprocess_query(&code_query); + let mask_query = preprocess_query(&mask_query); + let code_queries = + device_manager.htod_transfer_query(&code_query, &streams[0], 1, IRIS_CODE_LENGTH)?; + let mask_queries = + device_manager.htod_transfer_query(&mask_query, &streams[0], 1, MASK_CODE_LENGTH)?; + let code_query_sums = + upload_query_sums(&device_manager, &code_query, ROTATIONS, IRIS_CODE_LENGTH)?; + let mask_query_sums = + upload_query_sums(&device_manager, &mask_query, ROTATIONS, MASK_CODE_LENGTH)?; + device_manager.await_streams(&streams[0]); + + Ok(GpuScan { + device_manager, + code_engine, + mask_engine, + code_db, + mask_db, + db_sizes, + code_queries, + mask_queries, + code_query_sums, + mask_query_sums, + streams, + blass, + code_buffers, + mask_buffers, + _lifetime: std::marker::PhantomData, + }) +} + +fn build_cpu_pool(db_size: usize, query: &ArcIris) -> (LocalIrisWorkerPool, Vec) { + println!("loading CPU database ({db_size} records)"); + let mut store = SharedIrises::new( + HashMap::new(), + Arc::new(GaloisRingSharedIris::default_for_party(0)), + ); + store.reserve(db_size); + let mut vector_ids = Vec::with_capacity(db_size); + for _ in 0..db_size { + // A deep clone gives every logical record its own backing memory, so + // the scan cannot turn into an unrealistically cache-resident test. + vector_ids.push(store.append(Arc::new((**query).clone()))); + } + let store = store.to_arc(); + ( + LocalIrisWorkerPool::new_local(store, DistanceMode::MinRotation, 0), + vector_ids, + ) +} + +fn median(samples: &[Duration]) -> Duration { + let mut samples = samples.to_vec(); + samples.sort_unstable(); + samples[samples.len() / 2] +} + +fn print_result(backend: &str, db_size: usize, internal_rotations: usize, samples: &[Duration]) { + let median = median(samples).as_secs_f64(); + let mean = samples.iter().map(Duration::as_secs_f64).sum::() / samples.len() as f64; + let comparisons_per_second = db_size as f64 / median; + let rotation_pairs_per_second = db_size as f64 * internal_rotations as f64 / median; + println!( + "BENCH_RESULT backend={backend} batch_size=1 db_size={db_size} samples={} \ + median_seconds={median:.6} mean_seconds={mean:.6} comparisons_per_second={comparisons_per_second:.3} \ + internal_rotations={internal_rotations} rotation_pairs_per_second={rotation_pairs_per_second:.3}", + samples.len(), + ); +} + +fn measure( + backend: &str, + warmup_runs: usize, + measured_runs: usize, + mut run: impl FnMut() -> Duration, +) -> Vec { + for _ in 0..warmup_runs { + black_box(run()); + } + (0..measured_runs) + .map(|sample| { + let elapsed = run(); + println!( + "BENCH_SAMPLE backend={backend} run={sample} seconds={:.6}", + elapsed.as_secs_f64() + ); + elapsed + }) + .collect() +} + +fn main() -> Result<()> { + let db_size = env_usize("IRIS_MPC_DOT_BENCH_DB_SIZE", DEFAULT_DB_SIZE)?; + 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)?; + ensure!(db_size > 0, "DB size must be positive"); + ensure!(measured_runs > 0, "measured run count must be positive"); + + println!( + "BENCH_CONFIG batch_size=1 db_size={db_size} warmup_runs={warmup_runs} \ + measured_runs={measured_runs} chunk_size={PRODUCTION_DB_CHUNK_SIZE} \ + gpu_loading=page_locked_async_double_buffered scope=code_mask_dot_and_reduce" + ); + + let query = Arc::new(deterministic_iris(0)); + + let mut gpu = build_gpu_scan(db_size, &query)?; + println!("BENCH_GPU gpu_count={}", gpu.device_manager.device_count()); + let gpu_copy_samples = measure("gpu_copy_only", warmup_runs, measured_runs, || { + gpu.run_copy_only_once() + }); + gpu.prepare_resident_dot(); + let gpu_dot_samples = measure("gpu_resident_dot", warmup_runs, measured_runs, || { + gpu.run_resident_dot_once() + }); + let gpu_samples = measure("gpu_combined", warmup_runs, measured_runs, || { + gpu.run_combined_once() + }); + + let runtime = tokio::runtime::Runtime::new()?; + let (cpu, vector_ids) = build_cpu_pool(db_size, &query); + let query_id = QueryId::new(); + runtime.block_on(cpu.cache_queries(vec![(query_id, query.clone())]))?; + let query = QuerySpec::new(query_id); + + let run_cpu = |run: usize, report: bool| -> Result { + let ids = vector_ids.clone(); + let started = Instant::now(); + let output = runtime.block_on(cpu.compute_dot_products_full_rotations(query, ids))?; + let elapsed = started.elapsed(); + ensure!( + output.len() == 2 * ROTATIONS * db_size, + "CPU worker returned an unexpected dot-product result count" + ); + black_box(output); + if report { + println!( + "BENCH_SAMPLE backend=cpu run={run} seconds={:.6}", + elapsed.as_secs_f64() + ); + } + Ok(elapsed) + }; + for run in 0..warmup_runs { + black_box(run_cpu(run, false)?); + } + let cpu_samples = (0..measured_runs) + .map(|run| run_cpu(run, true)) + .collect::>>()?; + + print_result("cpu", db_size, ROTATIONS, &cpu_samples); + print_result("gpu_copy_only", db_size, 0, &gpu_copy_samples); + print_result("gpu_resident_dot", db_size, ROTATIONS, &gpu_dot_samples); + print_result("gpu_combined", db_size, ROTATIONS, &gpu_samples); + let copy_seconds = median(&gpu_copy_samples).as_secs_f64(); + let dot_seconds = median(&gpu_dot_samples).as_secs_f64(); + let combined_seconds = median(&gpu_samples).as_secs_f64(); + let transferred_bytes = db_size * 2 * (IRIS_CODE_LENGTH + MASK_CODE_LENGTH); + let copy_gbps = transferred_bytes as f64 / copy_seconds / 1e9; + let limiting_component = if copy_seconds >= dot_seconds { + "pci_copy" + } else { + "dot_compute" + }; + println!( + "GPU_BOTTLENECK limiting_component={limiting_component} copy_gbps={copy_gbps:.3} \ + copy_seconds={copy_seconds:.6} resident_dot_seconds={dot_seconds:.6} \ + combined_seconds={combined_seconds:.6} combined_over_max_component={:.3}", + combined_seconds / copy_seconds.max(dot_seconds), + ); + let speedup = median(&cpu_samples).as_secs_f64() / median(&gpu_samples).as_secs_f64(); + println!("BENCH_COMPARISON metric=logical_comparisons_per_second gpu_over_cpu={speedup:.3}"); + + Ok(()) +} diff --git a/iris-mpc-utils/src/client/options/types.rs b/iris-mpc-utils/src/client/options/types.rs index 0efff63021..4d9031a6b3 100644 --- a/iris-mpc-utils/src/client/options/types.rs +++ b/iris-mpc-utils/src/client/options/types.rs @@ -683,120 +683,107 @@ fn ground_truth_into_iter( operations.extend(cycle); } - for scheduled_operation in operations + for operation in operations .into_iter() .take(request_count.saturating_sub(batches.len())) { let label = format!("ground-truth-{sequence:04}"); sequence += 1; - // A deletion may empty the identity set. Replace any operation that - // needs a target with an enrollment until there is a live identity - // again. New enrollments deliberately do not consume target-selection - // randomness, keeping target draws tied to operations that use them. - let operation = if live.is_empty() { - GroundTruthOperation::NewUniqueness - } else { - scheduled_operation - }; - - let request = if matches!(operation, GroundTruthOperation::NewUniqueness) { - let identity_label = format!("{label}-enroll"); - let iris_pair = fresh_pair(); - let request = RequestOptions::new( - Some(&identity_label), - RequestPayloadOptions::Uniqueness { - iris_pair: Some(iris_pair), - insertion_layers: None, - }, - ) - .with_expected(json!({ "is_match": false })); - live.push(GroundTruthIdentity { - label: identity_label, - iris_pair, - }); - request - } else { - let target_idx = rng.gen_range(0..live.len()); - let target = live[target_idx].clone(); - match operation { - GroundTruthOperation::NewUniqueness => { - unreachable!("new uniqueness is handled before target selection") - } - GroundTruthOperation::DuplicateUniqueness => RequestOptions::new( - Some(&format!("{label}-duplicate")), + let target_idx = rng.gen_range(0..live.len()); + let target = live[target_idx].clone(); + + let request = match operation { + GroundTruthOperation::NewUniqueness => { + let identity_label = format!("{label}-enroll"); + let iris_pair = fresh_pair(); + let request = RequestOptions::new( + Some(&identity_label), RequestPayloadOptions::Uniqueness { - iris_pair: Some(target.iris_pair), + iris_pair: Some(iris_pair), insertion_layers: None, }, ) - .with_expected(json!({ "is_match": true })), - GroundTruthOperation::MirroredUniqueness => RequestOptions::new( - Some(&format!("{label}-mirror")), - RequestPayloadOptions::Mirrored { - iris_pair: Some(target.iris_pair), - insertion_layers: None, - }, - ), - GroundTruthOperation::ReauthMatch => RequestOptions::new( - Some(&format!("{label}-reauth-match")), - RequestPayloadOptions::Reauthorisation { - iris_pair: Some(target.iris_pair), + .with_expected(json!({ "is_match": false })); + live.push(GroundTruthIdentity { + label: identity_label, + iris_pair, + }); + request + } + GroundTruthOperation::DuplicateUniqueness => RequestOptions::new( + Some(&format!("{label}-duplicate")), + RequestPayloadOptions::Uniqueness { + iris_pair: Some(target.iris_pair), + insertion_layers: None, + }, + ) + .with_expected(json!({ "is_match": true })), + GroundTruthOperation::MirroredUniqueness => RequestOptions::new( + Some(&format!("{label}-mirror")), + RequestPayloadOptions::Mirrored { + iris_pair: Some(target.iris_pair), + insertion_layers: None, + }, + ), + GroundTruthOperation::ReauthMatch => RequestOptions::new( + Some(&format!("{label}-reauth-match")), + RequestPayloadOptions::Reauthorisation { + iris_pair: Some(target.iris_pair), + parent: Parent::Label(target.label), + }, + ) + .with_expected(json!({ "success": true })), + GroundTruthOperation::ReauthNonMatch => RequestOptions::new( + Some(&format!("{label}-reauth-non-match")), + RequestPayloadOptions::Reauthorisation { + iris_pair: Some(fresh_pair()), + parent: Parent::Label(target.label), + }, + ) + .with_expected(json!({ "success": false })), + GroundTruthOperation::ResetCheck => RequestOptions::new( + Some(&format!("{label}-reset-check")), + RequestPayloadOptions::ResetCheck { + iris_pair: Some(target.iris_pair), + }, + ), + GroundTruthOperation::RecoveryCheck => RequestOptions::new( + Some(&format!("{label}-recovery-check")), + RequestPayloadOptions::RecoveryCheck { + iris_pair: Some(target.iris_pair), + }, + ), + GroundTruthOperation::ResetUpdate => { + let iris_pair = fresh_pair(); + live[target_idx].iris_pair = iris_pair; + RequestOptions::new( + Some(&format!("{label}-reset-update")), + RequestPayloadOptions::ResetUpdate { + iris_pair: Some(iris_pair), parent: Parent::Label(target.label), }, ) - .with_expected(json!({ "success": true })), - GroundTruthOperation::ReauthNonMatch => RequestOptions::new( - Some(&format!("{label}-reauth-non-match")), - RequestPayloadOptions::Reauthorisation { - iris_pair: Some(fresh_pair()), + } + GroundTruthOperation::RecoveryUpdate => { + let iris_pair = fresh_pair(); + live[target_idx].iris_pair = iris_pair; + RequestOptions::new( + Some(&format!("{label}-recovery-update")), + RequestPayloadOptions::RecoveryUpdate { + iris_pair: Some(iris_pair), parent: Parent::Label(target.label), }, ) - .with_expected(json!({ "success": false })), - GroundTruthOperation::ResetCheck => RequestOptions::new( - Some(&format!("{label}-reset-check")), - RequestPayloadOptions::ResetCheck { - iris_pair: Some(target.iris_pair), - }, - ), - GroundTruthOperation::RecoveryCheck => RequestOptions::new( - Some(&format!("{label}-recovery-check")), - RequestPayloadOptions::RecoveryCheck { - iris_pair: Some(target.iris_pair), + } + GroundTruthOperation::IdentityDeletion => { + let deleted = live.swap_remove(target_idx); + RequestOptions::new( + Some(&format!("{label}-delete")), + RequestPayloadOptions::IdentityDeletion { + parent: Parent::Label(deleted.label), }, - ), - GroundTruthOperation::ResetUpdate => { - let iris_pair = fresh_pair(); - live[target_idx].iris_pair = iris_pair; - RequestOptions::new( - Some(&format!("{label}-reset-update")), - RequestPayloadOptions::ResetUpdate { - iris_pair: Some(iris_pair), - parent: Parent::Label(target.label), - }, - ) - } - GroundTruthOperation::RecoveryUpdate => { - let iris_pair = fresh_pair(); - live[target_idx].iris_pair = iris_pair; - RequestOptions::new( - Some(&format!("{label}-recovery-update")), - RequestPayloadOptions::RecoveryUpdate { - iris_pair: Some(iris_pair), - parent: Parent::Label(target.label), - }, - ) - } - GroundTruthOperation::IdentityDeletion => { - let deleted = live.swap_remove(target_idx); - RequestOptions::new( - Some(&format!("{label}-delete")), - RequestPayloadOptions::IdentityDeletion { - parent: Parent::Label(deleted.label), - }, - ) - .with_expected(json!({ "success": true })) - } + ) + .with_expected(json!({ "success": true })) } }; batches.push(vec![request]); @@ -1196,20 +1183,4 @@ mod tests { ]) ); } - - #[test] - fn ground_truth_schedule_recovers_when_deletion_empties_live_identities() { - for rng_seed in 0..64 { - let batches = RequestBatchOptions::GroundTruth { - request_count: 100, - rng_seed, - initial_uniqueness_count: 1, - } - .into_iter() - .collect::>(); - - assert_eq!(batches.len(), 100, "rng seed {rng_seed}"); - assert!(batches.iter().all(|batch| batch.len() == 1)); - } - } } diff --git a/iris-mpc/src/server/mod.rs b/iris-mpc/src/server/mod.rs index 6d018b59de..5872e08867 100644 --- a/iris-mpc/src/server/mod.rs +++ b/iris-mpc/src/server/mod.rs @@ -297,7 +297,6 @@ async fn server_main_with_search_mode(config: Config, search_mode: HawkSearchMod let results_persistence = ResultsPersistence { iris_store: iris_store.clone(), graph_store, - search_mode, ack_handle: hawk_actor.persistence_ack_handle(), }; let tx_results = start_results_thread( @@ -453,6 +452,25 @@ fn validate_max_batch_size(max_batch_size: usize, search_mode: HawkSearchMode) - Ok(()) } +#[cfg(test)] +mod config_tests { + use super::*; + + #[test] + fn linear_scan_requires_single_request_batches() { + assert!(validate_max_batch_size(1, HawkSearchMode::LinearScan).is_ok()); + + let error = validate_max_batch_size(2, HawkSearchMode::LinearScan) + .expect_err("linear scan must reject batches larger than one"); + assert!(error.to_string().contains("requires max_batch_size=1")); + } + + #[test] + fn hnsw_keeps_batched_requests() { + assert!(validate_max_batch_size(64, HawkSearchMode::Hnsw).is_ok()); + } +} + /// Returns initialized PostgreSQL clients for interacting /// with iris share and HNSW graph stores. async fn prepare_stores( @@ -786,8 +804,14 @@ async fn init_hawk_actor( "Initialize iris db: Loading from DB (parallelism: {})", parallelism, ); - let initializer: Box = - Box::new(LocalWorkerPoolInitializer::new_load_from_db( + // Exact-scan pools store irises in the mixed-plane layout on CPUs with + // the UMMLA kernel; HNSW pools keep plain ArcIris values. + let resident_layout = match search_mode { + HawkSearchMode::LinearScan => iris_mpc_cpu::protocol::shared_iris::preferred_scan_layout(), + HawkSearchMode::Hnsw => iris_mpc_cpu::protocol::shared_iris::ResidentLayout::U16, + }; + let initializer: Box = Box::new( + LocalWorkerPoolInitializer::new_load_from_db( hawk_args.party_index, HAWK_DISTANCE_MODE, hawk_args.numa, @@ -805,7 +829,9 @@ async fn init_hawk_actor( ampc_anon_stats::types::Eye::Right => 1, }), }, - )); + ) + .with_resident_layout(resident_layout), + ); let now = Instant::now(); let ct = shutdown_handler.get_network_cancellation_token(); @@ -893,7 +919,6 @@ async fn init_hawk_actor( struct ResultsPersistence { iris_store: Store, graph_store: GraphPg>, - search_mode: HawkSearchMode, ack_handle: IrisPersistenceAckHandle, } @@ -911,7 +936,6 @@ async fn start_results_thread( let ResultsPersistence { iris_store: store_bg, graph_store, - search_mode, ack_handle: persistence_ack_handle, } = persistence; let shutdown_handler_bg = Arc::clone(shutdown_handler); @@ -926,7 +950,6 @@ async fn start_results_thread( &graph_store, &sns_client_bg, &config_bg, - search_mode, &sns_attributes_maps.uniqueness_result_attributes, &sns_attributes_maps.reauth_result_attributes, &sns_attributes_maps.identity_deletion_result_attributes, @@ -1238,22 +1261,3 @@ async fn run_main_server_loop( } Ok(()) } - -#[cfg(test)] -mod config_tests { - use super::*; - - #[test] - fn linear_scan_requires_single_request_batches() { - assert!(validate_max_batch_size(1, HawkSearchMode::LinearScan).is_ok()); - - let error = validate_max_batch_size(2, HawkSearchMode::LinearScan) - .expect_err("linear scan must reject batches larger than one"); - assert!(error.to_string().contains("requires max_batch_size=1")); - } - - #[test] - fn hnsw_keeps_batched_requests() { - assert!(validate_max_batch_size(64, HawkSearchMode::Hnsw).is_ok()); - } -} diff --git a/iris-mpc/src/services/processors/job.rs b/iris-mpc/src/services/processors/job.rs index f850107691..efca31e2d5 100644 --- a/iris-mpc/src/services/processors/job.rs +++ b/iris-mpc/src/services/processors/job.rs @@ -17,7 +17,7 @@ use iris_mpc_common::helpers::sync::ModificationKey::{RequestId, RequestSerialId use iris_mpc_common::iris_db::get_dummy_shares_for_deletion; use iris_mpc_common::{job::ServerJobResult, SerialId}; use iris_mpc_cpu::execution::hawk_main::{ - iris_worker::IrisPersistenceAckHandle, GraphStore, HawkMutation, HawkSearchMode, + iris_worker::IrisPersistenceAckHandle, GraphStore, HawkMutation, }; use iris_mpc_store::{ExplicitVersionToken, Store, StoredIrisRef}; use itertools::{izip, Itertools}; @@ -37,29 +37,6 @@ fn one_indexed_ids_at(ids: &[Vec], index: usize) -> Option> { .map(|ids| ids.iter().map(|id| id + 1).collect()) } -fn linear_scan_one_indexed_nonempty_ids_at( - search_mode: HawkSearchMode, - ids: &[Vec], - index: usize, -) -> Option> { - if search_mode != HawkSearchMode::LinearScan { - return None; - } - ids.get(index) - .filter(|ids| !ids.is_empty()) - .map(|ids| ids.iter().map(|id| id + 1).collect()) -} - -fn linear_scan_count_at( - search_mode: HawkSearchMode, - counts: &[usize], - index: usize, -) -> Option { - (search_mode == HawkSearchMode::LinearScan) - .then(|| counts.get(index).copied()) - .flatten() -} - /// Processes a ServerJobResult, storing data in the database and sending result messages /// through SNS. #[allow(clippy::too_many_arguments)] @@ -71,7 +48,6 @@ pub async fn process_job_result( graph_store: &GraphStore, sns_client: &SNSClient, config: &Config, - search_mode: HawkSearchMode, uniqueness_result_attributes: &HashMap, reauth_result_attributes: &HashMap, identity_deletion_result_attributes: &HashMap, @@ -167,31 +143,22 @@ pub async fn process_job_result( .get(i) .filter(|rotations| !rotations.is_empty()) .cloned(), - linear_scan_one_indexed_nonempty_ids_at( - search_mode, - &full_face_mirror_match_ids, - i, - ), - linear_scan_one_indexed_nonempty_ids_at( - search_mode, - &full_face_mirror_partial_match_ids_left, - i, - ), - linear_scan_one_indexed_nonempty_ids_at( - search_mode, - &full_face_mirror_partial_match_ids_right, - i, - ), - linear_scan_count_at( - search_mode, - &full_face_mirror_partial_match_counters_left, - i, - ), - linear_scan_count_at( - search_mode, - &full_face_mirror_partial_match_counters_right, - i, - ), + full_face_mirror_match_ids + .get(i) + .filter(|ids| !ids.is_empty()) + .map(|ids| ids.iter().map(|x| x + 1).collect::>()), + full_face_mirror_partial_match_ids_left + .get(i) + .filter(|ids| !ids.is_empty()) + .map(|ids| ids.iter().map(|x| x + 1).collect::>()), + full_face_mirror_partial_match_ids_right + .get(i) + .filter(|ids| !ids.is_empty()) + .map(|ids| ids.iter().map(|x| x + 1).collect::>()), + (!full_face_mirror_partial_match_counters_left.is_empty()) + .then_some(full_face_mirror_partial_match_counters_left[i]), + (!full_face_mirror_partial_match_counters_right.is_empty()) + .then_some(full_face_mirror_partial_match_counters_right[i]), match full_face_mirror_attack_detected.is_empty() { false => full_face_mirror_attack_detected[i], true => false, @@ -676,34 +643,11 @@ pub async fn process_job_result( #[cfg(test)] mod tests { - use super::{ - linear_scan_count_at, linear_scan_one_indexed_nonempty_ids_at, one_indexed_ids_at, - }; - use iris_mpc_cpu::execution::hawk_main::HawkSearchMode; + use super::one_indexed_ids_at; #[test] fn hidden_partial_results_are_absent_instead_of_indexing_an_empty_outer_vector() { assert_eq!(one_indexed_ids_at(&[], 0), None); assert_eq!(one_indexed_ids_at(&[vec![0, 4]], 0), Some(vec![1, 5])); } - - #[test] - fn full_face_mirror_details_are_exposed_only_by_linear_scan() { - let ids = [vec![0, 4]]; - let counts = [7]; - - assert_eq!( - linear_scan_one_indexed_nonempty_ids_at(HawkSearchMode::Hnsw, &ids, 0), - None - ); - assert_eq!(linear_scan_count_at(HawkSearchMode::Hnsw, &counts, 0), None); - assert_eq!( - linear_scan_one_indexed_nonempty_ids_at(HawkSearchMode::LinearScan, &ids, 0), - Some(vec![1, 5]) - ); - assert_eq!( - linear_scan_count_at(HawkSearchMode::LinearScan, &counts, 0), - Some(7) - ); - } } diff --git a/scripts/analyze-linear-scan-server-benchmark.py b/scripts/analyze-linear-scan-server-benchmark.py new file mode 100755 index 0000000000..b0a8218a0c --- /dev/null +++ b/scripts/analyze-linear-scan-server-benchmark.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +"""Summarize production linear-scan cascade logs from all three parties. + +The normal and mirror searches run concurrently. For one logical request the +cluster therefore performs the sum of both orientations' comparisons, without +summing the work performed redundantly by the three MPC parties. + +The service starts a party's cascade when that party receives its SQS message. +When timestamps are present, report the fan-out skew separately and measure the +synchronized MPC interval from the last party's ingress until the last party +finishes. This prevents a slow SNS/SQS emulator fan-out from being mistaken for +linear-scan compute/network time while retaining an end-to-end rate that includes +that skew. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import json +import re +import statistics +from collections import defaultdict +from pathlib import Path + + +SUMMARY_MARKER = "LINEAR_SCAN_CASCADE_SUMMARY" +FIELD_RE = re.compile(r"([a-z_]+)=(\"[^\"]*\"|\S+)") +PARTY_RE = re.compile(r"(?:server|party)[-_]?([0-2])") +TIMESTAMP_RE = re.compile(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:\d{2})") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("logs", nargs="+", type=Path) + parser.add_argument( + "--warmup-requests", + type=int, + default=1, + help="discard this many requests from each party and orientation", + ) + parser.add_argument( + "--minimum-cps", + type=float, + help="exit unsuccessfully when median cluster comparisons/s is lower", + ) + parser.add_argument("--json", type=Path, help="also write the summary as JSON") + return parser.parse_args() + + +def party_from_path(path: Path) -> int: + match = PARTY_RE.search(path.name) + if match is None: + raise ValueError(f"cannot infer party id from log filename: {path}") + return int(match.group(1)) + + +def parse_fields(line: str) -> dict[str, str]: + try: + record = json.loads(line) + except json.JSONDecodeError: + record = None + if isinstance(record, dict) and record.get("message") == SUMMARY_MARKER: + return {name: str(value) for name, value in record.items()} + + fields = {} + for name, value in FIELD_RE.findall(line.split(SUMMARY_MARKER, 1)[1]): + fields[name] = value.strip('"').rstrip(",") + if match := TIMESTAMP_RE.search(line): + fields["timestamp"] = match.group(0) + return fields + + +def parse_timestamp(value: str) -> float: + # Python's ISO parser accepts microseconds; tracing emits nanoseconds. + value = re.sub(r"(\.\d{6})\d+", r"\1", value).replace("Z", "+00:00") + return dt.datetime.fromisoformat(value).timestamp() + + +def load_cascades(paths: list[Path]) -> dict[int, dict[str, list[dict[str, str]]]]: + cascades: dict[int, dict[str, list[dict[str, str]]]] = defaultdict( + lambda: defaultdict(list) + ) + for path in paths: + party = party_from_path(path) + with path.open(encoding="utf-8", errors="replace") as log: + for line in log: + if SUMMARY_MARKER not in line: + continue + fields = parse_fields(line) + orientation = fields.get("orientation") + if orientation not in {"normal", "mirror"}: + raise ValueError(f"missing/invalid orientation in {path}: {line.rstrip()}") + for required in ("total_comparisons", "elapsed_seconds"): + if required not in fields: + raise ValueError(f"missing {required} in {path}: {line.rstrip()}") + cascades[party][orientation].append(fields) + return cascades + + +def summarize( + cascades: dict[int, dict[str, list[dict[str, str]]]], warmups: int +) -> dict[str, object]: + if set(cascades) != {0, 1, 2}: + raise ValueError(f"expected logs for parties 0, 1, and 2; got {sorted(cascades)}") + + orientations = set(cascades[0]) + if not orientations or any(set(cascades[party]) != orientations for party in range(3)): + raise ValueError("parties have different orientation sets") + + retained: dict[int, dict[str, list[dict[str, str]]]] = defaultdict(dict) + sample_counts = set() + for party in range(3): + for orientation in orientations: + rows = cascades[party][orientation][warmups:] + retained[party][orientation] = rows + sample_counts.add(len(rows)) + if len(sample_counts) != 1: + raise ValueError(f"party/orientation sample counts differ: {sorted(sample_counts)}") + samples = sample_counts.pop() + if samples == 0: + raise ValueError("no samples remain after warm-up removal") + + synchronized_rates = [] + synchronized_durations = [] + end_to_end_rates = [] + end_to_end_durations = [] + ingress_skews = [] + comparisons_per_request = [] + for sample in range(samples): + total_comparisons = 0 + slowest_duration = 0.0 + party_starts: list[float] = [] + completion_times: list[float] = [] + timestamps_available = True + for orientation in sorted(orientations): + party_comparisons = { + int(retained[party][orientation][sample]["total_comparisons"]) + for party in range(3) + } + if len(party_comparisons) != 1: + raise ValueError( + f"comparison disagreement in sample {sample}, {orientation}: " + f"{sorted(party_comparisons)}" + ) + total_comparisons += party_comparisons.pop() + slowest_duration = max( + slowest_duration, + *( + float(retained[party][orientation][sample]["elapsed_seconds"]) + for party in range(3) + ), + ) + for party in range(3): + starts = [] + for orientation in orientations: + row = retained[party][orientation][sample] + timestamp = row.get("timestamp") + if timestamp is None: + timestamps_available = False + break + end = parse_timestamp(timestamp) + completion_times.append(end) + starts.append(end - float(row["elapsed_seconds"])) + if not timestamps_available: + break + # Both orientations are spawned for the same request. The first + # cascade to enter is the best observable party-ingress timestamp. + party_starts.append(min(starts)) + + if timestamps_available: + earliest_ingress = min(party_starts) + latest_ingress = max(party_starts) + completion = max(completion_times) + ingress_skew = latest_ingress - earliest_ingress + end_to_end_duration = completion - earliest_ingress + synchronized_duration = completion - latest_ingress + else: + ingress_skew = 0.0 + end_to_end_duration = slowest_duration + synchronized_duration = slowest_duration + + comparisons_per_request.append(total_comparisons) + ingress_skews.append(ingress_skew) + end_to_end_durations.append(end_to_end_duration) + synchronized_durations.append(synchronized_duration) + end_to_end_rates.append(total_comparisons / end_to_end_duration) + synchronized_rates.append(total_comparisons / synchronized_duration) + + result: dict[str, object] = { + "samples": samples, + "warmup_requests": warmups, + "orientations": sorted(orientations), + "comparisons_per_request": comparisons_per_request, + "ingress_skew_seconds": ingress_skews, + "median_ingress_skew_seconds": statistics.median(ingress_skews), + "max_ingress_skew_seconds": max(ingress_skews), + "end_to_end_elapsed_seconds": end_to_end_durations, + "end_to_end_comparisons_per_second": end_to_end_rates, + "median_end_to_end_comparisons_per_second": statistics.median(end_to_end_rates), + "synchronized_elapsed_seconds": synchronized_durations, + "synchronized_comparisons_per_second": synchronized_rates, + "median_comparisons_per_second": statistics.median(synchronized_rates), + "min_comparisons_per_second": min(synchronized_rates), + "max_comparisons_per_second": max(synchronized_rates), + } + return result + + +def main() -> int: + args = parse_args() + if args.warmup_requests < 0: + raise ValueError("--warmup-requests must be non-negative") + result = summarize(load_cascades(args.logs), args.warmup_requests) + print( + "REAL_SERVER_BENCH_RESULT " + f"samples={result['samples']} " + f"orientations={','.join(result['orientations'])} " + f"median_synchronized_comparisons_per_second=" + f"{result['median_comparisons_per_second']:.3f} " + f"min_synchronized_comparisons_per_second=" + f"{result['min_comparisons_per_second']:.3f} " + f"max_synchronized_comparisons_per_second=" + f"{result['max_comparisons_per_second']:.3f} " + f"median_end_to_end_comparisons_per_second=" + f"{result['median_end_to_end_comparisons_per_second']:.3f} " + f"median_ingress_skew_seconds={result['median_ingress_skew_seconds']:.6f} " + f"max_ingress_skew_seconds={result['max_ingress_skew_seconds']:.6f}" + ) + if args.json is not None: + args.json.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + if ( + args.minimum_cps is not None + and result["median_comparisons_per_second"] < args.minimum_cps + ): + print( + "REAL_SERVER_BENCH_BELOW_TARGET " + f"actual={result['median_comparisons_per_second']:.3f} " + f"minimum={args.minimum_cps:.3f}" + ) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/init-moto-linear-scan-benchmark.py b/scripts/init-moto-linear-scan-benchmark.py new file mode 100755 index 0000000000..426e5e331f --- /dev/null +++ b/scripts/init-moto-linear-scan-benchmark.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Create the Moto resources needed by the real iris-mpc server benchmark.""" + +from __future__ import annotations + +import argparse +import json + +import boto3 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--endpoint", required=True) + parser.add_argument("--region", default="us-east-1") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + common = { + "endpoint_url": args.endpoint, + "region_name": args.region, + "aws_access_key_id": "test", + "aws_secret_access_key": "test", + } + s3 = boto3.client("s3", **common) + sns = boto3.client("sns", **common) + sqs = boto3.client("sqs", **common) + secrets = boto3.client("secretsmanager", **common) + + for bucket in ( + "wf-dev-public-keys", + "wf-smpcv2-dev-sns-requests", + "wf-smpcv2-dev-sync-protocol", + "wf-smpcv2-dev-hnsw-performance-reports", + "wf-smpcv2-dev-hnsw-checkpoint", + ): + s3.create_bucket(Bucket=bucket) + s3.put_bucket_policy( + Bucket="wf-dev-public-keys", + Policy=json.dumps( + { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": "*", + "Action": "s3:GetObject", + "Resource": "arn:aws:s3:::wf-dev-public-keys/*", + } + ], + } + ), + ) + s3.put_object( + Bucket="wf-smpcv2-dev-sync-protocol", + Key="dev_deleted_serial_ids.json", + Body=json.dumps({"deleted_serial_ids": []}).encode(), + ) + + input_topic = sns.create_topic( + Name="iris-mpc-input.fifo", + Attributes={"FifoTopic": "true", "ContentBasedDeduplication": "true"}, + )["TopicArn"] + result_topic = sns.create_topic( + Name="iris-mpc-results.fifo", + Attributes={"FifoTopic": "true", "ContentBasedDeduplication": "true"}, + )["TopicArn"] + + queues = [] + for name in ( + "smpcv2-0-dev.fifo", + "smpcv2-1-dev.fifo", + "smpcv2-2-dev.fifo", + "iris-mpc-results-us-east-1.fifo", + ): + url = sqs.create_queue( + QueueName=name, + Attributes={ + "FifoQueue": "true", + "ContentBasedDeduplication": "true", + "VisibilityTimeout": "600", + }, + )["QueueUrl"] + arn = sqs.get_queue_attributes( + QueueUrl=url, AttributeNames=["QueueArn"] + )["Attributes"]["QueueArn"] + queues.append((url, arn)) + + for _, arn in queues[:3]: + sns.subscribe(TopicArn=input_topic, Protocol="sqs", Endpoint=arn) + sns.subscribe(TopicArn=result_topic, Protocol="sqs", Endpoint=queues[3][1]) + + for party in range(3): + secrets.create_secret( + Name=f"dev/iris-mpc/ecdh-private-key-{party}", + SecretString='{"private-key":""}', + ) + + print( + json.dumps( + { + "endpoint": args.endpoint, + "input_topic": input_topic, + "result_topic": result_topic, + "request_queues": [url for url, _ in queues[:3]], + "result_queue": queues[3][0], + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/moto-server-with-sns-sequence.py b/scripts/moto-server-with-sns-sequence.py index 4a4878c7fb..82347453d5 100644 --- a/scripts/moto-server-with-sns-sequence.py +++ b/scripts/moto-server-with-sns-sequence.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """Run Moto with the FIFO SNS SequenceNumber field emitted by AWS. -Moto 5.2.2 omits ``SequenceNumber`` from the SNS notification envelope sent +Moto 5.1.22 omits ``SequenceNumber`` from the SNS notification envelope sent to subscribed SQS queues. Production AWS includes it for FIFO topics and the iris-mpc server deliberately requires it. Keep the compatibility adjustment at the emulator boundary instead of weakening production message parsing. diff --git a/scripts/run-distributed-linear-scan-benchmark.sh b/scripts/run-distributed-linear-scan-benchmark.sh new file mode 100755 index 0000000000..927b28ef11 --- /dev/null +++ b/scripts/run-distributed-linear-scan-benchmark.sh @@ -0,0 +1,281 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Run the actual iris-mpc-linear-scan service across three remote hosts. Moto is +# only the AWS control plane; every request still traverses the production +# client, S3/SNS/SQS ingestion, server scheduler, TLS MPC network, persistence, +# and result publication paths. +# +# Usage: +# LINEAR_SCAN_BENCH_SSH_KEY=/path/to/key.pem \ +# scripts/run-distributed-linear-scan-benchmark.sh \ +# ubuntu@host0 ubuntu@host1 ubuntu@host2 [output-directory] +# +# Important overrides: +# LINEAR_SCAN_BENCH_DATABASE_SIZE=1048576 # >= 256 production 4K chunks +# LINEAR_SCAN_BENCH_REQUEST_COUNT=6 # first request is warm-up +# LINEAR_SCAN_BENCH_REQUEST_PARALLELISM=48 # tuned for r8g.24xlarge +# LINEAR_SCAN_BENCH_CONNECTION_PARALLELISM=16 +# LINEAR_SCAN_BENCH_PIPELINED_REQUESTS=1 # queue independent requests together +# LINEAR_SCAN_BENCH_REUSE_DB=1 # reuse the expensive seeded DB +# LINEAR_SCAN_BENCH_SKIP_BUILD=1 # reuse binaries already copied +# LINEAR_SCAN_BENCH_KEEP_RUNNING=1 # leave servers and Moto running +# LINEAR_SCAN_BENCH_NODE_ADDRESSES=a,b,c # override detected private IPs + +[[ $# -ge 3 && $# -le 4 ]] || { + sed -n '5,20p' "$0" >&2 + exit 2 +} + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(cd "${SCRIPT_DIR}/.." && pwd) +HOSTS=("$1" "$2" "$3") +OUTPUT_DIR=${4:-${PROJECT_ROOT}/target/linear-scan-real-server-benchmark} +DATABASE_SIZE=${LINEAR_SCAN_BENCH_DATABASE_SIZE:-1048576} +REQUEST_COUNT=${LINEAR_SCAN_BENCH_REQUEST_COUNT:-6} +WARMUP_REQUESTS=${LINEAR_SCAN_BENCH_WARMUP_REQUESTS:-1} +REQUEST_PARALLELISM=${LINEAR_SCAN_BENCH_REQUEST_PARALLELISM:-48} +CONNECTION_PARALLELISM=${LINEAR_SCAN_BENCH_CONNECTION_PARALLELISM:-16} +TOKIO_CORES=${LINEAR_SCAN_BENCH_TOKIO_CORES:-11} +CLIENT_RNG_SEED=${LINEAR_SCAN_BENCH_CLIENT_RNG_SEED:-8675309} +PIPELINED_REQUESTS=${LINEAR_SCAN_BENCH_PIPELINED_REQUESTS:-0} +REMOTE_RUN_DIR=${LINEAR_SCAN_BENCH_REMOTE_RUN_DIR:-/var/tmp/iris-mpc-real-server-bench} +COMMIT=$(git -C "$PROJECT_ROOT" rev-parse HEAD) +REMOTE_SOURCE=${LINEAR_SCAN_BENCH_REMOTE_SOURCE:-/var/tmp/iris-mpc-source-${COMMIT}} + +[[ ${DATABASE_SIZE} =~ ^[1-9][0-9]*$ ]] || { + echo "LINEAR_SCAN_BENCH_DATABASE_SIZE must be positive" >&2 + exit 2 +} +[[ ${REQUEST_COUNT} =~ ^[1-9][0-9]*$ ]] || { + echo "LINEAR_SCAN_BENCH_REQUEST_COUNT must be positive" >&2 + exit 2 +} +[[ ${CLIENT_RNG_SEED} =~ ^[0-9]+$ ]] || { + echo "LINEAR_SCAN_BENCH_CLIENT_RNG_SEED must be a non-negative integer" >&2 + exit 2 +} +[[ ${CONNECTION_PARALLELISM} =~ ^[1-9][0-9]*$ ]] || { + echo "LINEAR_SCAN_BENCH_CONNECTION_PARALLELISM must be positive" >&2 + exit 2 +} +[[ ${TOKIO_CORES} =~ ^[1-9][0-9]*$ ]] || { + echo "LINEAR_SCAN_BENCH_TOKIO_CORES must be positive" >&2 + exit 2 +} +[[ ${PIPELINED_REQUESTS} =~ ^[01]$ ]] || { + echo "LINEAR_SCAN_BENCH_PIPELINED_REQUESTS must be 0 or 1" >&2 + exit 2 +} +[[ ${WARMUP_REQUESTS} =~ ^[0-9]+$ && ${WARMUP_REQUESTS} -lt ${REQUEST_COUNT} ]] || { + echo "warm-up count must be non-negative and smaller than request count" >&2 + exit 2 +} +if [[ ${LINEAR_SCAN_BENCH_SKIP_SYNC:-0} != 1 ]] && \ + [[ -n $(git -C "$PROJECT_ROOT" status --short) ]]; then + echo "distributed benchmark requires a clean committed tree" >&2 + exit 1 +fi + +SSH_OPTIONS=(-o BatchMode=yes -o ConnectTimeout=10 -o ServerAliveInterval=30) +SCP_OPTIONS=(-o BatchMode=yes -o ConnectTimeout=10) +if [[ -n ${LINEAR_SCAN_BENCH_SSH_KEY:-} ]]; then + SSH_OPTIONS+=(-i "$LINEAR_SCAN_BENCH_SSH_KEY" -o IdentitiesOnly=yes) + SCP_OPTIONS+=(-i "$LINEAR_SCAN_BENCH_SSH_KEY" -o IdentitiesOnly=yes) +fi + +remote() { + local host=$1 quoted + shift + printf -v quoted '%q ' "$@" + ssh "${SSH_OPTIONS[@]}" "$host" "${quoted% }" +} + +remote_env() { + local host=$1 + shift + remote "$host" env \ + "LINEAR_SCAN_BENCH_RUN_DIR=${REMOTE_RUN_DIR}" \ + "LINEAR_SCAN_BENCH_DATABASE_SIZE=${DATABASE_SIZE}" \ + "LINEAR_SCAN_BENCH_REQUEST_COUNT=${REQUEST_COUNT}" \ + "LINEAR_SCAN_BENCH_REQUEST_PARALLELISM=${REQUEST_PARALLELISM}" \ + "LINEAR_SCAN_BENCH_CONNECTION_PARALLELISM=${CONNECTION_PARALLELISM}" \ + "LINEAR_SCAN_BENCH_TOKIO_CORES=${TOKIO_CORES}" \ + "LINEAR_SCAN_BENCH_CLIENT_RNG_SEED=${CLIENT_RNG_SEED}" \ + "LINEAR_SCAN_BENCH_PIPELINED_REQUESTS=${PIPELINED_REQUESTS}" \ + "LINEAR_SCAN_BENCH_NODE_HOSTNAMES=${NODE_HOSTNAMES_JSON}" \ + "LINEAR_SCAN_BENCH_AWS_ENDPOINT=${AWS_ENDPOINT}" \ + "LINEAR_SCAN_BENCH_SERVER_BINARY=${REMOTE_SOURCE}/target/release/iris-mpc-linear-scan" \ + "LINEAR_SCAN_BENCH_KEY_MANAGER_BINARY=${REMOTE_SOURCE}/target/release/key-manager" \ + "LINEAR_SCAN_BENCH_CLIENT_BINARY=${REMOTE_SOURCE}/target/release/service-client" \ + "LINEAR_SCAN_BENCH_IMAGE_NAME=real-server-benchmark-${COMMIT}" \ + "$@" +} + +mkdir -p "$OUTPUT_DIR" +LOCAL_TMP=$(mktemp -d "${TMPDIR:-/tmp}/iris-mpc-real-server-bench.XXXXXX") +SERVERS_STARTED=false +MOTO_STARTED=false +cleanup() { + local exit_code=$? + trap - EXIT INT TERM + if [[ ${LINEAR_SCAN_BENCH_KEEP_RUNNING:-0} != 1 ]]; then + if [[ ${SERVERS_STARTED} == true ]]; then + for host in "${HOSTS[@]}"; do + remote_env "$host" "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" \ + stop-server >/dev/null 2>&1 & + done + wait || true + fi + if [[ ${MOTO_STARTED} == true ]]; then + remote_env "${HOSTS[0]}" \ + "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" \ + stop-moto >/dev/null 2>&1 || true + fi + fi + rm -rf "$LOCAL_TMP" + exit "$exit_code" +} +trap cleanup EXIT INT TERM + +for host in "${HOSTS[@]}"; do + remote "$host" true +done + +if [[ -n ${LINEAR_SCAN_BENCH_NODE_ADDRESSES:-} ]]; then + IFS=',' read -r -a NODE_ADDRESSES <<<"$LINEAR_SCAN_BENCH_NODE_ADDRESSES" +else + NODE_ADDRESSES=() + for host in "${HOSTS[@]}"; do + NODE_ADDRESSES+=("$(remote "$host" sh -c "hostname -I | cut -d' ' -f1")") + done +fi +[[ ${#NODE_ADDRESSES[@]} -eq 3 ]] || { + echo "expected exactly three node addresses" >&2 + exit 1 +} +NODE_HOSTNAMES_JSON=$(python3 -c 'import json,sys; print(json.dumps(sys.argv[1:]))' \ + "${NODE_ADDRESSES[@]}") +AWS_ENDPOINT="http://${NODE_ADDRESSES[0]}:4566" +echo "REAL_SERVER_BENCH_TOPOLOGY commit=${COMMIT} nodes=${NODE_ADDRESSES[*]} moto=${AWS_ENDPOINT} request_parallelism=${REQUEST_PARALLELISM} connection_parallelism=${CONNECTION_PARALLELISM} tokio_cores=${TOKIO_CORES} client_rng_seed=${CLIENT_RNG_SEED} pipelined_requests=${PIPELINED_REQUESTS}" + +if [[ ${LINEAR_SCAN_BENCH_SKIP_SYNC:-0} != 1 ]]; then + git -C "$PROJECT_ROOT" archive --format=tar "$COMMIT" -o "${LOCAL_TMP}/source.tar" + for host in "${HOSTS[@]}"; do + remote "$host" mkdir -p "$REMOTE_SOURCE" + scp "${SCP_OPTIONS[@]}" "${LOCAL_TMP}/source.tar" "${host}:${REMOTE_SOURCE}/source.tar" + remote "$host" tar -xf "${REMOTE_SOURCE}/source.tar" -C "$REMOTE_SOURCE" + remote "$host" chmod +x \ + "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" \ + "${REMOTE_SOURCE}/scripts/moto-server-with-sns-sequence.py" \ + "${REMOTE_SOURCE}/scripts/init-moto-linear-scan-benchmark.py" + done +fi + +if [[ ${LINEAR_SCAN_BENCH_SKIP_BUILD:-0} != 1 ]]; then + remote "${HOSTS[0]}" bash -lc \ + "cd '$REMOTE_SOURCE' && RUSTFLAGS='--cfg aes_armv8 -C force-frame-pointers=yes -Ctarget-cpu=neoverse-v2 -Ctarget-feature=+lse' cargo build --release -p iris-mpc-bins --features aes_rng_prf --bin iris-mpc-linear-scan --bin key-manager --bin service-client" + for binary in iris-mpc-linear-scan key-manager service-client; do + scp "${SCP_OPTIONS[@]}" \ + "${HOSTS[0]}:${REMOTE_SOURCE}/target/release/${binary}" \ + "${LOCAL_TMP}/${binary}" + done + for host in "${HOSTS[@]}"; do + remote "$host" mkdir -p "${REMOTE_SOURCE}/target/release" + scp "${SCP_OPTIONS[@]}" "${LOCAL_TMP}/iris-mpc-linear-scan" \ + "${LOCAL_TMP}/key-manager" "${LOCAL_TMP}/service-client" \ + "${host}:${REMOTE_SOURCE}/target/release/" + remote "$host" chmod +x \ + "${REMOTE_SOURCE}/target/release/iris-mpc-linear-scan" \ + "${REMOTE_SOURCE}/target/release/key-manager" \ + "${REMOTE_SOURCE}/target/release/service-client" + done +fi + +TLS_DIR="${LOCAL_TMP}/tls" +mkdir -p "$TLS_DIR" +SAN=$(printf 'IP:%s,' "${NODE_ADDRESSES[@]}") +SAN=${SAN%,} +openssl req -x509 -newkey rsa:2048 -nodes -days 2 \ + -subj '/CN=iris-mpc-real-server-benchmark-ca' \ + -addext 'basicConstraints=critical,CA:TRUE' \ + -addext 'keyUsage=critical,keyCertSign,cRLSign' \ + -keyout "${TLS_DIR}/ca.key" -out "${TLS_DIR}/ca.crt" >/dev/null 2>&1 +openssl req -newkey rsa:2048 -nodes \ + -subj '/CN=iris-mpc-real-server-benchmark' \ + -addext "subjectAltName=${SAN}" \ + -keyout "${TLS_DIR}/tls.key" -out "${TLS_DIR}/tls.csr" >/dev/null 2>&1 +printf 'subjectAltName=%s\nextendedKeyUsage=serverAuth,clientAuth\n' "$SAN" \ + >"${TLS_DIR}/tls.ext" +openssl x509 -req -days 2 -in "${TLS_DIR}/tls.csr" \ + -CA "${TLS_DIR}/ca.crt" -CAkey "${TLS_DIR}/ca.key" -CAcreateserial \ + -extfile "${TLS_DIR}/tls.ext" -out "${TLS_DIR}/tls.crt" >/dev/null 2>&1 +for host in "${HOSTS[@]}"; do + remote "$host" mkdir -p "${REMOTE_RUN_DIR}/tls" + scp "${SCP_OPTIONS[@]}" "${TLS_DIR}/ca.crt" "${TLS_DIR}/tls.crt" \ + "${TLS_DIR}/tls.key" "${host}:${REMOTE_RUN_DIR}/tls/" +done + +remote_env "${HOSTS[0]}" "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" \ + start-moto +MOTO_STARTED=true +remote_env "${HOSTS[0]}" "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" \ + init-moto +remote_env "${HOSTS[0]}" "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" \ + rotate-keys + +RESET_DB=1 +[[ ${LINEAR_SCAN_BENCH_REUSE_DB:-0} == 1 ]] && RESET_DB=0 +for party in 0 1 2; do + remote_env "${HOSTS[$party]}" "LINEAR_SCAN_BENCH_RESET_DB=${RESET_DB}" \ + "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" prepare-db "$party" & +done +wait + +for party in 0 1 2; do + remote_env "${HOSTS[$party]}" \ + "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" start-server "$party" & +done +wait +SERVERS_STARTED=true + +for _ in $(seq 1 7200); do + all_ready=true + for party in 0 1 2; do + remote_env "${HOSTS[$party]}" \ + "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" status "$party" \ + >/dev/null 2>&1 || all_ready=false + done + [[ ${all_ready} == true ]] && break + sleep 1 +done +[[ ${all_ready} == true ]] || { + echo "servers did not become ready" >&2 + exit 1 +} +echo "REAL_SERVER_BENCH_SERVERS_READY database_size=${DATABASE_SIZE}" + +remote_env "${HOSTS[0]}" "${REMOTE_SOURCE}/scripts/run-distributed-linear-scan-node.sh" \ + run-client + +for party in 0 1 2; do + scp "${SCP_OPTIONS[@]}" \ + "${HOSTS[$party]}:${REMOTE_RUN_DIR}/server-${party}.log" \ + "${OUTPUT_DIR}/server-${party}.log" +done +scp "${SCP_OPTIONS[@]}" "${HOSTS[0]}:${REMOTE_RUN_DIR}/client.log" \ + "${OUTPUT_DIR}/client.log" + +ANALYZER_ARGS=( + "${SCRIPT_DIR}/analyze-linear-scan-server-benchmark.py" + --warmup-requests "$WARMUP_REQUESTS" + --json "${OUTPUT_DIR}/summary.json" +) +if [[ -n ${LINEAR_SCAN_BENCH_MINIMUM_CPS:-} ]]; then + ANALYZER_ARGS+=(--minimum-cps "$LINEAR_SCAN_BENCH_MINIMUM_CPS") +fi +python3 "${ANALYZER_ARGS[@]}" \ + "${OUTPUT_DIR}/server-0.log" \ + "${OUTPUT_DIR}/server-1.log" \ + "${OUTPUT_DIR}/server-2.log" +echo "REAL_SERVER_BENCH_COMPLETE output=${OUTPUT_DIR}" diff --git a/scripts/run-distributed-linear-scan-node.sh b/scripts/run-distributed-linear-scan-node.sh new file mode 100755 index 0000000000..1d76a49b8e --- /dev/null +++ b/scripts/run-distributed-linear-scan-node.sh @@ -0,0 +1,378 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Node-side lifecycle helper for the distributed real-server benchmark. The +# controller copies the same committed tree and TLS bundle to all three hosts, +# then invokes this script over SSH/SSM with the environment documented below. + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(cd "${SCRIPT_DIR}/.." && pwd) +RUN_DIR=${LINEAR_SCAN_BENCH_RUN_DIR:-/var/tmp/iris-mpc-real-server-bench} +POSTGRES_PORT=${LINEAR_SCAN_BENCH_POSTGRES_PORT:-55432} +DATABASE_SIZE=${LINEAR_SCAN_BENCH_DATABASE_SIZE:-1048576} +REQUEST_COUNT=${LINEAR_SCAN_BENCH_REQUEST_COUNT:-6} +REQUEST_PARALLELISM=${LINEAR_SCAN_BENCH_REQUEST_PARALLELISM:-48} +CONNECTION_PARALLELISM=${LINEAR_SCAN_BENCH_CONNECTION_PARALLELISM:-16} +TOKIO_CORES=${LINEAR_SCAN_BENCH_TOKIO_CORES:-11} +CLIENT_RNG_SEED=${LINEAR_SCAN_BENCH_CLIENT_RNG_SEED:-8675309} +PIPELINED_REQUESTS=${LINEAR_SCAN_BENCH_PIPELINED_REQUESTS:-0} +AUX_CPU_LIST="0-$((TOKIO_CORES - 1))" + +usage() { + echo "usage: $0 [party-id]" >&2 + echo " $0 " >&2 + exit 2 +} + +start_moto() { + : "${LINEAR_SCAN_BENCH_AWS_ENDPOINT:?set the externally reachable Moto endpoint}" + local pid_file="${RUN_DIR}/moto.pid" + local python=${LINEAR_SCAN_BENCH_MOTO_PYTHON:-${RUN_DIR}/moto-venv/bin/python} + mkdir -p "$RUN_DIR" + if [[ -f ${pid_file} ]] && pid_running "$(<"$pid_file")"; then + echo "Moto is already running: pid=$(<"$pid_file")" >&2 + exit 1 + fi + if [[ ! -x ${python} ]]; then + python3 -m venv "${RUN_DIR}/moto-venv" + python="${RUN_DIR}/moto-venv/bin/python" + fi + if ! "$python" -c 'import boto3, moto' >/dev/null 2>&1; then + "$python" -m pip install -q 'moto[server]==5.1.22' boto3 + fi + nohup taskset -c "$AUX_CPU_LIST" env \ + MOTO_ACCOUNT_ID=000000000000 S3_IGNORE_SUBDOMAIN_BUCKETNAME=true \ + "$python" "${PROJECT_ROOT}/scripts/moto-server-with-sns-sequence.py" \ + -H 0.0.0.0 -p 4566 >"${RUN_DIR}/moto.log" 2>&1 & + echo "$!" >"$pid_file" + for _ in $(seq 1 120); do + curl -fsS "http://127.0.0.1:4566/moto-api/" >/dev/null 2>&1 && { + echo "NODE_BENCH_MOTO_STARTED pid=$! endpoint=${LINEAR_SCAN_BENCH_AWS_ENDPOINT}" + return + } + pid_running "$!" || { + tail -100 "${RUN_DIR}/moto.log" >&2 + exit 1 + } + sleep 1 + done + echo "Moto did not become ready" >&2 + exit 1 +} + +stop_moto() { + local pid_file="${RUN_DIR}/moto.pid" + if [[ -f ${pid_file} ]]; then + local pid + pid=$(<"$pid_file") + pid_running "$pid" && kill "$pid" + rm -f "$pid_file" + fi + echo "NODE_BENCH_MOTO_STOPPED" +} + +init_moto() { + : "${LINEAR_SCAN_BENCH_AWS_ENDPOINT:?set the externally reachable Moto endpoint}" + local python=${LINEAR_SCAN_BENCH_MOTO_PYTHON:-${RUN_DIR}/moto-venv/bin/python} + [[ -x ${python} ]] || { + echo "missing Moto Python environment; run start-moto first" >&2 + exit 1 + } + "$python" "${PROJECT_ROOT}/scripts/init-moto-linear-scan-benchmark.py" \ + --endpoint "$LINEAR_SCAN_BENCH_AWS_ENDPOINT" +} + +rotate_keys() { + : "${LINEAR_SCAN_BENCH_AWS_ENDPOINT:?set the externally reachable Moto endpoint}" + local binary=${LINEAR_SCAN_BENCH_KEY_MANAGER_BINARY:-${PROJECT_ROOT}/target/release/key-manager} + [[ -x ${binary} ]] || { + echo "missing key-manager binary: ${binary}" >&2 + exit 1 + } + for party in 0 1 2; do + for _ in 1 2; do + AWS_ACCESS_KEY_ID=test AWS_SECRET_ACCESS_KEY=test AWS_REGION=us-east-1 \ + AWS_DEFAULT_REGION=us-east-1 AWS_ENDPOINT_URL="$LINEAR_SCAN_BENCH_AWS_ENDPOINT" \ + "$binary" --region us-east-1 \ + --endpoint-url "$LINEAR_SCAN_BENCH_AWS_ENDPOINT" \ + --node-id "$party" --env dev rotate \ + --public-key-bucket-name wf-dev-public-keys + done + done + echo "NODE_BENCH_KEYS_READY" +} + +run_client() { + : "${LINEAR_SCAN_BENCH_AWS_ENDPOINT:?set the externally reachable Moto endpoint}" + local binary=${LINEAR_SCAN_BENCH_CLIENT_BINARY:-${PROJECT_ROOT}/target/release/service-client} + local config="${RUN_DIR}/client.toml" + local aws_config="${RUN_DIR}/aws.toml" + local output="${RUN_DIR}/results.json" + [[ -x ${binary} ]] || { + echo "missing service-client binary: ${binary}" >&2 + exit 1 + } + mkdir -p "$RUN_DIR" + local batch_count=$REQUEST_COUNT + local batch_size=1 + if [[ ${PIPELINED_REQUESTS} == 1 ]]; then + # Publish independent requests together so the production server has a + # sustained queue. SMPC__MAX_BATCH_SIZE=1 still makes the server scan + # them serially; this only removes client-side S3/response idle gaps. + batch_count=1 + batch_size=$REQUEST_COUNT + fi + cat >"$config" <"$aws_config" <"${RUN_DIR}/client.log" 2>&1 + python3 - "$output" "$REQUEST_COUNT" <<'PY' +import json +import sys + +path, expected = sys.argv[1], int(sys.argv[2]) +with open(path, encoding="utf-8") as result_file: + records = json.load(result_file).get("records", []) +if len(records) != expected: + raise SystemExit(f"expected {expected} result records, got {len(records)}") +for index, record in enumerate(records): + responses = record.get("responses", []) + if len(responses) != 3: + raise SystemExit(f"record {index} has {len(responses)} party responses") +print(f"NODE_BENCH_CLIENT_COMPLETE records={len(records)}") +PY +} + +find_postgres_command() { + local name=$1 candidate + if command -v "$name" >/dev/null 2>&1; then + command -v "$name" + return + fi + for candidate in /usr/lib/postgresql/*/bin/"$name" /usr/pgsql-*/bin/"$name"; do + if [[ -x ${candidate} ]]; then + echo "$candidate" + return + fi + done + echo "missing PostgreSQL command: $name" >&2 + exit 1 +} + +pid_running() { + local pid=$1 state + kill -0 "$pid" >/dev/null 2>&1 || return 1 + state=$(ps -o stat= -p "$pid" 2>/dev/null || true) + [[ -n ${state} && ${state} != Z* ]] +} + +prepare_db() { + local party=$1 + local initdb pg_ctl createdb dropdb pg_data db_name + initdb=$(find_postgres_command initdb) + pg_ctl=$(find_postgres_command pg_ctl) + createdb=$(find_postgres_command createdb) + dropdb=$(find_postgres_command dropdb) + pg_data="${RUN_DIR}/postgres" + db_name="SMPC_bench_${party}" + mkdir -p "$RUN_DIR" + + if [[ ! -f ${pg_data}/PG_VERSION ]]; then + mkdir -p "$pg_data" + "$initdb" -D "$pg_data" -A trust -U postgres --no-locale >/dev/null + fi + if ! "$pg_ctl" -D "$pg_data" status >/dev/null 2>&1; then + taskset -c "$AUX_CPU_LIST" "$pg_ctl" -D "$pg_data" -l "${pg_data}/postgres.log" \ + -o "-h 127.0.0.1 -k ${pg_data} -p ${POSTGRES_PORT}" start >/dev/null + fi + + # Postgres is local only for the self-contained benchmark. Keep it off the + # production-equivalent dot cores, including a server reused from an older + # run and all of its currently live children. Future children inherit the + # postmaster's affinity. + local postgres_pid + postgres_pid=$(head -n 1 "${pg_data}/postmaster.pid") + taskset -pc "$AUX_CPU_LIST" "$postgres_pid" >/dev/null + while read -r child_pid; do + [[ -z ${child_pid} ]] || taskset -pc "$AUX_CPU_LIST" "$child_pid" >/dev/null + done < <(pgrep -P "$postgres_pid" || true) + + if [[ ${LINEAR_SCAN_BENCH_RESET_DB:-1} == 1 ]]; then + "$dropdb" -h 127.0.0.1 -p "$POSTGRES_PORT" -U postgres \ + --if-exists "$db_name" + "$createdb" -h 127.0.0.1 -p "$POSTGRES_PORT" -U postgres "$db_name" + else + "$createdb" -h 127.0.0.1 -p "$POSTGRES_PORT" -U postgres "$db_name" \ + 2>/dev/null || true + fi + echo "NODE_BENCH_DB_READY party=${party} database=${db_name} port=${POSTGRES_PORT}" +} + +start_server() { + local party=$1 + : "${LINEAR_SCAN_BENCH_NODE_HOSTNAMES:?set JSON array of the three private host addresses}" + : "${LINEAR_SCAN_BENCH_AWS_ENDPOINT:?set the shared Moto endpoint}" + + local binary=${LINEAR_SCAN_BENCH_SERVER_BINARY:-${PROJECT_ROOT}/target/release/iris-mpc-linear-scan} + local tls_dir=${LINEAR_SCAN_BENCH_TLS_DIR:-${RUN_DIR}/tls} + local db_name="SMPC_bench_${party}" + local pid_file="${RUN_DIR}/server.pid" + local log_file="${RUN_DIR}/server-${party}.log" + [[ -x ${binary} ]] || { + echo "missing server binary: ${binary}" >&2 + exit 1 + } + for file in "${tls_dir}/tls.key" "${tls_dir}/tls.crt" "${tls_dir}/ca.crt"; do + [[ -r ${file} ]] || { + echo "missing TLS file: ${file}" >&2 + exit 1 + } + done + if [[ -f ${pid_file} ]] && pid_running "$(<"$pid_file")"; then + echo "server is already running: pid=$(<"$pid_file")" >&2 + exit 1 + fi + + local endpoint=${LINEAR_SCAN_BENCH_AWS_ENDPOINT%/} + local max_db_size=$((DATABASE_SIZE + REQUEST_COUNT + 1024)) + local root_certs="[\"${tls_dir}/ca.crt\",\"${tls_dir}/ca.crt\",\"${tls_dir}/ca.crt\"]" + local image_name=${LINEAR_SCAN_BENCH_IMAGE_NAME:-real-server-benchmark} + mkdir -p "$RUN_DIR" + : >"$log_file" + + nohup env \ + AWS_ACCESS_KEY_ID=test \ + AWS_SECRET_ACCESS_KEY=test \ + AWS_REGION=us-east-1 \ + AWS_DEFAULT_REGION=us-east-1 \ + AWS_ENDPOINT_URL="$endpoint" \ + AWS_EC2_METADATA_DISABLED=true \ + RUST_LOG=${RUST_LOG:-info} \ + RUST_BACKTRACE=1 \ + RUST_MIN_STACK=104857600 \ + SMPC__ENVIRONMENT=dev \ + SMPC__PARTY_ID="$party" \ + SMPC__DATABASE__URL="postgres://postgres@127.0.0.1:${POSTGRES_PORT}/${db_name}" \ + SMPC__DATABASE__MIGRATE=true \ + SMPC__DATABASE__CREATE=true \ + SMPC__DATABASE__LOAD_PARALLELISM=8 \ + SMPC__AWS__REGION=us-east-1 \ + SMPC__AWS__ENDPOINT="$endpoint" \ + SMPC__PUBLIC_KEY_BASE_URL="${endpoint}/wf-dev-public-keys" \ + SMPC__REQUESTS_QUEUE_URL="${endpoint}/000000000000/smpcv2-${party}-dev.fifo" \ + SMPC__RESULTS_TOPIC_ARN="arn:aws:sns:us-east-1:000000000000:iris-mpc-results.fifo" \ + SMPC__SHARES_BUCKET_NAME=wf-smpcv2-dev-sns-requests \ + SMPC__GRAPH_CHECKPOINT_BUCKET_NAME=wf-smpcv2-dev-hnsw-checkpoint \ + SMPC__KMS_KEY_ARNS='["unused-0","unused-1","unused-2"]' \ + SMPC__FIXED_SHARED_SECRETS=true \ + SMPC__MAX_BATCH_SIZE=1 \ + SMPC__MAX_DB_SIZE="$max_db_size" \ + SMPC__INIT_DB_SIZE="$DATABASE_SIZE" \ + SMPC__CLEAR_DB_BEFORE_INIT=true \ + SMPC__FAKE_DB_SIZE=0 \ + SMPC__DISABLE_PERSISTENCE=false \ + SMPC__RETURN_PARTIAL_RESULTS=true \ + SMPC__ENABLE_REAUTH=true \ + SMPC__ENABLE_DELETION=true \ + SMPC__ENABLE_RESET=true \ + SMPC__ENABLE_RECOVERY=true \ + SMPC__LUC_ENABLED=true \ + SMPC__LUC_LOOKBACK_RECORDS=500 \ + SMPC__COLD_EYE_LFU_CACHE_RECORDS=4096 \ + SMPC__LUC_SERIAL_IDS_FROM_SMPC_REQUEST=true \ + SMPC__FULL_SCAN_SIDE=Left \ + SMPC__FULL_SCAN_SIDE_SWITCHING_ENABLED=false \ + SMPC__HAWK_REQUEST_PARALLELISM="$REQUEST_PARALLELISM" \ + SMPC__HAWK_CONNECTION_PARALLELISM="$CONNECTION_PARALLELISM" \ + SMPC__SEPARATE_TOKIO_CORES_PER_NODE="$TOKIO_CORES" \ + SMPC__SERVICE_PORTS='["4000","4001","4002"]' \ + SMPC__NODE_HOSTNAMES="$LINEAR_SCAN_BENCH_NODE_HOSTNAMES" \ + SMPC__SERVER_COORDINATION__NODE_HOSTNAMES="$LINEAR_SCAN_BENCH_NODE_HOSTNAMES" \ + SMPC__SERVER_COORDINATION__PARTY_ID="$party" \ + SMPC__SERVER_COORDINATION__HEALTHCHECK_PORTS='["13000","13000","13000"]' \ + SMPC__SERVER_COORDINATION__IMAGE_NAME="$image_name" \ + SMPC__SERVER_COORDINATION__HTTP_QUERY_RETRY_DELAY_MS=250 \ + SMPC__SERVER_COORDINATION__HEARTBEAT_INTERVAL_SECS=1 \ + SMPC__SERVER_COORDINATION__HEARTBEAT_INITIAL_RETRIES=3600 \ + SMPC__TLS__PRIVATE_KEY="${tls_dir}/tls.key" \ + SMPC__TLS__LEAF_CERT="${tls_dir}/tls.crt" \ + SMPC__TLS__ROOT_CERTS="$root_certs" \ + SMPC__SERVICE__SERVICE_NAME="iris-mpc-linear-scan-benchmark-${party}" \ + SMPC__SERVICE__METRICS__HOST=127.0.0.1 \ + SMPC__SERVICE__METRICS__PORT=8125 \ + SMPC__SERVICE__METRICS__QUEUE_SIZE=5000 \ + SMPC__SERVICE__METRICS__BUFFER_SIZE=1024 \ + SMPC__SERVICE__METRICS__PREFIX="linear-scan-benchmark-${party}" \ + "$binary" >>"$log_file" 2>&1 & + echo "$!" >"$pid_file" + echo "NODE_BENCH_SERVER_STARTED party=${party} pid=$! log=${log_file}" +} + +stop_server() { + local pid_file="${RUN_DIR}/server.pid" + if [[ ! -f ${pid_file} ]]; then + echo "NODE_BENCH_SERVER_STOPPED already=true" + return + fi + local pid + pid=$(<"$pid_file") + if pid_running "$pid"; then + kill "$pid" + for _ in $(seq 1 30); do + pid_running "$pid" || break + sleep 1 + done + pid_running "$pid" && kill -KILL "$pid" + fi + rm -f "$pid_file" + echo "NODE_BENCH_SERVER_STOPPED pid=${pid}" +} + +status() { + local party=$1 pid_file="${RUN_DIR}/server.pid" + if [[ -f ${pid_file} ]] && pid_running "$(<"$pid_file")"; then + curl -fsS "http://127.0.0.1:13000/ready" >/dev/null + echo "NODE_BENCH_SERVER_READY party=${party} pid=$(<"$pid_file")" + else + echo "NODE_BENCH_SERVER_NOT_RUNNING party=${party}" >&2 + return 1 + fi +} + +action=${1:-} +party=${2:-} +case "$action" in +prepare-db | start-server | status) + [[ ${party} =~ ^[0-2]$ ]] || usage + "${action//-/_}" "$party" + ;; +stop-server) + stop_server + ;; +start-moto | stop-moto | init-moto | rotate-keys | run-client) + "${action//-/_}" + ;; +*) usage ;; +esac diff --git a/scripts/run-native-ground-truth.sh b/scripts/run-native-ground-truth.sh index a54d3c7b3b..568f19b7d7 100755 --- a/scripts/run-native-ground-truth.sh +++ b/scripts/run-native-ground-truth.sh @@ -18,7 +18,7 @@ set -euo pipefail # # Required commands: cargo, curl, Python 3.10+, initdb, pg_ctl, createdb. # Set MOTO_PYTHON_BIN to a Python interpreter containing moto and boto3. -# Otherwise a temporary venv is populated with moto[server]==5.2.2 and boto3. +# Otherwise a temporary venv is populated with moto[server]==5.1.22 and boto3. # Set MOTO_BASE_PYTHON to select the interpreter used to create that venv. SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) @@ -209,11 +209,11 @@ elif [[ -z ${MOTO_BASE_PYTHON:-} ]] && python3 -c 'import boto3, moto' >/dev/nul else MOTO_BASE_PYTHON=${MOTO_BASE_PYTHON:-python3} "$MOTO_BASE_PYTHON" -c 'import sys; assert sys.version_info >= (3, 10)' || { - echo "Moto 5.2.2 requires Python 3.10 or newer; set MOTO_BASE_PYTHON" >&2 + echo "Moto 5.1.22 requires Python 3.10 or newer; set MOTO_BASE_PYTHON" >&2 exit 1 } "$MOTO_BASE_PYTHON" -m venv "$RUN_ROOT/moto-venv" - "$RUN_ROOT/moto-venv/bin/pip" install -q 'moto[server]==5.2.2' boto3 + "$RUN_ROOT/moto-venv/bin/pip" install -q 'moto[server]==5.1.22' boto3 MOTO_PYTHON="$RUN_ROOT/moto-venv/bin/python" fi diff --git a/scripts/tools/pin-network-irqs.sh b/scripts/tools/pin-network-irqs.sh new file mode 100755 index 0000000000..09f3ea701b --- /dev/null +++ b/scripts/tools/pin-network-irqs.sh @@ -0,0 +1,125 @@ +#!/usr/bin/env bash +set -euo pipefail + +# This script is intended for manual testing and recovery. For Kubernetes +# deployment, configure the node's existing irqbalance service during node +# provisioning instead of running this script from the application pod (the +# pod is unprivileged and hostNetwork does not grant access to host IRQ state). +# +# The benchmarked r8g.24xlarge layout reserves CPUs 0-10 for Tokio/TLS/network +# work and uses CPUs 11-95 for dot-product workers. Bake the following systemd +# drop-in into the node group's Launch Template user-data or AMI, then restart +# irqbalance: +# +# /etc/systemd/system/irqbalance.service.d/iris-mpc.conf +# [Service] +# Environment="IRQBALANCE_BANNED_CPULIST=11-95" +# +# This performs equivalently to pinning the ENA queue IRQs to CPUs 0-10 while +# retaining the host's standard irqbalance service. Adjust the CPU list if the +# instance type or SMPC__SEPARATE_TOKIO_CORES_PER_NODE changes. + +usage() { + echo "usage: sudo $0 [reserved-cpu-list]" >&2 + echo "example: sudo $0 pin ens66 0-7" >&2 + exit 2 +} + +[[ $# -ge 2 ]] || usage +action=$1 +interface=$2 +cpu_list=${3:-} +state_dir=${STATE_DIR:-/var/lib/iris-mpc} +state_file="${state_dir}/irq-affinity-${interface}.state" + +[[ ${EUID} -eq 0 ]] || { + echo "this script must run as root" >&2 + exit 1 +} +[[ -d "/sys/class/net/${interface}" ]] || { + echo "network interface ${interface} does not exist" >&2 + exit 1 +} + +expand_cpu_list() { + local part first last cpu + local -a expanded=() + IFS=',' read -ra parts <<<"$1" + for part in "${parts[@]}"; do + if [[ ${part} == *-* ]]; then + first=${part%-*} + last=${part#*-} + for ((cpu = first; cpu <= last; cpu++)); do + expanded+=("${cpu}") + done + else + expanded+=("${part}") + fi + done + printf '%s\n' "${expanded[@]}" +} + +case ${action} in +pin) + [[ -n ${cpu_list} ]] || usage + [[ ! -e ${state_file} ]] || { + echo "${state_file} already exists; restore before pinning again" >&2 + exit 1 + } + mapfile -t cpus < <(expand_cpu_list "${cpu_list}") + ((${#cpus[@]} > 0)) || { + echo "reserved CPU list is empty" >&2 + exit 1 + } + for cpu in "${cpus[@]}"; do + [[ -d "/sys/devices/system/cpu/cpu${cpu}" ]] || { + echo "CPU ${cpu} does not exist" >&2 + exit 1 + } + done + mapfile -t irqs < <( + awk -v queue_prefix="${interface}-Tx-Rx-" \ + 'index($0, queue_prefix) { gsub(":", "", $1); print $1 }' \ + /proc/interrupts + ) + ((${#irqs[@]} > 0)) || { + echo "no ${interface} Tx/Rx queue IRQs found" >&2 + exit 1 + } + + install -d "${state_dir}" + irqbalance_active=0 + if systemctl is-active --quiet irqbalance; then + irqbalance_active=1 + systemctl stop irqbalance + fi + printf 'irqbalance_active %s\n' "${irqbalance_active}" >"${state_file}" + for index in "${!irqs[@]}"; do + irq=${irqs[index]} + cpu_index=$((index % ${#cpus[@]})) + affinity=$(<"/proc/irq/${irq}/smp_affinity_list") + printf 'irq %s %s\n' "${irq}" "${affinity}" >>"${state_file}" + printf '%s\n' "${cpus[cpu_index]}" >"/proc/irq/${irq}/smp_affinity_list" + done + echo "pinned ${#irqs[@]} ${interface} queue IRQs across CPUs ${cpu_list}" + ;; +restore) + [[ -f ${state_file} ]] || { + echo "missing saved state ${state_file}" >&2 + exit 1 + } + irqbalance_active=0 + while read -r kind first second; do + case ${kind} in + irqbalance_active) irqbalance_active=${first} ;; + irq) printf '%s\n' "${second}" >"/proc/irq/${first}/smp_affinity_list" ;; + esac + done <"${state_file}" + if [[ ${irqbalance_active} == 1 ]]; then + systemctl start irqbalance + fi + rm "${state_file}" + echo "restored ${interface} IRQ affinities" + ;; +*) usage ;; +esac From 8ef678f46b193d673e26450e3290570bd50c6db0 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Wed, 19 Aug 2026 13:18:24 +0200 Subject: [PATCH 02/14] Prevent stale cold-eye prefetch reservations --- iris-mpc-cpu/src/execution/hawk_main.rs | 3 +- .../src/execution/hawk_main/search.rs | 71 +++++++++++++++---- iris-mpc-cpu/src/protocol/ops.rs | 10 +-- iris-mpc/src/server/mod.rs | 38 +++++----- 4 files changed, 86 insertions(+), 36 deletions(-) diff --git a/iris-mpc-cpu/src/execution/hawk_main.rs b/iris-mpc-cpu/src/execution/hawk_main.rs index 4074358eeb..7c91283d05 100644 --- a/iris-mpc-cpu/src/execution/hawk_main.rs +++ b/iris-mpc-cpu/src/execution/hawk_main.rs @@ -2204,7 +2204,8 @@ impl HawkHandle { let do_search = async |orient, precomputed_search: Option< SearchResults, - >| -> Result<_> { + >| + -> Result<_> { let search_queries = &request.queries(orient); let (luc_ids, request_types, forced_anon_stats_ids) = { // Choice of LEFT registry is arbitrary — both sides are in sync diff --git a/iris-mpc-cpu/src/execution/hawk_main/search.rs b/iris-mpc-cpu/src/execution/hawk_main/search.rs index 6f27deb5e3..a22c6aeac0 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/search.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/search.rs @@ -59,6 +59,9 @@ struct LinearScanChunk { #[derive(Clone)] struct LinearScanPrefetch { worker: Arc, + /// Candidate IDs already consumed by the concurrent known-candidate + /// stage. They must not reserve prefetch slots that no later stage reads. + excluded_ids: Arc>>, } #[derive(Clone, Default)] @@ -366,6 +369,12 @@ pub async fn linear_scan_cascade( LinearScanHooks { prefetch: Some(LinearScanPrefetch { worker: prefetch_worker.clone(), + excluded_ids: Arc::new( + known_second_stage_ids + .iter() + .map(|ids| ids.to_vec()) + .collect(), + ), }), progress: Some(first_eye_progress.clone()), }, @@ -526,6 +535,11 @@ fn collect_live_second_stage_ids( ids } +fn exclude_known_second_stage_ids(ids: &mut Vec, known_ids: &[VectorId]) { + debug_assert!(known_ids.windows(2).all(|pair| pair[0] < pair[1])); + ids.retain(|id| known_ids.binary_search(id).is_err()); +} + async fn linear_scan_eye( sessions: &BothEyes>, search_queries: &SearchQueries, @@ -608,7 +622,7 @@ async fn linear_scan_eye( .await?; if let Some(prefetch) = &hooks.prefetch { let chunk_ids = &candidate_ids[chunk.i_request][chunk.range.clone()]; - let prefetch_ids = collect_live_second_stage_ids( + let mut prefetch_ids = collect_live_second_stage_ids( chunk_ids, result .classified @@ -618,6 +632,10 @@ async fn linear_scan_eye( .map(|(id, _)| *id), &[], ); + exclude_known_second_stage_ids( + &mut prefetch_ids, + &prefetch.excluded_ids[chunk.i_request], + ); prefetch.worker.prefetch_irises(prefetch_ids).await?; } if let Some(progress) = &hooks.progress { @@ -1057,7 +1075,7 @@ async fn linear_scan_full_stage_paired( if let Some(prefetch) = &hooks.prefetch { // One union prefetch warms the cold eye for both // orientations' second-stage candidates. - let prefetch_ids = collect_live_second_stage_ids( + let mut prefetch_ids = collect_live_second_stage_ids( chunk_ids, plans.iter().flat_map(|plan| { plan.classified @@ -1068,6 +1086,10 @@ async fn linear_scan_full_stage_paired( }), &[], ); + exclude_known_second_stage_ids( + &mut prefetch_ids, + &prefetch.excluded_ids[chunk.i_request], + ); prefetch.worker.prefetch_irises(prefetch_ids).await?; } if let Some(progress) = &hooks.progress { @@ -1204,6 +1226,17 @@ pub async fn linear_scan_cascade_paired( .collect::>(), ) }); + let prefetch_excluded_ids = Arc::new( + (0..n_requests) + .map(|i_request| { + let mut ids = known_second_stage_ids_both[0][i_request].to_vec(); + ids.extend_from_slice(&known_second_stage_ids_both[1][i_request]); + ids.sort_unstable(); + ids.dedup(); + ids + }) + .collect(), + ); for orientation in orientations { tracing::info!( @@ -1235,6 +1268,7 @@ pub async fn linear_scan_cascade_paired( LinearScanHooks { prefetch: Some(LinearScanPrefetch { worker: prefetch_worker.clone(), + excluded_ids: prefetch_excluded_ids, }), progress: Some(first_eye_progress.clone()), }, @@ -1769,16 +1803,8 @@ async fn per_linear_scan_chunk_pair( classified, }; Ok([ - as_plan( - queries[0], - graph_stores.0.last_update_seq_no, - classified_a, - ), - as_plan( - queries[1], - graph_stores.1.last_update_seq_no, - classified_b, - ), + as_plan(queries[0], graph_stores.0.last_update_seq_no, classified_a), + as_plan(queries[1], graph_stores.1.last_update_seq_no, classified_b), ]) } @@ -2125,6 +2151,27 @@ mod tests { ); } + #[test] + fn prefetch_excludes_candidates_already_consumed_by_known_stage() { + let mut discovered = vec![ + VectorId::from_serial_id(1), + VectorId::from_serial_id(3), + VectorId::from_serial_id(5), + ]; + let known = [ + VectorId::from_serial_id(2), + VectorId::from_serial_id(3), + VectorId::from_serial_id(4), + ]; + + exclude_known_second_stage_ids(&mut discovered, &known); + + assert_eq!( + discovered, + vec![VectorId::from_serial_id(1), VectorId::from_serial_id(5)] + ); + } + #[tokio::test] async fn test_search() -> Result<()> { let actors = setup_hawk_actors().await?; diff --git a/iris-mpc-cpu/src/protocol/ops.rs b/iris-mpc-cpu/src/protocol/ops.rs index fb0dd9580f..165953ce5d 100644 --- a/iris-mpc-cpu/src/protocol/ops.rs +++ b/iris-mpc-cpu/src/protocol/ops.rs @@ -1331,14 +1331,16 @@ mod mixed_scan { let acc_c = &block[2]; let lane_scale = 1 + result_lane as u16; for (query, shares) in additive_shares.iter_mut().enumerate() { - let first = (acc_a[query].wrapping_add(acc_a[2 + query].wrapping_add(acc_c[query]) << 8) - as u16) - .wrapping_mul(lane_scale); + let first = + (acc_a[query].wrapping_add(acc_a[2 + query].wrapping_add(acc_c[query]) << 8) + as u16) + .wrapping_mul(lane_scale); let second = (acc_b[query] .wrapping_add(acc_b[2 + query].wrapping_add(acc_c[2 + query]) << 8) as u16) .wrapping_mul(lane_scale); - let first_idx = (base_target_idx + pair * 2) * ROTATIONS * 2 + rotation * 2 + result_lane; + let first_idx = + (base_target_idx + pair * 2) * ROTATIONS * 2 + rotation * 2 + result_lane; let second_idx = (base_target_idx + pair * 2 + 1) * ROTATIONS * 2 + rotation * 2 + result_lane; shares[first_idx].0 = shares[first_idx].0.wrapping_add(first); diff --git a/iris-mpc/src/server/mod.rs b/iris-mpc/src/server/mod.rs index 5872e08867..0c32383597 100644 --- a/iris-mpc/src/server/mod.rs +++ b/iris-mpc/src/server/mod.rs @@ -452,25 +452,6 @@ fn validate_max_batch_size(max_batch_size: usize, search_mode: HawkSearchMode) - Ok(()) } -#[cfg(test)] -mod config_tests { - use super::*; - - #[test] - fn linear_scan_requires_single_request_batches() { - assert!(validate_max_batch_size(1, HawkSearchMode::LinearScan).is_ok()); - - let error = validate_max_batch_size(2, HawkSearchMode::LinearScan) - .expect_err("linear scan must reject batches larger than one"); - assert!(error.to_string().contains("requires max_batch_size=1")); - } - - #[test] - fn hnsw_keeps_batched_requests() { - assert!(validate_max_batch_size(64, HawkSearchMode::Hnsw).is_ok()); - } -} - /// Returns initialized PostgreSQL clients for interacting /// with iris share and HNSW graph stores. async fn prepare_stores( @@ -1261,3 +1242,22 @@ async fn run_main_server_loop( } Ok(()) } + +#[cfg(test)] +mod config_tests { + use super::*; + + #[test] + fn linear_scan_requires_single_request_batches() { + assert!(validate_max_batch_size(1, HawkSearchMode::LinearScan).is_ok()); + + let error = validate_max_batch_size(2, HawkSearchMode::LinearScan) + .expect_err("linear scan must reject batches larger than one"); + assert!(error.to_string().contains("requires max_batch_size=1")); + } + + #[test] + fn hnsw_keeps_batched_requests() { + assert!(validate_max_batch_size(64, HawkSearchMode::Hnsw).is_ok()); + } +} From c476115e35bfd0a402d9f78aac412dd54243655b Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Wed, 19 Aug 2026 14:57:14 +0200 Subject: [PATCH 03/14] Keep ground-truth scheduling consistent --- iris-mpc-utils/src/client/options/types.rs | 199 ++++++++++++--------- 1 file changed, 114 insertions(+), 85 deletions(-) diff --git a/iris-mpc-utils/src/client/options/types.rs b/iris-mpc-utils/src/client/options/types.rs index 4d9031a6b3..0efff63021 100644 --- a/iris-mpc-utils/src/client/options/types.rs +++ b/iris-mpc-utils/src/client/options/types.rs @@ -683,107 +683,120 @@ fn ground_truth_into_iter( operations.extend(cycle); } - for operation in operations + for scheduled_operation in operations .into_iter() .take(request_count.saturating_sub(batches.len())) { let label = format!("ground-truth-{sequence:04}"); sequence += 1; - let target_idx = rng.gen_range(0..live.len()); - let target = live[target_idx].clone(); - - let request = match operation { - GroundTruthOperation::NewUniqueness => { - let identity_label = format!("{label}-enroll"); - let iris_pair = fresh_pair(); - let request = RequestOptions::new( - Some(&identity_label), - RequestPayloadOptions::Uniqueness { - iris_pair: Some(iris_pair), - insertion_layers: None, - }, - ) - .with_expected(json!({ "is_match": false })); - live.push(GroundTruthIdentity { - label: identity_label, - iris_pair, - }); - request - } - GroundTruthOperation::DuplicateUniqueness => RequestOptions::new( - Some(&format!("{label}-duplicate")), + // A deletion may empty the identity set. Replace any operation that + // needs a target with an enrollment until there is a live identity + // again. New enrollments deliberately do not consume target-selection + // randomness, keeping target draws tied to operations that use them. + let operation = if live.is_empty() { + GroundTruthOperation::NewUniqueness + } else { + scheduled_operation + }; + + let request = if matches!(operation, GroundTruthOperation::NewUniqueness) { + let identity_label = format!("{label}-enroll"); + let iris_pair = fresh_pair(); + let request = RequestOptions::new( + Some(&identity_label), RequestPayloadOptions::Uniqueness { - iris_pair: Some(target.iris_pair), + iris_pair: Some(iris_pair), insertion_layers: None, }, ) - .with_expected(json!({ "is_match": true })), - GroundTruthOperation::MirroredUniqueness => RequestOptions::new( - Some(&format!("{label}-mirror")), - RequestPayloadOptions::Mirrored { - iris_pair: Some(target.iris_pair), - insertion_layers: None, - }, - ), - GroundTruthOperation::ReauthMatch => RequestOptions::new( - Some(&format!("{label}-reauth-match")), - RequestPayloadOptions::Reauthorisation { - iris_pair: Some(target.iris_pair), - parent: Parent::Label(target.label), - }, - ) - .with_expected(json!({ "success": true })), - GroundTruthOperation::ReauthNonMatch => RequestOptions::new( - Some(&format!("{label}-reauth-non-match")), - RequestPayloadOptions::Reauthorisation { - iris_pair: Some(fresh_pair()), - parent: Parent::Label(target.label), - }, - ) - .with_expected(json!({ "success": false })), - GroundTruthOperation::ResetCheck => RequestOptions::new( - Some(&format!("{label}-reset-check")), - RequestPayloadOptions::ResetCheck { - iris_pair: Some(target.iris_pair), - }, - ), - GroundTruthOperation::RecoveryCheck => RequestOptions::new( - Some(&format!("{label}-recovery-check")), - RequestPayloadOptions::RecoveryCheck { - iris_pair: Some(target.iris_pair), - }, - ), - GroundTruthOperation::ResetUpdate => { - let iris_pair = fresh_pair(); - live[target_idx].iris_pair = iris_pair; - RequestOptions::new( - Some(&format!("{label}-reset-update")), - RequestPayloadOptions::ResetUpdate { - iris_pair: Some(iris_pair), - parent: Parent::Label(target.label), + .with_expected(json!({ "is_match": false })); + live.push(GroundTruthIdentity { + label: identity_label, + iris_pair, + }); + request + } else { + let target_idx = rng.gen_range(0..live.len()); + let target = live[target_idx].clone(); + match operation { + GroundTruthOperation::NewUniqueness => { + unreachable!("new uniqueness is handled before target selection") + } + GroundTruthOperation::DuplicateUniqueness => RequestOptions::new( + Some(&format!("{label}-duplicate")), + RequestPayloadOptions::Uniqueness { + iris_pair: Some(target.iris_pair), + insertion_layers: None, }, ) - } - GroundTruthOperation::RecoveryUpdate => { - let iris_pair = fresh_pair(); - live[target_idx].iris_pair = iris_pair; - RequestOptions::new( - Some(&format!("{label}-recovery-update")), - RequestPayloadOptions::RecoveryUpdate { - iris_pair: Some(iris_pair), + .with_expected(json!({ "is_match": true })), + GroundTruthOperation::MirroredUniqueness => RequestOptions::new( + Some(&format!("{label}-mirror")), + RequestPayloadOptions::Mirrored { + iris_pair: Some(target.iris_pair), + insertion_layers: None, + }, + ), + GroundTruthOperation::ReauthMatch => RequestOptions::new( + Some(&format!("{label}-reauth-match")), + RequestPayloadOptions::Reauthorisation { + iris_pair: Some(target.iris_pair), parent: Parent::Label(target.label), }, ) - } - GroundTruthOperation::IdentityDeletion => { - let deleted = live.swap_remove(target_idx); - RequestOptions::new( - Some(&format!("{label}-delete")), - RequestPayloadOptions::IdentityDeletion { - parent: Parent::Label(deleted.label), + .with_expected(json!({ "success": true })), + GroundTruthOperation::ReauthNonMatch => RequestOptions::new( + Some(&format!("{label}-reauth-non-match")), + RequestPayloadOptions::Reauthorisation { + iris_pair: Some(fresh_pair()), + parent: Parent::Label(target.label), }, ) - .with_expected(json!({ "success": true })) + .with_expected(json!({ "success": false })), + GroundTruthOperation::ResetCheck => RequestOptions::new( + Some(&format!("{label}-reset-check")), + RequestPayloadOptions::ResetCheck { + iris_pair: Some(target.iris_pair), + }, + ), + GroundTruthOperation::RecoveryCheck => RequestOptions::new( + Some(&format!("{label}-recovery-check")), + RequestPayloadOptions::RecoveryCheck { + iris_pair: Some(target.iris_pair), + }, + ), + GroundTruthOperation::ResetUpdate => { + let iris_pair = fresh_pair(); + live[target_idx].iris_pair = iris_pair; + RequestOptions::new( + Some(&format!("{label}-reset-update")), + RequestPayloadOptions::ResetUpdate { + iris_pair: Some(iris_pair), + parent: Parent::Label(target.label), + }, + ) + } + GroundTruthOperation::RecoveryUpdate => { + let iris_pair = fresh_pair(); + live[target_idx].iris_pair = iris_pair; + RequestOptions::new( + Some(&format!("{label}-recovery-update")), + RequestPayloadOptions::RecoveryUpdate { + iris_pair: Some(iris_pair), + parent: Parent::Label(target.label), + }, + ) + } + GroundTruthOperation::IdentityDeletion => { + let deleted = live.swap_remove(target_idx); + RequestOptions::new( + Some(&format!("{label}-delete")), + RequestPayloadOptions::IdentityDeletion { + parent: Parent::Label(deleted.label), + }, + ) + .with_expected(json!({ "success": true })) + } } }; batches.push(vec![request]); @@ -1183,4 +1196,20 @@ mod tests { ]) ); } + + #[test] + fn ground_truth_schedule_recovers_when_deletion_empties_live_identities() { + for rng_seed in 0..64 { + let batches = RequestBatchOptions::GroundTruth { + request_count: 100, + rng_seed, + initial_uniqueness_count: 1, + } + .into_iter() + .collect::>(); + + assert_eq!(batches.len(), 100, "rng seed {rng_seed}"); + assert!(batches.iter().all(|batch| batch.len() == 1)); + } + } } From 6294c902a9554240061668101be686b2e3da2d10 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Fri, 21 Aug 2026 11:02:21 +0200 Subject: [PATCH 04/14] Address fused-scan review findings - Port the GPU linear-scan bench to the resident-layout pool constructor; it still called the pre-layout LocalIrisWorkerPool::new_local and did not compile under --all-features --all-targets. - Restore the empty-request early return in search() and its test; the identity-update no-match search runs every batch and does not need a session per eye (the linear-scan no-match shortcut in per_session is back through the merge). - Lookahead dot tasks abort when their handle is dropped, so a lane error or a sibling lane failing try_join! cancels the in-flight chunk instead of leaving it running on the dot workers; the distance-mode and center rotation checks are made before spawning rather than inside the task. - Port the restored db-backed cold-eye test to the layout-aware pool API. Restored through the merge of the updated base: LinearScan-only gating of the full_face_mirror_* uniqueness result fields (HNSW keeps serializing null) and the dead linear-scan branch removal in per_session. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- .../src/execution/hawk_main/search.rs | 7 +- iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs | 64 +++++++++++++++---- iris-mpc-gpu/benches/linear_scan_dot.rs | 11 ++-- 3 files changed, 62 insertions(+), 20 deletions(-) diff --git a/iris-mpc-cpu/src/execution/hawk_main/search.rs b/iris-mpc-cpu/src/execution/hawk_main/search.rs index 6cd366e3ce..fe0cdcdc28 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/search.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/search.rs @@ -1089,10 +1089,13 @@ async fn linear_scan_full_stage_paired( // dot work then completes early and the stage drains on // thresholds alone with idle dot workers. const DOT_PIPELINE_DEPTH: usize = 1; + // Handles abort their task when dropped, so an error anywhere + // in this lane (or a sibling lane failing `try_join!`) also + // cancels the lookahead chunk instead of leaving it running. let mut pending_dots = std::collections::VecDeque::new(); if do_match { for chunk in batch.iter().take(DOT_PIPELINE_DEPTH) { - pending_dots.push_back(dispatch(&store_a, chunk)); + pending_dots.push_back(dispatch(&store_a, chunk)?); } } for (index, chunk) in batch.iter().enumerate() { @@ -1104,7 +1107,7 @@ async fn linear_scan_full_stage_paired( }; if do_match { if let Some(next) = batch.get(index + DOT_PIPELINE_DEPTH) { - pending_dots.push_back(dispatch(&store_a, next)); + pending_dots.push_back(dispatch(&store_a, next)?); } } let (query_a, query_b) = queries_for(chunk); diff --git a/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs b/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs index e8cb0bec48..202e73ea15 100644 --- a/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs +++ b/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs @@ -70,6 +70,38 @@ pub type Aby3Query = QuerySpec; pub type Aby3DistanceRef = DistanceShare; pub type RotationMatchIndices = Vec>; +/// Both orientations' additive dot-product shares for one chunk. +pub type PairDotContributions = [Vec>; 2]; + +/// A spawned task handle that aborts the task when dropped, so a lookahead +/// chunk cannot keep running detached after the lane that requested it has +/// failed or been cancelled. +pub struct AbortOnDropHandle(tokio::task::JoinHandle); + +impl AbortOnDropHandle { + fn new(handle: tokio::task::JoinHandle) -> Self { + Self(handle) + } +} + +impl Drop for AbortOnDropHandle { + fn drop(&mut self) { + self.0.abort(); + } +} + +impl std::future::Future for AbortOnDropHandle { + type Output = std::result::Result; + + fn poll( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll { + // `JoinHandle` is `Unpin`, so projecting through the wrapper is sound. + std::pin::Pin::new(&mut self.get_mut().0).poll(cx) + } +} + /// GPU-equivalent exact-scan classification for one chunk. Thresholds are /// evaluated directly for all 31 rotations; no secret minimum is computed. #[derive(Debug)] @@ -686,32 +718,36 @@ where /// worker-level target streaming is shared. This performs no network /// communication, so fusing and pipelining the dot passes is invisible to /// the MPC transcript. + /// + /// Configuration errors are reported before anything is spawned. The + /// returned handle aborts the task when dropped, so a lane that fails + /// while a lookahead chunk is in flight does not leave that chunk running + /// detached on the dot-product workers. pub fn spawn_full_rotation_dot_contributions_pair( &self, queries: [&Aby3Query; 2], vectors: &[VectorId], - ) -> tokio::task::JoinHandle>; 2]>> { - let mode = self.distance_fn.mode; + ) -> Result>> { + eyre::ensure!( + self.distance_fn.mode == DistanceMode::MinRotation, + "full-rotation scan requires min-rotation distance mode" + ); + for query in queries { + eyre::ensure!( + query.rotation == crate::execution::hawk_main::iris_worker::CENTER_ROTATION, + "full-rotation scan must start from the center query rotation" + ); + } let specs = [*queries[0], *queries[1]]; let workers = self.workers.clone(); let vectors = vectors.to_vec(); - tokio::spawn(async move { - eyre::ensure!( - mode == DistanceMode::MinRotation, - "full-rotation scan requires min-rotation distance mode" - ); - for query in specs { - eyre::ensure!( - query.rotation == crate::execution::hawk_main::iris_worker::CENTER_ROTATION, - "full-rotation scan must start from the center query rotation" - ); - } + Ok(AbortOnDropHandle::new(tokio::spawn(async move { metrics::counter!("distance_evaluations_total").increment(2 * vectors.len() as u64); metrics::histogram!("distance_evaluations_batch_size").record(vectors.len() as f64); workers .compute_dot_products_full_rotations_pair(specs, vectors) .await - }) + }))) } /// Check whether a batch of distances are matches at the given threshold. diff --git a/iris-mpc-gpu/benches/linear_scan_dot.rs b/iris-mpc-gpu/benches/linear_scan_dot.rs index 50314b3c10..c954718343 100644 --- a/iris-mpc-gpu/benches/linear_scan_dot.rs +++ b/iris-mpc-gpu/benches/linear_scan_dot.rs @@ -6,7 +6,7 @@ use iris_mpc_common::{ use iris_mpc_cpu::{ execution::hawk_main::iris_worker::{IrisWorkerPool, LocalIrisWorkerPool, QueryId, QuerySpec}, hawkers::{aby3::aby3_store::DistanceMode, shared_irises::SharedIrises}, - protocol::shared_iris::{ArcIris, GaloisRingSharedIris}, + protocol::shared_iris::{preferred_scan_layout, ArcIris, GaloisRingSharedIris, ResidentIris}, }; use iris_mpc_gpu::{ dot::share_db::{preprocess_query, ProcessedDatabase, ShareDB, SlicedProcessedDatabase}, @@ -415,20 +415,23 @@ fn build_gpu_scan(db_size: usize, query: &GaloisRingSharedIris) -> Result (LocalIrisWorkerPool, Vec) { println!("loading CPU database ({db_size} records)"); + // Same resident layout the production linear-scan pool selects on this + // CPU (mixed planes where the UMMLA kernel is available). + let layout = preferred_scan_layout(); let mut store = SharedIrises::new( HashMap::new(), - Arc::new(GaloisRingSharedIris::default_for_party(0)), + 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 _ in 0..db_size { // A deep clone gives every logical record its own backing memory, so // the scan cannot turn into an unrealistically cache-resident test. - vector_ids.push(store.append(Arc::new((**query).clone()))); + vector_ids.push(store.append(ResidentIris::from_arc(Arc::new((**query).clone()), layout))); } let store = store.to_arc(); ( - LocalIrisWorkerPool::new_local(store, DistanceMode::MinRotation, 0), + LocalIrisWorkerPool::new_local(store, layout, DistanceMode::MinRotation, 0), vector_ids, ) } From 05a07ce2734e977609c488e53a7d4e818306eaaf Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Fri, 21 Aug 2026 11:19:06 +0200 Subject: [PATCH 05/14] Assert shared worker pool in the paired linear scan The fused dot pass dispatches both orientations' queries through the normal-orientation store. That resolves the mirror query only because every session of an eye shares that eye's worker pool, into which HawkRequest::cache_into caches normal and mirror queries alike. Check the invariant per lane instead of relying on it silently. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- iris-mpc-cpu/src/execution/hawk_main/search.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/iris-mpc-cpu/src/execution/hawk_main/search.rs b/iris-mpc-cpu/src/execution/hawk_main/search.rs index fe0cdcdc28..7cd41f6a7d 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/search.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/search.rs @@ -1062,6 +1062,16 @@ async fn linear_scan_full_stage_paired( async move { let mut store_a = session_a.aby3_store.write().await; let mut store_b = session_b.aby3_store.write().await; + // The fused dot pass below is dispatched through `store_a` + // for both orientations' queries. That resolves the mirror + // query only because every session of one eye shares that + // eye's worker pool, into which `HawkRequest::cache_into` + // caches the normal and mirror queries alike. Make the + // assumption explicit rather than relying on it silently. + eyre::ensure!( + Arc::ptr_eq(&store_a.workers, &store_b.workers), + "paired linear scan requires both orientations' sessions to share one worker pool" + ); let graph_a = session_a.graph_store.clone().read_owned().await; let graph_b = session_b.graph_store.clone().read_owned().await; let mut results = Vec::with_capacity(batch.len()); From cca4fbdf9fa4b5e203d5ee0c06e586421a3782d8 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Fri, 21 Aug 2026 12:05:31 +0200 Subject: [PATCH 06/14] Refuse windowed dot products on mixed-plane residents The HNSW-style windowed dot products read resident records through ResidentIris::to_arc, which rebuilds a u16 iris per target when the resident layout is the mixed-plane scan layout. The exact scan streams planes directly and never takes that path, so production pools now refuse it instead of running a silent per-target reconstruction if a future caller reaches it. The cross-kernel parity tests that compare both kernels on the same mixed-plane data opt in explicitly. Test actors pick the resident layout by search mode exactly like the server does (HNSW keeps u16, the exact scan uses the CPU's scan layout), through a shared HawkActor::resident_layout_for. Also take the cached live VectorId list from the registry in the paired cascade, matching the single-orientation path. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- iris-mpc-cpu/src/execution/hawk_main.rs | 17 ++++++++-- .../src/execution/hawk_main/iris_worker.rs | 23 +++++++++++++ .../src/execution/hawk_main/search.rs | 16 ++------- iris-mpc-cpu/src/hawkers/aby3/test_utils.rs | 33 ++++++++++++------- iris-mpc-cpu/tests/full_rotation_dot.rs | 5 ++- iris-mpc/src/server/mod.rs | 5 +-- 6 files changed, 67 insertions(+), 32 deletions(-) diff --git a/iris-mpc-cpu/src/execution/hawk_main.rs b/iris-mpc-cpu/src/execution/hawk_main.rs index 41b53598c1..c4f6807b3e 100644 --- a/iris-mpc-cpu/src/execution/hawk_main.rs +++ b/iris-mpc-cpu/src/execution/hawk_main.rs @@ -668,7 +668,7 @@ impl HawkActor { HAWK_DISTANCE_MODE, args.numa, ) - .with_resident_layout(crate::protocol::shared_iris::preferred_scan_layout()), + .with_resident_layout(Self::resident_layout_for(search_mode)), ); let graph = [(); 2].map(|_| GraphMem::new()); Self::from_cli_with_initializer_and_graph( @@ -731,7 +731,7 @@ impl HawkActor { args.numa, iris_store, ) - .with_resident_layout(crate::protocol::shared_iris::preferred_scan_layout()), + .with_resident_layout(Self::resident_layout_for(search_mode)), ); Self::from_cli_with_initializer_and_graph( args, @@ -762,6 +762,19 @@ impl HawkActor { .await } + /// The resident iris layout the production server selects for a search + /// mode: the exact scan streams mixed planes where the UMMLA kernel is + /// available, while HNSW pools keep plain u16 irises for their windowed + /// dot products. + pub fn resident_layout_for( + search_mode: HawkSearchMode, + ) -> crate::protocol::shared_iris::ResidentLayout { + match search_mode { + HawkSearchMode::LinearScan => crate::protocol::shared_iris::preferred_scan_layout(), + HawkSearchMode::Hnsw => crate::protocol::shared_iris::ResidentLayout::U16, + } + } + async fn from_cli_with_initializer_and_graph( args: &HawkArgs, shutdown_ct: CancellationToken, diff --git a/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs b/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs index 9fb2d53fcc..dbea96c3c0 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs @@ -1211,6 +1211,12 @@ pub struct LocalIrisWorkerPool { /// rolling LUC window, a bounded frequency cache, and mutations awaiting a /// database commit; older explicit candidates are fetched sparsely. cold_storage: Option, + /// The HNSW-style windowed dot products read resident records through + /// `ResidentIris::to_arc`, which rebuilds a u16 iris per target when the + /// resident layout is the mixed-plane scan layout. The exact scan never + /// takes that path (it streams planes directly), so production pools + /// refuse it; the cross-kernel parity tests opt in explicitly. + windowed_ops_on_mixed_residents: bool, } #[derive(Clone)] @@ -1701,6 +1707,7 @@ impl LocalIrisWorkerPool { mode, party_id, cold_storage: None, + windowed_ops_on_mixed_residents: false, } } @@ -1783,9 +1790,18 @@ impl LocalIrisWorkerPool { state, prefetch_tx, }), + windowed_ops_on_mixed_residents: false, }) } + /// Allow windowed (HNSW-style) dot products against mixed-plane + /// residents. Each target is rebuilt as a u16 iris, so this is only for + /// tests and tools that compare the two kernels on the same data. + pub fn with_windowed_ops_on_mixed_residents(mut self) -> Self { + self.windowed_ops_on_mixed_residents = true; + self + } + /// Create a local worker pool for shard 0 with NUMA pinning. /// Standard construction for tests, benchmarks, and single-node tools. pub fn new_local( @@ -1968,6 +1984,8 @@ impl IrisWorkerPool for LocalIrisWorkerPool { let query_cache = self.query_cache.clone(); let mut inner = self.inner.clone(); let mode = self.mode; + let layout = self.layout; + let windowed_ops_on_mixed_residents = self.windowed_ops_on_mixed_residents; let pool = self.clone(); let is_cold = self.cold_storage.is_some(); Box::pin(async move { @@ -2009,6 +2027,11 @@ impl IrisWorkerPool for LocalIrisWorkerPool { return Ok(results); } + eyre::ensure!( + layout == ResidentLayout::U16 || windowed_ops_on_mixed_residents, + "windowed dot products are not served from mixed-plane residents: this path \ + rebuilds a u16 iris per target, and the exact scan does not use it" + ); match mode { DistanceMode::Simple => { let mut results = Vec::with_capacity(iris_batches.len()); diff --git a/iris-mpc-cpu/src/execution/hawk_main/search.rs b/iris-mpc-cpu/src/execution/hawk_main/search.rs index ce2b99cd65..4a36069b5b 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/search.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/search.rs @@ -1241,22 +1241,12 @@ pub async fn linear_scan_cascade_paired( assert_eq!(n_requests, forced_anon_stats_ids.len()); // Both eye registries and both orientation groups observe the same live - // VectorIds. Build the list once and share it across the fused stage. + // VectorIds. The registry caches this list between mutations, so the + // fused stage shares one allocation with every other user in the batch. let live_ids = { let vector_store = sessions_both[0][first_eye][0].aby3_store.read().await; let registry = vector_store.registry.read().await; - Arc::<[VectorId]>::from( - registry - .get_points() - .iter() - .enumerate() - .filter_map(|(serial_id, entry)| { - entry - .as_ref() - .map(|(version, ())| VectorId::new(serial_id as u32, *version)) - }) - .collect::>(), - ) + registry.live_vector_ids() }; let full_scan_ids = Arc::new(vec![live_ids.clone(); n_requests]); let first_eye_comparisons = live_ids.len() * n_requests; diff --git a/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs b/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs index 199b169660..1e88a394b4 100644 --- a/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs +++ b/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs @@ -86,12 +86,18 @@ pub async fn setup_local_aby3_players_with_preloaded_db( crate::protocol::shared_iris::ResidentIris::from_arc(iris, layout) }) .to_arc(); - let workers: Arc = Arc::new(LocalIrisWorkerPool::new_local( - resident_storage, - layout, - plain_store.distance_mode, - party_id, - )); + // Test stores use the scan layout of this CPU so the HNSW-style + // windowed ops double as a cross-kernel check; production pools + // refuse that combination. + let workers: Arc = Arc::new( + LocalIrisWorkerPool::new_local( + resident_storage, + layout, + plain_store.distance_mode, + party_id, + ) + .with_windowed_ops_on_mixed_residents(), + ); let registry = storage.data.try_read().unwrap().to_registry().to_arc(); Ok(Arc::new(Mutex::new(Aby3Store::new( registry, @@ -116,12 +122,15 @@ pub async fn setup_local_store_aby3_players(network_t: NetworkType) -> Result = Arc::new(LocalIrisWorkerPool::new_local( - storage.clone(), - layout, - TEST_DISTANCE_MODE, - party_id, - )); + let workers: Arc = Arc::new( + LocalIrisWorkerPool::new_local( + storage.clone(), + layout, + TEST_DISTANCE_MODE, + party_id, + ) + .with_windowed_ops_on_mixed_residents(), + ); let registry = storage.data.try_read().unwrap().to_registry().to_arc(); Ok(Arc::new(Mutex::new(Aby3Store::new( registry, diff --git a/iris-mpc-cpu/tests/full_rotation_dot.rs b/iris-mpc-cpu/tests/full_rotation_dot.rs index eb8436a613..20167d3124 100644 --- a/iris-mpc-cpu/tests/full_rotation_dot.rs +++ b/iris-mpc-cpu/tests/full_rotation_dot.rs @@ -65,7 +65,10 @@ fn run_test() -> Result<()> { )) }) .collect::>(); - let pool = LocalIrisWorkerPool::new_local(store.to_arc(), layout, DistanceMode::MinRotation, 0); + // The windowed comparison below is the cross-kernel oracle; production + // pools refuse it on mixed-plane residents, so opt in explicitly. + let pool = LocalIrisWorkerPool::new_local(store.to_arc(), layout, DistanceMode::MinRotation, 0) + .with_windowed_ops_on_mixed_residents(); let query_id = QueryId::new(); runtime.block_on(pool.cache_queries(vec![(query_id, Arc::new(deterministic_iris(0x5a5a)))]))?; diff --git a/iris-mpc/src/server/mod.rs b/iris-mpc/src/server/mod.rs index e5424f1d8c..5a8d0a0af5 100644 --- a/iris-mpc/src/server/mod.rs +++ b/iris-mpc/src/server/mod.rs @@ -836,10 +836,7 @@ async fn init_hawk_actor( ); // Exact-scan pools store irises in the mixed-plane layout on CPUs with // the UMMLA kernel; HNSW pools keep plain ArcIris values. - let resident_layout = match search_mode { - HawkSearchMode::LinearScan => iris_mpc_cpu::protocol::shared_iris::preferred_scan_layout(), - HawkSearchMode::Hnsw => iris_mpc_cpu::protocol::shared_iris::ResidentLayout::U16, - }; + let resident_layout = HawkActor::resident_layout_for(search_mode); let initializer: Box = Box::new( LocalWorkerPoolInitializer::new_load_from_db( hawk_args.party_index, From 7a2dc38dd47e07fc17709db9149552ae95449174 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Mon, 24 Aug 2026 11:24:49 +0200 Subject: [PATCH 07/14] Address kernel review comments - MixedPlaneIris keeps one share id; the code and mask shares of an iris always carry the same id (asserted in from_iris). - Use tokio-util's AbortOnDropHandle instead of a hand-rolled guard (enables the crate's "rt" feature). - Run the aarch64-only kernel parity tests on an arm64 runner in CI; the amd64 unit-test job compiles them out. - Build the resident representation before taking the store write lock in insert_irises. - Prove at compile time that every rotation amount is a multiple of 4, which the doubled-row mixed kernel's aligned-window-plus-phase addressing relies on. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- .github/workflows/run-unit-tests.yaml | 58 +++++++++++++++++++ Cargo.lock | 1 + iris-mpc-cpu/Cargo.toml | 2 +- .../src/execution/hawk_main/iris_worker.rs | 11 +++- iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs | 30 +--------- iris-mpc-cpu/src/protocol/ops.rs | 18 ++++++ iris-mpc-cpu/src/protocol/shared_iris.rs | 19 +++--- 7 files changed, 99 insertions(+), 40 deletions(-) diff --git a/.github/workflows/run-unit-tests.yaml b/.github/workflows/run-unit-tests.yaml index 0128ccab17..02d19db1ab 100644 --- a/.github/workflows/run-unit-tests.yaml +++ b/.github/workflows/run-unit-tests.yaml @@ -102,3 +102,61 @@ jobs: S3_TEST_ENDPOINT: http://localhost:4566 RUST_LOG: info RUST_BACKTRACE: 1 + + # The NEON/UMMLA exact-scan kernels and their scalar-parity tests are + # compiled only for aarch64, so the amd64 job above never executes them. + # Run the kernel-bearing crate on an arm64 runner to keep them covered. + kernel-tests-arm64: + timeout-minutes: 30 + runs-on: + labels: arm64-ubuntu-22.04-16core + permissions: + contents: read + + steps: + - name: Checkout code + uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 + + - name: Get all test, doc and src files that have changed + id: changed-files-yaml + uses: tj-actions/changed-files@24d32ffd492484c1d75e0c0b894501ddb9d30d62 + with: + files_yaml: | + src: + - Cargo.lock + - Cargo.toml + - rust-toolchain.toml + - iris-mpc-cpu/** + - iris-mpc-common/** + - .github/workflows/run-unit-tests.yaml + + - name: Cache Rust build + if: steps.changed-files-yaml.outputs.src_any_changed == 'true' + uses: actions/cache@9255dc7a253b0ccc959486e2bca901246202afeb + id: cache-rust + with: + path: | + ~/.cargo/registry + ~/.cargo/git + target + key: rust-build-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('**/Cargo.lock') }} + restore-keys: | + rust-build-${{ runner.os }}-${{ runner.arch }}- + + - name: Install Rust + if: steps.changed-files-yaml.outputs.src_any_changed == 'true' + run: | + command -v rustup >/dev/null || { + echo "::error::rustup is not available on this arm64 runner image" + exit 1 + } + rustup toolchain install 1.95.0 + rustup default 1.95.0 + + - name: Run kernel and scan tests + if: steps.changed-files-yaml.outputs.src_any_changed == 'true' + run: cargo test --release -p iris-mpc-cpu --lib -- --test-threads=1 + + - name: Run fused full-rotation parity test + if: steps.changed-files-yaml.outputs.src_any_changed == 'true' + run: cargo test --release -p iris-mpc-cpu --test full_rotation_dot -- --test-threads=1 diff --git a/Cargo.lock b/Cargo.lock index 1634dfecc0..561f8426ec 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6268,6 +6268,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/iris-mpc-cpu/Cargo.toml b/iris-mpc-cpu/Cargo.toml index 37f99a9f1d..a1fd5c821a 100644 --- a/iris-mpc-cpu/Cargo.toml +++ b/iris-mpc-cpu/Cargo.toml @@ -50,7 +50,7 @@ thiserror.workspace = true toml.workspace = true tokio.workspace = true tokio-stream = "0.1" -tokio-util = { workspace = true, features = ["io-util"] } +tokio-util = { workspace = true, features = ["io-util", "rt"] } metrics.workspace = true moka = { version = "0.12.15", features = ["sync"] } tracing.workspace = true diff --git a/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs b/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs index dbea96c3c0..53ff8c4612 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs @@ -2223,12 +2223,19 @@ impl IrisWorkerPool for LocalIrisWorkerPool { return Ok(0); } + // Build the resident representation before taking the lock; the + // mixed-plane interleave has no reason to run under it. + let resident = resolved + .into_iter() + .map(|(vector_id, iris)| (vector_id, ResidentIris::from_arc(iris, layout))) + .collect::>(); + // Write directly to the shared store (not via IrisPoolHandle::insert // which is fire-and-forget). HNSW insertion needs the iris to be // visible in the store immediately after this returns. let mut store = iris_store.data.write().await; - for (vector_id, iris) in resolved { - store.insert(vector_id, ResidentIris::from_arc(iris, layout)); + for (vector_id, iris) in resident { + store.insert(vector_id, iris); } Ok(store.set_hash.checksum()) }) diff --git a/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs b/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs index 73beda72e4..ab4ac38679 100644 --- a/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs +++ b/iris-mpc-cpu/src/hawkers/aby3/aby3_store.rs @@ -49,6 +49,7 @@ use std::{ sync::Arc, vec, }; +pub use tokio_util::task::AbortOnDropHandle; use tracing::instrument; mod distance_fn; @@ -73,35 +74,6 @@ pub type RotationMatchIndices = Vec>; /// Both orientations' additive dot-product shares for one chunk. pub type PairDotContributions = [Vec>; 2]; -/// A spawned task handle that aborts the task when dropped, so a lookahead -/// chunk cannot keep running detached after the lane that requested it has -/// failed or been cancelled. -pub struct AbortOnDropHandle(tokio::task::JoinHandle); - -impl AbortOnDropHandle { - fn new(handle: tokio::task::JoinHandle) -> Self { - Self(handle) - } -} - -impl Drop for AbortOnDropHandle { - fn drop(&mut self) { - self.0.abort(); - } -} - -impl std::future::Future for AbortOnDropHandle { - type Output = std::result::Result; - - fn poll( - self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll { - // `JoinHandle` is `Unpin`, so projecting through the wrapper is sound. - std::pin::Pin::new(&mut self.get_mut().0).poll(cx) - } -} - /// GPU-equivalent exact-scan classification for one chunk. Thresholds are /// evaluated directly for all 31 rotations; no secret minimum is computed. #[derive(Debug)] diff --git a/iris-mpc-cpu/src/protocol/ops.rs b/iris-mpc-cpu/src/protocol/ops.rs index 165953ce5d..c911f11c7b 100644 --- a/iris-mpc-cpu/src/protocol/ops.rs +++ b/iris-mpc-cpu/src/protocol/ops.rs @@ -188,6 +188,21 @@ impl PrerotatedQueryRowMajorView<'_, ROTATIONS> { amounts }; + /// Compile-time proof that every rotation amount is a multiple of 4. + /// The doubled-row mixed-plane kernel addresses each rotation as an + /// 8-byte-aligned window plus a 4-element phase + /// (`mixed_scan::query_window`), which is only sound under this property. + const ROTATION_AMOUNTS_ARE_MULTIPLES_OF_FOUR: () = { + let mut i = 0; + while i < ROTATIONS { + assert!( + Self::ROTATION_AMOUNTS[i] % 4 == 0, + "rotation amounts must be multiples of 4" + ); + i += 1; + } + }; + /// Rotate row directly into destination buffer (zero allocations). #[inline] fn rotate_row_into(src: &[u16], dst: &mut [u16], left_amount: usize) { @@ -1463,6 +1478,9 @@ mod mixed_scan { row: usize, rotation: usize, ) -> *const u8 { + // Referencing the const forces its compile-time evaluation for this + // ROTATIONS instantiation. + const { PrerotatedQueryRowMajorView::::ROTATION_AMOUNTS_ARE_MULTIPLES_OF_FOUR } let amount = PrerotatedQueryRowMajorView::::ROTATION_AMOUNTS[rotation]; let phase = (amount / 4) & 1; let aligned_amount = amount - phase * 4; diff --git a/iris-mpc-cpu/src/protocol/shared_iris.rs b/iris-mpc-cpu/src/protocol/shared_iris.rs index 3e4ae7bf14..37aa303556 100644 --- a/iris-mpc-cpu/src/protocol/shared_iris.rs +++ b/iris-mpc-cpu/src/protocol/shared_iris.rs @@ -133,10 +133,10 @@ impl GaloisRingSharedIris { /// original share is reconstructed exactly by [`MixedPlaneIris::to_iris`]. #[derive(Debug, Clone, PartialEq, Eq)] pub struct MixedPlaneIris { - /// `id` of the code share (party_id + 1), preserved for reconstruction. - code_id: usize, - /// `id` of the mask share (party_id + 1), preserved for reconstruction. - mask_id: usize, + /// Share `id` (party_id + 1) of the code and mask shares, preserved for + /// reconstruction. Both shares of one iris always carry the same id, so + /// one copy is enough. + share_id: usize, /// 16 rows x 1600 bytes. code: Box<[u8]>, /// 8 rows x 1600 bytes. @@ -164,13 +164,16 @@ fn unmix_planes(src: &[u8], dst: &mut [u16]) { impl MixedPlaneIris { pub fn from_iris(iris: &GaloisRingSharedIris) -> Self { + debug_assert_eq!( + iris.code.id, iris.mask.id, + "code and mask shares of one iris must carry the same share id" + ); let mut code = vec![0u8; iris.code.coefs.len() * 2].into_boxed_slice(); let mut mask = vec![0u8; iris.mask.coefs.len() * 2].into_boxed_slice(); mix_planes(&iris.code.coefs, &mut code); mix_planes(&iris.mask.coefs, &mut mask); Self { - code_id: iris.code.id, - mask_id: iris.mask.id, + share_id: iris.code.id, code, mask, } @@ -178,8 +181,8 @@ impl MixedPlaneIris { pub fn to_iris(&self) -> GaloisRingSharedIris { let mut iris = GaloisRingSharedIris::default_for_party(0); - iris.code.id = self.code_id; - iris.mask.id = self.mask_id; + iris.code.id = self.share_id; + iris.mask.id = self.share_id; unmix_planes(&self.code, &mut iris.code.coefs); unmix_planes(&self.mask, &mut iris.mask.coefs); iris From 348c5b77cf7b99d197acf518b9b2e898b55c2370 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Mon, 24 Aug 2026 11:26:39 +0200 Subject: [PATCH 08/14] Install a pinned rustup on the arm64 kernel-test runner The arm64 runner image has no Rust toolchain. Download a version-pinned rustup-init, verify its sha256, and execute it, rather than assuming rustup or piping a live script into a shell. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- .github/workflows/run-unit-tests.yaml | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/.github/workflows/run-unit-tests.yaml b/.github/workflows/run-unit-tests.yaml index 02d19db1ab..7699905ab5 100644 --- a/.github/workflows/run-unit-tests.yaml +++ b/.github/workflows/run-unit-tests.yaml @@ -143,13 +143,24 @@ jobs: restore-keys: | rust-build-${{ runner.os }}-${{ runner.arch }}- + # The arm64 runner image ships without rustup. Install a pinned + # rustup-init and verify its checksum before executing it (no + # pipe-to-shell), then install the workspace toolchain. - name: Install Rust if: steps.changed-files-yaml.outputs.src_any_changed == 'true' + env: + RUSTUP_VERSION: 1.29.0 + RUSTUP_SHA256: 9732d6c5e2a098d3521fca8145d826ae0aaa067ef2385ead08e6feac88fa5792 run: | - command -v rustup >/dev/null || { - echo "::error::rustup is not available on this arm64 runner image" - exit 1 - } + if ! command -v rustup >/dev/null && [ ! -x "$HOME/.cargo/bin/rustup" ]; then + curl -fsSL -o rustup-init "https://static.rust-lang.org/rustup/archive/${RUSTUP_VERSION}/aarch64-unknown-linux-gnu/rustup-init" + echo "${RUSTUP_SHA256} rustup-init" | sha256sum -c - + chmod +x rustup-init + ./rustup-init -y --default-toolchain none --profile minimal + rm rustup-init + fi + echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" + export PATH="$HOME/.cargo/bin:$PATH" rustup toolchain install 1.95.0 rustup default 1.95.0 From 4d8a894ecd9dcbeafe8ca91bd9736af8a81d76a1 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Mon, 24 Aug 2026 11:29:02 +0200 Subject: [PATCH 09/14] Install a C toolchain on the arm64 kernel-test runner Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- .github/workflows/run-unit-tests.yaml | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/.github/workflows/run-unit-tests.yaml b/.github/workflows/run-unit-tests.yaml index 7699905ab5..53fb87f5b0 100644 --- a/.github/workflows/run-unit-tests.yaml +++ b/.github/workflows/run-unit-tests.yaml @@ -143,9 +143,15 @@ jobs: restore-keys: | rust-build-${{ runner.os }}-${{ runner.arch }}- - # The arm64 runner image ships without rustup. Install a pinned - # rustup-init and verify its checksum before executing it (no - # pipe-to-shell), then install the workspace toolchain. + # The arm64 runner image is minimal: no C toolchain and no rustup. + - name: Install build dependencies + if: steps.changed-files-yaml.outputs.src_any_changed == 'true' + run: | + sudo apt-get update + sudo apt-get install -y build-essential pkg-config libssl-dev protobuf-compiler + + # Install a pinned rustup-init and verify its checksum before executing + # it (no pipe-to-shell), then install the workspace toolchain. - name: Install Rust if: steps.changed-files-yaml.outputs.src_any_changed == 'true' env: From 5c83d6823bca8335bcf96eb8357683d481fe5c46 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Mon, 24 Aug 2026 12:30:54 +0200 Subject: [PATCH 10/14] Unify the u16 NEON dot kernels into one const-generic function dot_product_nx4_u16:: replaces the hand-unrolled 1x4/4x4/6x4 kernels. The compiler fully unrolls the width loops: with target-cpu=neoverse-v2 the 6-wide instantiation compiles to the same 37-instruction hot loop (24 mla, 10 loads, no spills) as the hand version, and kernel and full-server benchmarks on r8g.24xlarge measure identically. The scalar-parity test now covers every width, and one function keeps the remainder tiles' schedules identical to the main tile by construction. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- iris-mpc-cpu/src/protocol/ops.rs | 306 ++++++------------------------- 1 file changed, 57 insertions(+), 249 deletions(-) diff --git a/iris-mpc-cpu/src/protocol/ops.rs b/iris-mpc-cpu/src/protocol/ops.rs index c911f11c7b..33ed66090a 100644 --- a/iris-mpc-cpu/src/protocol/ops.rs +++ b/iris-mpc-cpu/src/protocol/ops.rs @@ -546,7 +546,7 @@ fn accumulate_group_rows( &query_rows[(rotation_idx + 4) * ROW_SIZE..(rotation_idx + 5) * ROW_SIZE]; let query5 = &query_rows[(rotation_idx + 5) * ROW_SIZE..(rotation_idx + 6) * ROW_SIZE]; - let partials = dot_product_6x4_u16( + let partials = dot_product_nx4_u16( [query0, query1, query2, query3, query4, query5], [target0, target1, target2, target3], ); @@ -575,7 +575,7 @@ fn accumulate_group_rows( &query_rows[(rotation_idx + 2) * ROW_SIZE..(rotation_idx + 3) * ROW_SIZE]; let query3 = &query_rows[(rotation_idx + 3) * ROW_SIZE..(rotation_idx + 4) * ROW_SIZE]; - let partials = dot_product_4x4_u16( + let partials = dot_product_nx4_u16( [query0, query1, query2, query3], [target0, target1, target2, target3], ); @@ -597,8 +597,8 @@ fn accumulate_group_rows( while rotation_idx < ROTATIONS { let query_row = &query_rows[rotation_idx * ROW_SIZE..(rotation_idx + 1) * ROW_SIZE]; - let partials = - dot_product_1x4_u16(query_row, [target0, target1, target2, target3]); + let [partials] = + dot_product_nx4_u16([query_row], [target0, target1, target2, target3]); for (target_offset, partial) in partials.into_iter().enumerate() { let result_idx = (target_idx + target_offset) * ROTATIONS * 2 + rotation_idx * 2 @@ -674,144 +674,27 @@ fn accumulate_scalar_target( } } -/// One rotation against four targets: the remainder tile after the 6-wide -/// (and optional 4-wide) passes. Keeping the four targets in one pass reuses -/// each query vector load for four MLAs instead of evaluating four -/// independent dot products. -#[cfg(target_arch = "aarch64")] -#[inline(always)] -fn dot_product_1x4_u16(query: &[u16], targets: [&[u16]; 4]) -> [u16; 4] { - use std::arch::aarch64::{uint16x8_t, vaddvq_u16, vdupq_n_u16, vld1q_u16, vmlaq_u16}; - - debug_assert_eq!(query.len(), 800); - debug_assert!(targets.iter().all(|target| target.len() == 800)); - - // SAFETY: AArch64 guarantees Advanced SIMD. Each pointer is derived from - // an 800-element slice and the loop only issues unaligned-safe 8-lane - // loads at offsets 0..792. All arithmetic deliberately wraps in u16. - unsafe { - let zero = vdupq_n_u16(0); - let mut acc0: uint16x8_t = zero; - let mut acc1: uint16x8_t = zero; - let mut acc2: uint16x8_t = zero; - let mut acc3: uint16x8_t = zero; - - let mut idx = 0; - while idx < 800 { - let query0 = vld1q_u16(query.as_ptr().add(idx)); - let target0 = vld1q_u16(targets[0].as_ptr().add(idx)); - let target1 = vld1q_u16(targets[1].as_ptr().add(idx)); - let target2 = vld1q_u16(targets[2].as_ptr().add(idx)); - let target3 = vld1q_u16(targets[3].as_ptr().add(idx)); - acc0 = vmlaq_u16(acc0, query0, target0); - acc1 = vmlaq_u16(acc1, query0, target1); - acc2 = vmlaq_u16(acc2, query0, target2); - acc3 = vmlaq_u16(acc3, query0, target3); - idx += 8; - } - - [ - vaddvq_u16(acc0), - vaddvq_u16(acc1), - vaddvq_u16(acc2), - vaddvq_u16(acc3), - ] - } -} - -#[cfg(target_arch = "aarch64")] -#[inline(always)] -fn dot_product_4x4_u16(queries: [&[u16]; 4], targets: [&[u16]; 4]) -> [[u16; 4]; 4] { - use std::arch::aarch64::{uint16x8_t, vaddvq_u16, vdupq_n_u16, vld1q_u16, vmlaq_u16}; - - debug_assert!(queries.iter().all(|query| query.len() == 800)); - debug_assert!(targets.iter().all(|target| target.len() == 800)); - - // SAFETY: AArch64 guarantees Advanced SIMD. Each pointer is derived from - // an 800-element slice and the loop only issues unaligned-safe 8-lane - // loads at offsets 0..792. All arithmetic deliberately wraps in u16. - unsafe { - let zero = vdupq_n_u16(0); - let mut acc00: uint16x8_t = zero; - let mut acc01: uint16x8_t = zero; - let mut acc02: uint16x8_t = zero; - let mut acc03: uint16x8_t = zero; - let mut acc10: uint16x8_t = zero; - let mut acc11: uint16x8_t = zero; - let mut acc12: uint16x8_t = zero; - let mut acc13: uint16x8_t = zero; - let mut acc20: uint16x8_t = zero; - let mut acc21: uint16x8_t = zero; - let mut acc22: uint16x8_t = zero; - let mut acc23: uint16x8_t = zero; - let mut acc30: uint16x8_t = zero; - let mut acc31: uint16x8_t = zero; - let mut acc32: uint16x8_t = zero; - let mut acc33: uint16x8_t = zero; - - let mut idx = 0; - while idx < 800 { - let query0 = vld1q_u16(queries[0].as_ptr().add(idx)); - let query1 = vld1q_u16(queries[1].as_ptr().add(idx)); - let query2 = vld1q_u16(queries[2].as_ptr().add(idx)); - let query3 = vld1q_u16(queries[3].as_ptr().add(idx)); - let target0 = vld1q_u16(targets[0].as_ptr().add(idx)); - let target1 = vld1q_u16(targets[1].as_ptr().add(idx)); - let target2 = vld1q_u16(targets[2].as_ptr().add(idx)); - let target3 = vld1q_u16(targets[3].as_ptr().add(idx)); - - acc00 = vmlaq_u16(acc00, query0, target0); - acc01 = vmlaq_u16(acc01, query0, target1); - acc02 = vmlaq_u16(acc02, query0, target2); - acc03 = vmlaq_u16(acc03, query0, target3); - acc10 = vmlaq_u16(acc10, query1, target0); - acc11 = vmlaq_u16(acc11, query1, target1); - acc12 = vmlaq_u16(acc12, query1, target2); - acc13 = vmlaq_u16(acc13, query1, target3); - acc20 = vmlaq_u16(acc20, query2, target0); - acc21 = vmlaq_u16(acc21, query2, target1); - acc22 = vmlaq_u16(acc22, query2, target2); - acc23 = vmlaq_u16(acc23, query2, target3); - acc30 = vmlaq_u16(acc30, query3, target0); - acc31 = vmlaq_u16(acc31, query3, target1); - acc32 = vmlaq_u16(acc32, query3, target2); - acc33 = vmlaq_u16(acc33, query3, target3); - idx += 8; - } - - [ - [ - vaddvq_u16(acc00), - vaddvq_u16(acc01), - vaddvq_u16(acc02), - vaddvq_u16(acc03), - ], - [ - vaddvq_u16(acc10), - vaddvq_u16(acc11), - vaddvq_u16(acc12), - vaddvq_u16(acc13), - ], - [ - vaddvq_u16(acc20), - vaddvq_u16(acc21), - vaddvq_u16(acc22), - vaddvq_u16(acc23), - ], - [ - vaddvq_u16(acc30), - vaddvq_u16(acc31), - vaddvq_u16(acc32), - vaddvq_u16(acc33), - ], - ] - } -} - +/// `N` query rotations against four targets in one pass. +/// +/// One 8-lane block streams each target vector past all `N` query vectors, so +/// every target load is reused for `N` MLAs and the `N * 4` accumulators stay +/// in registers. The 6-wide instantiation is the main scan tile; 4 and 1 cover +/// the remainders of the 11/13/31-rotation schedules (6+4+1, 6+6+1, 6x5+1). +/// +/// The compiler fully unrolls the `N`-loops: on rustc 1.94/1.95 with +/// `target-cpu=neoverse-v2`, `N = 6` compiles to the same 37-instruction hot +/// loop (24 mla, 10 loads, no spills) as the previously hand-unrolled kernel, +/// and measures identically on r8g.24xlarge in kernel and full-server +/// benchmarks. The scalar-parity test below guards the correctness of every +/// width; keeping all widths in one function keeps their schedules identical +/// by construction. #[cfg(target_arch = "aarch64")] #[inline(always)] -fn dot_product_6x4_u16(queries: [&[u16]; 6], targets: [&[u16]; 4]) -> [[u16; 4]; 6] { - use std::arch::aarch64::{uint16x8_t, vaddvq_u16, vdupq_n_u16, vld1q_u16, vmlaq_u16}; +fn dot_product_nx4_u16( + queries: [&[u16]; N], + targets: [&[u16]; 4], +) -> [[u16; 4]; N] { + use std::arch::aarch64::{vaddvq_u16, vdupq_n_u16, vld1q_u16, vmlaq_u16}; debug_assert!(queries.iter().all(|query| query.len() == 800)); debug_assert!(targets.iter().all(|target| target.len() == 800)); @@ -821,115 +704,34 @@ fn dot_product_6x4_u16(queries: [&[u16]; 6], targets: [&[u16]; 4]) -> [[u16; 4]; // loads at offsets 0..792. All arithmetic deliberately wraps in u16. unsafe { let zero = vdupq_n_u16(0); - let mut acc00: uint16x8_t = zero; - let mut acc01: uint16x8_t = zero; - let mut acc02: uint16x8_t = zero; - let mut acc03: uint16x8_t = zero; - let mut acc10: uint16x8_t = zero; - let mut acc11: uint16x8_t = zero; - let mut acc12: uint16x8_t = zero; - let mut acc13: uint16x8_t = zero; - let mut acc20: uint16x8_t = zero; - let mut acc21: uint16x8_t = zero; - let mut acc22: uint16x8_t = zero; - let mut acc23: uint16x8_t = zero; - let mut acc30: uint16x8_t = zero; - let mut acc31: uint16x8_t = zero; - let mut acc32: uint16x8_t = zero; - let mut acc33: uint16x8_t = zero; - let mut acc40: uint16x8_t = zero; - let mut acc41: uint16x8_t = zero; - let mut acc42: uint16x8_t = zero; - let mut acc43: uint16x8_t = zero; - let mut acc50: uint16x8_t = zero; - let mut acc51: uint16x8_t = zero; - let mut acc52: uint16x8_t = zero; - let mut acc53: uint16x8_t = zero; + let mut acc = [[zero; 4]; N]; let mut idx = 0; while idx < 800 { - let query0 = vld1q_u16(queries[0].as_ptr().add(idx)); - let query1 = vld1q_u16(queries[1].as_ptr().add(idx)); - let query2 = vld1q_u16(queries[2].as_ptr().add(idx)); - let query3 = vld1q_u16(queries[3].as_ptr().add(idx)); - let query4 = vld1q_u16(queries[4].as_ptr().add(idx)); - let query5 = vld1q_u16(queries[5].as_ptr().add(idx)); - + let mut query = [zero; N]; + for (vector, source) in query.iter_mut().zip(&queries) { + *vector = vld1q_u16(source.as_ptr().add(idx)); + } // Stream one target at a time. Keeping all four target vectors - // live alongside 24 accumulators and 6 query vectors would exceed - // the architectural SIMD register file and force stack spills. - let target0 = vld1q_u16(targets[0].as_ptr().add(idx)); - acc00 = vmlaq_u16(acc00, query0, target0); - acc10 = vmlaq_u16(acc10, query1, target0); - acc20 = vmlaq_u16(acc20, query2, target0); - acc30 = vmlaq_u16(acc30, query3, target0); - acc40 = vmlaq_u16(acc40, query4, target0); - acc50 = vmlaq_u16(acc50, query5, target0); - - let target1 = vld1q_u16(targets[1].as_ptr().add(idx)); - acc01 = vmlaq_u16(acc01, query0, target1); - acc11 = vmlaq_u16(acc11, query1, target1); - acc21 = vmlaq_u16(acc21, query2, target1); - acc31 = vmlaq_u16(acc31, query3, target1); - acc41 = vmlaq_u16(acc41, query4, target1); - acc51 = vmlaq_u16(acc51, query5, target1); - - let target2 = vld1q_u16(targets[2].as_ptr().add(idx)); - acc02 = vmlaq_u16(acc02, query0, target2); - acc12 = vmlaq_u16(acc12, query1, target2); - acc22 = vmlaq_u16(acc22, query2, target2); - acc32 = vmlaq_u16(acc32, query3, target2); - acc42 = vmlaq_u16(acc42, query4, target2); - acc52 = vmlaq_u16(acc52, query5, target2); - - let target3 = vld1q_u16(targets[3].as_ptr().add(idx)); - acc03 = vmlaq_u16(acc03, query0, target3); - acc13 = vmlaq_u16(acc13, query1, target3); - acc23 = vmlaq_u16(acc23, query2, target3); - acc33 = vmlaq_u16(acc33, query3, target3); - acc43 = vmlaq_u16(acc43, query4, target3); - acc53 = vmlaq_u16(acc53, query5, target3); + // live alongside the accumulators and query vectors would exceed + // the architectural SIMD register file for the 6-wide tile and + // force stack spills. + for (lane, source) in targets.iter().enumerate() { + let target = vld1q_u16(source.as_ptr().add(idx)); + for (row, query) in acc.iter_mut().zip(&query) { + row[lane] = vmlaq_u16(row[lane], *query, target); + } + } idx += 8; } - [ - [ - vaddvq_u16(acc00), - vaddvq_u16(acc01), - vaddvq_u16(acc02), - vaddvq_u16(acc03), - ], - [ - vaddvq_u16(acc10), - vaddvq_u16(acc11), - vaddvq_u16(acc12), - vaddvq_u16(acc13), - ], - [ - vaddvq_u16(acc20), - vaddvq_u16(acc21), - vaddvq_u16(acc22), - vaddvq_u16(acc23), - ], - [ - vaddvq_u16(acc30), - vaddvq_u16(acc31), - vaddvq_u16(acc32), - vaddvq_u16(acc33), - ], - [ - vaddvq_u16(acc40), - vaddvq_u16(acc41), - vaddvq_u16(acc42), - vaddvq_u16(acc43), - ], - [ - vaddvq_u16(acc50), - vaddvq_u16(acc51), - vaddvq_u16(acc52), - vaddvq_u16(acc53), - ], - ] + let mut out = [[0u16; 4]; N]; + for (row, accs) in out.iter_mut().zip(&acc) { + for (value, acc) in row.iter_mut().zip(accs) { + *value = vaddvq_u16(*acc); + } + } + out } } @@ -2179,26 +1981,32 @@ mod tests { #[cfg(target_arch = "aarch64")] #[test] - fn dot_product_6x4_u16_matches_scalar_wrapping_reference() { + fn dot_product_nx4_u16_matches_scalar_wrapping_reference() { const ROW_SIZE: usize = PrerotatedQueryRowMajor::ROW_SIZE; - for seed in [0, 1, 42, u64::MAX] { + fn check_width(seed: u64) { let mut rng = AesRng::seed_from_u64(seed); - let queries: [Vec; 6] = + let queries: [Vec; N] = std::array::from_fn(|_| (0..ROW_SIZE).map(|_| rng.next_u32() as u16).collect()); let targets: [Vec; 4] = std::array::from_fn(|_| (0..ROW_SIZE).map(|_| rng.next_u32() as u16).collect()); - let query_refs: [&[u16]; 6] = std::array::from_fn(|idx| queries[idx].as_slice()); + let query_refs: [&[u16]; N] = std::array::from_fn(|idx| queries[idx].as_slice()); let target_refs: [&[u16]; 4] = std::array::from_fn(|idx| targets[idx].as_slice()); - let tiled = dot_product_6x4_u16(query_refs, target_refs); - let reference = std::array::from_fn(|query_idx| { + let tiled = dot_product_nx4_u16(query_refs, target_refs); + let reference: [[u16; 4]; N] = std::array::from_fn(|query_idx| { std::array::from_fn(|target_idx| { simple_dot_product(&queries[query_idx], &targets[target_idx]) }) }); - assert_eq!(tiled, reference, "6x4 tile mismatch for seed {seed}"); + assert_eq!(tiled, reference, "{N}x4 tile mismatch for seed {seed}"); + } + + for seed in [0, 1, 42, u64::MAX] { + check_width::<6>(seed); + check_width::<4>(seed); + check_width::<1>(seed); } } From 359a080a0a7a15fd99f72b1e476dc52d50f87d88 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Mon, 24 Aug 2026 12:41:41 +0200 Subject: [PATCH 11/14] Remove the arm64 kernel-test CI job The const-generic kernel was validated directly against the hand-written version on r8g.24xlarge (identical hot-loop codegen, kernel microbench and full three-party server benchmark within noise), so the dedicated arm64 CI job is not needed; the parity tests remain for on-demand runs on ARM hardware. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- .github/workflows/run-unit-tests.yaml | 75 --------------------------- 1 file changed, 75 deletions(-) diff --git a/.github/workflows/run-unit-tests.yaml b/.github/workflows/run-unit-tests.yaml index 53fb87f5b0..0128ccab17 100644 --- a/.github/workflows/run-unit-tests.yaml +++ b/.github/workflows/run-unit-tests.yaml @@ -102,78 +102,3 @@ jobs: S3_TEST_ENDPOINT: http://localhost:4566 RUST_LOG: info RUST_BACKTRACE: 1 - - # The NEON/UMMLA exact-scan kernels and their scalar-parity tests are - # compiled only for aarch64, so the amd64 job above never executes them. - # Run the kernel-bearing crate on an arm64 runner to keep them covered. - kernel-tests-arm64: - timeout-minutes: 30 - runs-on: - labels: arm64-ubuntu-22.04-16core - permissions: - contents: read - - steps: - - name: Checkout code - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 - - - name: Get all test, doc and src files that have changed - id: changed-files-yaml - uses: tj-actions/changed-files@24d32ffd492484c1d75e0c0b894501ddb9d30d62 - with: - files_yaml: | - src: - - Cargo.lock - - Cargo.toml - - rust-toolchain.toml - - iris-mpc-cpu/** - - iris-mpc-common/** - - .github/workflows/run-unit-tests.yaml - - - name: Cache Rust build - if: steps.changed-files-yaml.outputs.src_any_changed == 'true' - uses: actions/cache@9255dc7a253b0ccc959486e2bca901246202afeb - id: cache-rust - with: - path: | - ~/.cargo/registry - ~/.cargo/git - target - key: rust-build-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('**/Cargo.lock') }} - restore-keys: | - rust-build-${{ runner.os }}-${{ runner.arch }}- - - # The arm64 runner image is minimal: no C toolchain and no rustup. - - name: Install build dependencies - if: steps.changed-files-yaml.outputs.src_any_changed == 'true' - run: | - sudo apt-get update - sudo apt-get install -y build-essential pkg-config libssl-dev protobuf-compiler - - # Install a pinned rustup-init and verify its checksum before executing - # it (no pipe-to-shell), then install the workspace toolchain. - - name: Install Rust - if: steps.changed-files-yaml.outputs.src_any_changed == 'true' - env: - RUSTUP_VERSION: 1.29.0 - RUSTUP_SHA256: 9732d6c5e2a098d3521fca8145d826ae0aaa067ef2385ead08e6feac88fa5792 - run: | - if ! command -v rustup >/dev/null && [ ! -x "$HOME/.cargo/bin/rustup" ]; then - curl -fsSL -o rustup-init "https://static.rust-lang.org/rustup/archive/${RUSTUP_VERSION}/aarch64-unknown-linux-gnu/rustup-init" - echo "${RUSTUP_SHA256} rustup-init" | sha256sum -c - - chmod +x rustup-init - ./rustup-init -y --default-toolchain none --profile minimal - rm rustup-init - fi - echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" - export PATH="$HOME/.cargo/bin:$PATH" - rustup toolchain install 1.95.0 - rustup default 1.95.0 - - - name: Run kernel and scan tests - if: steps.changed-files-yaml.outputs.src_any_changed == 'true' - run: cargo test --release -p iris-mpc-cpu --lib -- --test-threads=1 - - - name: Run fused full-rotation parity test - if: steps.changed-files-yaml.outputs.src_any_changed == 'true' - run: cargo test --release -p iris-mpc-cpu --test full_rotation_dot -- --test-threads=1 From e198d501750f88f5fe99d01b58239db7d544ce10 Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Mon, 24 Aug 2026 13:41:33 +0200 Subject: [PATCH 12/14] Trim kernel doc comment Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- iris-mpc-cpu/src/protocol/ops.rs | 8 -------- 1 file changed, 8 deletions(-) diff --git a/iris-mpc-cpu/src/protocol/ops.rs b/iris-mpc-cpu/src/protocol/ops.rs index 33ed66090a..73508ceb5b 100644 --- a/iris-mpc-cpu/src/protocol/ops.rs +++ b/iris-mpc-cpu/src/protocol/ops.rs @@ -680,14 +680,6 @@ fn accumulate_scalar_target( /// every target load is reused for `N` MLAs and the `N * 4` accumulators stay /// in registers. The 6-wide instantiation is the main scan tile; 4 and 1 cover /// the remainders of the 11/13/31-rotation schedules (6+4+1, 6+6+1, 6x5+1). -/// -/// The compiler fully unrolls the `N`-loops: on rustc 1.94/1.95 with -/// `target-cpu=neoverse-v2`, `N = 6` compiles to the same 37-instruction hot -/// loop (24 mla, 10 loads, no spills) as the previously hand-unrolled kernel, -/// and measures identically on r8g.24xlarge in kernel and full-server -/// benchmarks. The scalar-parity test below guards the correctness of every -/// width; keeping all widths in one function keeps their schedules identical -/// by construction. #[cfg(target_arch = "aarch64")] #[inline(always)] fn dot_product_nx4_u16( From 62548dfcad45354b359f62cd469a9dc0306b8e4b Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Mon, 24 Aug 2026 14:02:11 +0200 Subject: [PATCH 13/14] Surface full-rotation scan targets missing from the resident store A missing exact-version target receives the max-distance sentinel and silently cannot match. The linear scan builds its candidate lists from the registry the resident store is updated in lockstep with, so this is a desync signal, not an expected state: count every occurrence in linear_scan_missing_resident_targets_total and log the first loudly. Counting happens once per kernel dispatch over the already-materialized target list, so it costs nothing next to the scan itself. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- .../src/execution/hawk_main/iris_worker.rs | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs b/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs index 53ff8c4612..cdcb676a82 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/iris_worker.rs @@ -807,6 +807,32 @@ fn worker_thread( } } +/// A full-rotation scan target that is absent from the resident store means +/// the caller's exact `(serial, version)` is not resident. The kernels +/// substitute the max-distance sentinel, so such a record silently cannot +/// match. The linear scan builds its candidate lists from the registry that +/// the resident store is updated in lockstep with, so a nonzero count here is +/// a registry/store desync, not an expected state: count every occurrence and +/// log the first loudly. +fn record_missing_resident_targets(missing: usize, total: usize) { + use std::sync::atomic::{AtomicBool, Ordering}; + + if missing == 0 { + return; + } + metrics::counter!("linear_scan_missing_resident_targets_total").increment(missing as u64); + static WARNED: AtomicBool = AtomicBool::new(false); + if !WARNED.swap(true, Ordering::Relaxed) { + tracing::warn!( + missing, + total, + "full-rotation scan targets are missing from the resident store; these records \ + receive the max-distance sentinel and cannot match (further occurrences are \ + counted in linear_scan_missing_resident_targets_total)" + ); + } +} + /// Full 31-rotation distances against resident targets, dispatching on the /// resident representation: mixed-plane targets use the UMMLA kernel, /// u16 targets the MLA kernel. Results are bit-identical between the two. @@ -815,6 +841,10 @@ where I: Iterator> + ExactSizeIterator, { let residents: Vec> = targets.collect(); + record_missing_resident_targets( + residents.iter().filter(|target| target.is_none()).count(), + residents.len(), + ); #[cfg(target_arch = "aarch64")] { @@ -855,6 +885,10 @@ where I: Iterator> + ExactSizeIterator, { let residents: Vec> = targets.collect(); + record_missing_resident_targets( + residents.iter().filter(|target| target.is_none()).count(), + residents.len(), + ); #[cfg(target_arch = "aarch64")] { From 2861087283d46b290e4cf92f4147f8c484ed3f1f Mon Sep 17 00:00:00 2001 From: Philipp Sippl Date: Mon, 24 Aug 2026 14:37:38 +0200 Subject: [PATCH 14/14] Gate the rotation-amount const assert to aarch64 Its only consumer is the aarch64-only mixed-plane kernel; on other targets clippy correctly flagged it as dead code. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XAbAnQyfsSrWH6UGcRM1FF --- iris-mpc-cpu/src/protocol/ops.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/iris-mpc-cpu/src/protocol/ops.rs b/iris-mpc-cpu/src/protocol/ops.rs index 73508ceb5b..72e5fa5367 100644 --- a/iris-mpc-cpu/src/protocol/ops.rs +++ b/iris-mpc-cpu/src/protocol/ops.rs @@ -192,6 +192,9 @@ impl PrerotatedQueryRowMajorView<'_, ROTATIONS> { /// The doubled-row mixed-plane kernel addresses each rotation as an /// 8-byte-aligned window plus a 4-element phase /// (`mixed_scan::query_window`), which is only sound under this property. + /// Its sole consumer is that aarch64-only kernel, so gate it to keep + /// non-aarch64 builds free of dead code. + #[cfg(target_arch = "aarch64")] const ROTATION_AMOUNTS_ARE_MULTIPLES_OF_FOUR: () = { let mut i = 0; while i < ROTATIONS {