diff --git a/.github/workflows/build-and-push-linear-scan-server.yaml b/.github/workflows/build-and-push-linear-scan-server.yaml index d56a83560..d7a1583ac 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/Cargo.lock b/Cargo.lock index 642eaab06..b46804c2a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6269,6 +6269,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 31fe30a35..a1fd5c821 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 @@ -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 31c91e319..d0fdce4ac 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 000000000..3f5c2703b --- /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 000000000..146312045 --- /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 c168fec47..c4f6807b3 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(Self::resident_layout_for(search_mode)), ); let graph = [(); 2].map(|_| GraphMem::new()); Self::from_cli_with_initializer_and_graph( @@ -729,7 +730,8 @@ impl HawkActor { HAWK_DISTANCE_MODE, args.numa, iris_store, - ), + ) + .with_resident_layout(Self::resident_layout_for(search_mode)), ); Self::from_cli_with_initializer_and_graph( args, @@ -760,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, @@ -2076,6 +2091,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; @@ -2178,8 +2282,15 @@ 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 @@ -2241,7 +2352,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) @@ -2311,9 +2424,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, @@ -2321,7 +2449,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 4b2f8681a..cdcb676a8 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,121 @@ fn worker_thread(ch: Receiver, iris_store: SharedIrisesRef, n } } +/// 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. +fn full_rotation_distance_resident<'a, I>(query: &ArcIris, targets: I) -> Vec> +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")] + { + 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(); + record_missing_resident_targets( + residents.iter().filter(|target| target.is_none()).count(), + residents.len(), + ); + + #[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 +1034,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 @@ -910,6 +1148,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,17 +1237,20 @@ 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. 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)] @@ -1355,6 +1603,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, @@ -1452,23 +1717,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") @@ -1480,7 +1728,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 { @@ -1488,10 +1737,11 @@ 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, + windowed_ops_on_mixed_residents: false, } } @@ -1502,7 +1752,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, @@ -1564,34 +1815,37 @@ 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, 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( - 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> { @@ -1599,7 +1853,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()); }; @@ -1764,6 +2018,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 { @@ -1805,6 +2061,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()); @@ -1832,7 +2093,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(); @@ -1859,6 +2120,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 }) } @@ -1920,6 +2230,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<_> = { @@ -1946,11 +2257,18 @@ 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 { + for (vector_id, iris) in resident { store.insert(vector_id, iris); } Ok(store.set_hash.checksum()) @@ -2038,6 +2356,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 { @@ -2050,9 +2369,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(()) }) @@ -2261,26 +2581,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 9fec04e4d..aeb07c0f0 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 0d111a469..4a36069b5 100644 --- a/iris-mpc-cpu/src/execution/hawk_main/search.rs +++ b/iris-mpc-cpu/src/execution/hawk_main/search.rs @@ -10,17 +10,26 @@ 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::collections::HashSet; -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 @@ -48,6 +57,60 @@ struct LinearScanChunk { range: std::ops::Range, } +#[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)] +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, @@ -91,11 +154,6 @@ const fn orientation_label(orientation: Orientation) -> &'static str { } } -#[derive(Clone)] -struct LinearScanPrefetch { - worker: Arc, -} - pub type SearchQueries = Arc>>>; pub type SearchResults = @@ -265,6 +323,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, @@ -277,16 +347,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 second-stage candidates regardless of - // the first-eye result, so their database reads can start before the - // resident-eye scan. Reservations are single-use and idempotent: a record - // hinted again by a first-eye chunk or by the other orientation is not - // reserved twice, and leftovers are released when the batch completes. - 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, @@ -297,11 +364,38 @@ 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(), + excluded_ids: Arc::new( + known_second_stage_ids + .iter() + .map(|ids| ids.to_vec()) + .collect(), + ), + }), + 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(); @@ -321,43 +415,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(); @@ -467,6 +572,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, @@ -474,7 +584,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); @@ -531,7 +641,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; @@ -547,9 +657,9 @@ 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( + let mut prefetch_ids = collect_live_second_stage_ids( chunk_ids, result .classified @@ -559,8 +669,15 @@ 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 { + progress.record(chunk.range.len()); + } results.push((chunk.i_request, chunk.i_chunk, result)); } Ok(results) @@ -576,7 +693,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)| { @@ -601,16 +838,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()) @@ -626,9 +887,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); @@ -638,7 +898,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, @@ -652,7 +912,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, @@ -666,7 +926,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, @@ -680,7 +940,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, @@ -699,66 +959,622 @@ 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" ); - - Ok(results) } -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); - target.saturated |= source.saturated; +/// 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, + }); + } } - merge_matches(&mut target.classified.matches, source.classified.matches); - merge_matches( - &mut target.classified.anon_stats_matches, - source.classified.anon_stats_matches, - ); - match ( - &mut target.classified.pre_extension, - source.classified.pre_extension, - ) { - (Some(target), Some(source)) => merge_matches(target, source), - (target @ None, source @ Some(_)) => *target = source, - _ => {} + 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); } - target.classified.linear_scan_supermatch_threshold = target - .classified - .linear_scan_supermatch_threshold - .or(source.classified.linear_scan_supermatch_threshold); - target - .classified - .partial_match_rotations - .append(&mut source.classified.partial_match_rotations); -} + 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); -#[instrument(level = "trace", target = "searcher::network", skip_all)] -async fn per_session( - session: &HawkSession, - search_queries: &SearchQueries, - search_ids: &SearchIds, - search_params: &SearchParams, - tx: UnboundedSender<(TaskId, HawkInsertPlan)>, - batch: Batch, -) -> Result<()> { - let inner = async { - // Linear scan does not build graph links for identity updates. The + 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; + // 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()); + // 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; + // 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)?); + } + } + 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 mut prefetch_ids = collect_live_second_stage_ids( + chunk_ids, + plans.iter().flat_map(|plan| { + plan.classified + .anon_stats_matches + .results + .iter() + .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 { + 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. 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; + 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; + 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::>(), + ) + }); + 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!( + 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(), + excluded_ids: prefetch_excluded_ids, + }), + 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); + target.saturated |= source.saturated; + } + + merge_matches(&mut target.classified.matches, source.classified.matches); + merge_matches( + &mut target.classified.anon_stats_matches, + source.classified.anon_stats_matches, + ); + match ( + &mut target.classified.pre_extension, + source.classified.pre_extension, + ) { + (Some(target), Some(source)) => merge_matches(target, source), + (target @ None, source @ Some(_)) => *target = source, + _ => {} + } + target.classified.linear_scan_supermatch_threshold = target + .classified + .linear_scan_supermatch_threshold + .or(source.classified.linear_scan_supermatch_threshold); + target + .classified + .partial_match_rotations + .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, + search_queries: &SearchQueries, + search_ids: &SearchIds, + search_params: &SearchParams, + tx: UnboundedSender<(TaskId, HawkInsertPlan)>, + 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 { @@ -845,47 +1661,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))), + extend_classified_from_thresholds( + &mut classified, + ids, + thresholds, + search_params.return_partial_results, ); - classified.anon_stats_matches.results.extend( - thresholds - .anon_stats_matches - .into_iter() - .map(|(vector, _rotation, distance)| (ids[vector], distance)), - ); - - 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) - }) - }), - ); - } } // The CUDA actor compares its per-eye match counters against @@ -916,6 +1718,125 @@ 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 @@ -1210,22 +2131,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); @@ -1275,6 +2180,43 @@ 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 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(()) + } + #[tokio::test] async fn test_search() -> Result<()> { let actors = setup_hawk_actors().await?; @@ -1313,7 +2255,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, @@ -1328,6 +2283,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 @@ -1375,6 +2344,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 527d5993c..34bdd51a1 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 { @@ -245,6 +261,7 @@ impl WorkerPoolInitializer for LocalWorkerPoolInitializer { LocalIrisWorkerPool::new( workers_handle[side].clone(), iris_stores[side].clone(), + layout, distance_mode, party_id, ) @@ -349,9 +366,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 7096338d7..ae506dc05 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 28a676021..ab4ac3867 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, @@ -47,6 +49,7 @@ use std::{ sync::Arc, vec, }; +pub use tokio_util::task::AbortOnDropHandle; use tracing::instrument; mod distance_fn; @@ -68,6 +71,9 @@ 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]; + /// GPU-equivalent exact-scan classification for one chunk. Thresholds are /// evaluated directly for all 31 rotations; no secret minimum is computed. #[derive(Debug)] @@ -585,12 +591,73 @@ where .await } + /// Full-rotation exact scan with the opened per-rotation match metadata + /// 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, + query: &Aby3Query, + vectors: &[VectorId], + ) -> Result<(Vec>, RotationMatchIndices)> { + if vectors.is_empty() { + return Ok((Vec::new(), Vec::new())); + } + let rotation_distances = self.full_rotation_distances(query, vectors).await?; + let distances = self + .oblivious_min_distance_batch(distance_fn::transpose_from_flat_with_rotations( + &rotation_distances, + ROTATIONS, + )) + .await?; + let rotation_match_bits = + D::lte_and_open(&mut self.session, &rotation_distances, Threshold::Match).await?; + let rotation_matches = rotation_match_bits + .chunks(ROTATIONS) + .map(|bits| { + bits.iter() + .enumerate() + .filter_map(|(rotation, &is_match)| is_match.then_some(rotation)) + .collect() + }) + .collect(); + Ok((distances, rotation_matches)) + } + + #[cfg(test)] + async fn full_rotation_distances( + &mut self, + query: &Aby3Query, + vectors: &[VectorId], + ) -> Result>> { + if vectors.is_empty() { + return Ok(Vec::new()); + } + let dot_shares = self.full_rotation_dot_shares(query, vectors).await?; + 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>> { // This scan is neither the simple nor the min-rotation distance: it // opens a threshold for each of the 31 rotations separately. What it // does rely on is the Hawk query layout, where the cached query @@ -602,11 +669,47 @@ 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. + /// + /// 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], + ) -> Result>> { + // See `full_rotation_dot_contributions`: only the center-rotation + // query layout matters, not the configured 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(); + 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. @@ -624,15 +727,14 @@ where } impl Aby3Store { - /// 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. + /// 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_with_forced_anon_stats( + pub async fn eval_distance_batch_full_rotation_thresholds( &mut self, query: &Aby3Query, vectors: &[VectorId], - forced_anon_stats_vectors: &[usize], ) -> Result { if vectors.is_empty() { return Ok(FullRotationThresholdResult { @@ -662,7 +764,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)) @@ -673,17 +775,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 @@ -778,6 +869,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 98a8eb13b..21c358494 100644 --- a/iris-mpc-cpu/src/hawkers/aby3/aby3_store/tests.rs +++ b/iris-mpc-cpu/src/hawkers/aby3/aby3_store/tests.rs @@ -49,6 +49,188 @@ fn gpu_prefilter_expands_candidate_records_to_all_rotations() { assert_eq!(expanded, expected); } +#[tokio::test(flavor = "multi_thread")] +async fn full_rotation_threshold_scan_matches_min_distance_reference() -> Result<()> { + let mut rng = AesRng::seed_from_u64(0x7468_7265_7368_6f6c); + let vectors_and_graphs = shared_random_setup(&mut rng, 2, 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)]; + let query = store.cache_query_from_store(&ids[0]).await?; + + let direct = store + .eval_distance_batch_full_rotation_thresholds(&query, &ids) + .await?; + let (minimums, reference_rotations) = store + .eval_distance_batch_full_rotations_with_rotation_matches(&query, &ids) + .await?; + let reference_anon = store.is_match_at(&minimums, Threshold::AnonStats).await?; + let reference_match = store.is_match_at(&minimums, Threshold::Match).await?; + + let direct_anon = (0..ids.len()) + .map(|vector| { + direct + .anon_stats_matches + .iter() + .any(|(matched_vector, _, _)| *matched_vector == vector) + }) + .collect::>(); + let direct_match = direct + .matches + .iter() + .map(Option::is_some) + .collect::>(); + assert_eq!(direct_anon, reference_anon); + assert_eq!(direct_match, reference_match); + assert_eq!(direct.match_rotations, reference_rotations); + Ok(()) + }); + + parallelize(tasks).await?; + 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 f8a45bc89..1e88a394b 100644 --- a/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs +++ b/iris-mpc-cpu/src/hawkers/aby3/test_utils.rs @@ -76,11 +76,28 @@ 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 workers: Arc = Arc::new(LocalIrisWorkerPool::new_local( - storage.clone(), - plain_store.distance_mode, - party_id, - )); + 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(); + // 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, @@ -99,12 +116,21 @@ pub async fn setup_local_store_aby3_players(network_t: NetworkType) -> Result::new_storage(None).to_arc(); - let workers: Arc = Arc::new(LocalIrisWorkerPool::new_local( - storage.clone(), - TEST_DISTANCE_MODE, - party_id, - )); + 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, + ) + .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/src/hawkers/shared_irises.rs b/iris-mpc-cpu/src/hawkers/shared_irises.rs index 9d55124ff..b40ba1c32 100644 --- a/iris-mpc-cpu/src/hawkers/shared_irises.rs +++ b/iris-mpc-cpu/src/hawkers/shared_irises.rs @@ -216,6 +216,25 @@ 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, + // Ids and versions are unchanged, so the cached list stays valid. + live_ids: self.live_ids, + } + } + /// 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 6e76c2bbf..72e5fa536 100644 --- a/iris-mpc-cpu/src/protocol/ops.rs +++ b/iris-mpc-cpu/src/protocol/ops.rs @@ -188,6 +188,24 @@ 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. + /// 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 { + 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) { @@ -458,23 +476,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 { @@ -490,7 +549,7 @@ fn accumulate_component_tiled_6x4( &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], ); @@ -519,7 +578,7 @@ fn accumulate_component_tiled_6x4( &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], ); @@ -541,10 +600,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_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 + result_lane; @@ -553,40 +611,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,99 +677,19 @@ fn accumulate_scalar_target( } } +/// `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). #[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), - ], - ] - } -} - -#[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)); @@ -710,115 +699,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 } } @@ -839,6 +747,998 @@ 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 { + // 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; + 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::*; @@ -1076,26 +1976,164 @@ 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); + } + } + + #[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"); + } } } diff --git a/iris-mpc-cpu/src/protocol/shared_iris.rs b/iris-mpc-cpu/src/protocol/shared_iris.rs index 65792e3e8..37aa30355 100644 --- a/iris-mpc-cpu/src/protocol/shared_iris.rs +++ b/iris-mpc-cpu/src/protocol/shared_iris.rs @@ -126,3 +126,151 @@ 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 { + /// 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. + 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 { + 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 { + share_id: iris.code.id, + code, + mask, + } + } + + pub fn to_iris(&self) -> GaloisRingSharedIris { + let mut iris = GaloisRingSharedIris::default_for_party(0); + 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 + } + + #[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 6a65c18c1..9d2ed5ffa 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 { @@ -155,11 +165,19 @@ async fn cold_eye_luc_window_rolls_forward_and_survives_persistence_ack() -> Res .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 { @@ -226,11 +244,19 @@ async fn cold_eye_version_miss_fails_foreground_fetch_after_prefetch() -> Result .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 { diff --git a/iris-mpc-cpu/tests/full_rotation_dot.rs b/iris-mpc-cpu/tests/full_rotation_dot.rs index b59cb861b..20167d312 100644 --- a/iris-mpc-cpu/tests/full_rotation_dot.rs +++ b/iris-mpc-cpu/tests/full_rotation_dot.rs @@ -46,14 +46,29 @@ 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); + // 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-gpu/Cargo.toml b/iris-mpc-gpu/Cargo.toml index ab9aa0358..078d3200f 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 000000000..c95471834 --- /dev/null +++ b/iris-mpc-gpu/benches/linear_scan_dot.rs @@ -0,0 +1,562 @@ +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::{preferred_scan_layout, ArcIris, GaloisRingSharedIris, ResidentIris}, +}; +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)"); + // 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(), + 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(ResidentIris::from_arc(Arc::new((**query).clone()), layout))); + } + let store = store.to_arc(); + ( + LocalIrisWorkerPool::new_local(store, layout, 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/src/server/mod.rs b/iris-mpc/src/server/mod.rs index 1baa18931..5a8d0a0af 100644 --- a/iris-mpc/src/server/mod.rs +++ b/iris-mpc/src/server/mod.rs @@ -834,8 +834,11 @@ 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 = HawkActor::resident_layout_for(search_mode); + let initializer: Box = Box::new( + LocalWorkerPoolInitializer::new_load_from_db( hawk_args.party_index, HAWK_DISTANCE_MODE, hawk_args.numa, @@ -853,7 +856,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(); diff --git a/scripts/analyze-linear-scan-server-benchmark.py b/scripts/analyze-linear-scan-server-benchmark.py new file mode 100755 index 000000000..b0a8218a0 --- /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 000000000..426e5e331 --- /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 4a4878c7f..82347453d 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 000000000..927b28ef1 --- /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 000000000..1d76a49b8 --- /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 a54d3c7b3..568f19b7d 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 000000000..09f3ea701 --- /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