From 90fc533dc2314cac2c269df8fc0bb06f01a2e664 Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Thu, 25 Jun 2026 00:30:51 +0530 Subject: [PATCH 1/5] Add staged (tiled_reduce_staged) multi-vector MaxSim kernel POC Introduce an experimental, generic cache-tiled reduction driver, `tiled_reduce_staged`, that computes multi-vector MaxSim/Chamfer for any element type and quantization by swapping four pluggable stages (StagedKernel, Postprocess, Reducer, StagedConvert) instead of forking the tiled loop nest per datatype. Validated with two instantiations: - f32: bit-identical to and on par with the hand-fused V3 kernel (selectable for A/B as MaxSimIsa::X86_64_V3_Staged). - 4-bit MinMax-quantized i8: a new datatype added by swapping Stage A + the Acc type + Stage B only; 1.5-4.1x over the per-pair SIMD reference. The driver owns all scratch, allocated from a caller-supplied ScopedAllocator; zero-allocation steady state is provided by a single-owner resettable bump arena (ResettableArena). A crisp design overview lives in the staged module README (diskann-quantization/.../kernels/staged/README.md). Also adds benchmark examples (multi-vector-{staged,quant,3way}.json) and the quantized multi-vector benchmark backend. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../example/multi-vector-3way.json | 47 + .../example/multi-vector-quant.json | 20 + .../example/multi-vector-staged.json | 41 + diskann-benchmark/src/inputs/multi_vector.rs | 82 ++ diskann-benchmark/src/multi_vector/mod.rs | 8 +- diskann-benchmark/src/multi_vector/quant.rs | 268 +++++ .../src/multi_vector/distance/factory.rs | 152 ++- .../src/multi_vector/distance/isa.rs | 9 +- .../src/multi_vector/distance/kernel.rs | 8 +- .../src/multi_vector/distance/kernels/mod.rs | 8 + .../distance/kernels/staged/README.md | 91 ++ .../distance/kernels/staged/arena.rs | 122 +++ .../distance/kernels/staged/driver.rs | 227 +++++ .../distance/kernels/staged/i8.rs | 962 ++++++++++++++++++ .../distance/kernels/staged/maxsim.rs | 155 +++ .../distance/kernels/staged/mod.rs | 428 ++++++++ .../distance/kernels/staged/v3.rs | 395 +++++++ .../src/multi_vector/distance/mod.rs | 5 + 18 files changed, 3021 insertions(+), 7 deletions(-) create mode 100644 diskann-benchmark/example/multi-vector-3way.json create mode 100644 diskann-benchmark/example/multi-vector-quant.json create mode 100644 diskann-benchmark/example/multi-vector-staged.json create mode 100644 diskann-benchmark/src/multi_vector/quant.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/staged/README.md create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/staged/arena.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/staged/driver.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/staged/i8.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/staged/maxsim.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/staged/mod.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/staged/v3.rs diff --git a/diskann-benchmark/example/multi-vector-3way.json b/diskann-benchmark/example/multi-vector-3way.json new file mode 100644 index 000000000..3a036d60e --- /dev/null +++ b/diskann-benchmark/example/multi-vector-3way.json @@ -0,0 +1,47 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-op", + "content": { + "element_type": "float32", + "isa": "reference", + "runs": [ + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 16, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 64, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 128, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 256, "loops_per_measurement": 25, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 512, "loops_per_measurement": 12, "num_measurements": 50 } + ] + } + }, + { + "type": "multi-vector-op", + "content": { + "element_type": "float32", + "isa": "x86-64-v3", + "runs": [ + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 16, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 64, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 128, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 256, "loops_per_measurement": 25, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 512, "loops_per_measurement": 12, "num_measurements": 50 } + ] + } + }, + { + "type": "multi-vector-op", + "content": { + "element_type": "float32", + "isa": "x86-64-v3-staged", + "runs": [ + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 16, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 64, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 128, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 256, "loops_per_measurement": 25, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 256, "dim": 512, "loops_per_measurement": 12, "num_measurements": 50 } + ] + } + } + ] +} diff --git a/diskann-benchmark/example/multi-vector-quant.json b/diskann-benchmark/example/multi-vector-quant.json new file mode 100644 index 000000000..ffc4131ac --- /dev/null +++ b/diskann-benchmark/example/multi-vector-quant.json @@ -0,0 +1,20 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-quant-op", + "content": { + "runs": [ + { "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 500, "num_measurements": 50 }, + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] + } + } + ] +} diff --git a/diskann-benchmark/example/multi-vector-staged.json b/diskann-benchmark/example/multi-vector-staged.json new file mode 100644 index 000000000..8cb5798dc --- /dev/null +++ b/diskann-benchmark/example/multi-vector-staged.json @@ -0,0 +1,41 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-op", + "content": { + "element_type": "float32", + "isa": "x86-64-v3", + "runs": [ + { "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 500, "num_measurements": 50 }, + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 264, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] + } + }, + { + "type": "multi-vector-op", + "content": { + "element_type": "float32", + "isa": "x86-64-v3-staged", + "runs": [ + { "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 500, "num_measurements": 50 }, + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 264, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] + } + } + ] +} diff --git a/diskann-benchmark/src/inputs/multi_vector.rs b/diskann-benchmark/src/inputs/multi_vector.rs index c74f9d232..5b2aaa833 100644 --- a/diskann-benchmark/src/inputs/multi_vector.rs +++ b/diskann-benchmark/src/inputs/multi_vector.rs @@ -26,6 +26,9 @@ pub(crate) enum BenchIsa { #[serde(rename = "x86-64-v3")] #[allow(non_camel_case_types)] X86_64_V3, + #[serde(rename = "x86-64-v3-staged")] + #[allow(non_camel_case_types)] + X86_64_V3_Staged, Neon, Scalar, Reference, @@ -37,6 +40,7 @@ impl std::fmt::Display for BenchIsa { let st = match self { Self::X86_64_V4 => "x86-64-v4", Self::X86_64_V3 => "x86-64-v3", + Self::X86_64_V3_Staged => "x86-64-v3-staged", Self::Neon => "neon", Self::Scalar => "scalar", Self::Reference => "reference", @@ -51,6 +55,7 @@ impl From for MaxSimIsa { match b { BenchIsa::X86_64_V4 => MaxSimIsa::X86_64_V4, BenchIsa::X86_64_V3 => MaxSimIsa::X86_64_V3, + BenchIsa::X86_64_V3_Staged => MaxSimIsa::X86_64_V3_Staged, BenchIsa::Neon => MaxSimIsa::Neon, BenchIsa::Scalar => MaxSimIsa::Scalar, BenchIsa::Reference => MaxSimIsa::Reference, @@ -149,3 +154,80 @@ impl std::fmt::Display for MultiVectorOp { Ok(()) } } + +/////////////////////////////// +// Multi-Vector Quantized Op // +/////////////////////////////// + +/// A 4-bit MinMax **quantized** multi-vector MaxSim A/B benchmark job: the +/// experimental staged integer kernel vs the scalar `MinMaxKernel` reference, +/// at identical shapes and quantization. +/// +/// The element type is implicitly f32 input → 4-bit MinMax codes, and the ISA is +/// fixed to V3/AVX2 (the only quantized staged kernel), so neither is a JSON +/// field. x86_64-only, like the kernel it drives. +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct MultiVectorQuantOp { + pub(crate) runs: Vec, +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl MultiVectorQuantOp { + pub(crate) const fn tag() -> &'static str { + "multi-vector-quant-op" + } +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl Input for MultiVectorQuantOp { + type Raw = Self; + + fn tag() -> &'static str { + Self::tag() + } + + fn from_raw(raw: Self::Raw, _checker: &mut Checker) -> anyhow::Result { + Ok(raw) + } + + fn serialize(&self) -> anyhow::Result { + Ok(serde_json::to_value(self)?) + } + + fn example() -> Self { + const NUM_DOC_VECTORS: NonZeroUsize = NonZeroUsize::new(64).unwrap(); + const DIM: NonZeroUsize = NonZeroUsize::new(128).unwrap(); + const LOOPS_PER_MEASUREMENT: NonZeroUsize = NonZeroUsize::new(50).unwrap(); + const NUM_MEASUREMENTS: NonZeroUsize = NonZeroUsize::new(20).unwrap(); + + let runs = vec![ + Run { + num_query_vectors: NonZeroUsize::new(32).unwrap(), + num_doc_vectors: NUM_DOC_VECTORS, + dim: DIM, + loops_per_measurement: LOOPS_PER_MEASUREMENT, + num_measurements: NUM_MEASUREMENTS, + }, + Run { + num_query_vectors: NonZeroUsize::new(64).unwrap(), + num_doc_vectors: NUM_DOC_VECTORS, + dim: DIM, + loops_per_measurement: LOOPS_PER_MEASUREMENT, + num_measurements: NUM_MEASUREMENTS, + }, + ]; + + Self { runs } + } +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl std::fmt::Display for MultiVectorQuantOp { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "Multi-Vector Quantized Operation (4-bit MinMax)\n")?; + write_field!(f, "tag", Self::tag())?; + write_field!(f, "number of runs", self.runs.len())?; + Ok(()) + } +} diff --git a/diskann-benchmark/src/multi_vector/mod.rs b/diskann-benchmark/src/multi_vector/mod.rs index dfad330af..a01285ba1 100644 --- a/diskann-benchmark/src/multi_vector/mod.rs +++ b/diskann-benchmark/src/multi_vector/mod.rs @@ -25,9 +25,15 @@ cfg_if::cfg_if! { if #[cfg(feature = "multi-vector")] { mod driver; mod kernels; + // The quantized A/B op drives the V3-only staged integer kernel. + #[cfg(target_arch = "x86_64")] + mod quant; pub(super) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> { - kernels::register(registry) + kernels::register(registry)?; + #[cfg(target_arch = "x86_64")] + quant::register(registry)?; + Ok(()) } } else { crate::utils::stub_impl!("multi-vector", inputs::multi_vector::MultiVectorOp); diff --git a/diskann-benchmark/src/multi_vector/quant.rs b/diskann-benchmark/src/multi_vector/quant.rs new file mode 100644 index 000000000..1d09ca88b --- /dev/null +++ b/diskann-benchmark/src/multi_vector/quant.rs @@ -0,0 +1,268 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! A/B benchmark for **4-bit MinMax quantized** multi-vector MaxSim: the +//! experimental *staged integer* kernel (block-transposed `i16` query + `u8` +//! doc codes, `vpmaddwd` accumulation, metadata postprocess) vs the scalar +//! [`MinMaxKernel`] reference — at identical shapes and identical quantization. +//! +//! Both paths consume the *same* random f32 multi-vectors quantized to 4-bit +//! MinMax (Null transform, scale 1.0), so the comparison isolates the distance +//! kernel. The build / quantize cost is excluded from the timing. +//! +//! x86_64 (V3/AVX2) only — the quantized staged kernel has no other backend. + +use std::io::Write; +use std::num::NonZeroUsize; + +use diskann_benchmark_runner::{ + benchmark::{FailureScore, MatchScore}, + utils::{fmt::Table, percentiles, MicroSeconds}, + Benchmark, Checkpoint, Output, Registry, +}; +use diskann_quantization::algorithms::transforms::NullTransform; +use diskann_quantization::algorithms::Transform; +use diskann_quantization::minmax::{MinMaxMeta, MinMaxQuantizer}; +use diskann_quantization::multi_vector::distance::{QuantStagedDocs, QuantStagedQuery}; +use diskann_quantization::multi_vector::{Defaulted, Mat, MatRef, MaxSim, QueryMatRef, Standard}; +use diskann_quantization::num::Positive; +use diskann_quantization::CompressInto; +use diskann_utils::ReborrowMut; +use diskann_vector::DistanceFunctionMut; +use serde::{Deserialize, Serialize}; + +use super::driver::Data; +use crate::inputs::multi_vector::{MultiVectorQuantOp, Run}; +use crate::utils::DisplayWrapper; + +// ───────────────────────────────────────────────────────────────────────── +// Kernel. +// ───────────────────────────────────────────────────────────────────────── + +#[derive(Debug)] +pub(super) struct QuantKernel; + +impl QuantKernel { + pub(super) const fn new() -> Self { + Self + } +} + +impl Benchmark for QuantKernel { + type Input = MultiVectorQuantOp; + type Output = Vec; + + fn try_match(&self, _from: &MultiVectorQuantOp) -> Result { + // The staged integer kernel requires AVX2 (V3). + if QuantStagedQuery::is_supported() { + Ok(MatchScore(0)) + } else { + Err(FailureScore(0)) + } + } + + fn run( + &self, + input: &MultiVectorQuantOp, + _: Checkpoint<'_>, + mut output: &mut dyn Output, + ) -> anyhow::Result { + writeln!(output, "{}", input)?; + let mut results = Vec::with_capacity(input.runs.len()); + for run in input.runs.iter() { + results.push(run_ab(run)?); + } + writeln!(output, "\n\n{}", DisplayWrapper(&*results))?; + Ok(results) + } + + fn description( + &self, + f: &mut std::fmt::Formatter<'_>, + input: Option<&MultiVectorQuantOp>, + ) -> std::fmt::Result { + match input { + None => writeln!(f, "- 4-bit MinMax quantized staged MaxSim (V3/AVX2)")?, + Some(_) => { + if !QuantStagedQuery::is_supported() { + writeln!(f, "\n - AVX2 (V3) unavailable on this CPU")?; + } + } + } + Ok(()) + } +} + +// ───────────────────────────────────────────────────────────────────────── +// A/B timing. +// ───────────────────────────────────────────────────────────────────────── + +/// Quantize an f32 multi-vector to 4-bit MinMax (Null transform, scale 1.0) — +/// the same quantizer both paths share so the codes + metadata are identical. +fn quantize(input: MatRef<'_, Standard>) -> Mat> { + let dim = input.vector_dim(); + let n = input.num_vectors(); + let q = MinMaxQuantizer::new( + Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())), + Positive::new(1.0).unwrap(), + ); + let mut out: Mat> = Mat::new(MinMaxMeta::new(n, dim), Defaulted).unwrap(); + q.compress_into(input, out.reborrow_mut()).unwrap(); + out +} + +/// Run `f` `loops_per_measurement` times per measurement, `num_measurements` +/// times, returning the per-measurement latencies and their percentiles. +fn measure(run: &Run, mut f: impl FnMut()) -> Series { + let mut latencies = Vec::with_capacity(run.num_measurements.get()); + for _ in 0..run.num_measurements.get() { + let start = std::time::Instant::now(); + for _ in 0..run.loops_per_measurement.get() { + f(); + } + latencies.push(start.elapsed().into()); + } + let percentiles = percentiles::compute_percentiles(&mut latencies).unwrap(); + Series { + latencies, + percentiles, + } +} + +/// Build both kernels for one shape and time them (build / quantize excluded). +fn run_ab(run: &Run) -> anyhow::Result { + let data = Data::::new(run)?; + + // Path A — staged integer kernel (quantizes internally at build time). + let mut query = QuantStagedQuery::build(data.queries.as_view()) + .ok_or_else(|| anyhow::anyhow!("AVX2 (V3) unavailable for the staged quantized kernel"))?; + let docs = QuantStagedDocs::build(data.docs.as_view()); + + // Path B — scalar MinMax reference over the same quantization. + let q_ref = quantize(data.queries.as_view()); + let d_ref = quantize(data.docs.as_view()); + + let nq = run.num_query_vectors.get(); + let mut scores = vec![0.0f32; nq]; + + // Launder BOTH the inputs and the output through `black_box` each iteration. + // Output-only `black_box` is not enough: the reference chain is `#[inline(always)]` + // end-to-end with loop-invariant inputs, so the optimizer could hoist/elide it out + // of the measured loop (the staged path is an opaque cross-crate call and cannot be), + // making the A/B asymmetric. Laundering the inputs forces both paths to re-run the + // full per-call work every iteration. + let staged = measure(run, || { + let docs = std::hint::black_box(&docs); + query.compute_max_sim(docs, &mut scores); + std::hint::black_box(&mut scores); + }); + + let reference = measure(run, || { + let q_ref = std::hint::black_box(&q_ref); + let d_ref = std::hint::black_box(&d_ref); + let query_ref: QueryMatRef<_> = q_ref.as_view().into(); + MaxSim::new(&mut scores).evaluate(query_ref, d_ref.as_view()); + std::hint::black_box(&mut scores); + }); + + Ok(QuantRunResult { + run: run.clone(), + staged, + reference, + }) +} + +// ───────────────────────────────────────────────────────────────────────── +// Result types. +// ───────────────────────────────────────────────────────────────────────── + +/// One timed series (per-measurement latencies + percentiles). +#[derive(Debug, Serialize, Deserialize)] +pub(super) struct Series { + latencies: Vec, + percentiles: percentiles::Percentiles, +} + +impl Series { + /// Minimum latency, in microseconds. + fn min_us(&self) -> f64 { + self.latencies + .iter() + .min() + .copied() + .unwrap_or(MicroSeconds::new(u64::MAX)) + .as_f64() + } +} + +/// Staged-vs-reference result for one shape. +#[derive(Debug, Serialize, Deserialize)] +pub(super) struct QuantRunResult { + pub(super) run: Run, + pub(super) staged: Series, + pub(super) reference: Series, +} + +impl QuantRunResult { + fn computations(&self) -> f64 { + (self.run.num_query_vectors.get() + * self.run.num_doc_vectors.get() + * self.run.loops_per_measurement.get()) as f64 + } +} + +impl std::fmt::Display for DisplayWrapper<'_, [QuantRunResult]> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.is_empty() { + return Ok(()); + } + + writeln!( + f, + "ns/IP = min time per (query, doc) inner-product call; \ + Speedup = reference / staged (>1 ⇒ staged faster)" + )?; + + let header = [ + "Q", + "D", + "Dim", + "Staged (ns/IP)", + "Reference (ns/IP)", + "Speedup", + ]; + let mut table = Table::new(header, self.len()); + + self.iter().enumerate().for_each(|(row, r)| { + let comps = r.computations(); + let staged = r.staged.min_us() / comps * 1000.0; + let reference = r.reference.min_us() / comps * 1000.0; + let speedup = if staged > 0.0 { + reference / staged + } else { + 0.0 + }; + + let mut row = table.row(row); + row.insert(r.run.num_query_vectors, 0); + row.insert(r.run.num_doc_vectors, 1); + row.insert(r.run.dim, 2); + row.insert(format!("{:.3}", staged), 3); + row.insert(format!("{:.3}", reference), 4); + row.insert(format!("{:.2}x", speedup), 5); + }); + + table.fmt(f) + } +} + +// ───────────────────────────────────────────────────────────────────────── +// Registration. +// ───────────────────────────────────────────────────────────────────────── + +pub(super) fn register(registry: &mut Registry) -> anyhow::Result<()> { + registry.register("multi-vector-quant-op", QuantKernel::new())?; + Ok(()) +} diff --git a/diskann-quantization/src/multi_vector/distance/factory.rs b/diskann-quantization/src/multi_vector/distance/factory.rs index 5dcd4b8cd..65a69691a 100644 --- a/diskann-quantization/src/multi_vector/distance/factory.rs +++ b/diskann-quantization/src/multi_vector/distance/factory.rs @@ -18,6 +18,8 @@ use super::isa::{MaxSimIsa, NotSupported}; use super::kernel::{Erase, MaxSimKernel}; use super::kernels::f16::F16Entry; use super::kernels::f32::F32Kernel; +#[cfg(target_arch = "x86_64")] +use super::kernels::staged::{F32StagedScratch, StagedF32Kernel, StagedRun}; use super::max_sim::{MaxSim, MaxSimError}; use crate::multi_vector::distance::QueryMatRef; use crate::multi_vector::{BlockTransposed, BlockTransposedRef, Mat, MatRef, Standard}; @@ -266,6 +268,92 @@ impl> } } +// ───────────────────────────────────────────────────────────────────────── +// Staged kernel (experimental) — selected by MaxSimIsa::X86_64_V3_Staged. +// Coexists with the fused `Prepared` path above for A/B benchmarking. +// ───────────────────────────────────────────────────────────────────────── + +/// Counterpart to [`Prepared`] for the staged f32 kernel. Owns a +/// [`F32StagedScratch`] (reset-arena + reused `state`) behind a `RefCell`: because +/// `compute_max_sim` takes `&self` and [`MaxSimKernel`] is `Send` but not `Sync`, +/// interior mutability lets each call reset+reuse the scratch with **zero heap +/// allocation** in steady state. +#[cfg(target_arch = "x86_64")] +#[derive(Debug)] +struct PreparedStaged { + arch: A, + prepared: BlockTransposed, + scratch: std::cell::RefCell, +} + +#[cfg(target_arch = "x86_64")] +impl MaxSimKernel for PreparedStaged +where + A: Architecture, + StagedF32Kernel: for<'a> diskann_wide::arch::Target3< + A, + (), + BlockTransposedRef<'a, f32, GROUP>, + MatRef<'a, Standard>, + StagedRun<'a>, + >, +{ + fn nrows(&self) -> usize { + self.prepared.nrows() + } + + fn compute_max_sim( + &self, + doc: MatRef<'_, Standard>, + scores: &mut [f32], + ) -> Result<(), MaxSimError> { + if scores.len() != self.nrows() { + return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows())); + } + if doc.num_vectors() == 0 { + scores.fill(f32::MAX); + return Ok(()); + } + let padded = self.prepared.padded_nrows(); + let nrows = self.prepared.nrows(); + // Reset + reuse the owned arena/state scratch (no per-call allocation). + self.scratch.borrow_mut().run(padded, |state, alloc| { + self.arch.run3( + StagedF32Kernel::, + self.prepared.reborrow(), + doc, + StagedRun { + state: &mut state[..], + alloc, + }, + ); + // Distance = negated max inner product (matches the fused path). + for (dst, &src) in scores.iter_mut().zip(state[..nrows].iter()) { + *dst = -src; + } + }); + Ok(()) + } +} + +#[cfg(target_arch = "x86_64")] +struct BuildAndEraseStaged(E); + +#[cfg(target_arch = "x86_64")] +impl> diskann_wide::arch::Target1>> + for BuildAndEraseStaged +{ + fn run(self, arch: V3, query: MatRef<'_, Standard>) -> E::Output { + let prepared = BlockTransposed::::from_matrix_view(query.as_matrix_view()); + let scratch = std::cell::RefCell::new(F32StagedScratch::new(prepared.padded_nrows())); + self.0.erase(PreparedStaged { + arch, + prepared, + scratch, + }) + } +} + // ───────────────────────────────────────────────────────────────────────── // MaxSimElement — sealed trait gating accepted element types. // ───────────────────────────────────────────────────────────────────────── @@ -325,11 +413,21 @@ impl MaxSimElement for f32 { })?; Ok(arch.run1(BuildAndErase(erase), query)) } + #[cfg(target_arch = "x86_64")] + MaxSimIsa::X86_64_V3_Staged => { + let arch = V3::new_checked().ok_or(NotSupported { + isa, + reason: "AVX2/FMA unavailable on this CPU", + })?; + Ok(arch.run1(BuildAndEraseStaged(erase), query)) + } #[cfg(not(target_arch = "x86_64"))] - MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => Err(NotSupported { - isa, - reason: "x86_64 target only", - }), + MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 | MaxSimIsa::X86_64_V3_Staged => { + Err(NotSupported { + isa, + reason: "x86_64 target only", + }) + } #[cfg(target_arch = "aarch64")] MaxSimIsa::Neon => { let arch = Neon::new_checked().ok_or(NotSupported { @@ -394,6 +492,10 @@ impl MaxSimElement for half::f16 { isa, reason: "aarch64 target only", }), + MaxSimIsa::X86_64_V3_Staged => Err(NotSupported { + isa, + reason: "x86-64-v3-staged supports f32 only", + }), MaxSimIsa::Reference => Ok(erase.erase(ReferenceKernel::::new(query))), } } @@ -538,6 +640,48 @@ mod tests { assert_eq!(kernel.nrows(), 5); } + /// Reusing one staged f32 `MaxSimKernel` across multiple `compute_max_sim` + /// calls (different doc counts) must match the reference each time — the + /// regression guard for the `PreparedStaged`-owned reset-arena + /// (`F32StagedScratch`): each call rewinds + reuses the scratch, so a stale or + /// aliased buffer would corrupt a later call. + #[cfg(target_arch = "x86_64")] + #[test] + fn staged_f32_arena_reuse() { + if diskann_wide::arch::x86_64::V3::new_checked().is_none() { + return; // No AVX2 on this host; the staged f32 path cannot build. + } + + const NQ: usize = 17; // exercises the A-panel row padding (17 -> 32) + const DIM: usize = 96; + let q_data = make_test_data::(NQ * DIM, DIM, DIM / 2); + let query = make_mat(&q_data, NQ, DIM); + + let staged = build_max_sim::(MaxSimIsa::X86_64_V3_Staged, query, BoxErase).unwrap(); + let reference = build_max_sim::(MaxSimIsa::Reference, query, BoxErase).unwrap(); + + // Distinct doc counts (multi-tile, single panel, remainder, one) reusing the + // same kernel — the arena is reset, never reallocated, between calls. + for (call, &nd) in [200usize, 3, 33, 1].iter().enumerate() { + let d_data = make_test_data::(nd * DIM, DIM, DIM + call); + let doc = make_mat(&d_data, nd, DIM); + + let mut got = vec![0.0f32; NQ]; + staged.compute_max_sim(doc, &mut got).unwrap(); + let mut want = vec![0.0f32; NQ]; + reference.compute_max_sim(doc, &mut want).unwrap(); + + for i in 0..NQ { + assert!( + (got[i] - want[i]).abs() <= 1e-4 * want[i].abs().max(1.0), + "call {call} (nd={nd}) row {i}: reused staged f32 {} != reference {}", + got[i], + want[i], + ); + } + } + } + fn check_size_mismatch(label: &str) where T: MaxSimElement + FromF32, diff --git a/diskann-quantization/src/multi_vector/distance/isa.rs b/diskann-quantization/src/multi_vector/distance/isa.rs index d295438bc..29ec6dffd 100644 --- a/diskann-quantization/src/multi_vector/distance/isa.rs +++ b/diskann-quantization/src/multi_vector/distance/isa.rs @@ -24,6 +24,10 @@ pub enum MaxSimIsa { X86_64_V3, /// x86_64 AVX-512. X86_64_V4, + /// Experimental staged-pipeline kernel (x86_64 AVX2+FMA). Coexists with + /// [`Self::X86_64_V3`] for A/B benchmarking; produces the same results via a + /// different (kernel → postprocess → reducer) micro-kernel structure. + X86_64_V3_Staged, /// AArch64 Neon. Neon, /// Non-SIMD reference fallback. Slow; serves as a correctness baseline. @@ -41,8 +45,10 @@ impl MaxSimIsa { Self::X86_64_V3 => diskann_wide::arch::x86_64::V3::new_checked().is_some(), #[cfg(target_arch = "x86_64")] Self::X86_64_V4 => diskann_wide::arch::x86_64::V4::new_checked().is_some(), + #[cfg(target_arch = "x86_64")] + Self::X86_64_V3_Staged => diskann_wide::arch::x86_64::V3::new_checked().is_some(), #[cfg(not(target_arch = "x86_64"))] - Self::X86_64_V3 | Self::X86_64_V4 => false, + Self::X86_64_V3 | Self::X86_64_V4 | Self::X86_64_V3_Staged => false, #[cfg(target_arch = "aarch64")] Self::Neon => diskann_wide::arch::aarch64::Neon::new_checked().is_some(), #[cfg(not(target_arch = "aarch64"))] @@ -58,6 +64,7 @@ impl std::fmt::Display for MaxSimIsa { Self::Scalar => "scalar", Self::X86_64_V3 => "x86-64-v3", Self::X86_64_V4 => "x86-64-v4", + Self::X86_64_V3_Staged => "x86-64-v3-staged", Self::Neon => "neon", Self::Reference => "reference", }; diff --git a/diskann-quantization/src/multi_vector/distance/kernel.rs b/diskann-quantization/src/multi_vector/distance/kernel.rs index b292def54..127c7f99e 100644 --- a/diskann-quantization/src/multi_vector/distance/kernel.rs +++ b/diskann-quantization/src/multi_vector/distance/kernel.rs @@ -6,7 +6,13 @@ use crate::multi_vector::{MatRef, MaxSimError, Standard}; /// Object-safe interface for computing per-query MaxSim scores. -pub trait MaxSimKernel: Send + Sync + std::fmt::Debug { +/// +/// `Send` (not `Sync`): a built kernel can be **moved** to a worker thread that +/// owns it (the "each search thread owns its distance computer" model), but is +/// not required to be shared by reference across threads. Dropping `Sync` is what +/// lets a kernel own interior-mutable per-call scratch (e.g. the staged f32 +/// kernel's reset-arena) under a `&self` method. +pub trait MaxSimKernel: Send + std::fmt::Debug { /// Number of query rows whose scores this kernel produces. fn nrows(&self) -> usize; diff --git a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs index 55108698d..61182315a 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs @@ -16,8 +16,16 @@ pub(super) mod f16; pub(super) mod f32; mod layouts; mod reduce; +// The staged kernel is V3 (x86_64) only; gate the whole module so its support +// code isn't dead on other architectures. +#[cfg(target_arch = "x86_64")] +pub(super) mod staged; mod tiled_reduce; +// Re-export the quantized staged kernel's public POC entry (x86_64 only). +#[cfg(target_arch = "x86_64")] +pub use staged::{QuantStagedDocs, QuantStagedQuery}; + // ── Tile budget ────────────────────────────────────────────────── /// Cache budgets fed to the tile planner. diff --git a/diskann-quantization/src/multi_vector/distance/kernels/staged/README.md b/diskann-quantization/src/multi_vector/distance/kernels/staged/README.md new file mode 100644 index 000000000..61c1389a8 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/staged/README.md @@ -0,0 +1,91 @@ +# Staged MaxSim kernel (`tiled_reduce_staged`) — design + +> **Status: experimental POC.** This module validates a design hypothesis; it is +> not yet the production MaxSim path. See *Scope* below. + +## Thesis + +One generic, cache-tiled reduction driver — [`tiled_reduce_staged`](driver.rs) — +can compute multi-vector MaxSim/Chamfer for **any element type and any +quantization** by swapping small pluggable *stages*, instead of forking the tiled +loop nest once per datatype. This module proves that claim with two +instantiations: + +- **f32** — bit-identical to, and on par with, the hand-fused V3 kernel. +- **4-bit MinMax-quantized `i8`** — a *new datatype* added by swapping two stages + only, with a real speedup over the per-pair SIMD reference. + +## Shape + +The driver owns tiling/blocking and walks query×doc tiles. Per tile it calls four +pluggable stages, defined in [`mod.rs`](mod.rs): + +| Stage | Trait | Role | +|-------|-------|------| +| A | `StagedKernel` | SIMD inner kernel; writes a per-pair accumulator `Acc` into a `partial` buffer | +| B | `Postprocess` | maps `Acc → Score` (e.g. dequantize) using optional per-call metadata (`scratch_len` + `apply`) | +| C | `Reducer` | folds per-doc scores into the running MaxSim (max) | +| — | `StagedConvert` | optional input conversion at tile load (identity for f32; future: on-the-fly quantize) | + +The driver allocates **all** scratch (`partial`, `scored`, conversion buffers) +from a caller-supplied `ScopedAllocator`. Callers size nothing. + +## What varies per axis (the generality proof) + +| Axis | f32 | 4-bit MinMax (`i8`) | +|------|-----|---------------------| +| Stage A kernel | `StagedF32Kernel` ([`v3.rs`](v3.rs)) | `StagedI8Kernel` ([`i8.rs`](i8.rs)) | +| `Acc` type | `f32` | `i32` | +| Stage B postprocess | `Identity` — `scratch_len = 0`, returns acc ([`maxsim.rs`](maxsim.rs)) | `MinMaxPostprocess` — `a·x + b` dequant → `f32` ([`i8.rs`](i8.rs)) | +| Stage C reducer | `MaxReducer` ([`maxsim.rs`](maxsim.rs)) | `MaxReducer` *(shared, unchanged)* | +| Convert | identity | identity (codes are pre-quantized) | + +Only **Stage A + the `Acc` type + Stage B** change between f32 and quantized; the +driver, tiling, and reducer are reused verbatim. That is the thesis. + +## Evidence + +- **f32 = parity.** The staged f32 path is bit-for-bit equal to the hand-fused V3 + kernel (it *is* the same math, restructured) and within ±1.7% throughput. + - Tests: `staged_matches_fused_v3` ([`v3.rs`](v3.rs)), + `staged_f32_arena_reuse` ([`../../factory.rs`](../factory.rs)). + - Benches: `example/multi-vector-staged.json` (fused vs staged sweep), + `example/multi-vector-3way.json` (reference vs fused vs staged). +- **Quantized = new datatype, real win.** 4-bit MinMax over `i8` codes was added + with **no driver change**; it is correct and runs **1.5–4.1×** faster than the + per-pair SIMD reference across a dim sweep. + - Tests: `staged_i8_matches_minmax_reference`, + `staged_i8_arena_reuse_across_calls`, `staged_i8_multi_tile_tiny_budget` + ([`i8.rs`](i8.rs)). + - Bench: `example/multi-vector-quant.json` (reference vs staged). + - The reference (`MinMaxKernel`) is **not** scalar: its per-pair inner product + over 4-bit codes is itself SIMD. The staged win comes from + fusion / block-transposition / tiling, not from SIMD-vs-scalar. + +## Scratch & allocation + +Driver-owned scratch comes from a passed `ScopedAllocator`. For a +zero-allocation steady state, callers reuse a single-owner resettable bump arena, +[`ResettableArena`](arena.rs): + +- the f32 kernel reuses one via `RefCell` + (`PreparedStaged` in [`../../factory.rs`](../factory.rs)); +- the quantized POC owns one in `QuantStagedQuery` ([`i8.rs`](i8.rs)) and `reset`s + it per call. + +`ResettableArena` is deliberately **not** `Clone`/`Sync`: `reset(&mut self)` is +sound only because the borrow checker forbids resetting while any +`ScopedAllocator` still borrows it. (The shared `BumpAllocator` is grow-only and +`Sync`, so it has no `reset`.) + +## Scope + +- The f32 staged kernel is wired end-to-end and selectable for A/B benchmarking as + `MaxSimIsa::X86_64_V3_Staged` ([`../../isa.rs`](../isa.rs)); it coexists with the + fused `X86_64_V3` path. +- The quantized path is a standalone POC entry (`QuantStagedQuery` / + `QuantStagedDocs`), **not** yet behind a `MaxSimIsa` variant or a productized + storage `Repr`. +- Deferred: a quantized storage `Repr`, folding the quantized path into + `MaxSimIsa` / the factory, V4 (AVX-512) Stage-A kernels, and richer reducers + (argmax / top-k). diff --git a/diskann-quantization/src/multi_vector/distance/kernels/staged/arena.rs b/diskann-quantization/src/multi_vector/distance/kernels/staged/arena.rs new file mode 100644 index 000000000..d7860661a --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/staged/arena.rs @@ -0,0 +1,122 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! POC-local single-threaded resettable bump arena for the staged quantized +//! kernel. +//! +//! Unlike [`BumpAllocator`](crate::alloc::BumpAllocator) — a `Clone`/`Send`/`Sync` +//! grow-only arena reclaimed only when its last clone drops — this is a +//! *single-owner, single-threaded* arena with O(1) +//! [`reset`](ResettableArena::reset). It is owned by one +//! [`QuantStagedQuery`](super::i8::QuantStagedQuery) and reused across +//! `compute_max_sim` calls, so the staged driver's per-call `partial` / `scored` +//! scratch allocates from it instead of the global heap (zero heap traffic in +//! steady state). +//! +//! `reset` is sound *because* the arena is not shareable: it takes `&mut self`, +//! and the borrow checker therefore guarantees no +//! [`ScopedAllocator`](crate::alloc::ScopedAllocator) borrowing this arena — hence +//! no outstanding allocation — is alive at the rewind point. `deallocate` is a +//! no-op; storage is reclaimed wholesale by `reset` or when the arena drops. + +use std::cell::{Cell, UnsafeCell}; +use std::ptr::NonNull; + +use crate::alloc::{AlignedAllocator, AllocatorCore, AllocatorError, Poly}; + +/// A single-owner, single-threaded resettable bump arena over an owned, +/// 64-byte-aligned byte buffer. Hands out aligned sub-slices by bumping a +/// non-atomic `head`; [`reset`](Self::reset) rewinds it. +pub(crate) struct ResettableArena { + /// Backing storage. `UnsafeCell` legitimizes handing out `*mut` ranges while + /// `allocate` holds only `&self` (mirrors `BumpAllocator`'s buffer). + buffer: Poly, AlignedAllocator>, + /// Bump cursor (non-atomic — this arena is single-threaded). + head: Cell, +} + +impl ResettableArena { + /// Allocate a fresh arena with room for `capacity` bytes, base-aligned to 64. + pub(crate) fn with_capacity(capacity: usize) -> Result { + let buffer = Poly::<[u8], _>::new_uninit_slice(capacity.max(1), AlignedAllocator::A64)?; + let (ptr, alloc) = Poly::into_raw(buffer); + + // SAFETY: `UnsafeCell<[u8]>` shares the layout of `[u8]`, and `MaybeUninit` + // is layout-compatible with `u8` (`u8` is valid for any bit pattern, and bytes + // are only read after the allocator hands them out and the caller writes them). + // `ptr` is non-null, having come from `Poly::into_raw`. + let buffer = unsafe { + Poly::from_raw( + NonNull::new_unchecked(ptr.as_ptr() as *mut UnsafeCell<[u8]>), + alloc, + ) + }; + + Ok(Self { + buffer, + head: Cell::new(0), + }) + } + + /// Rewind the arena, freeing every prior allocation in O(1). + /// + /// The `&mut self` receiver is load-bearing: it makes the borrow checker + /// forbid calling `reset` while any [`ScopedAllocator`](crate::alloc::ScopedAllocator) + /// borrowing this arena (and hence any live allocation) exists — which is what + /// makes rewinding the cursor sound. + pub(crate) fn reset(&mut self) { + self.head.set(0); + } + + /// Total capacity in bytes. + fn capacity(&self) -> usize { + self.buffer.get().len() + } + + /// Base pointer of the backing buffer. + fn base(&self) -> *mut u8 { + self.buffer.get().cast::() + } +} + +// SAFETY: on success `allocate` returns a slice of exactly `layout.size()` bytes +// whose base is aligned to at least `layout.align()` (the running offset is aligned +// up to `layout.align()` relative to the real base address); a request that cannot +// fit the fixed capacity returns an error. `deallocate` is a no-op — storage is +// reclaimed only by `reset` or on drop. +unsafe impl AllocatorCore for ResettableArena { + fn allocate(&self, layout: std::alloc::Layout) -> Result, AllocatorError> { + let base = self.base() as usize; + let head = self.head.get(); + // Align the current free address up to `layout.align()`, then reserve `size`. + let cur = base.checked_add(head).ok_or(AllocatorError)?; + let aligned = cur + .checked_next_multiple_of(layout.align()) + .ok_or(AllocatorError)?; + let pad = aligned - cur; + let new_head = head + .checked_add(pad) + .and_then(|h| h.checked_add(layout.size())) + .ok_or(AllocatorError)?; + if new_head > self.capacity() { + return Err(AllocatorError); + } + self.head.set(new_head); + + // SAFETY: `head + pad <= new_head <= capacity`, so the offset is in-bounds of + // the backing buffer and the range `[off, off + size)` lies within it. + let ptr = unsafe { self.base().add(head + pad) }; + NonNull::new(std::ptr::slice_from_raw_parts_mut(ptr, layout.size())).ok_or(AllocatorError) + } + + unsafe fn deallocate(&self, _ptr: NonNull<[u8]>, _layout: std::alloc::Layout) {} +} + +impl std::fmt::Debug for ResettableArena { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ResettableArena") + .field("capacity", &self.capacity()) + .field("head", &self.head.get()) + .finish() + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/staged/driver.rs b/diskann-quantization/src/multi_vector/distance/kernels/staged/driver.rs new file mode 100644 index 000000000..194b527a8 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/staged/driver.rs @@ -0,0 +1,227 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Staged tiling driver. +//! +//! Structurally identical to [`super::super::tiled_reduce`] (A-tile → B-tile), but +//! the inner work per A-panel is split into three stages: Stage A fills +//! `partial_buf` for the whole B-tile, Stage B ([`Postprocess::apply`]) turns the +//! raw `Acc` block into `Score`s, and Stage C ([`Reducer::fold_block`]) folds them +//! into the running state. The driver runs all three uniformly; for the identity +//! postprocess (`Acc == Score`) Stage B is `#[inline(always)]` and returns the +//! `partial_buf` pointer unchanged, so it compiles away (no `scored_buf`, no pass). +//! +//! Every B-tile is `≤ b_tile_rows` rows: a full tile has `b_panels_per_tile` +//! complete panels, and a final short tile may end in a `< B_PANEL` remainder +//! panel — both handled by one loop (`full_panels` + an optional `tail`), so there +//! is no separate "peeled tail" code path. + +use diskann_wide::Architecture; + +use super::super::TileBudget; +use super::super::layouts::Layout; +use super::{FoldCtx, Postprocess, Reducer, StagedConvert, StagedKernel, StagedPlan}; +use crate::alloc::{Poly, ScopedAllocator}; + +/// Run the staged loop. `state` (len `a_padded_nrows`) is the persistent running +/// reduction across all B-tiles (the caller's output buffer). All transient +/// scratch — `partial` (Stage A output), the Stage-B `scored` region, and the +/// conversion buffers — is allocated **internally** from `alloc`; the caller +/// sizes nothing and hands in only the allocator. +/// +/// # Safety +/// +/// * `a_ptr` valid for `a_padded_nrows * k` `LA::Element`; `a_padded_nrows` +/// a multiple of `SK::A_PANEL`. +/// * `b_ptr` valid for `b_nrows * k` `LB::Element`. +#[allow(clippy::too_many_arguments, clippy::expect_used)] +pub(super) unsafe fn tiled_reduce_staged( + arch: A, + ca: &LA, + cb: &LB, + post: &P, + a_ptr: *const LA::Element, + a_padded_nrows: usize, + b_ptr: *const LB::Element, + b_nrows: usize, + k: usize, + state: &mut [R::State], + alloc: ScopedAllocator<'_>, + budget: TileBudget, +) where + A: Architecture, + SK: StagedKernel, + P: Postprocess, + R: Reducer, + LA: StagedConvert, + LB: StagedConvert, +{ + let a_panel = SK::A_PANEL; + let b_panel = SK::B_PANEL; + + // Initialize the running reduction state. + for s in state[..a_padded_nrows].iter_mut() { + *s = R::init(); + } + + // Zero-dimensional vectors: every IP is 0. The caller fills the score for + // this degenerate case; here we just avoid the zero-stride tiling nest. + if k == 0 { + return; + } + + debug_assert_eq!( + a_padded_nrows % a_panel, + 0, + "a_padded_nrows must be a multiple of A_PANEL" + ); + + let acc_bytes = core::mem::size_of::(); + let a_row_bytes = k * core::mem::size_of::<::Element>(); + let b_row_bytes = k * core::mem::size_of::<::Element>(); + let plan = StagedPlan::new( + a_row_bytes, + b_row_bytes, + a_panel, + b_panel, + acc_bytes, + budget, + ); + + let a_tile_rows = a_panel * plan.a_panels_per_tile; + let b_tile_rows = b_panel * plan.b_panels_per_tile; + + let a_kern_panel_stride = a_panel * k; + let b_kern_panel_stride = b_panel * k; + + // Conversion staging buffers, also from the caller's allocator — 0-length + // (a no-op dangling allocation) for the identity conversions every current + // staged kernel uses. Sized by the staged-local `StagedConvert` contract, so + // the staged driver never touches the shared `ConvertTo` machinery. + let a_conv_len = ca.scratch_len(a_tile_rows.min(a_padded_nrows), k); + let mut a_conv = + Poly::<[::Element], _>::new_uninit_slice(a_conv_len, alloc) + .expect("a-side conversion scratch allocation"); + let b_conv_len = cb.scratch_len(b_tile_rows.min(b_nrows), k); + let mut b_conv = + Poly::<[::Element], _>::new_uninit_slice(b_conv_len, alloc) + .expect("b-side conversion scratch allocation"); + let a_conv_ptr = a_conv.as_mut_ptr().cast::<::Element>(); + let b_conv_ptr = b_conv.as_mut_ptr().cast::<::Element>(); + + // Internal scratch, allocated from the caller's allocator — the caller sizes + // nothing. `partial` is Stage A's output (the kernel declares its size via + // `StagedKernel::partial_len`); `scored` is Stage B's output, sized by the + // postprocess contract (a 0-length, no-op dangling allocation for the identity + // postprocess). Every B-tile is `≤ b_tile_rows` wide, so `b_tile_rows` is the + // exact upper bound on `valid_b_cols`. Both `Poly`s live to the end of the + // call, then free via `alloc` (a global free, or a no-op for a bump allocator). + let partial_len = SK::partial_len(k, budget); + let mut partial = Poly::<[SK::Acc], _>::new_uninit_slice(partial_len, alloc) + .expect("partial scratch allocation"); + let scored_len = post.scratch_len(a_panel, b_tile_rows); + let mut scored = Poly::<[P::Score], _>::new_uninit_slice(scored_len, alloc) + .expect("scored scratch allocation"); + + let partial_ptr = partial.as_mut_ptr().cast::(); + let scored_ptr = scored.as_mut_ptr().cast::(); + let state_ptr = state.as_mut_ptr(); + + // SAFETY: all pointer arithmetic stays within the respective allocations; + // this mirrors `super::super::tiled_reduce`'s established bounds. + unsafe { + let mut rows_done: usize = 0; + + // Loop 1: A tiles. + while rows_done < a_padded_nrows { + let tile_rows = a_tile_rows.min(a_padded_nrows - rows_done); + let pa_tile_src = a_ptr.add(rows_done * k); + let pr_tile = state_ptr.add(rows_done); + + let pa_tile = ca.convert(a_conv_ptr, arch, pa_tile_src, tile_rows, k); + let pa_tile_end = pa_tile.add(tile_rows * k); + + // Loop 2: B tiles. Each is `bt_rows = min(b_tile_rows, remaining)` — + // full tiles end on a panel boundary (`tail == 0`); the final short + // tile may carry a `< B_PANEL` remainder panel. + let mut pb_tile_src = b_ptr; + let mut b_row_offset = 0usize; + while b_row_offset < b_nrows { + let bt_rows = b_tile_rows.min(b_nrows - b_row_offset); + let pb_tile = cb.convert(b_conv_ptr, arch, pb_tile_src, bt_rows, k); + let full_panels = bt_rows / b_panel; + let tail = bt_rows % b_panel; + + // Loop 3: A micro-panels. + // + // Partial-buffer granularity: one A-panel (P_a = 1) against the + // whole B-tile (P_b = b_panels_per_tile). Loop 4 below runs ONLY + // Stage A, so the kernel stays hot in i-cache for the entire + // B-tile; Stage C then folds the whole block in one pass. See + // `StagedPlan` and docs/staged_multi_vector_kernel.md §5. + let mut pa_panel = pa_tile; + let mut pr_panel = pr_tile; + let mut a_row_offset = rows_done; + while pa_panel < pa_tile_end { + // Stage A: fill partial_buf for this A-panel across the B-tile + // (the full panels, then a `< B_PANEL` remainder panel if any). + let mut pb_panel = pb_tile; + let mut col = 0usize; + for _ in 0..full_panels { + SK::full_panel( + arch, + pa_panel, + pb_panel, + k, + partial_ptr.add(col * a_panel), + a_panel, + ); + pb_panel = pb_panel.add(b_kern_panel_stride); + col += b_panel; + } + if tail > 0 { + SK::partial_panel( + arch, + tail, + pa_panel, + pb_panel, + k, + partial_ptr.add(col * a_panel), + a_panel, + ); + } + + // Stage B (Acc -> Score): identity returns `partial_ptr` for + // free; a quantized post fills `scored` (the driver-allocated + // region), using the global (a_row_offset, b_row_offset) to + // index its metadata. Output column stride is contractually + // `a_panel`. + let scores = post.apply( + scored_ptr, + arch, + partial_ptr, + FoldCtx { + a_panel, + valid_b_cols: bt_rows, + b_stride: a_panel, + a_row_offset, + b_row_offset, + }, + ); + // Stage C: fold the scores into the running state (one fold + // per A-panel × B-tile — the widest, cheapest fold). + R::fold_block(arch, pr_panel, scores, a_panel, bt_rows, a_panel); + + pa_panel = pa_panel.add(a_kern_panel_stride); + pr_panel = pr_panel.add(a_panel); + a_row_offset += a_panel; + } + + pb_tile_src = pb_tile_src.add(bt_rows * k); + b_row_offset += bt_rows; + } + + rows_done += tile_rows; + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/staged/i8.rs b/diskann-quantization/src/multi_vector/distance/kernels/staged/i8.rs new file mode 100644 index 000000000..a51c5a6a6 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/staged/i8.rs @@ -0,0 +1,962 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! V3 (AVX2) **quantized** staged path: 4-bit MinMax MaxSim. +//! +//! This is the payoff the staged framework was built for — the first +//! **non-identity Stage B**. Stage A accumulates a raw *integer* dot product +//! (`Acc = i32`), Stage B ([`MinMaxPostprocess`]) turns each `i32` into the +//! finished MinMax inner product (`Score = f32`) using per-vector +//! scale/center/sum metadata, and Stage C ([`MaxReducer`](super::maxsim::MaxReducer)) +//! folds the `f32` scores exactly as in the f32 path — **unchanged**. +//! +//! # Stage A mirrors the f32 kernel (block-transposed + broadcast, no reduction) +//! +//! The integer micro-kernel is structurally the f32 [`store_microkernel`] with +//! `PACK = 2` integer MACs: +//! +//! * **Query** (`Left`): [`BlockTransposed`] — codes widened `u8→i16` +//! once at build time, then block-transposed with two K-columns interleaved per +//! row (`[r0_k0, r0_k1, r1_k0, r1_k1, …]`). One col-pair is `GROUP·PACK = 32` +//! `i16` = two `i16x16` halves (rows `0..8` / `8..16`). +//! * **Doc** (`Right`): [`RowMajor`] — codes stream in place; the kernel reads +//! each doc col's 2-K word, widens it, and broadcasts it as +//! `[d_k0, d_k1, d_k0, d_k1, …]`. +//! * **MAC**: `i32x8::dot_simd(query_half, doc_broadcast)` — i.e. `vpmaddwd` +//! (`_mm256_madd_epi16`), which sums each interleaved K-pair into one `i32` +//! lane. Each lane **is** one A-row's running dot for that doc col, so the +//! accumulators *are* the outputs — no per-pair horizontal reduction, exactly +//! the property block-transposition buys in the f32 kernel. +//! +//! 4-bit codes are `u8 ∈ [0, 15]` (not nibble-packed), widened to `i16 ∈ [0, 15]`. +//! A 2-K partial is `≤ 2·15·15 = 450` and the full dot over `dim ≤ 512` is +//! `≤ 15·15·512 ≈ 1.2e5`, so the `i32` accumulation **never overflows** — the +//! integer dot is *exact*. No new `diskann-wide` op is needed: the existing +//! [`SIMDDotProduct`] for `i32x8` is the right shape in broadcast config. +//! +//! # Even-K contract (zero driver change) +//! +//! The shared [driver](super::driver) derives the A-side physical row stride from +//! its `k` argument (`rows_done · k == block_offset`), which only matches a +//! block-transposed query when `k == padded_ncols`. So the entry point passes the +//! query's **padded** (even) column count as the driver `k` and requires the doc +//! to match (zero-padded to even). The padding column holds `0` on *both* sides, +//! so it contributes `0` to every dot — the IP is unchanged, and the kernel walks +//! exactly `k/2` full K-pairs with no odd-K tail branch. + +use std::num::NonZeroUsize; + +use diskann_utils::ReborrowMut; +use diskann_wide::arch::x86_64::V3; +use diskann_wide::{SIMDCast, SIMDDotProduct, SIMDMulAdd, SIMDReinterpret, SIMDVector}; + +use super::super::TileBudget; +use super::super::layouts; +use super::arena::ResettableArena; +use super::driver::tiled_reduce_staged; +use super::maxsim::MaxReducer; +use super::{FoldCtx, Postprocess, StagedKernel}; +use crate::CompressInto; +use crate::algorithms::Transform; +use crate::algorithms::transforms::NullTransform; +use crate::alloc::ScopedAllocator; +use crate::minmax::{MinMaxCompensation, MinMaxMeta, MinMaxQuantizer}; +use crate::multi_vector::{BlockTransposed, BlockTransposedRef, Defaulted, Mat, MatRef, Standard}; +use crate::num::Positive; + +diskann_wide::alias!(i16s = ::i16x16); +diskann_wide::alias!(i32s = ::i32x8); +diskann_wide::alias!(u32s = ::u32x8); +diskann_wide::alias!(f32s = ::f32x8); + +// ── Stage A: integer store-out micro-kernel ────────────────────── + +/// Zero-sized Stage-A kernel marker for the quantized (4-bit MinMax) staged path +/// with block size `GROUP`. +pub(crate) struct StagedI8Kernel; + +// SAFETY: `full_panel`/`partial_panel` read A_PANEL(16) i16 query rows × K +// (block-transposed, K padded to even) and UNROLL × K u8 doc elements, and write +// UNROLL columns of A_PANEL(16) i32 into `partial` at stride `partial_b_stride` — +// all within the bounds the `StagedKernel` contract guarantees. +unsafe impl StagedKernel for StagedI8Kernel<16> { + type Left = layouts::BlockTransposed; + type Right = layouts::RowMajor; + type Acc = i32; + const A_PANEL: usize = 16; + const B_PANEL: usize = 4; + + #[inline(always)] + unsafe fn full_panel( + arch: V3, + a: *const i16, + b: *const u8, + k: usize, + partial: *mut i32, + partial_b_stride: usize, + ) { + // SAFETY: pointer validity per the `StagedKernel` contract. + unsafe { + int_store_microkernel::<{ Self::B_PANEL }>(arch, a, b, k, partial, partial_b_stride) + } + } + + #[inline(always)] + unsafe fn partial_panel( + arch: V3, + remainder: usize, + a: *const i16, + b: *const u8, + k: usize, + partial: *mut i32, + partial_b_stride: usize, + ) { + // SAFETY: pointer validity per the `StagedKernel` contract. + unsafe { + match remainder { + 1 => int_store_microkernel::<1>(arch, a, b, k, partial, partial_b_stride), + 2 => int_store_microkernel::<2>(arch, a, b, k, partial, partial_b_stride), + 3 => int_store_microkernel::<3>(arch, a, b, k, partial, partial_b_stride), + _ => unreachable!( + "unexpected remainder {remainder} for B_PANEL={}", + Self::B_PANEL + ), + } + } + } +} + +/// V3 integer store-out micro-kernel: 16 A-rows × `UNROLL` B-rows. +/// +/// Mirrors [`super::v3::store_microkernel`] with `PACK = 2` integer MACs (see the +/// module docs). The epilogue stores each B-column's 16 A-row `i32` accumulators +/// into `partial` (A-major: column `j` at `partial + j*b_stride`, as two `i32x8` +/// halves) — identical contract to the f32 kernel, so Stage B / Stage C are +/// unchanged. +/// +/// # Safety +/// +/// 1. `a_packed` points to a block-transposed query block of `16 * k` `i16` +/// (`k` even — the padded column count). +/// 2. `b` points to `UNROLL` rows of `k` contiguous `u8` (`k` even). +/// 3. `partial` is valid for `UNROLL` columns of 16 `i32` at stride `b_stride`. +#[inline(always)] +unsafe fn int_store_microkernel( + arch: V3, + a_packed: *const i16, + b: *const u8, + k: usize, + partial: *mut i32, + b_stride: usize, +) { + let mut p0 = [i32s::default(arch); UNROLL]; + let mut p1 = [i32s::default(arch); UNROLL]; + let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); + + // One col-pair of the block-transposed query = GROUP·PACK = 32 i16 = two + // `i16x16` halves (rows 0..8 low, rows 8..16 high). + let a_pair_stride = 2 * i16s::LANES; + let a_half = i16s::LANES; + let pairs = k / 2; // `k` is even (the padded column count) ⇒ exact, no tail. + + for p in 0..pairs { + // SAFETY: precondition 1 — the query block has `pairs` col-pairs of 32 i16. + let (a0, a1) = unsafe { + ( + i16s::load_simd(arch, a_packed.add(a_pair_stride * p)), + i16s::load_simd(arch, a_packed.add(a_pair_stride * p + a_half)), + ) + }; + + for j in 0..UNROLL { + // SAFETY: precondition 2 — doc col j is `offsets[j]` in, and + // `2*p + 1 < k` because `pairs == k/2`. + let (d0, d1) = unsafe { + let base = 2 * p + offsets[j]; + ( + u32::from(b.add(base).read()), + u32::from(b.add(base + 1).read()), + ) + }; + // Broadcast the K-pair [d0, d1] across all 8 i32 lanes, reinterpreted + // as i16x16 = [d0, d1, d0, d1, …] — the shape `madd_epi16` consumes + // (pairing each query [k0, k1] with [d0, d1] into one i32 lane). + let packed = d0 | (d1 << 16); + let bcast: i16s = u32s::splat(arch, packed).reinterpret_simd(); + p0[j] = p0[j].dot_simd(a0, bcast); + p1[j] = p1[j].dot_simd(a1, bcast); + } + } + + for j in 0..UNROLL { + // SAFETY: precondition 3 — column j occupies [j*b_stride, j*b_stride+16) i32. + unsafe { + p0[j].store_simd(partial.add(j * b_stride)); + p1[j].store_simd(partial.add(j * b_stride + i32s::LANES)); + } + } +} + +// ── Stage B: integer code → MinMax inner product ───────────────── + +/// Stage B for 4-bit MinMax: convert each raw integer dot `⟨codes⟩` (the `i32` +/// `Acc` from Stage A) into the finished MinMax inner product +/// +/// ```text +/// IP = qm.a·dm.a·⟨codes⟩ + qm.n·dm.b + dm.n·qm.b + qm.b·dm.b·dim +/// ``` +/// +/// (the linear decomposition in `minmax::vectors`), emitting **+IP** so Stage C +/// folds `max` and the caller negates once at the end (`min distance = +/// max_a(-IP) = -max_a IP`). +/// +/// This is the first non-identity [`Postprocess`]: it reports a non-zero +/// [`scratch_len`](Postprocess::scratch_len) (so the driver allocates an f32 +/// region) and [`apply`] writes the converted scores into it, indexing its +/// per-vector metadata by the *global* row offsets in [`FoldCtx`]. +pub(crate) struct MinMaxPostprocess<'m> { + /// Per-query-vector metadata, indexed by `ctx.a_row_offset + i`. Length must + /// be `≥ padded_nrows` (padded rows carry default metadata; their scores are + /// computed but never read). + query_meta: &'m [MinMaxCompensation], + /// Per-doc-vector metadata, indexed by `ctx.b_row_offset + c`. Length `≥ nd`. + doc_meta: &'m [MinMaxCompensation], + /// The **logical** dimension (the `dim` term in the IP formula). The integer + /// dot is taken over the padded columns, but the extra column's codes are `0` + /// on both sides, so `⟨codes⟩` is unchanged. + dim: f32, +} + +impl<'m> MinMaxPostprocess<'m> { + pub(crate) fn new( + query_meta: &'m [MinMaxCompensation], + doc_meta: &'m [MinMaxCompensation], + dim: usize, + ) -> Self { + Self { + query_meta, + doc_meta, + dim: dim as f32, + } + } +} + +// SAFETY: `apply` reads exactly `ctx.valid_b_cols` columns of `ctx.a_panel` `i32` +// from `acc` at stride `ctx.b_stride`, writes only the corresponding +// `ctx.a_panel × ctx.valid_b_cols` region of `scratch` (the driver allocates +// `scratch_len(a_panel, max_b_cols) = a_panel · max_b_cols ≥ a_panel · +// valid_b_cols` `f32`), and returns a pointer into it. Metadata indices +// `a_row_offset + i < padded_nrows ≤ query_meta.len()` and `b_row_offset + c < nd +// ≤ doc_meta.len()` are in bounds by the entry-point's preconditions. +// +// V3-specific: the quantized kernel only runs on V3 (`StagedI8Kernel: StagedKernel`), +// so Stage B is implemented for V3 only and uses AVX2 for the score conversion. +unsafe impl Postprocess for MinMaxPostprocess<'_> { + type Acc = i32; + type Score = f32; + + #[inline] + fn scratch_len(&self, a_panel: usize, max_b_cols: usize) -> usize { + a_panel * max_b_cols + } + + #[inline] + unsafe fn apply( + &self, + scratch: *mut f32, + arch: V3, + acc: *const i32, + ctx: FoldCtx, + ) -> *const f32 { + let out = scratch; + + // Rewrite the per-(row i, col c) IP into a per-row-vector form whose only + // per-column inputs are three doc scalars (so the 16-row inner loop is a + // straight SIMD sweep): + // ip = dm.a·(qm.a·raw) + dm.b·qm.n + (dm.n + dm.b·dim)·qm.b + // = A_c·(qa·raw) + B_c·qn + C_c·qb, + // with A_c=dm.a, B_c=dm.b, C_c=dm.n+dm.b·dim, and qa/qn/qb the per-query-row + // metadata. This is the reference formula regrouped (within f32 rounding). + // The quantized kernel always uses A_PANEL = 16 = 2*LANES, so Stage B is a + // pure SIMD sweep — no scalar fallback for other panel widths. + let lanes = f32s::LANES; + debug_assert_eq!( + ctx.a_panel, + 2 * lanes, + "quantized Stage B expects A_PANEL == 2*LANES" + ); + + // Gather the per-A-row metadata into contiguous SoA arrays once (the + // source `MinMaxCompensation` is AoS, so a strided scalar gather), then + // hold them in registers across the whole B-column sweep. + let mut qa = [0.0f32; 16]; + let mut qb = [0.0f32; 16]; + let mut qn = [0.0f32; 16]; + for i in 0..16 { + let qm = self.query_meta[ctx.a_row_offset + i]; + qa[i] = qm.a; + qb[i] = qm.b; + qn[i] = qm.n; + } + // SAFETY: each array holds exactly 16 = 2·LANES f32. + let (qa0, qa1, qb0, qb1, qn0, qn1) = unsafe { + ( + f32s::load_simd(arch, qa.as_ptr()), + f32s::load_simd(arch, qa.as_ptr().add(lanes)), + f32s::load_simd(arch, qb.as_ptr()), + f32s::load_simd(arch, qb.as_ptr().add(lanes)), + f32s::load_simd(arch, qn.as_ptr()), + f32s::load_simd(arch, qn.as_ptr().add(lanes)), + ) + }; + + for c in 0..ctx.valid_b_cols { + let dm = self.doc_meta[ctx.b_row_offset + c]; + let a_c = f32s::splat(arch, dm.a); + let b_c = f32s::splat(arch, dm.b); + let c_c = f32s::splat(arch, dm.n + dm.b * self.dim); + let acc_col = c * ctx.b_stride; + let out_col = c * ctx.a_panel; + // SAFETY: `acc_col + 2·LANES ≤ valid_b_cols · b_stride`; the partial + // block is valid for that many i32, and `out_col + 2·LANES ≤ buf.len()`. + unsafe { + let raw0 = i32s::load_simd(arch, acc.add(acc_col)).simd_cast(); + let raw1 = i32s::load_simd(arch, acc.add(acc_col + lanes)).simd_cast(); + // a_c·(qa·raw) + (b_c·qn + c_c·qb) + let s0 = a_c.mul_add_simd(qa0 * raw0, b_c.mul_add_simd(qn0, c_c * qb0)); + let s1 = a_c.mul_add_simd(qa1 * raw1, b_c.mul_add_simd(qn1, c_c * qb1)); + s0.store_simd(out.add(out_col)); + s1.store_simd(out.add(out_col + lanes)); + } + } + scratch.cast_const() + } +} + +// ── Public POC entry: prepared 4-bit MinMax staged MaxSim ──────── + +/// Quantize an f32 multi-vector to 4-bit MinMax (Null transform, scale 1.0) — +/// the shared quantizer for the public query/doc builders so both sides decode +/// to comparable codes + metadata. +#[allow(clippy::expect_used)] // POC constructor: inputs are pre-validated by the caller. +fn quantize_minmax_4bit(input: MatRef<'_, Standard>) -> Mat> { + let dim = input.vector_dim(); + let n = input.num_vectors(); + let q = MinMaxQuantizer::new( + Transform::Null(NullTransform::new( + NonZeroUsize::new(dim).expect("dimension must be non-zero"), + )), + Positive::new(1.0).expect("1.0 is positive"), + ); + let mut out: Mat> = + Mat::new(MinMaxMeta::new(n, dim), Defaulted).expect("MinMaxMeta allocation"); + q.compress_into(input, out.reborrow_mut()) + .expect("input must be finite (no NaN)"); + out +} + +/// A prepared 4-bit MinMax **query** set for the staged MaxSim kernel (V3/AVX2). +/// +/// Built once from an f32 multi-vector; [`compute_max_sim`](Self::compute_max_sim) +/// is the per-document-set hot path. It owns a `ResettableArena` that backs the +/// staged driver's per-call `partial` / Stage-B scratch and a reused `state` +/// buffer, so **steady-state calls perform no heap allocation** (the arena is +/// reset, not reallocated, each call). +/// +/// This is a **standalone POC entry**: the quantized path is intentionally *not* +/// yet unified into [`MaxSimIsa`](crate::multi_vector::distance::MaxSimIsa) / +/// `build_max_sim`, which ties to a productized `QuantizedSoa` matrix `Repr` (see +/// [`QuantStagedDocs`] and `docs/staged_multi_vector_kernel.md`). +pub struct QuantStagedQuery { + /// Codes widened `u8→i16` and block-transposed (`PACK=2`) for Stage A. + query: BlockTransposed, + /// Per-vector metadata, padded to `query.padded_nrows()` (padded rows carry + /// default metadata; their scores are computed but never read). + meta: Vec, + /// Logical dimension. + dim: usize, + arch: V3, + /// Reusable running-reduction output (len `query.padded_nrows()`); the driver + /// re-initialises it each call, so it is reused, not reallocated. + state: Vec, + /// Reusable arena backing the driver's transient `partial` / `scored` scratch; + /// reset (not reallocated) at the top of every call. + arena: ResettableArena, +} + +impl QuantStagedQuery { + /// Quantize `query` to 4-bit MinMax and prepare the block-transposed layout + + /// the reusable scratch arena. Returns `None` if AVX2 (V3) is unavailable on + /// this host. + #[allow(clippy::expect_used)] // POC constructor: dims are valid by construction. + pub fn build(query: MatRef<'_, Standard>) -> Option { + let arch = V3::new_checked()?; + let dim = query.vector_dim(); + let nq = query.num_vectors(); + + let q_mat = quantize_minmax_4bit(query); + + let mut codes = vec![0i16; nq * dim]; + for r in 0..nq { + let row = q_mat.get_row(r).expect("row r < nq"); + for j in 0..dim { + codes[r * dim + j] = i16::from(row.vector().get(j).expect("col j < dim") as u8); + } + } + let code_view = MatRef::new(Standard::::new(nq, dim).expect("nq×dim i16"), &codes) + .expect("code slice length"); + let bt = BlockTransposed::::from_matrix_view(code_view.as_matrix_view()); + + let padded_nrows = bt.padded_nrows(); + let mut meta = vec![MinMaxCompensation::default(); padded_nrows]; + for (r, m) in meta.iter_mut().enumerate().take(nq) { + *m = q_mat.get_row(r).expect("row r < nq").meta(); + } + + // The driver allocates only `partial` + `scored` from the arena (the + // identity conversion buffers are zero-length). `StagedPlan` co-budgets so + // each is `<= l1_b`, so `2 * l1_b` is a provable upper bound for *any* k + // (no per-shape sizing needed); add one page of headroom for alignment. + let arena_bytes = 2 * TileBudget::default().l1_b + 4096; + let arena = ResettableArena::with_capacity(arena_bytes).expect("staged arena allocation"); + + Some(Self { + query: bt, + meta, + dim, + arch, + state: vec![f32::MIN; padded_nrows], + arena, + }) + } + + /// Whether this host supports the quantized staged kernel (requires AVX2 / + /// the `V3` ISA). Use this to gate before calling [`build`](Self::build). + pub fn is_supported() -> bool { + V3::new_checked().is_some() + } + + /// Number of (logical) query vectors. + pub fn num_vectors(&self) -> usize { + self.query.nrows() + } + + /// Compute the per-query **min distance** (`= -max_d IP`) against `docs`, + /// writing one score per query vector into `scores`. + /// + /// # Panics + /// + /// Panics if `scores.len() != self.num_vectors()` or the query and doc + /// logical dimensions differ. + #[allow(clippy::expect_used)] // doc view length is guaranteed by construction. + pub fn compute_max_sim(&mut self, docs: &QuantStagedDocs, scores: &mut [f32]) { + let nq = self.query.nrows(); + assert_eq!( + scores.len(), + nq, + "scores length {} must equal query vector count {nq}", + scores.len() + ); + assert_eq!( + self.dim, docs.dim, + "query dim {} != doc dim {}", + self.dim, docs.dim + ); + + let doc = MatRef::new( + Standard::::new(docs.nv, docs.padded_dim).expect("nv×padded_dim u8"), + &docs.codes, + ) + .expect("doc code slice length"); + let post = MinMaxPostprocess::new(&self.meta, &docs.meta, self.dim); + + // Rewind the arena so the driver's `partial` / `scored` reuse last call's + // storage (`&mut self` proves no prior allocation is still borrowing it), + // then hand it the arena and the reused `state` output. Steady-state: no + // heap allocation. + let padded = self.query.padded_nrows(); + self.arena.reset(); + max_ip_kernel_staged_i8( + self.arch, + self.query.as_view(), + doc, + &post, + &mut self.state[..padded], + ScopedAllocator::new(&self.arena), + TileBudget::default(), + ); + + for (s, &raw) in scores.iter_mut().zip(self.state.iter()) { + *s = -raw; // min distance = -(max inner product) + } + } +} + +/// A prepared 4-bit MinMax **document** set: the minimal "codes-together, +/// metadata-together" SoA the staged kernel streams. Stage A reads the contiguous +/// codes region (row-major `u8`, `padded_dim` per vector); Stage B reads the +/// contiguous metadata region (one [`MinMaxCompensation`] per vector). +/// +/// This is the doc-side storage the kernel was designed around — the interleaved +/// `MinMaxMeta` `Repr` (one codes+meta blob per row) *cannot* be streamed by a +/// kernel that needs the codes contiguous for SIMD and the metadata only in the +/// postprocess. +/// +/// # Productization (assessed, not built): a `QuantizedSoa` matrix `Repr` +/// +/// The natural next step is a matrix `Repr` that owns this layout in one +/// allocation, `[codes region | aligned metadata region]`, with `Row<'a> = (&'a +/// [u8], &'a MinMaxCompensation)`. It fits the existing `Repr`/`ReprOwned` +/// contract (`layout()` = `codes_bytes + pad + meta_bytes`; `get_row` splits the +/// two regions) and needs a `CompressInto` that emits SoA from `MinMaxQuantizer` +/// (today's emits the interleaved blob). That `Repr` is also the prerequisite for +/// unifying the quantized path into `MaxSimIsa`; this owning prototype keeps the +/// POC self-contained without committing to it. +pub struct QuantStagedDocs { + /// Row-major codes, `nv * padded_dim` `u8` (each `∈ [0, 15]` for 4-bit), with + /// the trailing (padded) column zeroed when `dim` is odd. + codes: Vec, + /// Per-vector metadata, `nv` entries. + meta: Vec, + /// Logical dimension (the IP formula's `dim`). + dim: usize, + /// Physical (even) column count `next_multiple_of(dim, 2)` — the doc row + /// stride, which must equal the query's padded column count. + padded_dim: usize, + /// Number of vectors. + nv: usize, +} + +impl QuantStagedDocs { + /// Quantize `docs` to 4-bit MinMax and pack the codes-together / + /// metadata-together SoA (codes zero-padded to an even `padded_dim`). + #[allow(clippy::expect_used)] // POC constructor: dims are valid by construction. + pub fn build(docs: MatRef<'_, Standard>) -> Self { + let dim = docs.vector_dim(); + let nv = docs.num_vectors(); + let padded_dim = dim.next_multiple_of(2); + + let d_mat = quantize_minmax_4bit(docs); + let mut codes = vec![0u8; nv * padded_dim]; + let mut meta = Vec::with_capacity(nv); + for r in 0..nv { + let row = d_mat.get_row(r).expect("row r < nv"); + for j in 0..dim { + codes[r * padded_dim + j] = row.vector().get(j).expect("col j < dim") as u8; + } + meta.push(row.meta()); + } + Self { + codes, + meta, + dim, + padded_dim, + nv, + } + } + + /// Number of document vectors. + pub fn num_vectors(&self) -> usize { + self.nv + } +} + +// ── Entry point ────────────────────────────────────────────────── + +/// Compute per-query-vector max MinMax inner product into `state` via the staged +/// quantized pipeline. `state` (len ≥ `query.padded_nrows()`) is the caller's +/// output, left holding the raw max-IP (the caller negates for min-distance). +/// Transient scratch (`partial`, Stage-B region) is allocated internally from +/// `alloc` — the caller sizes nothing. +/// +/// `query` is the block-transposed (`i16`, `GROUP=16`, `PACK=2`) widened codes; +/// `doc` is the row-major `u8` codes at stride `query.padded_ncols()` (even); and +/// `post` carries the query/doc metadata + logical `dim`. +/// +/// # Panics +/// +/// Panics if `state.len() < query.padded_nrows()` or `query.padded_ncols() != +/// doc.vector_dim()` (the even-K contract — see the module docs). +pub(crate) fn max_ip_kernel_staged_i8( + arch: V3, + query: BlockTransposedRef<'_, i16, 16, 2>, + doc: MatRef<'_, Standard>, + post: &MinMaxPostprocess<'_>, + state: &mut [f32], + alloc: ScopedAllocator<'_>, + budget: TileBudget, +) { + let padded = query.padded_nrows(); + // `k` is the *padded* (even) column count: the driver derives the A-side + // physical row stride from it, and it must match the doc stride. + let k = query.padded_ncols(); + if state.len() < padded || k != doc.vector_dim() { + max_ip_kernel_staged_i8_panic(state.len(), padded, k, doc.vector_dim()); + } + + let b_nrows = doc.num_vectors(); + + // Empty contraction: every IP reduces to the metadata-only terms. The POC + // does not exercise `dim == 0`; fill 0 and bail rather than enter the tiling + // nest with a zero stride (matches the f32 entry's degenerate guard). + if k == 0 { + state[..padded].fill(0.0); + return; + } + + let ca = layouts::BlockTransposed::::new(); + let cb = layouts::RowMajor::::new(); + + // SAFETY: + // - `query.as_ptr()` is valid for `padded * k` i16 (block-transposed, K padded + // to even == k), and `padded` is a multiple of GROUP == A_PANEL == 16. + // - `doc.as_slice()` is `b_nrows * k` contiguous u8 (k == doc.vector_dim()). + // - `state.len() >= padded` (checked); the driver allocates its scratch from + // `alloc`. + unsafe { + tiled_reduce_staged::, MinMaxPostprocess<'_>, MaxReducer, _, _>( + arch, + &ca, + &cb, + post, + query.as_ptr(), + padded, + doc.as_slice().as_ptr(), + b_nrows, + k, + &mut state[..padded], + alloc, + budget, + ); + } +} + +#[inline(never)] +#[cold] +#[allow(clippy::panic)] +fn max_ip_kernel_staged_i8_panic(state_len: usize, padded: usize, k: usize, doc_dim: usize) { + panic!( + "max_ip_kernel_staged_i8: precondition failed: \ + state.len()={state_len} (expected >= {padded}), \ + padded_ncols(k)={k}, doc.vector_dim()={doc_dim} (must be equal — even-K contract)" + ); +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroUsize; + + use diskann_utils::ReborrowMut; + use diskann_wide::arch::x86_64::V3; + + use super::super::super::TileBudget; + use super::{ + MinMaxPostprocess, QuantStagedDocs, QuantStagedQuery, int_store_microkernel, + max_ip_kernel_staged_i8, + }; + use crate::CompressInto; + use crate::algorithms::Transform; + use crate::algorithms::transforms::NullTransform; + use crate::alloc::ScopedAllocator; + use crate::minmax::{MinMaxCompensation, MinMaxMeta, MinMaxQuantizer}; + use crate::multi_vector::distance::{MaxSim, QueryMatRef}; + use crate::multi_vector::{BlockTransposed, Defaulted, Mat, MatRef, Standard}; + use crate::num::Positive; + use diskann_vector::DistanceFunctionMut; + + const NBITS: usize = 4; + + fn quantizer(dim: usize) -> MinMaxQuantizer { + MinMaxQuantizer::new( + Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())), + Positive::new(1.0).unwrap(), + ) + } + + /// Pseudo-random f32 in roughly `[-1, 1]`, deterministic per `(seed, idx)`. + fn rnd(seed: u64, idx: usize) -> f32 { + let x = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(idx as u64) + .wrapping_mul(1442695040888963407); + ((x >> 33) as f32 / (1u64 << 31) as f32) - 1.0 + } + + fn quantize(q: &MinMaxQuantizer, data: &[f32], n: usize, dim: usize) -> Mat> { + let input = MatRef::new(Standard::::new(n, dim).unwrap(), data).unwrap(); + let mut out: Mat> = Mat::new(MinMaxMeta::new(n, dim), Defaulted).unwrap(); + q.compress_into(input, out.reborrow_mut()).unwrap(); + out + } + + /// Extract `(codes_u8 [nv × padded_dim, zero-padded], meta [nv])` from a + /// quantized doc matrix — the minimal doc-side SoA. + fn doc_soa(mat: &Mat>, dim: usize) -> (Vec, Vec) { + let nv = mat.num_vectors(); + let padded_dim = dim.next_multiple_of(2); + let mut codes = vec![0u8; nv * padded_dim]; + let mut meta = Vec::with_capacity(nv); + for r in 0..nv { + let row = mat.get_row(r).unwrap(); + for j in 0..dim { + codes[r * padded_dim + j] = row.vector().get(j).unwrap() as u8; + } + meta.push(row.meta()); + } + (codes, meta) + } + + /// Extract `(codes_i16 [nq × dim], meta padded to padded_nrows)` for the query + /// side — i16 widening + row padding feeds `BlockTransposed`. + fn query_arrays( + mat: &Mat>, + dim: usize, + padded_nrows: usize, + ) -> (Vec, Vec) { + let nq = mat.num_vectors(); + let mut codes = vec![0i16; nq * dim]; + let mut meta = vec![MinMaxCompensation::default(); padded_nrows]; + for r in 0..nq { + let row = mat.get_row(r).unwrap(); + for j in 0..dim { + codes[r * dim + j] = i16::from(row.vector().get(j).unwrap() as u8); + } + meta[r] = row.meta(); + } + (codes, meta) + } + + /// (nq, nd, dim): every B-remainder class (`nd ∈ {1,5,6,7,8}`), A-panel + /// remainder (`17`), multi-tile B (`1250`), and `dim ∈ {64,128,256}`. + const CASES: &[(usize, usize, usize)] = &[ + (1, 1, 64), + (1, 5, 64), + (5, 1, 128), + (16, 4, 64), + (16, 5, 128), + (16, 6, 64), + (16, 7, 256), + (16, 8, 128), + (17, 9, 64), + (32, 16, 256), + (8, 1250, 128), + (64, 1250, 64), + // Odd dims exercise the even-K contract (padded_dim = dim + 1, the trailing + // column zero-padded on both sides), end-to-end against the reference. + (5, 3, 63), + (17, 9, 65), + (8, 33, 127), + (16, 7, 1), + ]; + + /// The public quantized staged path must match the scalar MinMax `MaxSim` + /// reference within tolerance (Stage A's integer dot is exact; only Stage B's + /// f32 accumulation order differs from the reference's). Exercises the full + /// public API: [`QuantStagedQuery`]/[`QuantStagedDocs`] build + compute. + #[test] + fn staged_i8_matches_minmax_reference() { + if V3::new_checked().is_none() { + return; // No AVX2 on this host. + } + + for &(nq, nd, dim) in CASES { + let q_data: Vec = (0..nq * dim).map(|i| rnd(1, i)).collect(); + let d_data: Vec = (0..nd * dim).map(|i| rnd(2, i)).collect(); + + // ── Path A: the public quantized staged kernel. ── + let q_f32 = MatRef::new(Standard::::new(nq, dim).unwrap(), &q_data).unwrap(); + let d_f32 = MatRef::new(Standard::::new(nd, dim).unwrap(), &d_data).unwrap(); + let mut query = QuantStagedQuery::build(q_f32).unwrap(); + let docs = QuantStagedDocs::build(d_f32); + let mut got = vec![0.0f32; nq]; + query.compute_max_sim(&docs, &mut got); + + // ── Path B: scalar MinMax MaxSim reference (identical quantization). ── + let q = quantizer(dim); + let q_mat = quantize(&q, &q_data, nq, dim); + let d_mat = quantize(&q, &d_data, nd, dim); + let query_ref: QueryMatRef<_> = q_mat.as_view().into(); + let mut ref_scores = vec![0.0f32; nq]; + MaxSim::new(&mut ref_scores).evaluate(query_ref, d_mat.as_view()); + + for i in 0..nq { + assert!( + (got[i] - ref_scores[i]).abs() <= 1e-4 * ref_scores[i].abs().max(1.0), + "({nq},{nd},{dim}) row {i}: staged-i8 min-dist {} != reference {}", + got[i], + ref_scores[i], + ); + } + } + } + + /// Reusing a single [`QuantStagedQuery`] across multiple `compute_max_sim` + /// calls (different doc sets / counts) must give correct results every time — + /// the regression guard for the [`ResettableArena`](super::super::arena::ResettableArena) + /// reset path: each call rewinds and re-fills the shared scratch, so a stale + /// or aliased buffer would corrupt the second/third call. + #[test] + fn staged_i8_arena_reuse_across_calls() { + if V3::new_checked().is_none() { + return; // No AVX2 on this host. + } + + const NQ: usize = 17; // exercises the A-panel row padding (17 -> 32) + const DIM: usize = 128; + let q_data: Vec = (0..NQ * DIM).map(|i| rnd(5, i)).collect(); + let q_f32 = MatRef::new(Standard::::new(NQ, DIM).unwrap(), &q_data).unwrap(); + let mut query = QuantStagedQuery::build(q_f32).unwrap(); + + let quant = quantizer(DIM); + let q_mat = quantize(&quant, &q_data, NQ, DIM); + + // Distinct doc counts (multi-tile, single panel, remainder) reusing the + // same query — the arena is reset, never reallocated, between calls. + for (call, &nd) in [251usize, 3, 64, 1].iter().enumerate() { + let d_data: Vec = (0..nd * DIM).map(|i| rnd(6 + call as u64, i)).collect(); + let d_f32 = MatRef::new(Standard::::new(nd, DIM).unwrap(), &d_data).unwrap(); + let docs = QuantStagedDocs::build(d_f32); + + let mut got = vec![0.0f32; NQ]; + query.compute_max_sim(&docs, &mut got); + + let d_mat = quantize(&quant, &d_data, nd, DIM); + let query_ref: QueryMatRef<_> = q_mat.as_view().into(); + let mut ref_scores = vec![0.0f32; NQ]; + MaxSim::new(&mut ref_scores).evaluate(query_ref, d_mat.as_view()); + + for i in 0..NQ { + assert!( + (got[i] - ref_scores[i]).abs() <= 1e-4 * ref_scores[i].abs().max(1.0), + "call {call} (nd={nd}) row {i}: reused staged-i8 {} != reference {}", + got[i], + ref_scores[i], + ); + } + } + } + + /// Isolate Stage A: the raw `i32` partial it stores must equal the brute-force + /// integer code dot `⟨codes_q, codes_d⟩` exactly (no float, no metadata). + #[test] + fn stage_a_integer_dot_exact() { + let Some(arch) = V3::new_checked() else { + return; + }; + + for &dim in &[64usize, 128, 130, 256] { + let padded_dim = dim.next_multiple_of(2); + let q = quantizer(dim); + // Exactly one A-panel (16 rows) × one B-panel (4 cols). + let q_data: Vec = (0..16 * dim).map(|i| rnd(3, i)).collect(); + let d_data: Vec = (0..4 * dim).map(|i| rnd(4, i)).collect(); + let q_mat = quantize(&q, &q_data, 16, dim); + let d_mat = quantize(&q, &d_data, 4, dim); + + let (d_codes, _) = doc_soa(&d_mat, dim); + let q_i16 = { + let bt = BlockTransposed::::new(16, dim); + let (c, _) = query_arrays(&q_mat, dim, bt.padded_nrows()); + c + }; + let q_mat_view = MatRef::new(Standard::::new(16, dim).unwrap(), &q_i16).unwrap(); + let bt = BlockTransposed::::from_matrix_view(q_mat_view.as_matrix_view()); + + let mut partial = vec![0i32; 16 * 4]; + // SAFETY: `bt` has exactly one block (16 rows) at `as_ptr()`; `d_codes` + // is 4 rows × padded_dim u8; `partial` is 4 cols × 16 i32 at stride 16. + unsafe { + int_store_microkernel::<4>( + arch, + bt.as_ptr(), + d_codes.as_ptr(), + padded_dim, + partial.as_mut_ptr(), + 16, + ); + } + + // Brute-force ⟨codes⟩ over the logical dim (codes are u8 ∈ [0,15]). + for i in 0..16 { + let qr = q_mat.get_row(i).unwrap(); + for jcol in 0..4 { + let dr = d_mat.get_row(jcol).unwrap(); + let expect: i32 = (0..dim) + .map(|d| { + i32::from(qr.vector().get(d).unwrap() as u8) + * i32::from(dr.vector().get(d).unwrap() as u8) + }) + .sum(); + assert_eq!( + partial[jcol * 16 + i], + expect, + "dim={dim} A-major partial[col {jcol}, row {i}] != brute force" + ); + } + } + } + } + + /// Drive the internal entry with a deliberately tiny cache budget so the + /// planner clamps to one A-panel and one B-panel per tile. With `nq > 16` and + /// `nd > 4` this forces **multiple A-tiles and multiple B-tiles**, exercising + /// the cross-tile `a_row_offset`/`b_row_offset` carry that the default-budget + /// reference test (one giant A-tile) never reaches — yet still matching the + /// scalar reference. + #[test] + fn staged_i8_multi_tile_tiny_budget() { + let Some(arch) = V3::new_checked() else { + return; + }; + + // l2_a / l1_b of 1 clamp `a_panels_per_tile` / `b_panels_per_tile` to 1 + // (both `.max(1)` in `StagedPlan::new`): a_tile_rows = 16, b_tile_rows = 4. + let budget = TileBudget { l2_a: 1, l1_b: 1 }; + + for &(nq, nd, dim) in &[(48usize, 22usize, 64usize), (33, 37, 128), (35, 19, 65)] { + let q = quantizer(dim); + let q_data: Vec = (0..nq * dim).map(|i| rnd(5, i)).collect(); + let d_data: Vec = (0..nd * dim).map(|i| rnd(6, i)).collect(); + let q_mat = quantize(&q, &q_data, nq, dim); + let d_mat = quantize(&q, &d_data, nd, dim); + + let padded_dim = dim.next_multiple_of(2); + let (d_codes, d_meta) = doc_soa(&d_mat, dim); + let doc = MatRef::new(Standard::::new(nd, padded_dim).unwrap(), &d_codes).unwrap(); + + let q_i16 = { + let probe = BlockTransposed::::new(nq, dim); + let (c, _) = query_arrays(&q_mat, dim, probe.padded_nrows()); + c + }; + let q_view = MatRef::new(Standard::::new(nq, dim).unwrap(), &q_i16).unwrap(); + let query_bt = BlockTransposed::::from_matrix_view(q_view.as_matrix_view()); + let (_, q_meta) = query_arrays(&q_mat, dim, query_bt.padded_nrows()); + let post = MinMaxPostprocess::new(&q_meta, &d_meta, dim); + + let mut state = vec![f32::MIN; query_bt.padded_nrows()]; + max_ip_kernel_staged_i8( + arch, + query_bt.as_view(), + doc, + &post, + &mut state, + ScopedAllocator::global(), + budget, + ); + + let query_ref: QueryMatRef<_> = q_mat.as_view().into(); + let mut ref_scores = vec![0.0f32; nq]; + MaxSim::new(&mut ref_scores).evaluate(query_ref, d_mat.as_view()); + + for i in 0..nq { + let got = -state[i]; + assert!( + (got - ref_scores[i]).abs() <= 1e-4 * ref_scores[i].abs().max(1.0), + "({nq},{nd},{dim}) row {i}: tiny-budget staged-i8 {got} != reference {}", + ref_scores[i], + ); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/staged/maxsim.rs b/diskann-quantization/src/multi_vector/distance/kernels/staged/maxsim.rs new file mode 100644 index 000000000..18cdfa88c --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/staged/maxsim.rs @@ -0,0 +1,155 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Stage B / Stage C impls for MaxSim, plus the owned reset-arena scratch for the +//! `MaxSimKernel` (f32) path. +//! +//! The Stage-A kernel and the V3 entry point live in [`super::v3`]; the +//! Stage-C reducer's SIMD fold is V3-specific and also lives there. + +use core::marker::PhantomData; + +use diskann_wide::Architecture; + +use super::super::TileBudget; +use super::arena::ResettableArena; +use super::{FoldCtx, Postprocess}; +use crate::alloc::ScopedAllocator; + +// ── Stage B: identity ──────────────────────────────────────────── + +/// Identity postprocess: the raw inner product (`Acc`) *is* the `Score`. +/// +/// Reports [`scratch_len`](Postprocess::scratch_len) 0 and +/// [`apply`](Postprocess::apply) returns its input pointer unchanged, so the +/// driver folds `partial_buf` directly — no `scored_buf`, no extra pass, no +/// boolean flag. (Mirrors [`ConvertTo`](super::super::layouts::ConvertTo)'s +/// zero-cost identity blanket impl.) +pub(super) struct Identity(PhantomData); + +impl Identity { + pub(super) fn new() -> Self { + Self(PhantomData) + } +} + +// SAFETY: identity reads nothing beyond `acc`, writes nothing (`scratch_len` +// is 0), and returns exactly `acc` — `Score == Acc == T`, already A-major at the +// fixed output stride `a_panel`. The returned pointer carries the caller's +// validity for `acc` unchanged. +unsafe impl Postprocess for Identity { + type Acc = T; + type Score = T; + + #[inline(always)] + fn scratch_len(&self, _a_panel: usize, _max_b_cols: usize) -> usize { + 0 + } + + #[inline(always)] + unsafe fn apply(&self, _scratch: *mut T, _arch: A, acc: *const T, _ctx: FoldCtx) -> *const T { + acc + } +} + +/// MaxSim reducer: per-A-row running maximum of the inner products. The +/// `Reducer` impl (a register-resident `max_simd` sweep) is V3-specific and +/// lives in [`super::v3`]. +pub(super) struct MaxReducer; + +// ── Dispatch carrier ───────────────────────────────────────────── + +/// `state` (the output running-reduction) + `alloc` (the caller's allocator the +/// driver carves its internal `partial`/`scored` scratch from) bundled to cross +/// the [`Target3`](diskann_wide::arch::Target3) dispatch boundary. Only the +/// allocator crosses here — the driver allocates the scratch buffers itself, so +/// the caller hands in nothing but `state` and `alloc`. +pub(crate) struct StagedRun<'a> { + pub(crate) state: &'a mut [f32], + pub(crate) alloc: ScopedAllocator<'a>, +} + +// ── Owned reset-arena scratch for the f32 `MaxSimKernel` path ───── + +/// Per-kernel reusable scratch for the staged f32 [`MaxSimKernel`](super::super::MaxSimKernel) +/// path: the running-reduction `state` output plus a [`ResettableArena`] backing +/// the driver's transient `partial`/`scored`. Owned by `PreparedStaged` (behind a +/// `RefCell`, since `compute_max_sim` is `&self`) so steady-state calls allocate +/// nothing — the arena is reset, not reallocated, each call. +/// +/// This is the f32 counterpart to the i8 path's `QuantStagedQuery`-owned arena; +/// it lives here (not in `factory`) because sizing reads [`TileBudget`], which is +/// private to the `kernels` module tree. +#[derive(Debug)] +pub(crate) struct F32StagedScratch { + state: Vec, + arena: ResettableArena, +} + +impl F32StagedScratch { + /// Build scratch for a query of `padded` rows. The arena is sized once to the + /// provable `2·l1_b` ceiling (`StagedPlan` co-budgets `partial`/`scored` each + /// `≤ l1_b`), so it never needs per-shape sizing or reallocation. + #[allow(clippy::expect_used)] // POC: 72 KB arena, OOM is not a recoverable case here. + pub(crate) fn new(padded: usize) -> Self { + let arena_bytes = 2 * TileBudget::default().l1_b + 4096; + Self { + state: vec![f32::MIN; padded], + arena: ResettableArena::with_capacity(arena_bytes).expect("f32 staged arena"), + } + } + + /// Reset the arena, ensure `state` covers `padded` rows, then run `f` with the + /// `state` slice (the driver re-initialises it) and a `ScopedAllocator` over + /// the arena. + pub(crate) fn run( + &mut self, + padded: usize, + f: impl FnOnce(&mut [f32], ScopedAllocator<'_>) -> R, + ) -> R { + if self.state.len() < padded { + self.state.resize(padded, f32::MIN); + } + // Split the borrow so `state` (mut) and `arena` (shared, via the allocator) + // are simultaneously live as disjoint fields. + let Self { state, arena } = self; + arena.reset(); + f(&mut state[..padded], ScopedAllocator::new(arena)) + } +} + +#[cfg(test)] +mod tests { + use diskann_wide::arch::Scalar; + + use super::*; + + /// The identity postprocess returns its input pointer unchanged (the + /// zero-cost identity), so the driver folds `partial_buf` directly. + #[test] + fn identity_returns_source_pointer() { + let acc = [1.0f32, 2.0, 3.0, 4.0]; + let id = Identity::::new(); + let ctx = FoldCtx { + a_panel: 2, + valid_b_cols: 2, + b_stride: 2, + a_row_offset: 0, + b_row_offset: 0, + }; + // SAFETY: identity ignores `scratch` (its `scratch_len` is 0) and returns + // its input pointer unchanged. UFCS pins `A = Scalar` (the impl is blanket + // over `A`); the real driver pins `A = V3` via turbofish, so production + // never needs this. + let out = unsafe { + as Postprocess>::apply( + &id, + core::ptr::null_mut(), + Scalar::new(), + acc.as_ptr(), + ctx, + ) + }; + assert_eq!(out, acc.as_ptr()); + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/staged/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/staged/mod.rs new file mode 100644 index 000000000..e151ec1cc --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/staged/mod.rs @@ -0,0 +1,428 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Experimental *staged* multi-vector distance kernel. +//! +//! The production kernel (`super::tiled_reduce` + `super::f32`) fuses three +//! concerns into one micro-kernel epilogue: inner-product accumulation, +//! cross-row reduction, and merge into the per-A-row score scratch. This module +//! is a parallel, *separately selectable* kernel that splits those concerns into +//! three independently-pluggable stages so future work (quantized distances, +//! other reductions) can reuse the tiling loop without forking the micro-kernel: +//! +//! * **Stage A — [`StagedKernel`]**: pure SIMD math. Writes a raw `Acc` block +//! (`f32` here) into a per-A-panel `partial_buf`; no reduction, no merge. +//! * **Stage B — [`Postprocess`]**: `Acc` → `Score`, modeled on +//! [`ConvertTo`](super::layouts::ConvertTo). `apply` returns a read pointer and +//! the identity impl returns its input unchanged (reporting `scratch_len` 0), so +//! the driver runs it uniformly and identity stays zero-cost — no extra memory +//! pass, no boolean flag. +//! * **Stage C — [`Reducer`]**: owns the per-A-row `State` and folds `Score` +//! blocks into it. +//! +//! The traits keep raw-pointer methods (no generic methods) so a future +//! `&dyn Postprocess` / `&dyn Reducer` switched at the per-(A-panel, B-tile) +//! boundary stays possible without a monomorphization blow-up; we keep them +//! static-generic for now. +//! +//! Scope: single K-segment (no fractured-K), f32 + V3 (AVX2/FMA) only — an +//! apples-to-apples A/B against the fused V3 kernel. + +use diskann_wide::Architecture; + +use super::TileBudget; +use super::layouts::Layout; + +pub(super) mod arena; +pub(super) mod driver; +pub(super) mod i8; +pub(super) mod maxsim; +pub(super) mod v3; + +pub(crate) use maxsim::{F32StagedScratch, StagedRun}; +pub(crate) use v3::StagedF32Kernel; +// Public POC entry for the quantized (4-bit MinMax) staged kernel. +pub use i8::{QuantStagedDocs, QuantStagedQuery}; + +// ── Stage A: kernel ────────────────────────────────────────────── + +/// Stage A micro-kernel. Computes an `A_PANEL × B_PANEL` block of raw `Acc` +/// accumulators and **writes them out** to a partial buffer — unlike +/// [`super::Kernel`], it performs no cross-row reduction or scratch merge. +/// +/// # Safety +/// +/// Implementors must respect the per-method pointer contracts. +pub(super) unsafe trait StagedKernel { + /// Layout consumed by the A (left / query) side. + type Left: Layout; + /// Layout consumed by the B (right / document) side. + type Right: Layout; + /// Raw accumulator element written into `partial_buf` (`f32` for MaxSim). + type Acc: Copy; + + /// A rows processed per invocation (= the query `BlockTransposed` GROUP). + const A_PANEL: usize; + /// B rows processed per full invocation. + const B_PANEL: usize; + + /// Number of `Acc` elements the per-(A-panel, B-tile) `partial_buf` needs at + /// contraction dim `k` under `budget`. The single source of truth for the + /// partial size formula: the driver sizes its internal allocation from this, + /// derived from the panel geometry + the element sizes the kernel already + /// knows. + fn partial_len(k: usize, budget: TileBudget) -> usize { + let a_elem = core::mem::size_of::<::Element>(); + let b_elem = core::mem::size_of::<::Element>(); + let acc = core::mem::size_of::(); + StagedPlan::new( + k * a_elem, + k * b_elem, + Self::A_PANEL, + Self::B_PANEL, + acc, + budget, + ) + .partial_len(Self::A_PANEL, Self::B_PANEL) + } + + /// Write a full `A_PANEL × B_PANEL` block into `partial`. + /// + /// `partial` points at the first B-column of this panel; column `j` + /// (`0..B_PANEL`) and its `A_PANEL` rows occupy `partial[j*partial_b_stride + /// ..][..A_PANEL]` (A-major: `partial_b_stride == A_PANEL`). + /// + /// # Safety + /// + /// * `a` valid for `A_PANEL * k` `Left::Element`. + /// * `b` valid for `B_PANEL * k` `Right::Element`. + /// * `partial` valid for `B_PANEL` columns of `A_PANEL` `Acc` at stride + /// `partial_b_stride`. + unsafe fn full_panel( + arch: A, + a: *const ::Element, + b: *const ::Element, + k: usize, + partial: *mut Self::Acc, + partial_b_stride: usize, + ); + + /// Like [`Self::full_panel`] but writes only `remainder` (`1..B_PANEL`) + /// B-columns. + /// + /// # Safety + /// + /// As [`Self::full_panel`], with `b` valid for `remainder * k` + /// `Right::Element` and only `remainder` columns written. + unsafe fn partial_panel( + arch: A, + remainder: usize, + a: *const ::Element, + b: *const ::Element, + k: usize, + partial: *mut Self::Acc, + partial_b_stride: usize, + ); +} + +// ── Stage B: postprocess ───────────────────────────────────────── + +/// Per-(A-panel, B-tile) context passed to [`Postprocess::apply`]. +/// +/// All fields are cheap `Copy` scalars the driver already tracks. The identity +/// postprocess ignores them (they vanish after inlining); a metadata-bearing +/// postprocess (e.g. the quantized path) uses the global +/// `a_row_offset`/`b_row_offset` to index its per-vector metadata. +#[derive(Debug, Clone, Copy)] +pub(super) struct FoldCtx { + /// Rows in this A-panel (`== A_PANEL`). + pub(super) a_panel: usize, + /// Valid B columns in this block (`≤` the tile width). + pub(super) valid_b_cols: usize, + /// `Acc` column stride within the partial block (`== a_panel`). + pub(super) b_stride: usize, + /// Global index of this A-panel's first row. + pub(super) a_row_offset: usize, + /// Global index of this B-tile's first row. + pub(super) b_row_offset: usize, +} + +/// Stage B: convert one A-major block of raw `Acc` accumulators into finished +/// `Score`s, returning a read pointer to the scores. +/// +/// The driver sizes the staging region from [`scratch_len`](Self::scratch_len) +/// and allocates it (from the caller's allocator), then hands it to +/// [`apply`](Self::apply) as a raw `*mut Score`. The identity impl +/// ([`Identity`](maxsim::Identity)) reports `scratch_len == 0` and returns its +/// input pointer unchanged, so the driver runs one uniform path — no boolean, no +/// branch — and the identity case is zero-cost. A non-identity impl — e.g. the +/// quantized [`MinMaxPostprocess`](i8::MinMaxPostprocess), which turns raw `i32` +/// integer dot products into f32 MinMax scores using its own captured per-vector +/// metadata — writes into `scratch` and returns a pointer into it. +/// +/// # Safety +/// +/// Implementors must ensure that [`apply`](Self::apply): +/// - reads at most `ctx.valid_b_cols` columns of `ctx.a_panel` `Acc` from `acc` +/// at column stride `ctx.b_stride` (never the stale padded remainder columns); +/// - writes only within the `scratch` region it was given; +/// - returns a `*const Score` valid for `ctx.valid_b_cols` columns of +/// `ctx.a_panel` `Score` at the **fixed output column stride `ctx.a_panel`**. +/// (Identity returns `acc`, already at stride `a_panel`.) +pub(super) unsafe trait Postprocess { + /// Raw accumulator type produced by Stage A. + type Acc: Copy; + /// Finished score type consumed by Stage C. + type Score: Copy; + + /// Number of `Score` elements [`apply`](Self::apply) needs for one A-panel + /// against a B-tile of up to `max_b_cols` columns. `0` for the identity + /// postprocess (no staging — `apply` returns `acc`); the driver allocates + /// exactly this many `Score`s from the caller's allocator and passes the + /// region to `apply`. + fn scratch_len(&self, a_panel: usize, max_b_cols: usize) -> usize; + + /// Convert the `ctx.a_panel × ctx.valid_b_cols` A-major `Acc` block at `acc` + /// (column stride `ctx.b_stride`) into `Score`s, returning a read pointer. + /// Output is A-major at the fixed column stride `ctx.a_panel`; the identity + /// impl returns `acc` unchanged. Metadata-bearing impls index their own + /// per-vector metadata by `ctx.a_row_offset` / `ctx.b_row_offset`. + /// + /// # Safety + /// + /// * `acc` is valid for `ctx.valid_b_cols` columns of `ctx.a_panel` `Acc` at + /// stride `ctx.b_stride`. + /// * `scratch` is valid+writable for `scratch_len(ctx.a_panel, max_b_cols)` + /// `Score` with `max_b_cols ≥ ctx.valid_b_cols` (dangling when that is `0`). + unsafe fn apply( + &self, + scratch: *mut Self::Score, + arch: A, + acc: *const Self::Acc, + ctx: FoldCtx, + ) -> *const Self::Score; +} + +// ── Stage C: reducer ───────────────────────────────────────────── + +/// Stage C owns the per-A-row reduction `State` and folds `Score` blocks into +/// it. `Max` here; richer `State` shapes (argmax `(f32,u32)`, top-k) and an +/// `Output`/`finalize` step are follow-on work. +pub(super) trait Reducer { + /// Score element folded in (matches [`Postprocess::Score`]). + type Score: Copy; + /// Per-A-row running state (the score scratch element). + type State: Copy; + + /// Identity state for an A-row before any B-rows are seen. + fn init() -> Self::State; + + /// Fold an `A_PANEL × valid_b_cols` block of `Score` (read from + /// `partial_buf`) into `state[0..a_panel]`, in place. + /// + /// Column `c` (`0..valid_b_cols`), row `i` (`0..a_panel`) is at + /// `scores[c*b_stride + i]`. Only `valid_b_cols` columns are read — the + /// padded remainder columns hold stale data and **must not** be folded. + /// + /// # Safety + /// + /// * `state` valid+writable for `a_panel` `State`. + /// * `scores` valid for `valid_b_cols` columns of `a_panel` `Score` at + /// stride `b_stride`. + unsafe fn fold_block( + arch: A, + state: *mut Self::State, + scores: *const Self::Score, + a_panel: usize, + valid_b_cols: usize, + b_stride: usize, + ); +} + +// ── Stage conversion: StagedConvert ────────────────────────────── + +/// Staged-local tile conversion from layout `Self` to layout `To` — the staged +/// path's self-contained replacement for the shared +/// [`ConvertTo`](super::layouts::ConvertTo). Same role (convert a tile of source +/// data into the kernel's element type), but inverted ownership: instead of +/// owning a `Buffer`, the impl reports a [`scratch_len`](Self::scratch_len) and +/// the **driver** allocates that staging region from the caller's allocator and +/// hands it to [`convert`](Self::convert). The blanket identity impl reports `0` +/// and returns `src` unchanged, so identity conversions cost nothing and the +/// staged driver never touches the shared `ConvertTo` machinery. +/// +/// # Safety +/// +/// Implementors must ensure [`convert`](Self::convert) reads at most `rows * k` +/// source elements, writes only within the `scratch` region it was given, and +/// returns a pointer valid for `rows * k` `To::Element`. +pub(super) unsafe trait StagedConvert: Layout { + /// Number of `To::Element` the driver must allocate to convert up to + /// `max_tile_rows × k`. `0` for identity (no conversion — `convert` returns + /// `src`, ignoring `scratch`). + fn scratch_len(&self, max_tile_rows: usize, k: usize) -> usize; + + /// Convert `rows × k` `Self::Element` at `src` into `To::Element`, writing + /// into `scratch` (the driver-allocated region of `scratch_len(..)`), and + /// returning a read pointer. The identity impl returns `src` unchanged. + /// + /// # Safety + /// + /// * `src` points to `rows * k` valid `Self::Element`. + /// * `scratch` is valid+writable for `scratch_len(max_tile_rows, k)` + /// `To::Element` with `max_tile_rows ≥ rows` (dangling when that is `0`). + unsafe fn convert( + &self, + scratch: *mut To::Element, + arch: A, + src: *const Self::Element, + rows: usize, + k: usize, + ) -> *const To::Element; +} + +/// Identity conversion: every layout converts to itself at zero cost (no +/// scratch, returns `src`). Mirrors the shared `ConvertTo` blanket identity. +// SAFETY: identity reads nothing beyond `src`, writes nothing (`scratch_len` is +// 0), and returns exactly `src`, valid for the caller's lifetime. +unsafe impl StagedConvert for L { + fn scratch_len(&self, _max_tile_rows: usize, _k: usize) -> usize { + 0 + } + + unsafe fn convert( + &self, + _scratch: *mut L::Element, + _arch: A, + src: *const L::Element, + _rows: usize, + _k: usize, + ) -> *const L::Element { + src + } +} + +// ── Planner ────────────────────────────────────────────────────── + +/// Tile-panel counts for the staged loop. +/// +/// Encodes the **partial-buffer granularity** decision: the partial buffer is +/// always **one A-panel** (`P_a = 1`) wide in the A direction, against **as many +/// B-panels as co-fit L1** (`P_b = b_panels_per_tile`) in the B direction. +/// +/// `P_a = 1` because extra A-panels cut no B reads (B is re-streamed per A-panel +/// regardless) — they only enlarge `partial_buf`. `P_b = co-fit` rather than +/// `1×1` because total partial traffic is invariant to the fold granularity, so +/// the widest fold minimizes fold-call / `state`-reload overhead, maximizes the +/// SIMD sweep, and keeps Stage A (kernel) and Stage C (reduce) as contiguous +/// non-interleaved phases. See `docs/staged_multi_vector_kernel.md` §5. +#[derive(Debug, Clone, Copy)] +pub(super) struct StagedPlan { + pub(super) a_panels_per_tile: usize, + pub(super) b_panels_per_tile: usize, +} + +impl StagedPlan { + /// Choose `a_panels_per_tile` / `b_panels_per_tile` from the cache budgets, + /// the panel sizes, and the partial-buffer footprint. + /// + /// **L2** holds the A-tile (it is reused across every B-tile): + /// + /// ```text + /// a_panels_per_tile · A_PANEL · a_row_bytes ≤ l2_a + /// ``` + /// + /// **L1** holds *three* things at once during Stage A / Stage C, so the + /// planner co-budgets all three against `l1_b` (a usable fraction of L1) + /// rather than letting each independently claim the whole budget: + /// + /// * one A micro-panel — `A_PANEL · a_row_bytes` (re-read per B-panel); + /// * the B-tile data — `B_TILE_ROWS · b_row_bytes` (re-read per A-panel); + /// * `partial_buf` — `A_PANEL · B_TILE_ROWS · acc_bytes`. + /// + /// Each B-row added to the tile therefore costs `b_row_bytes` of document + /// data **plus** `A_PANEL · acc_bytes` of partial scratch, so we keep the + /// largest `B_TILE_ROWS` satisfying + /// + /// ```text + /// A_PANEL·a_row_bytes + B_TILE_ROWS·(b_row_bytes + A_PANEL·acc_bytes) ≤ l1_b + /// ``` + /// + /// This bounds `partial_buf + B-tile` together. (For very large `k` the A + /// micro-panel alone can approach `l1_b`; then `b_panels_per_tile` clamps to + /// 1 and the A-panel is the limit — inherent, and identical to the fused + /// kernel's behaviour.) + pub(super) fn new( + a_row_bytes: usize, + b_row_bytes: usize, + a_panel: usize, + b_panel: usize, + acc_bytes: usize, + budget: TileBudget, + ) -> Self { + let a_row_bytes = a_row_bytes.max(1); + let b_row_bytes = b_row_bytes.max(1); + + // L2: the A-tile is reused across all B-tiles, so size it to L2. + let a_panels_per_tile = (budget.l2_a / (a_row_bytes * a_panel)).max(1); + + // L1: co-budget the A micro-panel + B-tile data + partial_buf. Each + // B-row costs its document data plus one A_PANEL-tall partial column. + let a_panel_bytes = a_panel * a_row_bytes; + let bytes_per_b_row = b_row_bytes + a_panel * acc_bytes; + let b_tile_budget = budget.l1_b.saturating_sub(a_panel_bytes); + let b_panels_per_tile = ((b_tile_budget / bytes_per_b_row) / b_panel).max(1); + + Self { + a_panels_per_tile, + b_panels_per_tile, + } + } + + /// `partial_buf` capacity (in `Acc` elements): one A-panel (`P_a = 1`) wide, + /// covering a full B-tile (`b_panels_per_tile · B_PANEL` rows). The driver + /// caps every B-tile at `b_tile_rows`, so this is the exact upper bound. + pub(super) fn partial_len(&self, a_panel: usize, b_panel: usize) -> usize { + a_panel * self.b_panels_per_tile * b_panel + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Mirror of `L1_CACHE` in `super::super::TileBudget::default`. + const L1_CACHE_BYTES: usize = 48_000; + + /// The co-budget must keep the inner-loop L1 working set — one A + /// micro-panel + the B-tile data + `partial_buf` — within real L1 for every + /// realistic `k`. The previous design budgeted `partial_buf` and the B-tile + /// independently against `l1_b` and overflowed at small/moderate `k` (e.g. + /// k=16 placed ~70 KB into a 48 KB L1); this pins the fix. + #[test] + fn l1_working_set_fits_for_all_k() { + const A_PANEL: usize = 16; + const B_PANEL: usize = 4; + const ACC: usize = 4; // f32 + + // k up to 512: beyond ~768 the A micro-panel alone exceeds L1, which is + // inherent (the fused kernel hits the same wall) and not a planner bug. + for &k in &[1usize, 2, 4, 8, 16, 32, 64, 128, 256, 384, 512] { + let row = k * 4; // f32 element + let plan = StagedPlan::new(row, row, A_PANEL, B_PANEL, ACC, TileBudget::default()); + let b_tile_rows = plan.b_panels_per_tile * B_PANEL; + + let a_panel_bytes = A_PANEL * row; + let b_data_bytes = b_tile_rows * row; + let partial_bytes = A_PANEL * b_tile_rows * ACC; + let working_set = a_panel_bytes + b_data_bytes + partial_bytes; + + assert!( + working_set <= L1_CACHE_BYTES, + "k={k}: L1 working set {working_set} B (a_panel={a_panel_bytes}, \ + b_data={b_data_bytes}, partial={partial_bytes}) exceeds L1 {L1_CACHE_BYTES} B", + ); + assert!(plan.a_panels_per_tile >= 1 && plan.b_panels_per_tile >= 1); + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/staged/v3.rs b/diskann-quantization/src/multi_vector/distance/kernels/staged/v3.rs new file mode 100644 index 000000000..e09328aa9 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/staged/v3.rs @@ -0,0 +1,395 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! V3 (AVX2+FMA) staged path: the store-out micro-kernel (Stage A), the +//! register-resident `max_simd` reducer (Stage C), and the f32 entry point. +//! +//! The Stage-A k-loop is byte-identical to the fused `super::super::f32::v3` +//! 16×4 kernel, so per-(query, doc) inner products are bit-identical. Only the +//! epilogue differs: instead of reducing the `UNROLL` accumulators and merging +//! into the score scratch, it **stores** them into `partial_buf` (A-major), +//! deferring the reduction to Stage C. +//! +//! The whole module is V3-specific; the experiment targets V3 only. + +use diskann_wide::arch::Target3; +use diskann_wide::arch::x86_64::V3; +use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDVector}; + +use super::super::TileBudget; +use super::super::layouts::{self, DescribeLayout, Layout}; +use super::driver::tiled_reduce_staged; +use super::maxsim::{Identity, MaxReducer, StagedRun}; +use super::{Reducer, StagedConvert, StagedKernel}; +use crate::alloc::ScopedAllocator; +use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; + +diskann_wide::alias!(f32s = ::f32x8); + +/// Zero-sized Stage-A kernel marker for the f32 staged path with block size +/// `GROUP`. +pub(crate) struct StagedF32Kernel; + +// SAFETY: `full_panel`/`partial_panel` read A_PANEL(16) * k A elements and +// UNROLL * k B elements, and write UNROLL columns of A_PANEL(16) f32 into +// `partial` at stride `partial_b_stride` — all within the bounds the +// `StagedKernel` contract guarantees. +unsafe impl StagedKernel for StagedF32Kernel<16> { + type Left = layouts::BlockTransposed; + type Right = layouts::RowMajor; + type Acc = f32; + const A_PANEL: usize = 16; + const B_PANEL: usize = 4; + + #[inline(always)] + unsafe fn full_panel( + arch: V3, + a: *const f32, + b: *const f32, + k: usize, + partial: *mut f32, + partial_b_stride: usize, + ) { + // SAFETY: pointer validity per the `StagedKernel` contract. + unsafe { store_microkernel::<{ Self::B_PANEL }>(arch, a, b, k, partial, partial_b_stride) } + } + + #[inline(always)] + unsafe fn partial_panel( + arch: V3, + remainder: usize, + a: *const f32, + b: *const f32, + k: usize, + partial: *mut f32, + partial_b_stride: usize, + ) { + // SAFETY: pointer validity per the `StagedKernel` contract. + unsafe { + match remainder { + 1 => store_microkernel::<1>(arch, a, b, k, partial, partial_b_stride), + 2 => store_microkernel::<2>(arch, a, b, k, partial, partial_b_stride), + 3 => store_microkernel::<3>(arch, a, b, k, partial, partial_b_stride), + _ => unreachable!( + "unexpected remainder {remainder} for B_PANEL={}", + Self::B_PANEL + ), + } + } + } +} + +/// V3 store-out micro-kernel: 16 A-rows × `UNROLL` B-rows. +/// +/// The accumulation loop matches `super::super::f32::v3::f32_microkernel` +/// exactly (two `f32x8` register tiles, FMA, same splat/stride/unroll order). +/// The epilogue stores each B-column's 16 A-row accumulators into `partial` +/// (A-major: column `j` at `partial + j*b_stride`, as two `f32x8` halves). +/// +/// # Safety +/// +/// 1. `a_packed` points to `16 * k` contiguous `f32`. +/// 2. `b` points to `UNROLL` rows of `k` contiguous `f32`. +/// 3. `partial` is valid for `UNROLL` columns of 16 `f32` at stride `b_stride`. +#[inline(always)] +unsafe fn store_microkernel( + arch: V3, + a_packed: *const f32, + b: *const f32, + k: usize, + partial: *mut f32, + b_stride: usize, +) { + let mut p0 = [f32s::default(arch); UNROLL]; + let mut p1 = [f32s::default(arch); UNROLL]; + let offsets: [usize; UNROLL] = core::array::from_fn(|i| k * i); + + let a_stride = 2 * f32s::LANES; + let a_stride_half = f32s::LANES; + + for i in 0..k { + // SAFETY: preconditions 1 and 2; i < k and j < UNROLL. + unsafe { + let a0 = f32s::load_simd(arch, a_packed.add(a_stride * i)); + let a1 = f32s::load_simd(arch, a_packed.add(a_stride * i + a_stride_half)); + + for j in 0..UNROLL { + let bj = f32s::splat(arch, b.add(i + offsets[j]).read_unaligned()); + p0[j] = a0.mul_add_simd(bj, p0[j]); + p1[j] = a1.mul_add_simd(bj, p1[j]); + } + } + } + + for j in 0..UNROLL { + // SAFETY: precondition 3; column j occupies [j*b_stride, j*b_stride+16). + unsafe { + p0[j].store_simd(partial.add(j * b_stride)); + p1[j].store_simd(partial.add(j * b_stride + a_stride_half)); + } + } +} + +// ── Stage C: V3 SIMD max reducer ───────────────────────────────── + +impl Reducer for MaxReducer { + type Score = f32; + type State = f32; + + #[inline(always)] + fn init() -> f32 { + f32::MIN + } + + #[inline(always)] + unsafe fn fold_block( + arch: V3, + state: *mut f32, + scores: *const f32, + a_panel: usize, + valid_b_cols: usize, + b_stride: usize, + ) { + let lanes = f32s::LANES; + + // The V3 staged kernel always folds a full A_PANEL = 16 = 2*LANES block: + // two register-resident accumulators sweep the valid B-columns of + // `partial_buf` in a single pass — the same access pattern the fused + // kernel uses, just hoisted out of the inner B-loop. + debug_assert_eq!( + a_panel, + 2 * lanes, + "V3 MaxReducer expects A_PANEL == 2*LANES" + ); + + // SAFETY: `state` is writable for 16; `scores` is valid for `valid_b_cols` + // columns of 16 f32 at `b_stride`; only the valid columns are read (never + // the stale padded remainder). + unsafe { + let mut a0 = f32s::load_simd(arch, state); + let mut a1 = f32s::load_simd(arch, state.add(lanes)); + for c in 0..valid_b_cols { + let col = scores.add(c * b_stride); + a0 = a0.max_simd(f32s::load_simd(arch, col)); + a1 = a1.max_simd(f32s::load_simd(arch, col.add(lanes))); + } + a0.store_simd(state); + a1.store_simd(state.add(lanes)); + } + } +} + +// ── Entry point ────────────────────────────────────────────────── + +/// Compute per-A-row max inner product (block-transposed A query, row-major B +/// doc) into `state` via the staged pipeline. `state` (len ≥ `padded_nrows`) is +/// the caller's output, left holding the raw max-IP (the caller negates). +/// Transient scratch (`partial`, Stage-B region) is allocated internally from +/// `alloc` — the caller sizes nothing. +/// +/// # Panics +/// +/// Panics if `state.len() < a.padded_nrows()` or `a.ncols() != b.vector_dim()`. +pub(crate) fn max_ip_kernel_staged( + arch: V3, + a: BlockTransposedRef<'_, f32, GROUP>, + b: MatRef<'_, Standard>, + state: &mut [f32], + alloc: ScopedAllocator<'_>, + budget: TileBudget, +) where + StagedF32Kernel: StagedKernel, + layouts::BlockTransposed: StagedConvert as StagedKernel>::Left> + + Layout, + layouts::RowMajor: StagedConvert as StagedKernel>::Right> + + Layout, +{ + let padded = a.padded_nrows(); + if state.len() < padded || a.ncols() != b.vector_dim() { + max_ip_kernel_staged_panic(state.len(), padded, a.ncols(), b.vector_dim()); + } + + // A_PANEL must equal GROUP for block-transposed layout correctness. + const { assert!( as StagedKernel>::A_PANEL == GROUP) } + + let k = a.ncols(); + let b_nrows = b.num_vectors(); + + // Empty contraction: every IP is 0 ⇒ max-IP is 0. Callers guarantee + // b_nrows > 0 (the zero-doc case is short-circuited before reaching here). + if k == 0 { + state[..padded].fill(0.0); + return; + } + + let ca = a.layout(); + let cb = b.layout(); + let post = Identity::::new(); + + // SAFETY: + // - a.as_ptr() is valid for padded * k f32, and padded is a multiple of + // GROUP == A_PANEL (const-asserted above). + // - b.as_slice() is num_vectors * vector_dim contiguous f32. + // - state.len() >= padded (checked); the driver allocates its scratch from + // `alloc`. + unsafe { + tiled_reduce_staged::, Identity, MaxReducer, _, _>( + arch, + &ca, + &cb, + &post, + a.as_ptr(), + padded, + b.as_slice().as_ptr(), + b_nrows, + k, + &mut state[..padded], + alloc, + budget, + ); + } +} + +#[inline(never)] +#[cold] +#[allow(clippy::panic)] +fn max_ip_kernel_staged_panic(state_len: usize, padded: usize, a_ncols: usize, b_dim: usize) { + panic!( + "max_ip_kernel_staged: precondition failed: \ + state.len()={state_len} (expected >= {padded}), \ + a.ncols()={a_ncols}, b.vector_dim()={b_dim}" + ); +} + +// ── Dispatch glue ──────────────────────────────────────────────── + +impl + Target3, MatRef<'_, Standard>, StagedRun<'_>> + for StagedF32Kernel +where + StagedF32Kernel: StagedKernel, + layouts::BlockTransposed: + StagedConvert>::Left> + Layout, + layouts::RowMajor: + StagedConvert>::Right> + Layout, +{ + #[inline(always)] + fn run( + self, + arch: V3, + lhs: BlockTransposedRef<'_, f32, GROUP>, + rhs: MatRef<'_, Standard>, + scratch: StagedRun<'_>, + ) { + max_ip_kernel_staged( + arch, + lhs, + rhs, + scratch.state, + scratch.alloc, + TileBudget::default(), + ); + } +} + +#[cfg(test)] +mod tests { + use diskann_wide::arch::x86_64::V3; + + use super::super::super::TileBudget; + use super::super::super::f32::max_ip_kernel; + use super::max_ip_kernel_staged; + use crate::alloc::ScopedAllocator; + use crate::multi_vector::{BlockTransposed, MatRef, Standard}; + + // (a_nrows, b_nrows, dim): degenerate, zero-dim, zero-doc, prime k, A/B-panel + // boundaries and every B-remainder class for V3 (B_PANEL=4), plus multi-tile. + const CASES: &[(usize, usize, usize)] = &[ + (1, 1, 4), + (1, 5, 8), + (5, 1, 8), + (5, 3, 5), + (3, 2, 0), // zero dim + (3, 0, 4), // zero docs + (7, 7, 32), + (2, 3, 128), + (16, 4, 64), // one A-panel, no B remainder + (17, 4, 64), // A-panel remainder + (16, 5, 8), // B remainder = 1 + (16, 6, 32), // B remainder = 2 + (16, 7, 32), // B remainder = 3 + (16, 8, 32), + (32, 5, 16), + (48, 3, 16), + (8, 32, 128), + (64, 32, 128), + (32, 16, 256), + (64, 1250, 512), // multi-tile B + ]; + + fn naive(a: &[f32], a_nrows: usize, b: &[f32], b_nrows: usize, k: usize) -> Vec { + (0..a_nrows) + .map(|i| { + (0..b_nrows) + .map(|j| (0..k).map(|d| a[i * k + d] * b[j * k + d]).sum::()) + .fold(f32::MIN, f32::max) + }) + .collect() + } + + /// The staged kernel must be **bit-identical** to the fused V3 kernel (same + /// k-loop ⇒ same per-pair IP; `max` is order-independent), and within + /// tolerance of the naive reference. + #[test] + fn staged_matches_fused_v3() { + let Some(arch) = V3::new_checked() else { + // No AVX2/FMA on this host; the staged V3 path cannot run. + return; + }; + + for &(a_nrows, b_nrows, dim) in CASES { + let a_data: Vec = (0..a_nrows * dim).map(|i| (i % 13 + 1) as f32).collect(); + let b_data: Vec = (0..b_nrows * dim).map(|i| (i % 7 + 1) as f32).collect(); + + let a_mat = MatRef::new(Standard::new(a_nrows, dim).unwrap(), &a_data).unwrap(); + let a_bt = BlockTransposed::::from_matrix_view(a_mat.as_matrix_view()); + let b_mat = MatRef::new(Standard::new(b_nrows, dim).unwrap(), &b_data).unwrap(); + + let mut fused = vec![f32::MIN; a_bt.padded_nrows()]; + max_ip_kernel::( + arch, + a_bt.as_view(), + b_mat, + &mut fused, + TileBudget::default(), + ); + + let mut state = vec![f32::MIN; a_bt.padded_nrows()]; + max_ip_kernel_staged::<16>( + arch, + a_bt.as_view(), + b_mat, + &mut state, + ScopedAllocator::global(), + TileBudget::default(), + ); + + let expected = naive(&a_data, a_nrows, &b_data, b_nrows, dim); + for i in 0..a_nrows { + assert_eq!( + state[i].to_bits(), + fused[i].to_bits(), + "staged != fused at row {i} for ({a_nrows},{b_nrows},{dim}): staged={}, fused={}", + state[i], + fused[i], + ); + assert!( + (state[i] - expected[i]).abs() < 1e-6 * expected[i].abs().max(1.0), + "staged != naive at row {i} for ({a_nrows},{b_nrows},{dim}): staged={}, naive={}", + state[i], + expected[i], + ); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/mod.rs b/diskann-quantization/src/multi_vector/distance/mod.rs index ef336161c..bdbc9e83c 100644 --- a/diskann-quantization/src/multi_vector/distance/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/mod.rs @@ -53,3 +53,8 @@ pub use fallback::QueryMatRef; pub use isa::{MaxSimIsa, NotSupported}; pub use kernel::{BoxErase, Erase, MaxSimKernel}; pub use max_sim::{Chamfer, MaxSim, MaxSimError}; + +/// Standalone POC entry for the 4-bit MinMax *staged* multi-vector MaxSim kernel +/// (V3/AVX2 only) — not yet unified into [`MaxSimIsa`]/[`build_max_sim`]. +#[cfg(target_arch = "x86_64")] +pub use kernels::{QuantStagedDocs, QuantStagedQuery}; From 817c10893ecea0f169289bc2de8b6235b3ad7394 Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Tue, 21 Jul 2026 04:43:09 +0530 Subject: [PATCH 2/5] Add tiler experimentation --- .../example/multi-vector-tiled-f16.json | 20 + diskann-benchmark/src/inputs/multi_vector.rs | 78 +++ diskann-benchmark/src/multi_vector/mod.rs | 5 + diskann-benchmark/src/multi_vector/quant.rs | 37 +- .../src/multi_vector/tiled_f16.rs | 259 ++++++++++ .../src/multi_vector/distance/kernels/mod.rs | 7 + .../distance/kernels/tiler/arena.rs | 90 ++++ .../distance/kernels/tiler/f16.rs | 304 ++++++++++++ .../distance/kernels/tiler/leaves.rs | 207 ++++++++ .../distance/kernels/tiler/minmax.rs | 445 ++++++++++++++++++ .../distance/kernels/tiler/mod.rs | 355 ++++++++++++++ .../distance/kernels/tiler/tilers.rs | 360 ++++++++++++++ .../src/multi_vector/distance/mod.rs | 12 + 13 files changed, 2173 insertions(+), 6 deletions(-) create mode 100644 diskann-benchmark/example/multi-vector-tiled-f16.json create mode 100644 diskann-benchmark/src/multi_vector/tiled_f16.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/tiler/arena.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/tiler/f16.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/tiler/leaves.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/tiler/minmax.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/tiler/mod.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/tiler/tilers.rs diff --git a/diskann-benchmark/example/multi-vector-tiled-f16.json b/diskann-benchmark/example/multi-vector-tiled-f16.json new file mode 100644 index 000000000..4ef913aa0 --- /dev/null +++ b/diskann-benchmark/example/multi-vector-tiled-f16.json @@ -0,0 +1,20 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-tiled-f16-op", + "content": { + "runs": [ + { "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 500, "num_measurements": 50 }, + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] + } + } + ] +} diff --git a/diskann-benchmark/src/inputs/multi_vector.rs b/diskann-benchmark/src/inputs/multi_vector.rs index 5b2aaa833..d431a0909 100644 --- a/diskann-benchmark/src/inputs/multi_vector.rs +++ b/diskann-benchmark/src/inputs/multi_vector.rs @@ -231,3 +231,81 @@ impl std::fmt::Display for MultiVectorQuantOp { Ok(()) } } + +/////////////////////////////// +// Multi-Vector Tiled f16 Op // +/////////////////////////////// + +/// An **f16** multi-vector MaxSim A/B benchmark job: the coarse tiler's f16 path +/// (build-time f16→f32 widen + f32 store kernel) vs the production `f16.rs` +/// preprocess path (per-tile f16→f32 convert + fused f32 kernel). +/// +/// Not apples-to-apples — the tiler is strip-based and converts once in build; the +/// reference is fused and converts per tile inside the timed loop. Read the ratio as +/// a ceiling, not a pure abstraction delta. Element type is f16 and the ISA is fixed +/// to V3/AVX2, so neither is a JSON field. x86_64-only. +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct MultiVectorTiledF16Op { + pub(crate) runs: Vec, +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl MultiVectorTiledF16Op { + pub(crate) const fn tag() -> &'static str { + "multi-vector-tiled-f16-op" + } +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl Input for MultiVectorTiledF16Op { + type Raw = Self; + + fn tag() -> &'static str { + Self::tag() + } + + fn from_raw(raw: Self::Raw, _checker: &mut Checker) -> anyhow::Result { + Ok(raw) + } + + fn serialize(&self) -> anyhow::Result { + Ok(serde_json::to_value(self)?) + } + + fn example() -> Self { + const NUM_DOC_VECTORS: NonZeroUsize = NonZeroUsize::new(64).unwrap(); + const DIM: NonZeroUsize = NonZeroUsize::new(128).unwrap(); + const LOOPS_PER_MEASUREMENT: NonZeroUsize = NonZeroUsize::new(50).unwrap(); + const NUM_MEASUREMENTS: NonZeroUsize = NonZeroUsize::new(20).unwrap(); + + let runs = vec![ + Run { + num_query_vectors: NonZeroUsize::new(32).unwrap(), + num_doc_vectors: NUM_DOC_VECTORS, + dim: DIM, + loops_per_measurement: LOOPS_PER_MEASUREMENT, + num_measurements: NUM_MEASUREMENTS, + }, + Run { + num_query_vectors: NonZeroUsize::new(64).unwrap(), + num_doc_vectors: NUM_DOC_VECTORS, + dim: DIM, + loops_per_measurement: LOOPS_PER_MEASUREMENT, + num_measurements: NUM_MEASUREMENTS, + }, + ]; + + Self { runs } + } +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl std::fmt::Display for MultiVectorTiledF16Op { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "Multi-Vector Tiled f16 Operation\n")?; + write_field!(f, "tag", Self::tag())?; + write_field!(f, "number of runs", self.runs.len())?; + Ok(()) + } +} diff --git a/diskann-benchmark/src/multi_vector/mod.rs b/diskann-benchmark/src/multi_vector/mod.rs index a01285ba1..f1f0e8f3b 100644 --- a/diskann-benchmark/src/multi_vector/mod.rs +++ b/diskann-benchmark/src/multi_vector/mod.rs @@ -28,11 +28,16 @@ cfg_if::cfg_if! { // The quantized A/B op drives the V3-only staged integer kernel. #[cfg(target_arch = "x86_64")] mod quant; + // The f16 A/B op: coarse tiler vs the f16.rs preprocess path (V3-only). + #[cfg(target_arch = "x86_64")] + mod tiled_f16; pub(super) fn register_benchmarks(registry: &mut Registry) -> anyhow::Result<()> { kernels::register(registry)?; #[cfg(target_arch = "x86_64")] quant::register(registry)?; + #[cfg(target_arch = "x86_64")] + tiled_f16::register(registry)?; Ok(()) } } else { diff --git a/diskann-benchmark/src/multi_vector/quant.rs b/diskann-benchmark/src/multi_vector/quant.rs index 1d09ca88b..efe981e6c 100644 --- a/diskann-benchmark/src/multi_vector/quant.rs +++ b/diskann-benchmark/src/multi_vector/quant.rs @@ -25,7 +25,9 @@ use diskann_benchmark_runner::{ use diskann_quantization::algorithms::transforms::NullTransform; use diskann_quantization::algorithms::Transform; use diskann_quantization::minmax::{MinMaxMeta, MinMaxQuantizer}; -use diskann_quantization::multi_vector::distance::{QuantStagedDocs, QuantStagedQuery}; +use diskann_quantization::multi_vector::distance::{ + QuantStagedDocs, QuantStagedQuery, QuantTiledDocs, QuantTiledQuery, +}; use diskann_quantization::multi_vector::{Defaulted, Mat, MatRef, MaxSim, QueryMatRef, Standard}; use diskann_quantization::num::Positive; use diskann_quantization::CompressInto; @@ -140,7 +142,13 @@ fn run_ab(run: &Run) -> anyhow::Result { .ok_or_else(|| anyhow::anyhow!("AVX2 (V3) unavailable for the staged quantized kernel"))?; let docs = QuantStagedDocs::build(data.docs.as_view()); - // Path B — scalar MinMax reference over the same quantization. + // Path B — the coarse tiler rebuild (same math + tiling, `Tiler`/`Tile` handles, + // kernel pinned on the panel types). + let mut tiled_query = QuantTiledQuery::build(data.queries.as_view()) + .ok_or_else(|| anyhow::anyhow!("AVX2 (V3) unavailable for the tiled quantized kernel"))?; + let tiled_docs = QuantTiledDocs::build(data.docs.as_view()); + + // Path C — scalar MinMax reference over the same quantization. let q_ref = quantize(data.queries.as_view()); let d_ref = quantize(data.docs.as_view()); @@ -159,6 +167,14 @@ fn run_ab(run: &Run) -> anyhow::Result { std::hint::black_box(&mut scores); }); + // Timed adjacent to `staged` so the staged-vs-tiled ratio is trustworthy + // despite this box's cross-run clock variance. + let tiled = measure(run, || { + let docs = std::hint::black_box(&tiled_docs); + tiled_query.compute_max_sim(docs, &mut scores); + std::hint::black_box(&mut scores); + }); + let reference = measure(run, || { let q_ref = std::hint::black_box(&q_ref); let d_ref = std::hint::black_box(&d_ref); @@ -170,6 +186,7 @@ fn run_ab(run: &Run) -> anyhow::Result { Ok(QuantRunResult { run: run.clone(), staged, + tiled, reference, }) } @@ -197,11 +214,12 @@ impl Series { } } -/// Staged-vs-reference result for one shape. +/// Staged-vs-tiled-vs-reference result for one shape. #[derive(Debug, Serialize, Deserialize)] pub(super) struct QuantRunResult { pub(super) run: Run, pub(super) staged: Series, + pub(super) tiled: Series, pub(super) reference: Series, } @@ -222,7 +240,8 @@ impl std::fmt::Display for DisplayWrapper<'_, [QuantRunResult]> { writeln!( f, "ns/IP = min time per (query, doc) inner-product call; \ - Speedup = reference / staged (>1 ⇒ staged faster)" + Tiled/Staged = tiled ÷ staged (1.00 ⇒ zero abstraction overhead, \ + >1 ⇒ tiled slower); Speedup = reference ÷ staged" )?; let header = [ @@ -230,6 +249,8 @@ impl std::fmt::Display for DisplayWrapper<'_, [QuantRunResult]> { "D", "Dim", "Staged (ns/IP)", + "Tiled (ns/IP)", + "Tiled/Staged", "Reference (ns/IP)", "Speedup", ]; @@ -238,7 +259,9 @@ impl std::fmt::Display for DisplayWrapper<'_, [QuantRunResult]> { self.iter().enumerate().for_each(|(row, r)| { let comps = r.computations(); let staged = r.staged.min_us() / comps * 1000.0; + let tiled = r.tiled.min_us() / comps * 1000.0; let reference = r.reference.min_us() / comps * 1000.0; + let overhead = if staged > 0.0 { tiled / staged } else { 0.0 }; let speedup = if staged > 0.0 { reference / staged } else { @@ -250,8 +273,10 @@ impl std::fmt::Display for DisplayWrapper<'_, [QuantRunResult]> { row.insert(r.run.num_doc_vectors, 1); row.insert(r.run.dim, 2); row.insert(format!("{:.3}", staged), 3); - row.insert(format!("{:.3}", reference), 4); - row.insert(format!("{:.2}x", speedup), 5); + row.insert(format!("{:.3}", tiled), 4); + row.insert(format!("{:.2}x", overhead), 5); + row.insert(format!("{:.3}", reference), 6); + row.insert(format!("{:.2}x", speedup), 7); }); table.fmt(f) diff --git a/diskann-benchmark/src/multi_vector/tiled_f16.rs b/diskann-benchmark/src/multi_vector/tiled_f16.rs new file mode 100644 index 000000000..32c7c7eb4 --- /dev/null +++ b/diskann-benchmark/src/multi_vector/tiled_f16.rs @@ -0,0 +1,259 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! A/B benchmark for **f16** multi-vector MaxSim: the coarse tiler's f16 path +//! (per-tile f16→f32 widen into a reused buffer, then an f32 store kernel + identity +//! postprocess reuse the tiled pipeline) vs the production `f16.rs` **preprocess** +//! path (per-tile f16→f32 `ConvertTo` + the fused f32 micro-kernel), reached via +//! [`build_max_sim`]. +//! +//! Both convert per tile now, so the ratio mostly isolates one structural +//! difference: the tiler is *strip-based* (the kernel stores an A-major strip, then a +//! separate reduce pass), while the reference is *fused* (the kernel maxes straight +//! into state, no strip). Expect near-parity. +//! +//! x86_64 (V3/AVX2) only. + +use std::io::Write; + +use diskann_benchmark_runner::{ + benchmark::{FailureScore, MatchScore}, + utils::{fmt::Table, percentiles, MicroSeconds}, + Benchmark, Checkpoint, Output, Registry, +}; +use diskann_quantization::multi_vector::distance::{QuantTiledF16Docs, QuantTiledF16Query}; +use diskann_quantization::multi_vector::{ + build_max_sim, BoxErase, Mat, MatRef, MaxSimIsa, Standard, +}; +use serde::{Deserialize, Serialize}; + +use super::driver::Data; +use crate::inputs::multi_vector::{MultiVectorTiledF16Op, Run}; +use crate::utils::DisplayWrapper; + +// ───────────────────────────────────────────────────────────────────────── +// Kernel. +// ───────────────────────────────────────────────────────────────────────── + +#[derive(Debug)] +pub(super) struct TiledF16Kernel; + +impl TiledF16Kernel { + pub(super) const fn new() -> Self { + Self + } +} + +impl Benchmark for TiledF16Kernel { + type Input = MultiVectorTiledF16Op; + type Output = Vec; + + fn try_match(&self, _from: &MultiVectorTiledF16Op) -> Result { + if QuantTiledF16Query::is_supported() { + Ok(MatchScore(0)) + } else { + Err(FailureScore(0)) + } + } + + fn run( + &self, + input: &MultiVectorTiledF16Op, + _: Checkpoint<'_>, + mut output: &mut dyn Output, + ) -> anyhow::Result { + writeln!(output, "{}", input)?; + let mut results = Vec::with_capacity(input.runs.len()); + for run in input.runs.iter() { + results.push(run_ab(run)?); + } + writeln!(output, "\n\n{}", DisplayWrapper(&*results))?; + Ok(results) + } + + fn description( + &self, + f: &mut std::fmt::Formatter<'_>, + input: Option<&MultiVectorTiledF16Op>, + ) -> std::fmt::Result { + match input { + None => writeln!(f, "- f16 tiler vs f16.rs preprocess (V3/AVX2)")?, + Some(_) => { + if !QuantTiledF16Query::is_supported() { + writeln!(f, "\n - AVX2 (V3) unavailable on this CPU")?; + } + } + } + Ok(()) + } +} + +// ───────────────────────────────────────────────────────────────────────── +// A/B timing. +// ───────────────────────────────────────────────────────────────────────── + +/// Run `f` `loops_per_measurement` times per measurement, `num_measurements` +/// times, returning the per-measurement latencies and their percentiles. +fn measure(run: &Run, mut f: impl FnMut()) -> Series { + let mut latencies = Vec::with_capacity(run.num_measurements.get()); + for _ in 0..run.num_measurements.get() { + let start = std::time::Instant::now(); + for _ in 0..run.loops_per_measurement.get() { + f(); + } + latencies.push(start.elapsed().into()); + } + let percentiles = percentiles::compute_percentiles(&mut latencies).unwrap(); + Series { + latencies, + percentiles, + } +} + +/// Narrow an f32 fixture to f16 (row-major, same shape) so both f16 paths see +/// identical bits. +fn narrow(src: MatRef<'_, Standard>) -> Mat> { + let (n, dim) = (src.num_vectors(), src.vector_dim()); + let s = src.as_slice(); + let mut i = 0; + Mat::from_fn(Standard::new(n, dim).expect("n×dim"), || { + let v = diskann_wide::cast_f32_to_f16(s[i]); + i += 1; + v + }) +} + +/// Build both f16 kernels for one shape and time them (build / convert excluded). +fn run_ab(run: &Run) -> anyhow::Result { + // f16 fixtures, generated as f32 then narrowed so both paths see identical bits. + let data = Data::::new(run)?; + let q_f16 = narrow(data.queries.as_view()); + let d_f16 = narrow(data.docs.as_view()); + + // Path A — the coarse tiler's f16 path (per-tile f16→f32 into a reused buffer). + let mut tiled_query = QuantTiledF16Query::build(q_f16.as_view()) + .ok_or_else(|| anyhow::anyhow!("AVX2 (V3) unavailable for the tiled f16 kernel"))?; + let tiled_docs = QuantTiledF16Docs::build(d_f16.as_view()); + + // Path B — the production preprocess path (per-tile f16→f32 + fused f32 kernel). + let preprocess = + build_max_sim::(MaxSimIsa::X86_64_V3, q_f16.as_view(), BoxErase)?; + + let nq = run.num_query_vectors.get(); + let mut scores = vec![0.0f32; nq]; + + let tiled = measure(run, || { + let docs = std::hint::black_box(&tiled_docs); + tiled_query.compute_max_sim(docs, &mut scores); + std::hint::black_box(&mut scores); + }); + + let preprocess = measure(run, || { + let doc = std::hint::black_box(d_f16.as_view()); + preprocess + .compute_max_sim(doc, &mut scores) + .expect("scores.len() == nrows by construction"); + std::hint::black_box(&mut scores); + }); + + Ok(TiledF16RunResult { + run: run.clone(), + tiled, + preprocess, + }) +} + +// ───────────────────────────────────────────────────────────────────────── +// Result types. +// ───────────────────────────────────────────────────────────────────────── + +/// One timed series (per-measurement latencies + percentiles). +#[derive(Debug, Serialize, Deserialize)] +pub(super) struct Series { + latencies: Vec, + percentiles: percentiles::Percentiles, +} + +impl Series { + fn min_us(&self) -> f64 { + self.latencies + .iter() + .min() + .copied() + .unwrap_or(MicroSeconds::new(u64::MAX)) + .as_f64() + } +} + +/// Tiled-vs-preprocess result for one shape. +#[derive(Debug, Serialize, Deserialize)] +pub(super) struct TiledF16RunResult { + pub(super) run: Run, + pub(super) tiled: Series, + pub(super) preprocess: Series, +} + +impl TiledF16RunResult { + fn computations(&self) -> f64 { + (self.run.num_query_vectors.get() + * self.run.num_doc_vectors.get() + * self.run.loops_per_measurement.get()) as f64 + } +} + +impl std::fmt::Display for DisplayWrapper<'_, [TiledF16RunResult]> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.is_empty() { + return Ok(()); + } + + writeln!( + f, + "ns/IP = min time per (query, doc) inner-product call; \ + Tiled/Preprocess = tiled ÷ preprocess (>1 ⇒ tiler slower). Both convert \ + f16→f32 per tile; the residual is strip-based (tiler) vs fused (preprocess)." + )?; + + let header = [ + "Q", + "D", + "Dim", + "Tiled (ns/IP)", + "Preprocess (ns/IP)", + "Tiled/Preprocess", + ]; + let mut table = Table::new(header, self.len()); + + self.iter().enumerate().for_each(|(row, r)| { + let comps = r.computations(); + let tiled = r.tiled.min_us() / comps * 1000.0; + let preprocess = r.preprocess.min_us() / comps * 1000.0; + let ratio = if preprocess > 0.0 { + tiled / preprocess + } else { + 0.0 + }; + + let mut row = table.row(row); + row.insert(r.run.num_query_vectors, 0); + row.insert(r.run.num_doc_vectors, 1); + row.insert(r.run.dim, 2); + row.insert(format!("{:.3}", tiled), 3); + row.insert(format!("{:.3}", preprocess), 4); + row.insert(format!("{:.2}x", ratio), 5); + }); + + table.fmt(f) + } +} + +// ───────────────────────────────────────────────────────────────────────── +// Registration. +// ───────────────────────────────────────────────────────────────────────── + +pub(super) fn register(registry: &mut Registry) -> anyhow::Result<()> { + registry.register("multi-vector-tiled-f16-op", TiledF16Kernel::new())?; + Ok(()) +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs index 61182315a..ad270c574 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs @@ -20,12 +20,19 @@ mod reduce; // code isn't dead on other architectures. #[cfg(target_arch = "x86_64")] pub(super) mod staged; +// Coarse Tiler/Tile rebuild — panels selected by type, kernel pinned on the pair (V3 only). mod tiled_reduce; +#[cfg(target_arch = "x86_64")] +pub(super) mod tiler; // Re-export the quantized staged kernel's public POC entry (x86_64 only). #[cfg(target_arch = "x86_64")] pub use staged::{QuantStagedDocs, QuantStagedQuery}; +// Re-export the coarse Tiler-based quantized POC entry (x86_64 only). +#[cfg(target_arch = "x86_64")] +pub use tiler::{QuantTiledDocs, QuantTiledF16Docs, QuantTiledF16Query, QuantTiledQuery}; + // ── Tile budget ────────────────────────────────────────────────── /// Cache budgets fed to the tile planner. diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiler/arena.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiler/arena.rs new file mode 100644 index 000000000..4314d98be --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiler/arena.rs @@ -0,0 +1,90 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Single-owner resettable bump arena (copy of `staged::arena`, kept local so the +//! `tiler` experiment doesn't reach into its sibling). + +use std::cell::{Cell, UnsafeCell}; +use std::ptr::NonNull; + +use crate::alloc::{AlignedAllocator, AllocatorCore, AllocatorError, Poly}; + +/// Single-owner, single-threaded resettable bump over a 64-byte-aligned buffer. +pub(crate) struct ResettableArena { + buffer: Poly, AlignedAllocator>, + head: Cell, +} + +impl ResettableArena { + pub(crate) fn with_capacity(capacity: usize) -> Result { + let buffer = Poly::<[u8], _>::new_uninit_slice(capacity.max(1), AlignedAllocator::A64)?; + let (ptr, alloc) = Poly::into_raw(buffer); + + // SAFETY: `UnsafeCell<[u8]>` shares the layout of `[u8]`; `MaybeUninit` is + // layout-compatible with `u8` (bytes are only read after being handed out and + // written). `ptr` is non-null, from `Poly::into_raw`. + let buffer = unsafe { + Poly::from_raw( + NonNull::new_unchecked(ptr.as_ptr() as *mut UnsafeCell<[u8]>), + alloc, + ) + }; + + Ok(Self { + buffer, + head: Cell::new(0), + }) + } + + /// Rewind in O(1). `&mut self` makes the borrow checker forbid resetting while any + /// [`ScopedAllocator`](crate::alloc::ScopedAllocator) borrowing this arena is live. + pub(crate) fn reset(&mut self) { + self.head.set(0); + } + + fn capacity(&self) -> usize { + self.buffer.get().len() + } + + fn base(&self) -> *mut u8 { + self.buffer.get().cast::() + } +} + +// SAFETY: `allocate` returns exactly `layout.size()` bytes aligned to `layout.align()` +// within the fixed buffer, or errors. `deallocate` is a no-op — storage is reclaimed by +// `reset` or on drop. +unsafe impl AllocatorCore for ResettableArena { + fn allocate(&self, layout: std::alloc::Layout) -> Result, AllocatorError> { + let base = self.base() as usize; + let head = self.head.get(); + let cur = base.checked_add(head).ok_or(AllocatorError)?; + let aligned = cur + .checked_next_multiple_of(layout.align()) + .ok_or(AllocatorError)?; + let pad = aligned - cur; + let new_head = head + .checked_add(pad) + .and_then(|h| h.checked_add(layout.size())) + .ok_or(AllocatorError)?; + if new_head > self.capacity() { + return Err(AllocatorError); + } + self.head.set(new_head); + + // SAFETY: `head + pad <= new_head <= capacity`, so the range is in-bounds. + let ptr = unsafe { self.base().add(head + pad) }; + NonNull::new(std::ptr::slice_from_raw_parts_mut(ptr, layout.size())).ok_or(AllocatorError) + } + + unsafe fn deallocate(&self, _ptr: NonNull<[u8]>, _layout: std::alloc::Layout) {} +} + +impl std::fmt::Debug for ResettableArena { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ResettableArena") + .field("capacity", &self.capacity()) + .field("head", &self.head.get()) + .finish() + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiler/f16.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiler/f16.rs new file mode 100644 index 000000000..c8aca6b9d --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiler/f16.rs @@ -0,0 +1,304 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! f16 entry for the coarse driver, via **per-tile** f16→f32 conversion. +//! +//! The query is stored block-transposed *as f16*; docs stay row-major f16. Each +//! tile is widened f16→f32 into a small buffer the walk **reuses** across tiles +//! (allocated once from the arena) — see [`QueryConvertWalk`]/[`DocConvertWalk`]. +//! The widened f32 tiles feed the same f32 store kernel → [`Identity`] postprocess → +//! `Max` reducer as a native-f32 walk would. No 2× whole-matrix f32 copy; the +//! lending [`TileWalk`] contract makes the buffer reuse sound. + +use core::mem::size_of; + +use diskann_wide::arch::x86_64::V3; + +use super::arena::ResettableArena; +use super::minmax::Max; +use super::tilers::{A_PANEL, B_PANEL, DPanel, DocConvertWalk, QPanel, QueryConvertWalk}; +use super::{Accumulate, Identity, Kernel, Plan, StripMut, TileBudget, drive, leaves, zeroed}; +use crate::alloc::{Poly, ScopedAllocator}; +use crate::multi_vector::{BlockTransposed, MatRef, Standard}; + +// ── Stage A (f32 store kernel) ─────────────────────────────────── + +pub(crate) struct F32Kernel; + +impl Kernel for F32Kernel { + type Acc = f32; + const A_PANEL: usize = A_PANEL; + const B_PANEL: usize = B_PANEL; +} + +impl<'a, 'b> Accumulate, DPanel<'b, f32>> for F32Kernel { + fn accumulate( + &self, + arch: V3, + a: QPanel<'a, f32>, + b: DPanel<'b, f32>, + mut out: StripMut<'_, f32>, + ) { + // SAFETY: `a` is a 16×k block-transposed f32 block; `b` is B_PANEL rows of k f32; + // `out` is B_PANEL columns of 16 f32 at stride 16. + unsafe { + leaves::f32_store_microkernel::( + arch, + a.as_ptr(), + b.as_ptr(), + a.k(), + out.as_mut_ptr(), + A_PANEL, + ); + } + } + + fn accumulate_tail( + &self, + arch: V3, + a: QPanel<'a, f32>, + b: DPanel<'b, f32>, + mut out: StripMut<'_, f32>, + ) { + let (ap, bp, op) = (a.as_ptr(), b.as_ptr(), out.as_mut_ptr()); + // SAFETY: as `accumulate`, with a runtime width `b.rows()` in 1..B_PANEL. + unsafe { + match b.rows() { + 3 => leaves::f32_store_microkernel::<3>(arch, ap, bp, a.k(), op, A_PANEL), + 2 => leaves::f32_store_microkernel::<2>(arch, ap, bp, a.k(), op, A_PANEL), + 1 => leaves::f32_store_microkernel::<1>(arch, ap, bp, a.k(), op, A_PANEL), + other => unreachable!("tail width {other} out of 1..{B_PANEL}"), + } + } + } +} + +// ── Public entry ───────────────────────────────────────────────── + +/// A prepared f16 query set, stored block-transposed as f16 (widened per tile). +pub struct QuantTiledF16Query { + query: BlockTransposed, + dim: usize, + arch: V3, + state: Vec, + arena: ResettableArena, +} + +impl QuantTiledF16Query { + /// `None` if AVX2 (V3) is unavailable. + #[allow(clippy::expect_used)] + pub fn build(query: MatRef<'_, Standard>) -> Option { + let arch = V3::new_checked()?; + let dim = query.vector_dim(); + let query = BlockTransposed::::from_matrix_view(query.as_matrix_view()); + + let k = query.padded_ncols(); + let padded = query.padded_nrows(); + + // Size the arena for the reused convert buffers (one A-tile = the whole query + // at most; one B-tile) plus the driver's `partial` strip, at the default + // budget. `Identity` postprocess needs no `scored` scratch. + let plan = Plan::new( + k * size_of::(), + k * size_of::(), + A_PANEL, + B_PANEL, + size_of::(), + TileBudget::default(), + ); + let b_tile_rows = B_PANEL * plan.b_panels; + let f32s = padded * k + b_tile_rows * k + A_PANEL * b_tile_rows; + let arena = ResettableArena::with_capacity(f32s * size_of::() + 8192) + .expect("arena allocation"); + + Some(Self { + query, + dim, + arch, + state: vec![f32::MIN; padded], + arena, + }) + } + + pub fn is_supported() -> bool { + V3::new_checked().is_some() + } + + pub fn num_vectors(&self) -> usize { + self.query.nrows() + } + + /// Per-query max inner product (the MaxSim similarity) against `docs`. + /// + /// # Panics + /// + /// If `scores.len() != self.num_vectors()` or the logical dims differ. + pub fn compute_max_sim(&mut self, docs: &QuantTiledF16Docs, scores: &mut [f32]) { + self.compute(docs, scores, TileBudget::default()); + } + + #[allow(clippy::expect_used)] + fn compute(&mut self, docs: &QuantTiledF16Docs, scores: &mut [f32], budget: TileBudget) { + let nq = self.query.nrows(); + assert_eq!(scores.len(), nq, "scores length must equal query count"); + assert_eq!(self.dim, docs.dim, "query dim != doc dim"); + + let k = self.query.padded_ncols(); + let padded = self.query.padded_nrows(); + + self.arena.reset(); + let plan = Plan::new( + k * size_of::(), + k * size_of::(), + A_PANEL, + B_PANEL, + size_of::(), + budget, + ); + + // Reused per-tile convert buffers, each sized to its largest single tile. + let q_src = self.query.as_slice(); + let qbuf_len = q_src.len().min(A_PANEL * plan.a_panels * k).max(1); + let dbuf_len = docs.codes.len().min(B_PANEL * plan.b_panels * k).max(1); + let alloc = ScopedAllocator::new(&self.arena); + let mut qbuf_poly = Poly::<[f32], _>::new_uninit_slice(qbuf_len, alloc).expect("q convert"); + let mut dbuf_poly = Poly::<[f32], _>::new_uninit_slice(dbuf_len, alloc).expect("d convert"); + let qbuf = zeroed(&mut qbuf_poly, qbuf_len); + let dbuf = zeroed(&mut dbuf_poly, dbuf_len); + + let a_walk = QueryConvertWalk::new(q_src, k, plan.a_panels, qbuf); + let b_walk = DocConvertWalk::new(&docs.codes, k, plan.b_panels, dbuf); + drive( + self.arch, + a_walk, + b_walk, + &F32Kernel, + &Identity, + &Max, + &mut self.state[..padded], + alloc, + ); + + scores.copy_from_slice(&self.state[..nq]); + } +} + +/// A prepared f16 document set, kept row-major as f16 (widened per tile). +pub struct QuantTiledF16Docs { + codes: Vec, + dim: usize, + nv: usize, +} + +impl QuantTiledF16Docs { + pub fn build(docs: MatRef<'_, Standard>) -> Self { + let (nv, dim) = (docs.num_vectors(), docs.vector_dim()); + let codes = docs.as_slice().to_vec(); + Self { codes, dim, nv } + } + + pub fn num_vectors(&self) -> usize { + self.nv + } +} + +#[cfg(test)] +mod tests { + use super::*; + use diskann_wide::{cast_f16_to_f32, cast_f32_to_f16}; + + fn rnd(seed: u64, idx: usize) -> f32 { + let x = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(idx as u64) + .wrapping_mul(1442695040888963407); + ((x >> 33) as f32 / (1u64 << 31) as f32) - 1.0 + } + + /// Naive f16 max-IP reference: widen every f16 to f32, dot, take the max doc. + fn reference(q: &[half::f16], nq: usize, d: &[half::f16], nd: usize, dim: usize) -> Vec { + (0..nq) + .map(|i| { + (0..nd) + .map(|j| { + (0..dim) + .map(|c| { + cast_f16_to_f32(q[i * dim + c]) * cast_f16_to_f32(d[j * dim + c]) + }) + .sum::() + }) + .fold(f32::MIN, f32::max) + }) + .collect() + } + + const CASES: &[(usize, usize, usize)] = &[ + (1, 1, 64), + (5, 3, 5), + (16, 4, 64), + (16, 5, 128), + (16, 7, 256), + (17, 9, 65), + (32, 16, 256), + (64, 1250, 64), + (8, 33, 127), + ]; + + #[test] + fn tiled_f16_matches_reference() { + if V3::new_checked().is_none() { + return; + } + for &(nq, nd, dim) in CASES { + let q: Vec = (0..nq * dim).map(|i| cast_f32_to_f16(rnd(1, i))).collect(); + let d: Vec = (0..nd * dim).map(|i| cast_f32_to_f16(rnd(2, i))).collect(); + + let q_mat = MatRef::new(Standard::::new(nq, dim).unwrap(), &q).unwrap(); + let d_mat = MatRef::new(Standard::::new(nd, dim).unwrap(), &d).unwrap(); + let mut query = QuantTiledF16Query::build(q_mat).unwrap(); + let docs = QuantTiledF16Docs::build(d_mat); + let mut got = vec![0.0f32; nq]; + query.compute_max_sim(&docs, &mut got); + + let want = reference(&q, nq, &d, nd, dim); + for i in 0..nq { + assert!( + (got[i] - want[i]).abs() <= 1e-3 * want[i].abs().max(1.0), + "({nq},{nd},{dim}) row {i}: tiled-f16 {} != reference {}", + got[i], + want[i], + ); + } + } + } + + /// Tiny cache budget forces multiple A- and B-tiles. + #[test] + fn tiled_f16_multi_tile_tiny_budget() { + if V3::new_checked().is_none() { + return; + } + let budget = TileBudget { l2_a: 1, l1_b: 1 }; + for &(nq, nd, dim) in &[(48usize, 22usize, 64usize), (33, 37, 128), (35, 19, 65)] { + let q: Vec = (0..nq * dim).map(|i| cast_f32_to_f16(rnd(3, i))).collect(); + let d: Vec = (0..nd * dim).map(|i| cast_f32_to_f16(rnd(4, i))).collect(); + + let q_mat = MatRef::new(Standard::::new(nq, dim).unwrap(), &q).unwrap(); + let d_mat = MatRef::new(Standard::::new(nd, dim).unwrap(), &d).unwrap(); + let mut query = QuantTiledF16Query::build(q_mat).unwrap(); + let docs = QuantTiledF16Docs::build(d_mat); + let mut got = vec![0.0f32; nq]; + query.compute(&docs, &mut got, budget); + + let want = reference(&q, nq, &d, nd, dim); + for i in 0..nq { + assert!( + (got[i] - want[i]).abs() <= 1e-3 * want[i].abs().max(1.0), + "({nq},{nd},{dim}) row {i}: tiny-budget tiled-f16 {} != reference {}", + got[i], + want[i], + ); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiler/leaves.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiler/leaves.rs new file mode 100644 index 000000000..324476e49 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiler/leaves.rs @@ -0,0 +1,207 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! The three SIMD leaves, copied verbatim from `staged`. Keeping the inner math +//! byte-identical is the point: any A/B against `staged` measures only the +//! abstraction, not a different kernel. These are the only `unsafe` here outside +//! the scratch bridge. +//! +//! All three assume the intermediate strip is A-major with column stride +//! `A_PANEL = 16 = 2·LANES`. + +use diskann_wide::arch::x86_64::V3; +use diskann_wide::{SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDMulAdd, SIMDReinterpret, SIMDVector}; + +use crate::minmax::MinMaxCompensation; + +diskann_wide::alias!(i16s = ::i16x16); +diskann_wide::alias!(i32s = ::i32x8); +diskann_wide::alias!(u32s = ::u32x8); +diskann_wide::alias!(f32s = ::f32x8); + +/// Stage A — integer store-out micro-kernel: 16 A-rows × `UNROLL` B-cols. +/// +/// # Safety +/// +/// 1. `a_packed` points to a `16 × k` block-transposed `i16` block (`k` even). +/// 2. `b` points to `UNROLL` rows of `k` contiguous `u8` (`k` even). +/// 3. `partial` is valid for `UNROLL` columns of 16 `i32` at stride `b_stride`. +#[inline(always)] +pub(super) unsafe fn int_store_microkernel( + arch: V3, + a_packed: *const i16, + b: *const u8, + k: usize, + partial: *mut i32, + b_stride: usize, +) { + let mut p0 = [i32s::default(arch); UNROLL]; + let mut p1 = [i32s::default(arch); UNROLL]; + let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); + + let a_pair_stride = 2 * i16s::LANES; + let a_half = i16s::LANES; + let pairs = k / 2; + + for p in 0..pairs { + // SAFETY: precondition 1 — the query block has `pairs` col-pairs of 32 i16. + let (a0, a1) = unsafe { + ( + i16s::load_simd(arch, a_packed.add(a_pair_stride * p)), + i16s::load_simd(arch, a_packed.add(a_pair_stride * p + a_half)), + ) + }; + + for j in 0..UNROLL { + // SAFETY: precondition 2 — doc col j is `offsets[j]` in, `2*p+1 < k`. + let (d0, d1) = unsafe { + let base = 2 * p + offsets[j]; + ( + u32::from(b.add(base).read()), + u32::from(b.add(base + 1).read()), + ) + }; + let packed = d0 | (d1 << 16); + let bcast: i16s = u32s::splat(arch, packed).reinterpret_simd(); + p0[j] = p0[j].dot_simd(a0, bcast); + p1[j] = p1[j].dot_simd(a1, bcast); + } + } + + for j in 0..UNROLL { + // SAFETY: precondition 3 — column j occupies [j*b_stride, j*b_stride+16) i32. + unsafe { + p0[j].store_simd(partial.add(j * b_stride)); + p1[j].store_simd(partial.add(j * b_stride + i32s::LANES)); + } + } +} + +/// Stage A (f32) — store-out micro-kernel: 16 A-rows × `UNROLL` B-cols of f32 IP. +/// +/// Mirrors [`int_store_microkernel`] but for f32 in/out with no packing; the max +/// reduction is deferred to the `Max` reducer (via [`fold_strip`]). +/// +/// # Safety +/// +/// 1. `a_packed` points to a `16 × k` block-transposed `f32` block (`PACK = 1`). +/// 2. `b` points to `UNROLL` rows of `k` contiguous `f32`. +/// 3. `partial` is valid for `UNROLL` columns of 16 `f32` at stride `b_stride`. +#[inline(always)] +pub(super) unsafe fn f32_store_microkernel( + arch: V3, + a_packed: *const f32, + b: *const f32, + k: usize, + partial: *mut f32, + b_stride: usize, +) { + let mut p0 = [f32s::default(arch); UNROLL]; + let mut p1 = [f32s::default(arch); UNROLL]; + let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); + + let a_stride = 2 * f32s::LANES; + let a_half = f32s::LANES; + + for i in 0..k { + // SAFETY: precondition 1 — the query block has `k` columns of 16 f32. + let (a0, a1) = unsafe { + ( + f32s::load_simd(arch, a_packed.add(a_stride * i)), + f32s::load_simd(arch, a_packed.add(a_stride * i + a_half)), + ) + }; + for j in 0..UNROLL { + // SAFETY: precondition 2 — doc col j is `offsets[j]` in, `i < k`. + let bj = unsafe { f32s::splat(arch, b.add(i + offsets[j]).read_unaligned()) }; + p0[j] = a0.mul_add_simd(bj, p0[j]); + p1[j] = a1.mul_add_simd(bj, p1[j]); + } + } + + for j in 0..UNROLL { + // SAFETY: precondition 3 — column j occupies [j*b_stride, j*b_stride+16) f32. + unsafe { + p0[j].store_simd(partial.add(j * b_stride)); + p1[j].store_simd(partial.add(j * b_stride + f32s::LANES)); + } + } +} + +/// Stage B — 4-bit MinMax dequant of one 16×`cols` A-major `i32` strip into f32. +/// +/// # Safety +/// +/// `acc` valid for `cols` columns of 16 `i32` (stride 16); `out` writable for the +/// same shape of `f32`; `q_meta.len() >= 16`; `d_meta.len() >= cols`. +#[inline(always)] +pub(super) unsafe fn score_strip( + arch: V3, + acc: *const i32, + out: *mut f32, + cols: usize, + q_meta: &[MinMaxCompensation], + d_meta: &[MinMaxCompensation], + dim: f32, +) { + let lanes = f32s::LANES; + + let mut qa = [0.0f32; 16]; + let mut qb = [0.0f32; 16]; + let mut qn = [0.0f32; 16]; + for i in 0..16 { + let qm = q_meta[i]; + qa[i] = qm.a; + qb[i] = qm.b; + qn[i] = qm.n; + } + // SAFETY: each array holds exactly 16 = 2·LANES f32. + let (qa0, qa1, qb0, qb1, qn0, qn1) = unsafe { + ( + f32s::load_simd(arch, qa.as_ptr()), + f32s::load_simd(arch, qa.as_ptr().add(lanes)), + f32s::load_simd(arch, qb.as_ptr()), + f32s::load_simd(arch, qb.as_ptr().add(lanes)), + f32s::load_simd(arch, qn.as_ptr()), + f32s::load_simd(arch, qn.as_ptr().add(lanes)), + ) + }; + + for (c, dm) in d_meta.iter().enumerate().take(cols) { + let a_c = f32s::splat(arch, dm.a); + let b_c = f32s::splat(arch, dm.b); + let c_c = f32s::splat(arch, dm.n + dm.b * dim); + let col = c * 16; + // SAFETY: `col + 2·LANES <= cols*16`; `acc`/`out` valid for that many. + unsafe { + let raw0 = i32s::load_simd(arch, acc.add(col)).simd_cast(); + let raw1 = i32s::load_simd(arch, acc.add(col + lanes)).simd_cast(); + let s0 = a_c.mul_add_simd(qa0 * raw0, b_c.mul_add_simd(qn0, c_c * qb0)); + let s1 = a_c.mul_add_simd(qa1 * raw1, b_c.mul_add_simd(qn1, c_c * qb1)); + s0.store_simd(out.add(col)); + s1.store_simd(out.add(col + lanes)); + } + } +} + +/// Stage C — fold a 16×`cols` A-major f32 score strip into the 16-wide running max. +/// +/// # Safety +/// +/// `state` writable for 16 `f32`; `scores` valid for `cols` columns of 16 `f32`. +#[inline(always)] +pub(super) unsafe fn fold_strip(arch: V3, state: *mut f32, scores: *const f32, cols: usize) { + let lanes = f32s::LANES; + // SAFETY: `state` writable for 16; `scores` valid for `cols` columns of 16. + unsafe { + let mut a0 = f32s::load_simd(arch, state); + let mut a1 = f32s::load_simd(arch, state.add(lanes)); + for c in 0..cols { + let col = scores.add(c * 16); + a0 = a0.max_simd(f32s::load_simd(arch, col)); + a1 = a1.max_simd(f32s::load_simd(arch, col.add(lanes))); + } + a0.store_simd(state); + a1.store_simd(state.add(lanes)); + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiler/minmax.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiler/minmax.rs new file mode 100644 index 000000000..415570252 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiler/minmax.rs @@ -0,0 +1,445 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! 4-bit MinMax instantiation of the coarse driver: identity [`QueryWalk`] / +//! [`DocWalk`] (panels borrow the source), the `I8Kernel` (Stage A), `MinMax` +//! postprocess (Stage B), `Max` reducer (Stage C), and the standalone +//! `QuantTiled{Query,Docs}` entry. Everything else is generic in [`super`]. + +use core::mem::size_of; +use std::num::NonZeroUsize; + +use diskann_utils::ReborrowMut; +use diskann_wide::arch::x86_64::V3; + +use super::arena::ResettableArena; +use super::tilers::{A_PANEL, B_PANEL, DPanel, DocWalk, QPanel, QueryWalk}; +use super::{ + Accumulate, BlockCtx, Kernel, Plan, Postprocess, Reducer, Strip, StripMut, TileBudget, drive, + leaves, +}; +use crate::CompressInto; +use crate::algorithms::Transform; +use crate::algorithms::transforms::NullTransform; +use crate::alloc::ScopedAllocator; +use crate::minmax::{MinMaxCompensation, MinMaxMeta, MinMaxQuantizer}; +use crate::multi_vector::{BlockTransposed, Defaulted, Mat, MatRef, Standard}; +use crate::num::Positive; + +// ── Stage A ────────────────────────────────────────────────────── + +pub(crate) struct I8Kernel; + +impl Kernel for I8Kernel { + type Acc = i32; + const A_PANEL: usize = A_PANEL; + const B_PANEL: usize = B_PANEL; +} + +impl<'a, 'b> Accumulate, DPanel<'b, u8>> for I8Kernel { + fn accumulate( + &self, + arch: V3, + a: QPanel<'a, i16>, + b: DPanel<'b, u8>, + mut out: StripMut<'_, i32>, + ) { + // SAFETY: `a` is a 16×k block-transposed i16 block; `b` is B_PANEL rows of k u8; + // `out` is B_PANEL columns of 16 i32 at stride 16 (`k` even). + unsafe { + leaves::int_store_microkernel::( + arch, + a.as_ptr(), + b.as_ptr(), + a.k(), + out.as_mut_ptr(), + A_PANEL, + ); + } + } + + fn accumulate_tail( + &self, + arch: V3, + a: QPanel<'a, i16>, + b: DPanel<'b, u8>, + mut out: StripMut<'_, i32>, + ) { + let (ap, bp, op) = (a.as_ptr(), b.as_ptr(), out.as_mut_ptr()); + // SAFETY: as `accumulate`, with a runtime width `b.rows()` in 1..B_PANEL. + unsafe { + match b.rows() { + 3 => leaves::int_store_microkernel::<3>(arch, ap, bp, a.k(), op, A_PANEL), + 2 => leaves::int_store_microkernel::<2>(arch, ap, bp, a.k(), op, A_PANEL), + 1 => leaves::int_store_microkernel::<1>(arch, ap, bp, a.k(), op, A_PANEL), + other => unreachable!("tail width {other} out of 1..{B_PANEL}"), + } + } + } +} + +// ── Stage B ────────────────────────────────────────────────────── + +/// 4-bit MinMax dequant. Rewrites each raw integer dot into the MinMax inner product +/// using per-vector `a`/`b`/`n` metadata, indexed by the block offsets. +pub(crate) struct MinMax<'m> { + query_meta: &'m [MinMaxCompensation], + doc_meta: &'m [MinMaxCompensation], + dim: f32, +} + +impl Postprocess for MinMax<'_> { + type Score = f32; + + fn scratch_len(&self, cols: usize) -> usize { + A_PANEL * cols + } + + fn apply<'s>( + &self, + arch: V3, + acc: Strip<'s, i32>, + scratch: StripMut<'s, f32>, + ctx: BlockCtx, + ) -> Strip<'s, f32> { + let cols = acc.cols(); + let q = &self.query_meta[ctx.a_row_offset..ctx.a_row_offset + A_PANEL]; + let d = &self.doc_meta[ctx.b_row_offset..ctx.b_row_offset + cols]; + let StripMut { data, rows } = scratch; + // SAFETY: `acc` is `cols` cols of 16 i32; `data` is writable for `A_PANEL*cols` + // f32; `q.len() == 16`, `d.len() == cols`. + unsafe { + leaves::score_strip(arch, acc.as_ptr(), data.as_mut_ptr(), cols, q, d, self.dim); + } + Strip { + data: &data[..rows * cols], + rows, + } + } +} + +// ── Stage C ────────────────────────────────────────────────────── + +pub(crate) struct Max; + +impl Reducer for Max { + type State = f32; + const A_PANEL: usize = A_PANEL; + + fn init() -> f32 { + f32::MIN + } + + fn fold(&self, arch: V3, state: &mut [f32], scores: Strip, _first_col: usize) { + // SAFETY: `state` is A_PANEL=16 f32; `scores` is `cols` cols of 16 f32. + unsafe { leaves::fold_strip(arch, state.as_mut_ptr(), scores.as_ptr(), scores.cols()) } + } +} + +// ── Public entry ───────────────────────────────────────────────── + +/// Quantize an f32 multi-vector to 4-bit MinMax (Null transform, scale 1.0). +#[allow(clippy::expect_used)] +fn quantize(input: MatRef<'_, Standard>) -> Mat> { + let (n, dim) = (input.num_vectors(), input.vector_dim()); + let q = MinMaxQuantizer::new( + Transform::Null(NullTransform::new( + NonZeroUsize::new(dim).expect("dimension must be non-zero"), + )), + Positive::new(1.0).expect("1.0 is positive"), + ); + let mut out: Mat> = + Mat::new(MinMaxMeta::new(n, dim), Defaulted).expect("MinMaxMeta allocation"); + q.compress_into(input, out.reborrow_mut()) + .expect("input must be finite"); + out +} + +/// A prepared 4-bit MinMax query set for the coarse tiled driver (V3/AVX2). +/// Standalone POC entry, mirroring `QuantStagedQuery` but built on `drive`. +pub struct QuantTiledQuery { + query: BlockTransposed, + meta: Vec, + dim: usize, + arch: V3, + state: Vec, + arena: ResettableArena, +} + +impl QuantTiledQuery { + /// `None` if AVX2 (V3) is unavailable. + #[allow(clippy::expect_used)] + pub fn build(query: MatRef<'_, Standard>) -> Option { + let arch = V3::new_checked()?; + let (nq, dim) = (query.num_vectors(), query.vector_dim()); + let q_mat = quantize(query); + + let mut codes = vec![0i16; nq * dim]; + for r in 0..nq { + let row = q_mat.get_row(r).expect("row < nq"); + for j in 0..dim { + codes[r * dim + j] = i16::from(row.vector().get(j).expect("col < dim") as u8); + } + } + let view = MatRef::new(Standard::::new(nq, dim).expect("nq×dim"), &codes) + .expect("code slice"); + let query = BlockTransposed::::from_matrix_view(view.as_matrix_view()); + + let padded = query.padded_nrows(); + let mut meta = vec![MinMaxCompensation::default(); padded]; + for (r, m) in meta.iter_mut().enumerate().take(nq) { + *m = q_mat.get_row(r).expect("row < nq").meta(); + } + + // `partial` and `scored` each fit `l1_b`, so `2·l1_b` bounds the arena for any + // k; a page of headroom covers alignment. + let arena = ResettableArena::with_capacity(2 * TileBudget::default().l1_b + 4096) + .expect("arena allocation"); + + Some(Self { + query, + meta, + dim, + arch, + state: vec![f32::MIN; padded], + arena, + }) + } + + pub fn is_supported() -> bool { + V3::new_checked().is_some() + } + + pub fn num_vectors(&self) -> usize { + self.query.nrows() + } + + /// Per-query min distance (`= -max_d IP`) against `docs`. + /// + /// # Panics + /// + /// If `scores.len() != self.num_vectors()` or the logical dims differ. + pub fn compute_max_sim(&mut self, docs: &QuantTiledDocs, scores: &mut [f32]) { + self.compute(docs, scores, TileBudget::default()); + } + + fn compute(&mut self, docs: &QuantTiledDocs, scores: &mut [f32], budget: TileBudget) { + let nq = self.query.nrows(); + assert_eq!(scores.len(), nq, "scores length must equal query count"); + assert_eq!(self.dim, docs.dim, "query dim != doc dim"); + + let k = self.query.padded_ncols(); + let padded = self.query.padded_nrows(); + + self.arena.reset(); + let plan = Plan::new( + k * size_of::(), + k * size_of::(), + A_PANEL, + B_PANEL, + size_of::(), + budget, + ); + let a_walk = QueryWalk::new(self.query.as_slice(), k, plan.a_panels); + let b_walk = DocWalk::new(&docs.codes, k, plan.b_panels); + let post = MinMax { + query_meta: &self.meta, + doc_meta: &docs.meta, + dim: self.dim as f32, + }; + drive( + self.arch, + a_walk, + b_walk, + &I8Kernel, + &post, + &Max, + &mut self.state[..padded], + ScopedAllocator::new(&self.arena), + ); + + for (s, &raw) in scores.iter_mut().zip(self.state.iter()) { + *s = -raw; + } + } +} + +/// A prepared 4-bit MinMax document set (codes-together / metadata-together SoA). +pub struct QuantTiledDocs { + codes: Vec, + meta: Vec, + dim: usize, + nv: usize, +} + +impl QuantTiledDocs { + #[allow(clippy::expect_used)] + pub fn build(docs: MatRef<'_, Standard>) -> Self { + let (nv, dim) = (docs.num_vectors(), docs.vector_dim()); + let padded_dim = dim.next_multiple_of(2); + let d_mat = quantize(docs); + + let mut codes = vec![0u8; nv * padded_dim]; + let mut meta = Vec::with_capacity(nv); + for r in 0..nv { + let row = d_mat.get_row(r).expect("row < nv"); + for j in 0..dim { + codes[r * padded_dim + j] = row.vector().get(j).expect("col < dim") as u8; + } + meta.push(row.meta()); + } + Self { + codes, + meta, + dim, + nv, + } + } + + pub fn num_vectors(&self) -> usize { + self.nv + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::multi_vector::distance::{MaxSim, QueryMatRef}; + use diskann_vector::DistanceFunctionMut; + + fn rnd(seed: u64, idx: usize) -> f32 { + let x = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(idx as u64) + .wrapping_mul(1442695040888963407); + ((x >> 33) as f32 / (1u64 << 31) as f32) - 1.0 + } + + #[allow(clippy::expect_used)] + fn reference(q: &[f32], nq: usize, d: &[f32], nd: usize, dim: usize) -> Vec { + let quantizer = MinMaxQuantizer::new( + Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())), + Positive::new(1.0).unwrap(), + ); + let quantize = |data: &[f32], n: usize| -> Mat> { + let input = MatRef::new(Standard::::new(n, dim).unwrap(), data).unwrap(); + let mut out: Mat> = Mat::new(MinMaxMeta::new(n, dim), Defaulted).unwrap(); + quantizer.compress_into(input, out.reborrow_mut()).unwrap(); + out + }; + let q_mat = quantize(q, nq); + let d_mat = quantize(d, nd); + let query: QueryMatRef<_> = q_mat.as_view().into(); + let mut out = vec![0.0f32; nq]; + MaxSim::new(&mut out).evaluate(query, d_mat.as_view()); + out + } + + /// (nq, nd, dim): B-remainder classes, A-panel remainder (17), multi-tile B, and + /// the odd-dim even-K contract. + const CASES: &[(usize, usize, usize)] = &[ + (1, 1, 64), + (5, 1, 128), + (16, 4, 64), + (16, 5, 128), + (16, 6, 64), + (16, 7, 256), + (17, 9, 64), + (32, 16, 256), + (64, 1250, 64), + (5, 3, 63), + (17, 9, 65), + (8, 33, 127), + ]; + + #[test] + fn tiled_i8_matches_minmax_reference() { + if V3::new_checked().is_none() { + return; + } + for &(nq, nd, dim) in CASES { + let q_data: Vec = (0..nq * dim).map(|i| rnd(1, i)).collect(); + let d_data: Vec = (0..nd * dim).map(|i| rnd(2, i)).collect(); + + let q_f32 = MatRef::new(Standard::::new(nq, dim).unwrap(), &q_data).unwrap(); + let d_f32 = MatRef::new(Standard::::new(nd, dim).unwrap(), &d_data).unwrap(); + let mut query = QuantTiledQuery::build(q_f32).unwrap(); + let docs = QuantTiledDocs::build(d_f32); + let mut got = vec![0.0f32; nq]; + query.compute_max_sim(&docs, &mut got); + + let want = reference(&q_data, nq, &d_data, nd, dim); + for i in 0..nq { + assert!( + (got[i] - want[i]).abs() <= 1e-4 * want[i].abs().max(1.0), + "({nq},{nd},{dim}) row {i}: tiled-i8 {} != reference {}", + got[i], + want[i], + ); + } + } + } + + /// Tiny cache budget clamps the planner to one A-panel and one B-panel per tile, + /// forcing multiple A-tiles and B-tiles — exercising the cross-tile offset carry + /// the default-budget cases (one A-tile) never reach. + #[test] + fn tiled_i8_multi_tile_tiny_budget() { + if V3::new_checked().is_none() { + return; + } + let budget = TileBudget { l2_a: 1, l1_b: 1 }; + for &(nq, nd, dim) in &[(48usize, 22usize, 64usize), (33, 37, 128), (35, 19, 65)] { + let q_data: Vec = (0..nq * dim).map(|i| rnd(3, i)).collect(); + let d_data: Vec = (0..nd * dim).map(|i| rnd(4, i)).collect(); + + let q_f32 = MatRef::new(Standard::::new(nq, dim).unwrap(), &q_data).unwrap(); + let d_f32 = MatRef::new(Standard::::new(nd, dim).unwrap(), &d_data).unwrap(); + let mut query = QuantTiledQuery::build(q_f32).unwrap(); + let docs = QuantTiledDocs::build(d_f32); + let mut got = vec![0.0f32; nq]; + query.compute(&docs, &mut got, budget); + + let want = reference(&q_data, nq, &d_data, nd, dim); + for i in 0..nq { + assert!( + (got[i] - want[i]).abs() <= 1e-4 * want[i].abs().max(1.0), + "({nq},{nd},{dim}) row {i}: tiny-budget tiled-i8 {} != reference {}", + got[i], + want[i], + ); + } + } + } + + /// Arena reuse across differently-sized doc sets stays correct. + #[test] + fn tiled_i8_arena_reuse() { + if V3::new_checked().is_none() { + return; + } + const NQ: usize = 17; + const DIM: usize = 128; + let q_data: Vec = (0..NQ * DIM).map(|i| rnd(5, i)).collect(); + let q_f32 = MatRef::new(Standard::::new(NQ, DIM).unwrap(), &q_data).unwrap(); + let mut query = QuantTiledQuery::build(q_f32).unwrap(); + + for (call, &nd) in [251usize, 3, 64, 1].iter().enumerate() { + let d_data: Vec = (0..nd * DIM).map(|i| rnd(6 + call as u64, i)).collect(); + let d_f32 = MatRef::new(Standard::::new(nd, DIM).unwrap(), &d_data).unwrap(); + let docs = QuantTiledDocs::build(d_f32); + let mut got = vec![0.0f32; NQ]; + query.compute_max_sim(&docs, &mut got); + + let want = reference(&q_data, NQ, &d_data, nd, DIM); + for i in 0..NQ { + assert!( + (got[i] - want[i]).abs() <= 1e-4 * want[i].abs().max(1.0), + "call {call} (nd={nd}) row {i}: tiled-i8 {} != reference {}", + got[i], + want[i], + ); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiler/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiler/mod.rs new file mode 100644 index 000000000..789c0c155 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiler/mod.rs @@ -0,0 +1,355 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Coarse tiled MaxSim. A [`TileWalk`] streams materialized [`Tile`]s — lending, so a +//! converting walk can reuse one buffer — and each tile yields panels the kernel +//! binds on by type. Three compute seams — [`Accumulate`], [`Postprocess`], +//! [`Reducer`] — fold panels into per-A-row state, which *is* the output. +//! Instantiated for 4-bit MinMax i8 ([`minmax`]) and f16 ([`f16`]); the driver and +//! seams stay generic. + +use core::mem::MaybeUninit; + +use crate::alloc::{AllocatorCore, Poly, ScopedAllocator}; + +use super::TileBudget; + +mod arena; +mod f16; +mod leaves; +mod minmax; +mod tilers; + +pub use f16::{QuantTiledF16Docs, QuantTiledF16Query}; +pub use minmax::{QuantTiledDocs, QuantTiledQuery}; + +// ── Tile planning (copy of `staged::StagedPlan`) ───────────────── + +/// Panel counts per tile. `a_panels` A-panels sit resident in L2; as many B-panels +/// as co-fit L1 alongside one A-panel and the partial strip. +#[derive(Clone, Copy)] +struct Plan { + a_panels: usize, + b_panels: usize, +} + +impl Plan { + fn new( + a_row_bytes: usize, + b_row_bytes: usize, + a_panel: usize, + b_panel: usize, + acc_bytes: usize, + budget: TileBudget, + ) -> Self { + let a_row_bytes = a_row_bytes.max(1); + let b_row_bytes = b_row_bytes.max(1); + let a_panels = (budget.l2_a / (a_row_bytes * a_panel)).max(1); + let a_panel_bytes = a_panel * a_row_bytes; + let per_b_row = b_row_bytes + a_panel * acc_bytes; + let b_budget = budget.l1_b.saturating_sub(a_panel_bytes); + let b_panels = ((b_budget / per_b_row) / b_panel).max(1); + Self { a_panels, b_panels } + } +} + +// ── Intermediate strips ────────────────────────────────────────── + +/// An A-major intermediate (rows = A-panel, cols = `len / rows`), read-only. +pub(crate) struct Strip<'a, T> { + data: &'a [T], + rows: usize, +} +/// An A-major intermediate, writable. Also the kernel's per-call output sub-view. +pub(crate) struct StripMut<'a, T> { + data: &'a mut [T], + rows: usize, +} + +impl Strip<'_, T> { + fn as_ptr(&self) -> *const T { + self.data.as_ptr() + } + fn cols(&self) -> usize { + if self.rows == 0 { + 0 + } else { + self.data.len() / self.rows + } + } +} +impl StripMut<'_, T> { + fn as_mut_ptr(&mut self) -> *mut T { + self.data.as_mut_ptr() + } +} + +/// Per-block context for [`Postprocess`]: the global row offsets a metadata-bearing +/// stage indexes by. +#[derive(Clone, Copy)] +pub(crate) struct BlockCtx { + pub a_row_offset: usize, + pub b_row_offset: usize, +} + +// ── Data side ──────────────────────────────────────────────────── + +/// Misuse guard for [`TileAt`]'s implicit-bounds parameter: `Bounds` and `Sealed` +/// are private, so no downstream impl can override the defaulted parameter with a +/// type that drops the `Self: 'a` implied bound. +mod sealed { + pub trait Sealed {} + pub struct Bounds(#[allow(dead_code)] T); + impl Sealed for Bounds {} +} + +/// The per-lifetime half of [`TileWalk`]: at lifetime `'a`, the walk yields a tile +/// borrowing `'a`. The defaulted `B = Bounds<&'a Self>` carries the `Self: 'a` +/// implied bound through the well-formedness of `&'a Self`, which is what makes the +/// `for<'a>` bound on the driver's kernel provable (a plain GAT `where Self: 'a` +/// collapses to `'static` under that HRTB on stable). +pub(crate) trait TileAt<'a, B: sealed::Sealed = sealed::Bounds<&'a Self>> { + type Tile: Tile; +} + +/// A **lending** walk over a source: `next` reborrows `&mut self`, so a materialized +/// tile (and its panels) can borrow a buffer the walk **reuses** on the following +/// `next` — the borrow checker forbids overwriting it while a tile is still live. +/// `reset` rewinds for a re-walk (the driver re-walks B once per A-tile) without +/// reallocating the buffer. +pub(crate) trait TileWalk: for<'a> TileAt<'a> { + fn next(&mut self) -> Option<>::Tile>; + fn reset(&mut self); + + /// Rows in a full tile — sizes the driver's inter-stage scratch. + fn max_tile_rows(&self) -> usize; +} + +/// One materialized tile: its global offset, its row count, and its panels split +/// into full panels plus an optional short tail. The borrow into the walker's buffer +/// lives on the concrete type (e.g. `QMat<'a, T>`), surfaced via [`TileAt::Tile`], so +/// the trait itself needs no lifetime. +pub(crate) trait Tile { + type Panel: Copy; + + /// Global A-row (queries) or B-col (docs) where this tile starts. + fn offset(&self) -> usize; + /// Rows in this tile. + fn rows(&self) -> usize; + /// Full, fixed-size panels. + fn panels(&self) -> impl Iterator + '_; + /// The short trailing panel, if the row count isn't a whole number of panels. + fn tail(&self) -> Option; +} + +// ── Compute side ───────────────────────────────────────────────── + +/// Stage A, datatype-independent facts: the accumulator type and the panel sizes. +/// Split from [`Accumulate`] so `K::Acc` stays unambiguous under the driver's +/// `for<'a, 'b>` bound over the (lifetime-carrying) panel types. +pub(crate) trait Kernel { + type Acc: Copy; + const A_PANEL: usize; + const B_PANEL: usize; +} + +/// Stage A — datatype axis. One A-panel × one B-panel → an A-major `Acc` block. +/// Pinned on the `(A, B)` panel pair as type parameters, so the walks' panel types +/// select the kernel with no `Convert::To = _` join. `accumulate` is the fixed-width +/// hot path; `accumulate_tail` handles a `1..B_PANEL` remainder. +pub(crate) trait Accumulate: Kernel { + fn accumulate(&self, arch: Arch, a: A, b: B, out: StripMut<'_, Self::Acc>); + fn accumulate_tail(&self, arch: Arch, a: A, b: B, out: StripMut<'_, Self::Acc>); +} + +/// Stage B — quantization axis. `Acc` strip → `Score` strip. Identity returns `acc` +/// (`Score = Acc`); a metadata stage writes `scratch` and returns it. +pub(crate) trait Postprocess { + type Score: Copy; + + fn scratch_len(&self, cols: usize) -> usize; + fn apply<'s>( + &self, + arch: Arch, + acc: Strip<'s, Acc>, + scratch: StripMut<'s, Self::Score>, + ctx: BlockCtx, + ) -> Strip<'s, Self::Score>; +} + +/// Stage C — reducer axis. Fold a `Score` strip into per-A-row `State`; `State` *is* +/// the output (the caller interprets it). `first_col` = global B offset for argmax. +pub(crate) trait Reducer { + type State: Copy; + const A_PANEL: usize; + + fn init() -> Self::State; + fn fold(&self, arch: Arch, state: &mut [Self::State], scores: Strip, first_col: usize); +} + +/// The zero-cost [`Postprocess`]: `Score = Acc`, returns the accumulator strip +/// untouched (no scratch, no metadata). For kernels whose `Acc` is already the score. +pub(crate) struct Identity; + +impl Postprocess for Identity { + type Score = Acc; + + fn scratch_len(&self, _cols: usize) -> usize { + 0 + } + fn apply<'s>( + &self, + _arch: Arch, + acc: Strip<'s, Acc>, + _scratch: StripMut<'s, Acc>, + _ctx: BlockCtx, + ) -> Strip<'s, Acc> { + acc + } +} + +// ── Scratch: uninit alloc → zeroed slice ───────────────────────── + +/// Marker for element types where all-zero is a valid value, so a zeroed allocation +/// is a sound `&mut [T]`. +pub(crate) trait ZeroInit: Copy {} +impl ZeroInit for i16 {} +impl ZeroInit for i32 {} +impl ZeroInit for f32 {} +impl ZeroInit for u8 {} + +fn zeroed( + poly: &mut Poly<[MaybeUninit], A>, + len: usize, +) -> &mut [T] { + let ptr = poly.as_mut_ptr().cast::(); + // SAFETY: the poly owns `len` `T`-sized slots; `T: ZeroInit` ⇒ all-zero is a valid + // `T`, so zeroing initializes every element and the slice is sound as `&mut [T]`. + unsafe { + core::ptr::write_bytes(ptr, 0, len); + core::slice::from_raw_parts_mut(ptr, len) + } +} + +// ── Driver ─────────────────────────────────────────────────────── + +/// Per-A-row reduction into `state` (len ≥ padded A rows) via the seams. The walks +/// carry the plan (baked tile sizes); the driver reads only the B-tile width to size +/// the scratch it allocates from `alloc`, so the caller sizes nothing. B is re-walked +/// (`reset`) once per A-tile — for a convert walk that re-runs its per-tile transform, +/// which under the default budget is a single A-tile (one pass). +#[allow(clippy::too_many_arguments, clippy::expect_used)] +pub(super) fn drive( + arch: Arch, + mut a_walk: AW, + mut b_walk: BW, + kernel: &K, + post: &P, + reducer: &R, + state: &mut [R::State], + alloc: ScopedAllocator<'_>, +) where + Arch: Copy, + AW: TileWalk, + BW: TileWalk, + K: Kernel + + for<'a, 'b> Accumulate< + Arch, + <>::Tile as Tile>::Panel, + <>::Tile as Tile>::Panel, + >, + P: Postprocess, + R: Reducer, + K::Acc: ZeroInit, + P::Score: ZeroInit, +{ + const { assert!(K::A_PANEL == R::A_PANEL) } + let a_panel = K::A_PANEL; + let b_panel = K::B_PANEL; + + for s in state.iter_mut() { + *s = R::init(); + } + + let b_tile_rows = b_walk.max_tile_rows(); + let strip_len = a_panel * b_tile_rows; + let scored_len = post.scratch_len(b_tile_rows); + + let mut partial_poly = + Poly::<[K::Acc], _>::new_uninit_slice(strip_len, alloc).expect("partial scratch"); + let mut scored_poly = + Poly::<[P::Score], _>::new_uninit_slice(scored_len, alloc).expect("scored scratch"); + let partial = zeroed(&mut partial_poly, strip_len); + let scored = zeroed(&mut scored_poly, scored_len); + + while let Some(a_mat) = a_walk.next() { + // The driver has no A-tail path (it never calls `a_mat.tail()`), so a partial + // A-panel would be silently dropped by `panels()`. Block-transposed padding + // guarantees whole panels; assert it so an unpadded A source fails loudly. + debug_assert_eq!( + a_mat.rows() % a_panel, + 0, + "A walk must yield whole A-panels; the driver has no A-tail path" + ); + let a_tile_off = a_mat.offset(); + b_walk.reset(); + while let Some(b_mat) = b_walk.next() { + let (cols, b_off) = (b_mat.rows(), b_mat.offset()); + let w = a_panel * cols; + let full_w = a_panel * (cols - cols % b_panel); + + for (i, a) in a_mat.panels().enumerate() { + let a_off = a_tile_off + i * a_panel; + + // Stage A: fill the partial strip. Full B-panels pair with equal-width + // output blocks; the short tail dispatches its width in the kernel. + for (b, obuf) in b_mat + .panels() + .zip(partial[..full_w].chunks_mut(a_panel * b_panel)) + { + kernel.accumulate( + arch, + a, + b, + StripMut { + data: obuf, + rows: a_panel, + }, + ); + } + if let Some(b) = b_mat.tail() { + let obuf = &mut partial[full_w..w]; + kernel.accumulate_tail( + arch, + a, + b, + StripMut { + data: obuf, + rows: a_panel, + }, + ); + } + + // Stage B: Acc → Score (identity returns the partial strip untouched). + let scores = post.apply( + arch, + Strip { + data: &partial[..w], + rows: a_panel, + }, + StripMut { + data: &mut scored[..post.scratch_len(cols)], + rows: a_panel, + }, + BlockCtx { + a_row_offset: a_off, + b_row_offset: b_off, + }, + ); + + // Stage C: fold into this A-panel's state slice. + reducer.fold(arch, &mut state[a_off..a_off + a_panel], scores, b_off); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiler/tilers.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiler/tilers.rs new file mode 100644 index 000000000..d27e77873 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiler/tilers.rs @@ -0,0 +1,360 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Block-transposed A / row-major B walks, generic over the element type. Two +//! flavours share one tile/panel shape: +//! +//! - **Identity** ([`QueryWalk`] / [`DocWalk`]): tiles reborrow the source (`i8` +//! path — `i16` / `u8`, and any pre-materialised `f32`). +//! - **Convert** ([`QueryConvertWalk`] / [`DocConvertWalk`]): each tile widens its +//! `f16` source into a **reused** `f32` buffer (allocated once from the arena) and +//! hands back the *same* `f32` tile the identity-`f32` walk would — so one `f32` +//! kernel serves both. The lending [`TileWalk`] contract makes the reuse sound: a +//! live tile borrows the buffer, so `next` (which overwrites it) can't be called +//! until the tile is dropped. + +use super::{Tile, TileAt, TileWalk}; + +/// Rows per A-panel (the block-transposed group and the reducer's state block). +pub(crate) const A_PANEL: usize = 16; +/// Rows per full B-panel (the kernel's micro-panel / max unroll). +pub(crate) const B_PANEL: usize = 4; + +// ── Panels ─────────────────────────────────────────────────────── + +/// One block-transposed query block: `A_PANEL` rows × `k` `T`. +pub(crate) struct QPanel<'a, T> { + data: &'a [T], + k: usize, +} +/// One row-major doc panel: `rows` (`1..=B_PANEL`) × `k` `T`. The contraction `k` +/// travels with the query panel, so it isn't repeated here. +pub(crate) struct DPanel<'a, T> { + data: &'a [T], + rows: usize, +} + +impl Clone for QPanel<'_, T> { + fn clone(&self) -> Self { + *self + } +} +impl Copy for QPanel<'_, T> {} +impl Clone for DPanel<'_, T> { + fn clone(&self) -> Self { + *self + } +} +impl Copy for DPanel<'_, T> {} + +impl QPanel<'_, T> { + pub(crate) fn as_ptr(&self) -> *const T { + self.data.as_ptr() + } + pub(crate) fn k(&self) -> usize { + self.k + } +} +impl DPanel<'_, T> { + pub(crate) fn as_ptr(&self) -> *const T { + self.data.as_ptr() + } + pub(crate) fn rows(&self) -> usize { + self.rows + } +} + +// ── Materialized tiles ─────────────────────────────────────────── + +/// One A-tile: a run of whole `A_PANEL`-row blocks. +pub(crate) struct QMat<'a, T> { + data: &'a [T], + offset: usize, + k: usize, +} +/// One B-tile: a run of docs (last panel possibly short). +pub(crate) struct DMat<'a, T> { + data: &'a [T], + offset: usize, + k: usize, +} + +impl<'a, T> Tile for QMat<'a, T> { + type Panel = QPanel<'a, T>; + + fn offset(&self) -> usize { + self.offset + } + fn rows(&self) -> usize { + if self.k == 0 { + 0 + } else { + self.data.len() / self.k + } + } + fn panels(&self) -> impl Iterator> + '_ { + let (data, k) = (self.data, self.k); + let block = A_PANEL * k; + let n = if block == 0 { 0 } else { data.len() / block }; + (0..n).map(move |p| QPanel { + data: &data[p * block..(p + 1) * block], + k, + }) + } + fn tail(&self) -> Option> { + None // the query is padded to a whole number of A_PANEL blocks + } +} + +impl<'a, T> Tile for DMat<'a, T> { + type Panel = DPanel<'a, T>; + + fn offset(&self) -> usize { + self.offset + } + fn rows(&self) -> usize { + if self.k == 0 { + 0 + } else { + self.data.len() / self.k + } + } + fn panels(&self) -> impl Iterator> + '_ { + let (data, k) = (self.data, self.k); + let full = self.rows() / B_PANEL; + (0..full).map(move |p| DPanel { + data: &data[p * B_PANEL * k..(p + 1) * B_PANEL * k], + rows: B_PANEL, + }) + } + fn tail(&self) -> Option> { + let (data, k) = (self.data, self.k); + let rem = self.rows() % B_PANEL; + (rem > 0).then(|| DPanel { + data: &data[(self.rows() - rem) * k..], + rows: rem, + }) + } +} + +// ── Identity walks (reborrow the source) ───────────────────────── + +/// Block-transposed query source, walked `tile_panels` `A_PANEL`-row blocks at a time. +pub(crate) struct QueryWalk<'s, T> { + src: &'s [T], + k: usize, + tile_panels: usize, + cur: usize, + ti: usize, +} +/// Row-major doc source, walked `tile_panels` `B_PANEL`-row panels at a time. +pub(crate) struct DocWalk<'s, T> { + src: &'s [T], + k: usize, + tile_panels: usize, + cur: usize, + ti: usize, +} + +impl<'s, T> QueryWalk<'s, T> { + pub(crate) fn new(src: &'s [T], k: usize, tile_panels: usize) -> Self { + Self { + src, + k, + tile_panels, + cur: 0, + ti: 0, + } + } +} +impl<'s, T> DocWalk<'s, T> { + pub(crate) fn new(src: &'s [T], k: usize, tile_panels: usize) -> Self { + Self { + src, + k, + tile_panels, + cur: 0, + ti: 0, + } + } +} + +impl<'a, 's, T> TileAt<'a> for QueryWalk<'s, T> { + type Tile = QMat<'a, T>; +} +impl<'s, T> TileWalk for QueryWalk<'s, T> { + fn next(&mut self) -> Option> { + let src = self.src; + if self.cur >= src.len() { + return None; + } + let span = (self.tile_panels * A_PANEL * self.k).max(1); + let (start, ti) = (self.cur, self.ti); + let end = (start + span).min(src.len()); + self.cur = end; + self.ti += 1; + Some(QMat { + data: &src[start..end], + offset: ti * self.tile_panels * A_PANEL, + k: self.k, + }) + } + fn reset(&mut self) { + self.cur = 0; + self.ti = 0; + } + fn max_tile_rows(&self) -> usize { + self.tile_panels * A_PANEL + } +} + +impl<'a, 's, T> TileAt<'a> for DocWalk<'s, T> { + type Tile = DMat<'a, T>; +} +impl<'s, T> TileWalk for DocWalk<'s, T> { + fn next(&mut self) -> Option> { + let src = self.src; + if self.cur >= src.len() { + return None; + } + let span = (self.tile_panels * B_PANEL * self.k).max(1); + let (start, ti) = (self.cur, self.ti); + let end = (start + span).min(src.len()); + self.cur = end; + self.ti += 1; + Some(DMat { + data: &src[start..end], + offset: ti * self.tile_panels * B_PANEL, + k: self.k, + }) + } + fn reset(&mut self) { + self.cur = 0; + self.ti = 0; + } + fn max_tile_rows(&self) -> usize { + self.tile_panels * B_PANEL + } +} + +// ── Convert walks (widen f16 → f32 into a reused buffer) ────────── + +/// Block-transposed `f16` query, widened per tile into `buf` (`≥ max_tile_rows·k`). +pub(crate) struct QueryConvertWalk<'s, 'buf> { + src: &'s [half::f16], + k: usize, + tile_panels: usize, + cur: usize, + ti: usize, + buf: &'buf mut [f32], +} +/// Row-major `f16` docs, widened per tile into `buf` (`≥ max_tile_rows·k`). +pub(crate) struct DocConvertWalk<'s, 'buf> { + src: &'s [half::f16], + k: usize, + tile_panels: usize, + cur: usize, + ti: usize, + buf: &'buf mut [f32], +} + +impl<'s, 'buf> QueryConvertWalk<'s, 'buf> { + pub(crate) fn new( + src: &'s [half::f16], + k: usize, + tile_panels: usize, + buf: &'buf mut [f32], + ) -> Self { + Self { + src, + k, + tile_panels, + cur: 0, + ti: 0, + buf, + } + } +} +impl<'s, 'buf> DocConvertWalk<'s, 'buf> { + pub(crate) fn new( + src: &'s [half::f16], + k: usize, + tile_panels: usize, + buf: &'buf mut [f32], + ) -> Self { + Self { + src, + k, + tile_panels, + cur: 0, + ti: 0, + buf, + } + } +} + +impl<'a, 's, 'buf> TileAt<'a> for QueryConvertWalk<'s, 'buf> { + type Tile = QMat<'a, f32>; +} +impl<'s, 'buf> TileWalk for QueryConvertWalk<'s, 'buf> { + fn next(&mut self) -> Option> { + let (src, k) = (self.src, self.k); + if self.cur >= src.len() { + return None; + } + let span = (self.tile_panels * A_PANEL * k).max(1); + let (start, ti) = (self.cur, self.ti); + let end = (start + span).min(src.len()); + let len = end - start; + for i in 0..len { + self.buf[i] = diskann_wide::cast_f16_to_f32(src[start + i]); + } + self.cur = end; + self.ti += 1; + Some(QMat { + data: &self.buf[..len], + offset: ti * self.tile_panels * A_PANEL, + k, + }) + } + fn reset(&mut self) { + self.cur = 0; + self.ti = 0; + } + fn max_tile_rows(&self) -> usize { + self.tile_panels * A_PANEL + } +} + +impl<'a, 's, 'buf> TileAt<'a> for DocConvertWalk<'s, 'buf> { + type Tile = DMat<'a, f32>; +} +impl<'s, 'buf> TileWalk for DocConvertWalk<'s, 'buf> { + fn next(&mut self) -> Option> { + let (src, k) = (self.src, self.k); + if self.cur >= src.len() { + return None; + } + let span = (self.tile_panels * B_PANEL * k).max(1); + let (start, ti) = (self.cur, self.ti); + let end = (start + span).min(src.len()); + let len = end - start; + for i in 0..len { + self.buf[i] = diskann_wide::cast_f16_to_f32(src[start + i]); + } + self.cur = end; + self.ti += 1; + Some(DMat { + data: &self.buf[..len], + offset: ti * self.tile_panels * B_PANEL, + k, + }) + } + fn reset(&mut self) { + self.cur = 0; + self.ti = 0; + } + fn max_tile_rows(&self) -> usize { + self.tile_panels * B_PANEL + } +} diff --git a/diskann-quantization/src/multi_vector/distance/mod.rs b/diskann-quantization/src/multi_vector/distance/mod.rs index 39c667d82..74fef8f5f 100644 --- a/diskann-quantization/src/multi_vector/distance/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/mod.rs @@ -60,3 +60,15 @@ pub use projected_eigen::ProjectedEigen; /// (V3/AVX2 only) — not yet unified into [`MaxSimIsa`]/[`build_max_sim`]. #[cfg(target_arch = "x86_64")] pub use kernels::{QuantStagedDocs, QuantStagedQuery}; + +/// Standalone POC entry for the coarse `Tiler`/`Tile` rebuild (V3/AVX2 only) — the +/// 4-bit MinMax kernel where tilers own movement and the kernel is pinned on the +/// A/B panel types. The i8 A/B baseline is `QuantStaged*`. +#[cfg(target_arch = "x86_64")] +pub use kernels::{QuantTiledDocs, QuantTiledQuery}; + +/// Standalone POC entry for the coarse tiler's **f16** path (V3/AVX2 only) — each +/// tile is widened f16→f32 into a reused buffer, then an f32 store kernel + identity +/// postprocess reuse the same tiled pipeline. +#[cfg(target_arch = "x86_64")] +pub use kernels::{QuantTiledF16Docs, QuantTiledF16Query}; From d9daec5712726e2486a1f8093f26c7826414138b Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Fri, 31 Jul 2026 05:50:27 +0530 Subject: [PATCH 3/5] Add paneled end to end design --- .../example/multi-vector-f16.json | 71 +++ .../example/multi-vector-paneled-f32.json | 20 + diskann-benchmark/src/inputs/multi_vector.rs | 81 ++++ diskann-benchmark/src/multi_vector/mod.rs | 5 + .../src/multi_vector/paneled_f32.rs | 273 ++++++++++++ diskann-benchmark/src/multi_vector/quant.rs | 80 ++-- .../src/multi_vector/block_transposed.rs | 151 +++++++ .../src/multi_vector/distance/kernels/mod.rs | 8 + .../distance/kernels/paneled/arena.rs | 90 ++++ .../distance/kernels/paneled/float.rs | 282 ++++++++++++ .../distance/kernels/paneled/leaves.rs | 234 ++++++++++ .../distance/kernels/paneled/minmax.rs | 414 ++++++++++++++++++ .../distance/kernels/paneled/mod.rs | 245 +++++++++++ .../distance/kernels/paneled/strip.rs | 100 +++++ .../distance/kernels/paneled/views.rs | 197 +++++++++ .../src/multi_vector/distance/mod.rs | 6 + .../src/multi_vector/matrix.rs | 48 ++ 17 files changed, 2279 insertions(+), 26 deletions(-) create mode 100644 diskann-benchmark/example/multi-vector-f16.json create mode 100644 diskann-benchmark/example/multi-vector-paneled-f32.json create mode 100644 diskann-benchmark/src/multi_vector/paneled_f32.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/paneled/arena.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs diff --git a/diskann-benchmark/example/multi-vector-f16.json b/diskann-benchmark/example/multi-vector-f16.json new file mode 100644 index 000000000..23f269d6d --- /dev/null +++ b/diskann-benchmark/example/multi-vector-f16.json @@ -0,0 +1,71 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-op", + "content": { "element_type": "float32", "isa": "x86-64-v3", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] } }, + { + "type": "multi-vector-op", + "content": { "element_type": "float32", "isa": "reference", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] } }, + { + "type": "multi-vector-op", + "content": { "element_type": "float16", "isa": "x86-64-v3", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] } }, + { + "type": "multi-vector-op", + "content": { "element_type": "float16", "isa": "x86-64-v3-f16-direct-v2", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] } }, + { + "type": "multi-vector-op", + "content": { "element_type": "float16", "isa": "x86-64-v3-f16-direct-v3", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] } }, + { + "type": "multi-vector-op", + "content": { "element_type": "float16", "isa": "reference", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 50 } + ] } } + ] +} diff --git a/diskann-benchmark/example/multi-vector-paneled-f32.json b/diskann-benchmark/example/multi-vector-paneled-f32.json new file mode 100644 index 000000000..4a281f532 --- /dev/null +++ b/diskann-benchmark/example/multi-vector-paneled-f32.json @@ -0,0 +1,20 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-paneled-f32-op", + "content": { + "runs": [ + { "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 500, "num_measurements": 30 }, + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 30 }, + { "num_query_vectors": 32, "num_doc_vectors": 128, "dim": 384, "loops_per_measurement": 20, "num_measurements": 30 }, + { "num_query_vectors": 32, "num_doc_vectors": 16, "dim": 256, "loops_per_measurement": 200, "num_measurements": 30 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 30 }, + { "num_query_vectors": 64, "num_doc_vectors": 1250, "dim": 512, "loops_per_measurement": 2, "num_measurements": 20 }, + { "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 128, "loops_per_measurement": 200, "num_measurements": 30 }, + { "num_query_vectors": 32, "num_doc_vectors": 32, "dim": 512, "loops_per_measurement": 50, "num_measurements": 30 } + ] + } + } + ] +} diff --git a/diskann-benchmark/src/inputs/multi_vector.rs b/diskann-benchmark/src/inputs/multi_vector.rs index d431a0909..efe5e5068 100644 --- a/diskann-benchmark/src/inputs/multi_vector.rs +++ b/diskann-benchmark/src/inputs/multi_vector.rs @@ -232,6 +232,87 @@ impl std::fmt::Display for MultiVectorQuantOp { } } +/////////////////////////////// +// Multi-Vector Paneled f32 // +/////////////////////////////// + +/// An **f32** multi-vector MaxSim A/B benchmark job: the paneled rebuild vs the +/// production block-transposed fused V3 kernel vs the non-SIMD reference, at +/// identical shapes over identical data. +/// +/// Not fully apples-to-apples — the paneled path pre-materializes its doc side once +/// at build (excluded from timing), while the fused kernel takes a `MatRef` per call. +/// The tiler is absent by construction: it has no f32 instantiation (only 4-bit +/// MinMax and f16). +/// +/// Element type is f32 and the ISA is fixed to V3/AVX2, so neither is a JSON field. +/// x86_64-only. +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct MultiVectorPaneledF32Op { + pub(crate) runs: Vec, +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl MultiVectorPaneledF32Op { + pub(crate) const fn tag() -> &'static str { + "multi-vector-paneled-f32-op" + } +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl Input for MultiVectorPaneledF32Op { + type Raw = Self; + + fn tag() -> &'static str { + Self::tag() + } + + fn from_raw(raw: Self::Raw, _checker: &mut Checker) -> anyhow::Result { + Ok(raw) + } + + fn serialize(&self) -> anyhow::Result { + Ok(serde_json::to_value(self)?) + } + + fn example() -> Self { + const NUM_DOC_VECTORS: NonZeroUsize = NonZeroUsize::new(64).unwrap(); + const DIM: NonZeroUsize = NonZeroUsize::new(128).unwrap(); + const LOOPS_PER_MEASUREMENT: NonZeroUsize = NonZeroUsize::new(50).unwrap(); + const NUM_MEASUREMENTS: NonZeroUsize = NonZeroUsize::new(20).unwrap(); + + let runs = vec![ + Run { + num_query_vectors: NonZeroUsize::new(32).unwrap(), + num_doc_vectors: NUM_DOC_VECTORS, + dim: DIM, + loops_per_measurement: LOOPS_PER_MEASUREMENT, + num_measurements: NUM_MEASUREMENTS, + }, + Run { + num_query_vectors: NonZeroUsize::new(64).unwrap(), + num_doc_vectors: NUM_DOC_VECTORS, + dim: DIM, + loops_per_measurement: LOOPS_PER_MEASUREMENT, + num_measurements: NUM_MEASUREMENTS, + }, + ]; + + Self { runs } + } +} + +#[cfg(all(feature = "multi-vector", target_arch = "x86_64"))] +impl std::fmt::Display for MultiVectorPaneledF32Op { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "Multi-Vector Paneled f32 Operation\n")?; + write_field!(f, "tag", Self::tag())?; + write_field!(f, "number of runs", self.runs.len())?; + Ok(()) + } +} + /////////////////////////////// // Multi-Vector Tiled f16 Op // /////////////////////////////// diff --git a/diskann-benchmark/src/multi_vector/mod.rs b/diskann-benchmark/src/multi_vector/mod.rs index f1f0e8f3b..c93538395 100644 --- a/diskann-benchmark/src/multi_vector/mod.rs +++ b/diskann-benchmark/src/multi_vector/mod.rs @@ -28,6 +28,9 @@ cfg_if::cfg_if! { // The quantized A/B op drives the V3-only staged integer kernel. #[cfg(target_arch = "x86_64")] mod quant; + // The f32 A/B op: paneled vs the production fused V3 kernel vs reference. + #[cfg(target_arch = "x86_64")] + mod paneled_f32; // The f16 A/B op: coarse tiler vs the f16.rs preprocess path (V3-only). #[cfg(target_arch = "x86_64")] mod tiled_f16; @@ -37,6 +40,8 @@ cfg_if::cfg_if! { #[cfg(target_arch = "x86_64")] quant::register(registry)?; #[cfg(target_arch = "x86_64")] + paneled_f32::register(registry)?; + #[cfg(target_arch = "x86_64")] tiled_f16::register(registry)?; Ok(()) } diff --git a/diskann-benchmark/src/multi_vector/paneled_f32.rs b/diskann-benchmark/src/multi_vector/paneled_f32.rs new file mode 100644 index 000000000..d038f4768 --- /dev/null +++ b/diskann-benchmark/src/multi_vector/paneled_f32.rs @@ -0,0 +1,273 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! A/B benchmark for **f32** multi-vector MaxSim across three paths at identical +//! shapes over identical data: +//! +//! - **Paneled** — the paneled rebuild (views own their panel decomposition, one +//! `Drain` seam), driven through `PaneledF32Query`. +//! - **Fused** — the production block-transposed V3 kernel via the public factory +//! (`MaxSimIsa::X86_64_V3`). +//! - **Reference** — the non-SIMD `MaxSimIsa::Reference` baseline. +//! +//! The coarse tiler is deliberately absent: it has no f32 instantiation (only 4-bit +//! MinMax and f16), so there is nothing to time. +//! +//! # Reading the numbers +//! +//! Not perfectly apples-to-apples. The paneled path pre-materializes its doc side +//! once in `build` (excluded from the timing), while the fused kernel is handed a +//! `MatRef` per call. Treat `Paneled/Fused` as a ceiling on the paneled structure's +//! win, not a pure abstraction delta. +//! +//! x86_64 (V3/AVX2) only. + +use std::io::Write; + +use diskann_benchmark_runner::{ + benchmark::{FailureScore, MatchScore}, + utils::{fmt::Table, percentiles, MicroSeconds}, + Benchmark, Checkpoint, Output, Registry, +}; +use diskann_quantization::multi_vector::distance::{PaneledF32Docs, PaneledF32Query}; +use diskann_quantization::multi_vector::{build_max_sim, BoxErase, MaxSimIsa}; +use serde::{Deserialize, Serialize}; + +use super::driver::Data; +use crate::inputs::multi_vector::{MultiVectorPaneledF32Op, Run}; +use crate::utils::DisplayWrapper; + +// ───────────────────────────────────────────────────────────────────────── +// Kernel. +// ───────────────────────────────────────────────────────────────────────── + +#[derive(Debug)] +pub(super) struct PaneledF32Kernel; + +impl PaneledF32Kernel { + pub(super) const fn new() -> Self { + Self + } +} + +impl Benchmark for PaneledF32Kernel { + type Input = MultiVectorPaneledF32Op; + type Output = Vec; + + fn try_match(&self, _from: &MultiVectorPaneledF32Op) -> Result { + if PaneledF32Query::is_supported() { + Ok(MatchScore(0)) + } else { + Err(FailureScore(0)) + } + } + + fn run( + &self, + input: &MultiVectorPaneledF32Op, + _: Checkpoint<'_>, + mut output: &mut dyn Output, + ) -> anyhow::Result { + writeln!(output, "{}", input)?; + let mut results = Vec::with_capacity(input.runs.len()); + for run in input.runs.iter() { + results.push(run_ab(run)?); + } + writeln!(output, "\n\n{}", DisplayWrapper(&*results))?; + Ok(results) + } + + fn description( + &self, + f: &mut std::fmt::Formatter<'_>, + input: Option<&MultiVectorPaneledF32Op>, + ) -> std::fmt::Result { + match input { + None => writeln!(f, "- f32 MaxSim, paneled / fused V3 / reference (V3/AVX2)")?, + Some(_) => { + if !PaneledF32Query::is_supported() { + writeln!(f, "\n - AVX2 (V3) unavailable on this CPU")?; + } + } + } + Ok(()) + } +} + +// ───────────────────────────────────────────────────────────────────────── +// A/B timing. +// ───────────────────────────────────────────────────────────────────────── + +/// Run `f` `loops_per_measurement` times per measurement, `num_measurements` +/// times, returning the per-measurement latencies and their percentiles. +fn measure(run: &Run, mut f: impl FnMut()) -> Series { + let mut latencies = Vec::with_capacity(run.num_measurements.get()); + for _ in 0..run.num_measurements.get() { + let start = std::time::Instant::now(); + for _ in 0..run.loops_per_measurement.get() { + f(); + } + latencies.push(start.elapsed().into()); + } + let percentiles = percentiles::compute_percentiles(&mut latencies).unwrap(); + Series { + latencies, + percentiles, + } +} + +/// Build all three paths for one shape and time them (build cost excluded). +fn run_ab(run: &Run) -> anyhow::Result { + let data = Data::::new(run)?; + + // Path A — the paneled rebuild. + let mut paneled_query = PaneledF32Query::build(data.queries.as_view()) + .ok_or_else(|| anyhow::anyhow!("AVX2 (V3) unavailable for the paneled f32 kernel"))?; + let paneled_docs = PaneledF32Docs::build(data.docs.as_view()); + + // Path B / C — the production factory kernels over the same query matrix. + let fused_kernel = + build_max_sim::(MaxSimIsa::X86_64_V3, data.queries.as_view(), BoxErase)?; + let ref_kernel = + build_max_sim::(MaxSimIsa::Reference, data.queries.as_view(), BoxErase)?; + + let nq = run.num_query_vectors.get(); + let mut scores = vec![0.0f32; nq]; + let doc_view = data.docs.as_view(); + + // Launder inputs *and* output through `black_box` each iteration, matching the + // quantized A/B: the factory kernels are opaque cross-crate calls, but the + // paneled path is in-crate and could otherwise be hoisted out of the loop. + let paneled = measure(run, || { + let docs = std::hint::black_box(&paneled_docs); + paneled_query.compute_max_sim(docs, &mut scores); + std::hint::black_box(&mut scores); + }); + + // Timed adjacent to `paneled` so the ratio survives cross-run clock variance. + let fused = measure(run, || { + let doc_view = std::hint::black_box(doc_view); + fused_kernel + .compute_max_sim(doc_view, &mut scores) + .expect("scores.len() == kernel.nrows() by construction"); + std::hint::black_box(&mut scores); + }); + + let reference = measure(run, || { + let doc_view = std::hint::black_box(doc_view); + ref_kernel + .compute_max_sim(doc_view, &mut scores) + .expect("scores.len() == kernel.nrows() by construction"); + std::hint::black_box(&mut scores); + }); + + Ok(F32RunResult { + run: run.clone(), + paneled, + fused, + reference, + }) +} + +// ───────────────────────────────────────────────────────────────────────── +// Result types. +// ───────────────────────────────────────────────────────────────────────── + +/// One timed series (per-measurement latencies + percentiles). +#[derive(Debug, Serialize, Deserialize)] +pub(super) struct Series { + latencies: Vec, + percentiles: percentiles::Percentiles, +} + +impl Series { + /// Minimum latency, in microseconds. + fn min_us(&self) -> f64 { + self.latencies + .iter() + .min() + .copied() + .unwrap_or(MicroSeconds::new(u64::MAX)) + .as_f64() + } +} + +/// Paneled-vs-fused-vs-reference result for one shape. +#[derive(Debug, Serialize, Deserialize)] +pub(super) struct F32RunResult { + pub(super) run: Run, + pub(super) paneled: Series, + pub(super) fused: Series, + pub(super) reference: Series, +} + +impl F32RunResult { + fn computations(&self) -> f64 { + (self.run.num_query_vectors.get() + * self.run.num_doc_vectors.get() + * self.run.loops_per_measurement.get()) as f64 + } +} + +impl std::fmt::Display for DisplayWrapper<'_, [F32RunResult]> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.is_empty() { + return Ok(()); + } + + writeln!( + f, + "ns/IP = min time per (query, doc) inner-product call. \ + Panel/Fused < 1 ⇒ paneled faster than the production V3 kernel. \ + Speedup = reference ÷ paneled." + )?; + + let header = [ + "Q", + "D", + "Dim", + "Paneled", + "Fused V3", + "Panel/Fused", + "Reference", + "Ref/Panel", + ]; + let mut table = Table::new(header, self.len()); + + self.iter().enumerate().for_each(|(row, r)| { + let comps = r.computations(); + let paneled = r.paneled.min_us() / comps * 1000.0; + let fused = r.fused.min_us() / comps * 1000.0; + let reference = r.reference.min_us() / comps * 1000.0; + let vs_fused = if fused > 0.0 { paneled / fused } else { 0.0 }; + let speedup = if paneled > 0.0 { + reference / paneled + } else { + 0.0 + }; + + let mut row = table.row(row); + row.insert(r.run.num_query_vectors, 0); + row.insert(r.run.num_doc_vectors, 1); + row.insert(r.run.dim, 2); + row.insert(format!("{:.3}", paneled), 3); + row.insert(format!("{:.3}", fused), 4); + row.insert(format!("{:.2}x", vs_fused), 5); + row.insert(format!("{:.3}", reference), 6); + row.insert(format!("{:.2}x", speedup), 7); + }); + + table.fmt(f) + } +} + +// ───────────────────────────────────────────────────────────────────────── +// Registration. +// ───────────────────────────────────────────────────────────────────────── + +pub(super) fn register(registry: &mut Registry) -> anyhow::Result<()> { + registry.register("multi-vector-paneled-f32-op", PaneledF32Kernel::new())?; + Ok(()) +} diff --git a/diskann-benchmark/src/multi_vector/quant.rs b/diskann-benchmark/src/multi_vector/quant.rs index efe981e6c..d51520ee2 100644 --- a/diskann-benchmark/src/multi_vector/quant.rs +++ b/diskann-benchmark/src/multi_vector/quant.rs @@ -3,16 +3,21 @@ * Licensed under the MIT license. */ -//! A/B benchmark for **4-bit MinMax quantized** multi-vector MaxSim: the -//! experimental *staged integer* kernel (block-transposed `i16` query + `u8` -//! doc codes, `vpmaddwd` accumulation, metadata postprocess) vs the scalar -//! [`MinMaxKernel`] reference — at identical shapes and identical quantization. +//! A/B benchmark for **4-bit MinMax quantized** multi-vector MaxSim across four +//! paths at identical shapes and identical quantization: //! -//! Both paths consume the *same* random f32 multi-vectors quantized to 4-bit -//! MinMax (Null transform, scale 1.0), so the comparison isolates the distance -//! kernel. The build / quantize cost is excluded from the timing. +//! - **Staged** — block-transposed `i16` query + `u8` doc codes, `vpmaddwd` +//! accumulation, metadata postprocess as a separate stage. +//! - **Tiled** — the coarse tiler rebuild (accumulate → postprocess → reduce). +//! - **Paneled** — the paneled rebuild (views own their panel decomposition; one +//! `Drain` seam fuses dequant + reduce). +//! - **Reference** — the scalar [`MinMaxKernel`] baseline. //! -//! x86_64 (V3/AVX2) only — the quantized staged kernel has no other backend. +//! All four consume the *same* random f32 multi-vectors quantized to 4-bit MinMax +//! (Null transform, scale 1.0), so the comparison isolates the distance kernel. The +//! build / quantize cost is excluded from the timing. +//! +//! x86_64 (V3/AVX2) only — the quantized experimental kernels have no other backend. use std::io::Write; use std::num::NonZeroUsize; @@ -26,7 +31,8 @@ use diskann_quantization::algorithms::transforms::NullTransform; use diskann_quantization::algorithms::Transform; use diskann_quantization::minmax::{MinMaxMeta, MinMaxQuantizer}; use diskann_quantization::multi_vector::distance::{ - QuantStagedDocs, QuantStagedQuery, QuantTiledDocs, QuantTiledQuery, + PaneledQuantDocs, PaneledQuantQuery, QuantStagedDocs, QuantStagedQuery, QuantTiledDocs, + QuantTiledQuery, }; use diskann_quantization::multi_vector::{Defaulted, Mat, MatRef, MaxSim, QueryMatRef, Standard}; use diskann_quantization::num::Positive; @@ -86,7 +92,10 @@ impl Benchmark for QuantKernel { input: Option<&MultiVectorQuantOp>, ) -> std::fmt::Result { match input { - None => writeln!(f, "- 4-bit MinMax quantized staged MaxSim (V3/AVX2)")?, + None => writeln!( + f, + "- 4-bit MinMax quantized MaxSim, staged / tiled / paneled / reference (V3/AVX2)" + )?, Some(_) => { if !QuantStagedQuery::is_supported() { writeln!(f, "\n - AVX2 (V3) unavailable on this CPU")?; @@ -148,7 +157,13 @@ fn run_ab(run: &Run) -> anyhow::Result { .ok_or_else(|| anyhow::anyhow!("AVX2 (V3) unavailable for the tiled quantized kernel"))?; let tiled_docs = QuantTiledDocs::build(data.docs.as_view()); - // Path C — scalar MinMax reference over the same quantization. + // Path C — the paneled rebuild (views own their panel decomposition, one `Drain` + // seam fusing dequant + reduce). + let mut paneled_query = PaneledQuantQuery::build(data.queries.as_view()) + .ok_or_else(|| anyhow::anyhow!("AVX2 (V3) unavailable for the paneled quantized kernel"))?; + let paneled_docs = PaneledQuantDocs::build(data.docs.as_view()); + + // Path D — scalar MinMax reference over the same quantization. let q_ref = quantize(data.queries.as_view()); let d_ref = quantize(data.docs.as_view()); @@ -175,6 +190,12 @@ fn run_ab(run: &Run) -> anyhow::Result { std::hint::black_box(&mut scores); }); + let paneled = measure(run, || { + let docs = std::hint::black_box(&paneled_docs); + paneled_query.compute_max_sim(docs, &mut scores); + std::hint::black_box(&mut scores); + }); + let reference = measure(run, || { let q_ref = std::hint::black_box(&q_ref); let d_ref = std::hint::black_box(&d_ref); @@ -187,6 +208,7 @@ fn run_ab(run: &Run) -> anyhow::Result { run: run.clone(), staged, tiled, + paneled, reference, }) } @@ -214,12 +236,13 @@ impl Series { } } -/// Staged-vs-tiled-vs-reference result for one shape. +/// Staged-vs-tiled-vs-paneled-vs-reference result for one shape. #[derive(Debug, Serialize, Deserialize)] pub(super) struct QuantRunResult { pub(super) run: Run, pub(super) staged: Series, pub(super) tiled: Series, + pub(super) paneled: Series, pub(super) reference: Series, } @@ -239,20 +262,22 @@ impl std::fmt::Display for DisplayWrapper<'_, [QuantRunResult]> { writeln!( f, - "ns/IP = min time per (query, doc) inner-product call; \ - Tiled/Staged = tiled ÷ staged (1.00 ⇒ zero abstraction overhead, \ - >1 ⇒ tiled slower); Speedup = reference ÷ staged" + "ns/IP = min time per (query, doc) inner-product call. \ + Ratios are vs Staged (1.00 ⇒ parity, >1 ⇒ slower). \ + Speedup = reference ÷ paneled." )?; let header = [ "Q", "D", "Dim", - "Staged (ns/IP)", - "Tiled (ns/IP)", - "Tiled/Staged", - "Reference (ns/IP)", - "Speedup", + "Staged", + "Tiled", + "Paneled", + "Tiled/St", + "Panel/St", + "Reference", + "Ref/Panel", ]; let mut table = Table::new(header, self.len()); @@ -260,10 +285,11 @@ impl std::fmt::Display for DisplayWrapper<'_, [QuantRunResult]> { let comps = r.computations(); let staged = r.staged.min_us() / comps * 1000.0; let tiled = r.tiled.min_us() / comps * 1000.0; + let paneled = r.paneled.min_us() / comps * 1000.0; let reference = r.reference.min_us() / comps * 1000.0; - let overhead = if staged > 0.0 { tiled / staged } else { 0.0 }; - let speedup = if staged > 0.0 { - reference / staged + let ratio = |x: f64| if staged > 0.0 { x / staged } else { 0.0 }; + let speedup = if paneled > 0.0 { + reference / paneled } else { 0.0 }; @@ -274,9 +300,11 @@ impl std::fmt::Display for DisplayWrapper<'_, [QuantRunResult]> { row.insert(r.run.dim, 2); row.insert(format!("{:.3}", staged), 3); row.insert(format!("{:.3}", tiled), 4); - row.insert(format!("{:.2}x", overhead), 5); - row.insert(format!("{:.3}", reference), 6); - row.insert(format!("{:.2}x", speedup), 7); + row.insert(format!("{:.3}", paneled), 5); + row.insert(format!("{:.2}x", ratio(tiled)), 6); + row.insert(format!("{:.2}x", ratio(paneled)), 7); + row.insert(format!("{:.3}", reference), 8); + row.insert(format!("{:.2}x", speedup), 9); }); table.fmt(f) diff --git a/diskann-quantization/src/multi_vector/block_transposed.rs b/diskann-quantization/src/multi_vector/block_transposed.rs index 6fe9a1315..81249986d 100644 --- a/diskann-quantization/src/multi_vector/block_transposed.rs +++ b/diskann-quantization/src/multi_vector/block_transposed.rs @@ -840,6 +840,58 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedRef<'a, } } + /// Return the backing elements of `block` — exactly `GROUP * padded_ncols()` of + /// them — or `None` if `block >= num_blocks()`. + /// + /// Unlike [`block()`](Self::block) this accepts the zero-padded remainder block, + /// so a kernel can walk `0..num_blocks()` uniformly instead of special-casing the + /// tail. The padding rows are part of the slice; use [`remainder()`](Self::remainder) + /// to find how many of the last block's rows are logical. + #[inline] + pub fn block_slice(&self, block: usize) -> Option<&'a [T]> { + if block >= self.num_blocks() { + return None; + } + // SAFETY: `block < num_blocks()` checked above, and `block_ptr_unchecked` + // guarantees `GROUP * padded_ncols()` valid elements even for the remainder. + unsafe { + Some(std::slice::from_raw_parts( + self.block_ptr_unchecked(block), + self.data.repr().block_stride(), + )) + } + } + + /// A view over `count` consecutive blocks starting at `start` — a cache-sized + /// tile that is a block-transposed matrix in its own right. + /// + /// `count` is clipped to the blocks available, so a walk can request a fixed tile + /// size and get a short final tile. Returns `None` once `start` reaches + /// [`num_blocks()`](Self::num_blocks), which ends such a walk. + /// + /// The sub-view reports its own *logical* row count, so a tile containing the + /// parent's partial block reports that block as its remainder rather than + /// claiming the padding rows are real. + #[inline] + #[allow(clippy::expect_used)] + pub fn block_range(&self, start: usize, count: usize) -> Option { + let count = count.min(self.num_blocks().checked_sub(start)?); + if count == 0 { + return None; + } + let nrows = (self.nrows() - start * GROUP).min(count * GROUP); + let repr = BlockTransposedRepr::::new(nrows, self.ncols()) + .expect("sub-view of a valid matrix cannot overflow"); + // SAFETY: `start + count <= num_blocks()`, so the range lies inside the + // backing allocation; `repr.storage_len()` is exactly `count` block strides. + let data: &[T] = unsafe { + std::slice::from_raw_parts(self.block_ptr_unchecked(start), repr.storage_len()) + }; + Some(Self { + data: MatRef::new(repr, data).expect("slice matches repr"), + }) + } + /// Retrieve the value at the logical `(row, col)`. /// /// # Panics @@ -894,6 +946,8 @@ impl<'a, T: Copy, const GROUP: usize, const PACK: usize> BlockTransposedMut<'a, delegate_to_ref!(#[allow(clippy::missing_safety_doc)] unsafe pub fn block_ptr_unchecked(&self, block: usize) -> *const T); delegate_to_ref!(#[allow(clippy::expect_used)] pub fn block(&self, block: usize) -> MatrixView<'_, T>); delegate_to_ref!(#[allow(clippy::expect_used)] pub fn remainder_block(&self) -> Option>); + delegate_to_ref!(pub fn block_slice(&self, block: usize) -> Option<&[T]>); + delegate_to_ref!(#[allow(clippy::expect_used)] pub fn block_range(&self, start: usize, count: usize) -> Option>); delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T); /// Group size (blocking factor `GROUP`). @@ -1042,6 +1096,8 @@ impl BlockTransposed *const T); delegate_to_ref!(#[allow(clippy::expect_used)] pub fn block(&self, block: usize) -> MatrixView<'_, T>); delegate_to_ref!(#[allow(clippy::expect_used)] pub fn remainder_block(&self) -> Option>); + delegate_to_ref!(pub fn block_slice(&self, block: usize) -> Option<&[T]>); + delegate_to_ref!(#[allow(clippy::expect_used)] pub fn block_range(&self, start: usize, count: usize) -> Option>); delegate_to_ref!(pub fn get_element(&self, row: usize, col: usize) -> T); /// Group size (blocking factor `GROUP`). @@ -1931,6 +1987,101 @@ mod tests { } } + /// `block_range` partitions the blocks exactly: every tiling reproduces the + /// parent's blocks in order, and each tile reports honest logical rows. + fn test_block_range( + nrows: usize, + ncols: usize, + gen_element: fn(usize) -> T, + ) { + let mut data = Matrix::new(T::default(), nrows, ncols); + data.as_mut_slice() + .iter_mut() + .enumerate() + .for_each(|(i, d)| *d = gen_element(i)); + let bt = BlockTransposed::::from_strided(data.as_view().into()); + let v = bt.as_view(); + + let all: Vec<&[T]> = (0..v.num_blocks()) + .map(|b| v.block_slice(b).unwrap()) + .collect(); + for s in &all { + assert_eq!(s.len(), GROUP * v.padded_ncols()); + } + assert!(v.block_slice(v.num_blocks()).is_none()); + + for tile in 1..v.num_blocks() + 3 { + let mut cur = 0; + while let Some(t) = v.block_range(cur, tile) { + assert!(t.num_blocks() <= tile); + assert_eq!(t.padded_ncols(), v.padded_ncols()); + assert_eq!( + t.nrows(), + (v.nrows() - cur * GROUP).min(t.num_blocks() * GROUP), + "tile must report its own logical rows, not the padding" + ); + for b in 0..t.num_blocks() { + assert_eq!( + t.block_slice(b).unwrap(), + all[cur + b], + "tile {cur} block {b}" + ); + } + cur += t.num_blocks(); + } + assert_eq!( + cur, + v.num_blocks(), + "tile size {tile} must cover every block (nrows={nrows}, ncols={ncols})" + ); + } + assert!(v.block_range(v.num_blocks(), 1).is_none()); + assert!(v.block_range(v.num_blocks() + 1, 1).is_none()); + assert!(v.block_range(0, 0).is_none()); + } + + #[test] + fn test_block_range_group16() { + for nrows in [0, 1, 15, 16, 17, 31, 32, 33, 64] { + for ncols in [1, 2, 5] { + test_block_range::(nrows, ncols, gen_f32); + } + } + } + + #[test] + fn test_block_range_group8_pack_agnostic() { + for nrows in [0, 1, 7, 8, 9, 17] { + for ncols in [1, 3, 4] { + test_block_range::(nrows, ncols, gen_i32); + } + } + } + + /// The sub-views must outlive the temporary the delegate creates, so this is a + /// lifetime test as much as a behavioural one. + #[test] + fn test_block_sub_views_delegate_to_owning_types() { + let mut data = Matrix::new(0.0f32, 33, 5); + data.as_mut_slice() + .iter_mut() + .enumerate() + .for_each(|(i, d)| *d = gen_f32(i)); + let mut bt = BlockTransposed::::from_strided(data.as_view().into()); + let expected = bt.as_view().block_slice(1).unwrap().to_vec(); + + assert_eq!(bt.block_slice(1).unwrap(), expected); + assert!(bt.block_slice(bt.num_blocks()).is_none()); + let tile = bt.block_range(1, 2).unwrap(); + assert_eq!((tile.num_blocks(), tile.nrows()), (2, 17)); + assert_eq!(tile.block_slice(0).unwrap(), expected); + assert!(bt.block_range(bt.num_blocks(), 1).is_none()); + + let m = bt.as_view_mut(); + assert_eq!(m.block_slice(1).unwrap(), expected); + assert_eq!(m.block_range(1, 2).unwrap().nrows(), 17); + } + // ════════════════════════════════════════════════════════════════ // Focused tests (not part of the unified parameterized test) // ════════════════════════════════════════════════════════════════ diff --git a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs index ad270c574..3aeb64b3c 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs @@ -25,6 +25,10 @@ mod tiled_reduce; #[cfg(target_arch = "x86_64")] pub(super) mod tiler; +// Paneled rebuild — views own their panel decomposition, one `Drain` seam (V3 only). +#[cfg(target_arch = "x86_64")] +pub(super) mod paneled; + // Re-export the quantized staged kernel's public POC entry (x86_64 only). #[cfg(target_arch = "x86_64")] pub use staged::{QuantStagedDocs, QuantStagedQuery}; @@ -33,6 +37,10 @@ pub use staged::{QuantStagedDocs, QuantStagedQuery}; #[cfg(target_arch = "x86_64")] pub use tiler::{QuantTiledDocs, QuantTiledF16Docs, QuantTiledF16Query, QuantTiledQuery}; +// Re-export the paneled POC entries (x86_64 only). +#[cfg(target_arch = "x86_64")] +pub use paneled::{PaneledF32Docs, PaneledF32Query, PaneledQuantDocs, PaneledQuantQuery}; + // ── Tile budget ────────────────────────────────────────────────── /// Cache budgets fed to the tile planner. diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/arena.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/arena.rs new file mode 100644 index 000000000..e59a5cd7f --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/arena.rs @@ -0,0 +1,90 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Single-owner resettable bump arena (copy of `staged::arena`, kept local so the +//! `paneled` experiment doesn't reach into its siblings). + +use std::cell::{Cell, UnsafeCell}; +use std::ptr::NonNull; + +use crate::alloc::{AlignedAllocator, AllocatorCore, AllocatorError, Poly}; + +/// Single-owner, single-threaded resettable bump over a 64-byte-aligned buffer. +pub(crate) struct ResettableArena { + buffer: Poly, AlignedAllocator>, + head: Cell, +} + +impl ResettableArena { + pub(crate) fn with_capacity(capacity: usize) -> Result { + let buffer = Poly::<[u8], _>::new_uninit_slice(capacity.max(1), AlignedAllocator::A64)?; + let (ptr, alloc) = Poly::into_raw(buffer); + + // SAFETY: `UnsafeCell<[u8]>` shares the layout of `[u8]`; `MaybeUninit` is + // layout-compatible with `u8` (bytes are only read after being handed out and + // written). `ptr` is non-null, from `Poly::into_raw`. + let buffer = unsafe { + Poly::from_raw( + NonNull::new_unchecked(ptr.as_ptr() as *mut UnsafeCell<[u8]>), + alloc, + ) + }; + + Ok(Self { + buffer, + head: Cell::new(0), + }) + } + + /// Rewind in O(1). `&mut self` makes the borrow checker forbid resetting while any + /// [`ScopedAllocator`](crate::alloc::ScopedAllocator) borrowing this arena is live. + pub(crate) fn reset(&mut self) { + self.head.set(0); + } + + fn capacity(&self) -> usize { + self.buffer.get().len() + } + + fn base(&self) -> *mut u8 { + self.buffer.get().cast::() + } +} + +// SAFETY: `allocate` returns exactly `layout.size()` bytes aligned to `layout.align()` +// within the fixed buffer, or errors. `deallocate` is a no-op — storage is reclaimed by +// `reset` or on drop. +unsafe impl AllocatorCore for ResettableArena { + fn allocate(&self, layout: std::alloc::Layout) -> Result, AllocatorError> { + let base = self.base() as usize; + let head = self.head.get(); + let cur = base.checked_add(head).ok_or(AllocatorError)?; + let aligned = cur + .checked_next_multiple_of(layout.align()) + .ok_or(AllocatorError)?; + let pad = aligned - cur; + let new_head = head + .checked_add(pad) + .and_then(|h| h.checked_add(layout.size())) + .ok_or(AllocatorError)?; + if new_head > self.capacity() { + return Err(AllocatorError); + } + self.head.set(new_head); + + // SAFETY: `head + pad <= new_head <= capacity`, so the range is in-bounds. + let ptr = unsafe { self.base().add(head + pad) }; + NonNull::new(std::ptr::slice_from_raw_parts_mut(ptr, layout.size())).ok_or(AllocatorError) + } + + unsafe fn deallocate(&self, _ptr: NonNull<[u8]>, _layout: std::alloc::Layout) {} +} + +impl std::fmt::Debug for ResettableArena { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ResettableArena") + .field("capacity", &self.capacity()) + .field("head", &self.head.get()) + .finish() + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs new file mode 100644 index 000000000..c490b4dc5 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs @@ -0,0 +1,282 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! f32 instantiation: block-transposed A, row-major B, running max. The accumulator is +//! already the score, so [`RawMax`] is a bare reduction — the degenerate [`Drain`]. + +use core::mem::size_of; + +use diskann_wide::arch::x86_64::V3; + +use super::arena::ResettableArena; +use super::views::{A_PANEL, B_PANEL, DPanel, DTail, DocWalk, QPanel, QueryWalk}; +use super::{ + Accumulate, At, Block, Drain, Plan, Short, Strip, StripRef, TileBudget, drive, leaves, +}; +use crate::alloc::{Poly, ScopedAllocator}; +use crate::multi_vector::{BlockTransposed, Mat, MatRef, Standard}; + +// ── Kernel ─────────────────────────────────────────────────────── + +pub(crate) struct F32Kernel; + +impl<'a, 'b, 'x, const R: usize, const N: usize> + Accumulate, DPanel<'b, f32, N>, Block<'x, f32, R, N>> for F32Kernel +{ + fn accumulate( + &self, + arch: V3, + a: QPanel<'a, f32, R>, + b: DPanel<'b, f32, N>, + mut out: Block<'x, f32, R, N>, + ) { + // SAFETY: `a` is an R×k block-transposed f32 block; `b` is N rows of k f32; + // `out` is N columns of R f32 at stride R. + unsafe { + leaves::f32_store_microkernel::( + arch, + a.as_ptr(), + b.as_ptr(), + a.k(), + out.as_mut_ptr(), + ); + } + } +} + +impl<'a, 'b, 'x, const R: usize, const N: usize> + Accumulate, DTail<'b, f32>, Short>> for F32Kernel +{ + fn accumulate( + &self, + arch: V3, + a: QPanel<'a, f32, R>, + b: DTail<'b, f32>, + mut out: Short>, + ) { + debug_assert_eq!(out.0.cols(), b.rows()); + let (ap, bp, op, k) = (a.as_ptr(), b.as_ptr(), out.0.as_mut_ptr(), a.k()); + // SAFETY: as the full-width impl, with a runtime width in 1..N. + unsafe { + match b.rows() { + 3 => leaves::f32_store_microkernel::<3, R>(arch, ap, bp, k, op), + 2 => leaves::f32_store_microkernel::<2, R>(arch, ap, bp, k, op), + 1 => leaves::f32_store_microkernel::<1, R>(arch, ap, bp, k, op), + other => unreachable!("tail width {other} out of 1..{N}"), + } + } + } +} + +// ── Drain ──────────────────────────────────────────────────────── + +/// Running max over an output it owns, padded to whole A-panels by the caller. +pub(crate) struct RawMax<'o> { + out: &'o mut [f32], +} + +impl<'o> RawMax<'o> { + fn new(out: &'o mut [f32]) -> Self { + out.fill(f32::MIN); + Self { out } + } +} + +impl Drain> for RawMax<'_> { + fn drain(&mut self, arch: V3, acc: StripRef<'_, f32, R>, at: At) { + let out = &mut self.out[at.a_panel * R..][..R]; + // SAFETY: `out` is R f32; `acc` is `cols` columns of R f32. + unsafe { leaves::fold_strip::(arch, out.as_mut_ptr(), acc.as_ptr(), acc.cols()) } + } +} + +// ── Public entry ───────────────────────────────────────────────── + +/// A prepared f32 query set for the paneled driver (V3/AVX2). +pub struct PaneledF32Query { + query: BlockTransposed, + dim: usize, + arch: V3, + state: Vec, + arena: ResettableArena, +} + +impl PaneledF32Query { + /// `None` if AVX2 (V3) is unavailable. + #[allow(clippy::expect_used)] + pub fn build(query: MatRef<'_, Standard>) -> Option { + let arch = V3::new_checked()?; + let dim = query.vector_dim(); + let query = BlockTransposed::::from_matrix_view(query.as_matrix_view()); + let padded = query.padded_nrows(); + + // The planner keeps `A_PANEL · acc_bytes · b_tile_rows` inside `l1_b`, so one + // page of headroom bounds the single strip for any k. + let arena = ResettableArena::with_capacity(TileBudget::default().l1_b + 4096) + .expect("arena allocation"); + + Some(Self { + query, + dim, + arch, + state: vec![f32::MIN; padded], + arena, + }) + } + + pub fn is_supported() -> bool { + V3::new_checked().is_some() + } + + pub fn num_vectors(&self) -> usize { + self.query.nrows() + } + + /// Per-query max inner product (the MaxSim similarity) against `docs`. + /// + /// # Panics + /// + /// If `scores.len() != self.num_vectors()` or the logical dims differ. + pub fn compute_max_sim(&mut self, docs: &PaneledF32Docs, scores: &mut [f32]) { + self.compute(docs, scores, TileBudget::default()); + } + + #[allow(clippy::expect_used)] + fn compute(&mut self, docs: &PaneledF32Docs, scores: &mut [f32], budget: TileBudget) { + let nq = self.query.nrows(); + assert_eq!(scores.len(), nq, "scores length must equal query count"); + assert_eq!(self.dim, docs.data.vector_dim(), "query dim != doc dim"); + + let k = self.query.padded_ncols(); + let padded = self.query.padded_nrows(); + assert_eq!(docs.data.vector_dim(), k, "doc row stride must equal k"); + + self.arena.reset(); + let row_bytes = k * size_of::(); + let plan = Plan::::new(row_bytes, row_bytes, size_of::(), budget); + let strip_len = plan.strip_len(); + let mut buf = + Poly::<[f32], _>::new_uninit_slice(strip_len, ScopedAllocator::new(&self.arena)) + .expect("strip fits the arena"); + let mut scratch = Strip::::from_uninit(&mut buf, strip_len); + drive::<_, _, _, _, Strip<'_, f32, A_PANEL, B_PANEL>, _>( + self.arch, + QueryWalk::new(self.query.as_view(), plan.a_panels), + DocWalk::::new(docs.data.as_view(), plan.b_panels), + &F32Kernel, + &mut scratch, + &mut RawMax::new(&mut self.state[..padded]), + ); + + scores.copy_from_slice(&self.state[..nq]); + } +} + +/// A prepared f32 document set, kept row-major. +pub struct PaneledF32Docs { + data: Mat>, +} + +impl PaneledF32Docs { + pub fn build(docs: MatRef<'_, Standard>) -> Self { + let mut src = docs.as_slice().iter().copied(); + Self { + data: Mat::from_fn(*docs.repr(), || src.next().unwrap_or_default()), + } + } + + pub fn num_vectors(&self) -> usize { + self.data.num_vectors() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rnd(seed: u64, idx: usize) -> f32 { + let x = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(idx as u64) + .wrapping_mul(1442695040888963407); + ((x >> 33) as f32 / (1u64 << 31) as f32) - 1.0 + } + + fn reference(q: &[f32], nq: usize, d: &[f32], nd: usize, dim: usize) -> Vec { + (0..nq) + .map(|i| { + (0..nd) + .map(|j| { + (0..dim) + .map(|c| q[i * dim + c] * d[j * dim + c]) + .sum::() + }) + .fold(f32::MIN, f32::max) + }) + .collect() + } + + /// (nq, nd, dim): every B-remainder class, an A-panel remainder (17), a + /// multi-B-tile case, and odd dims. + const CASES: &[(usize, usize, usize)] = &[ + (1, 1, 64), + (5, 3, 5), + (16, 4, 64), + (16, 5, 128), + (16, 6, 64), + (16, 7, 256), + (17, 9, 65), + (32, 16, 256), + (64, 1250, 64), + (8, 33, 127), + ]; + + #[allow(clippy::expect_used)] + fn run(nq: usize, nd: usize, dim: usize, seed: u64, budget: TileBudget) { + let q: Vec = (0..nq * dim).map(|i| rnd(seed, i)).collect(); + let d: Vec = (0..nd * dim).map(|i| rnd(seed + 1, i)).collect(); + + let q_mat = MatRef::new(Standard::::new(nq, dim).expect("nq×dim"), q.as_slice()) + .expect("query slice"); + let d_mat = MatRef::new(Standard::::new(nd, dim).expect("nd×dim"), d.as_slice()) + .expect("doc slice"); + + let mut query = PaneledF32Query::build(q_mat).expect("V3 checked by caller"); + let docs = PaneledF32Docs::build(d_mat); + let mut got = vec![0.0f32; nq]; + query.compute(&docs, &mut got, budget); + + let want = reference(&q, nq, &d, nd, dim); + for i in 0..nq { + assert!( + (got[i] - want[i]).abs() <= 1e-4 * want[i].abs().max(1.0), + "({nq},{nd},{dim}) row {i}: paneled-f32 {} != reference {}", + got[i], + want[i], + ); + } + } + + #[test] + fn paneled_f32_matches_reference() { + if V3::new_checked().is_none() { + return; + } + for &(nq, nd, dim) in CASES { + run(nq, nd, dim, 1, TileBudget::default()); + } + } + + /// A tiny budget clamps the planner to one panel per tile, forcing multiple A- and + /// B-tiles — the cross-tile offset carry the default budget never reaches. + #[test] + fn paneled_f32_multi_tile_tiny_budget() { + if V3::new_checked().is_none() { + return; + } + let budget = TileBudget { l2_a: 1, l1_b: 1 }; + for &(nq, nd, dim) in &[(48usize, 22usize, 64usize), (33, 37, 128), (35, 19, 65)] { + run(nq, nd, dim, 3, budget); + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs new file mode 100644 index 000000000..993bd33fb --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs @@ -0,0 +1,234 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! SIMD leaves. The store-out micro-kernels are byte-identical to `tiler`'s, so an A/B +//! between the two experiments measures only the abstraction; [`score_fold_strip`] is +//! the one that differs — it fuses `tiler`'s separate dequant and max passes. +//! +//! Each leaf const-asserts the A-panel width it was previously assuming against its own +//! register shape, so a caller that disagrees fails to compile. This is the only guard +//! past the raw pointers, where the type-level widths stop. + +use diskann_wide::arch::x86_64::V3; +use diskann_wide::{SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDMulAdd, SIMDReinterpret, SIMDVector}; + +use crate::minmax::MinMaxCompensation; + +diskann_wide::alias!(i16s = ::i16x16); +diskann_wide::alias!(i32s = ::i32x8); +diskann_wide::alias!(u32s = ::u32x8); +diskann_wide::alias!(f32s = ::f32x8); + +/// Integer store-out micro-kernel: `AROWS` A-rows × `UNROLL` B-rows. +/// +/// # Safety +/// +/// 1. `a_packed` points to an `AROWS × k` block-transposed `i16` block (`k` even). +/// 2. `b` points to `UNROLL` rows of `k` contiguous `u8` (`k` even). +/// 3. `partial` is valid for `UNROLL` columns of `AROWS` `i32` at stride `AROWS`. +#[inline(always)] +pub(super) unsafe fn int_store_microkernel( + arch: V3, + a_packed: *const i16, + b: *const u8, + k: usize, + partial: *mut i32, +) { + const { + assert!( + AROWS == 2 * i32s::LANES, + "leaf emits exactly two i32 registers of rows" + ); + assert!( + AROWS == i16s::LANES, + "leaf loads AROWS i16 per A column-pair half" + ); + } + let mut p0 = [i32s::default(arch); UNROLL]; + let mut p1 = [i32s::default(arch); UNROLL]; + let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); + + let a_pair_stride = 2 * AROWS; + let a_half = AROWS; + let pairs = k / 2; + + for p in 0..pairs { + // SAFETY: precondition 1 — the A block has `pairs` col-pairs of 2·AROWS i16. + let (a0, a1) = unsafe { + ( + i16s::load_simd(arch, a_packed.add(a_pair_stride * p)), + i16s::load_simd(arch, a_packed.add(a_pair_stride * p + a_half)), + ) + }; + + for j in 0..UNROLL { + // SAFETY: precondition 2 — B row j is `offsets[j]` in, `2*p+1 < k`. + let (d0, d1) = unsafe { + let base = 2 * p + offsets[j]; + ( + u32::from(b.add(base).read()), + u32::from(b.add(base + 1).read()), + ) + }; + let packed = d0 | (d1 << 16); + let bcast: i16s = u32s::splat(arch, packed).reinterpret_simd(); + p0[j] = p0[j].dot_simd(a0, bcast); + p1[j] = p1[j].dot_simd(a1, bcast); + } + } + + for j in 0..UNROLL { + // SAFETY: precondition 3 — column j occupies [j*AROWS, j*AROWS+AROWS) i32. + unsafe { + p0[j].store_simd(partial.add(j * AROWS)); + p1[j].store_simd(partial.add(j * AROWS + i32s::LANES)); + } + } +} + +/// f32 store-out micro-kernel: `AROWS` A-rows × `UNROLL` B-rows of inner product. +/// +/// # Safety +/// +/// 1. `a_packed` points to an `AROWS × k` block-transposed `f32` block (`PACK = 1`). +/// 2. `b` points to `UNROLL` rows of `k` contiguous `f32`. +/// 3. `partial` is valid for `UNROLL` columns of `AROWS` `f32` at stride `AROWS`. +#[inline(always)] +pub(super) unsafe fn f32_store_microkernel( + arch: V3, + a_packed: *const f32, + b: *const f32, + k: usize, + partial: *mut f32, +) { + const { + assert!( + AROWS == 2 * f32s::LANES, + "leaf emits exactly two f32 registers of rows" + ) + } + let mut p0 = [f32s::default(arch); UNROLL]; + let mut p1 = [f32s::default(arch); UNROLL]; + let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); + + let a_stride = AROWS; + let a_half = f32s::LANES; + + for i in 0..k { + // SAFETY: precondition 1 — the A block has `k` columns of AROWS f32. + let (a0, a1) = unsafe { + ( + f32s::load_simd(arch, a_packed.add(a_stride * i)), + f32s::load_simd(arch, a_packed.add(a_stride * i + a_half)), + ) + }; + for j in 0..UNROLL { + // SAFETY: precondition 2 — B row j is `offsets[j]` in, `i < k`. + let bj = unsafe { f32s::splat(arch, b.add(i + offsets[j]).read_unaligned()) }; + p0[j] = a0.mul_add_simd(bj, p0[j]); + p1[j] = a1.mul_add_simd(bj, p1[j]); + } + } + + for j in 0..UNROLL { + // SAFETY: precondition 3 — column j occupies [j*AROWS, j*AROWS+AROWS) f32. + unsafe { + p0[j].store_simd(partial.add(j * AROWS)); + p1[j].store_simd(partial.add(j * AROWS + f32s::LANES)); + } + } +} + +/// Fold an `R`×`cols` A-major f32 strip into the `R`-wide running max. +/// +/// # Safety +/// +/// `state` writable for `R` `f32`; `acc` valid for `cols` columns of `R` `f32`. +#[inline(always)] +pub(super) unsafe fn fold_strip( + arch: V3, + state: *mut f32, + acc: *const f32, + cols: usize, +) { + const { assert!(R == 2 * f32s::LANES, "fold reads two f32 registers of rows") } + let lanes = f32s::LANES; + // SAFETY: `state` writable for R; `acc` valid for `cols` columns of R. + unsafe { + let mut m0 = f32s::load_simd(arch, state); + let mut m1 = f32s::load_simd(arch, state.add(lanes)); + for c in 0..cols { + let col = acc.add(c * R); + m0 = m0.max_simd(f32s::load_simd(arch, col)); + m1 = m1.max_simd(f32s::load_simd(arch, col.add(lanes))); + } + m0.store_simd(state); + m1.store_simd(state.add(lanes)); + } +} + +/// 4-bit MinMax dequant of an `R`×`cols` A-major `i32` strip, folded straight into the +/// running max — the score never reaches memory. +/// +/// # Safety +/// +/// `acc` valid for `cols` columns of `R` `i32` (stride `R`); `state` writable for `R` +/// `f32`; `q_meta.len() >= R`; `d_meta.len() >= cols`. +#[inline(always)] +pub(super) unsafe fn score_fold_strip( + arch: V3, + acc: *const i32, + state: *mut f32, + cols: usize, + q_meta: &[MinMaxCompensation], + d_meta: &[MinMaxCompensation], + dim: f32, +) { + const { assert!(R == 2 * f32s::LANES, "fold reads two f32 registers of rows") } + let lanes = f32s::LANES; + + let mut qa = [0.0f32; R]; + let mut qb = [0.0f32; R]; + let mut qn = [0.0f32; R]; + for i in 0..R { + let qm = q_meta[i]; + qa[i] = qm.a; + qb[i] = qm.b; + qn[i] = qm.n; + } + // SAFETY: each array holds exactly R = 2·LANES f32; `state` writable for R. + let (qa0, qa1, qb0, qb1, qn0, qn1, mut m0, mut m1) = unsafe { + ( + f32s::load_simd(arch, qa.as_ptr()), + f32s::load_simd(arch, qa.as_ptr().add(lanes)), + f32s::load_simd(arch, qb.as_ptr()), + f32s::load_simd(arch, qb.as_ptr().add(lanes)), + f32s::load_simd(arch, qn.as_ptr()), + f32s::load_simd(arch, qn.as_ptr().add(lanes)), + f32s::load_simd(arch, state), + f32s::load_simd(arch, state.add(lanes)), + ) + }; + + for (c, dm) in d_meta.iter().enumerate().take(cols) { + let a_c = f32s::splat(arch, dm.a); + let b_c = f32s::splat(arch, dm.b); + let c_c = f32s::splat(arch, dm.n + dm.b * dim); + let col = c * R; + // SAFETY: `col + 2·LANES <= cols*R`; `acc` valid for that many i32. + unsafe { + let raw0 = i32s::load_simd(arch, acc.add(col)).simd_cast(); + let raw1 = i32s::load_simd(arch, acc.add(col + lanes)).simd_cast(); + let s0 = a_c.mul_add_simd(qa0 * raw0, b_c.mul_add_simd(qn0, c_c * qb0)); + let s1 = a_c.mul_add_simd(qa1 * raw1, b_c.mul_add_simd(qn1, c_c * qb1)); + m0 = m0.max_simd(s0); + m1 = m1.max_simd(s1); + } + } + + // SAFETY: `state` writable for R f32. + unsafe { + m0.store_simd(state); + m1.store_simd(state.add(lanes)); + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs new file mode 100644 index 000000000..0b1d16ba2 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs @@ -0,0 +1,414 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! 4-bit MinMax instantiation. The interesting half is [`MinMaxMax`]: dequant needs +//! per-vector metadata indexed by [`At`] and the reduction needs the dequantized +//! score, so both ride in one [`Drain`] and the score never reaches memory. + +use core::mem::size_of; +use std::num::NonZeroUsize; + +use diskann_utils::ReborrowMut; +use diskann_wide::arch::x86_64::V3; + +use super::arena::ResettableArena; +use super::views::{A_PANEL, B_PANEL, DPanel, DTail, DocWalk, QPanel, QueryWalk}; +use super::{ + Accumulate, At, Block, Drain, Plan, Short, Strip, StripRef, TileBudget, drive, leaves, +}; +use crate::CompressInto; +use crate::algorithms::Transform; +use crate::algorithms::transforms::NullTransform; +use crate::alloc::{Poly, ScopedAllocator}; +use crate::minmax::{MinMaxCompensation, MinMaxMeta, MinMaxQuantizer}; +use crate::multi_vector::{BlockTransposed, Defaulted, Mat, MatRef, Standard}; +use crate::num::Positive; + +// ── Kernel ─────────────────────────────────────────────────────── + +pub(crate) struct I8Kernel; + +impl<'a, 'b, 'x, const R: usize, const N: usize> + Accumulate, DPanel<'b, u8, N>, Block<'x, i32, R, N>> for I8Kernel +{ + fn accumulate( + &self, + arch: V3, + a: QPanel<'a, i16, R>, + b: DPanel<'b, u8, N>, + mut out: Block<'x, i32, R, N>, + ) { + // SAFETY: `a` is an R×k block-transposed i16 block; `b` is N rows of k u8; + // `out` is N columns of R i32 at stride R (`k` even). + unsafe { + leaves::int_store_microkernel::( + arch, + a.as_ptr(), + b.as_ptr(), + a.k(), + out.as_mut_ptr(), + ); + } + } +} + +impl<'a, 'b, 'x, const R: usize, const N: usize> + Accumulate, DTail<'b, u8>, Short>> for I8Kernel +{ + fn accumulate( + &self, + arch: V3, + a: QPanel<'a, i16, R>, + b: DTail<'b, u8>, + mut out: Short>, + ) { + debug_assert_eq!(out.0.cols(), b.rows()); + let (ap, bp, op, k) = (a.as_ptr(), b.as_ptr(), out.0.as_mut_ptr(), a.k()); + // SAFETY: as the full-width impl, with a runtime width in 1..N. + unsafe { + match b.rows() { + 3 => leaves::int_store_microkernel::<3, R>(arch, ap, bp, k, op), + 2 => leaves::int_store_microkernel::<2, R>(arch, ap, bp, k, op), + 1 => leaves::int_store_microkernel::<1, R>(arch, ap, bp, k, op), + other => unreachable!("tail width {other} out of 1..{N}"), + } + } + } +} + +// ── Drain ──────────────────────────────────────────────────────── + +/// Fused 4-bit MinMax dequant + running max: rewrites each raw integer dot into the +/// MinMax inner product using per-vector `a`/`b`/`n` metadata, then folds it straight +/// into the output it owns. +pub(crate) struct MinMaxMax<'m, 'o> { + query_meta: &'m [MinMaxCompensation], + doc_meta: &'m [MinMaxCompensation], + out: &'o mut [f32], + dim: f32, +} + +impl<'m, 'o> MinMaxMax<'m, 'o> { + fn new( + query_meta: &'m [MinMaxCompensation], + doc_meta: &'m [MinMaxCompensation], + out: &'o mut [f32], + dim: f32, + ) -> Self { + out.fill(f32::MIN); + Self { + query_meta, + doc_meta, + out, + dim, + } + } +} + +impl Drain> for MinMaxMax<'_, '_> { + fn drain(&mut self, arch: V3, acc: StripRef<'_, i32, R>, at: At) { + let cols = acc.cols(); + let lo = at.a_panel * R; + let q = &self.query_meta[lo..lo + R]; + let d = &self.doc_meta[at.b_row..at.b_row + cols]; + let dim = self.dim; + let out = &mut self.out[lo..][..R]; + // SAFETY: `acc` is `cols` columns of R i32; `out` is R writable f32; + // `q.len() == R`; `d.len() == cols`. + unsafe { + leaves::score_fold_strip::(arch, acc.as_ptr(), out.as_mut_ptr(), cols, q, d, dim); + } + } +} + +// ── Public entry ───────────────────────────────────────────────── + +/// Quantize an f32 multi-vector to 4-bit MinMax (Null transform, scale 1.0). +#[allow(clippy::expect_used)] +fn quantize(input: MatRef<'_, Standard>) -> Mat> { + let (n, dim) = (input.num_vectors(), input.vector_dim()); + let q = MinMaxQuantizer::new( + Transform::Null(NullTransform::new( + NonZeroUsize::new(dim).expect("dimension must be non-zero"), + )), + Positive::new(1.0).expect("1.0 is positive"), + ); + let mut out: Mat> = + Mat::new(MinMaxMeta::new(n, dim), Defaulted).expect("MinMaxMeta allocation"); + q.compress_into(input, out.reborrow_mut()) + .expect("input must be finite"); + out +} + +/// A prepared 4-bit MinMax query set for the paneled driver (V3/AVX2). +pub struct PaneledQuantQuery { + query: BlockTransposed, + meta: Vec, + dim: usize, + arch: V3, + state: Vec, + arena: ResettableArena, +} + +impl PaneledQuantQuery { + /// `None` if AVX2 (V3) is unavailable. + #[allow(clippy::expect_used)] + pub fn build(query: MatRef<'_, Standard>) -> Option { + let arch = V3::new_checked()?; + let (nq, dim) = (query.num_vectors(), query.vector_dim()); + let q_mat = quantize(query); + + let mut codes = vec![0i16; nq * dim]; + for r in 0..nq { + let row = q_mat.get_row(r).expect("row < nq"); + for j in 0..dim { + codes[r * dim + j] = i16::from(row.vector().get(j).expect("col < dim") as u8); + } + } + let view = MatRef::new(Standard::::new(nq, dim).expect("nq×dim"), &codes) + .expect("code slice"); + let query = BlockTransposed::::from_matrix_view(view.as_matrix_view()); + + let padded = query.padded_nrows(); + let mut meta = vec![MinMaxCompensation::default(); padded]; + for (r, m) in meta.iter_mut().enumerate().take(nq) { + *m = q_mat.get_row(r).expect("row < nq").meta(); + } + + // The planner keeps `A_PANEL · acc_bytes · b_tile_rows` inside `l1_b`, so one + // page of headroom bounds the single strip for any k. + let arena = ResettableArena::with_capacity(TileBudget::default().l1_b + 4096) + .expect("arena allocation"); + + Some(Self { + query, + meta, + dim, + arch, + state: vec![f32::MIN; padded], + arena, + }) + } + + pub fn is_supported() -> bool { + V3::new_checked().is_some() + } + + pub fn num_vectors(&self) -> usize { + self.query.nrows() + } + + /// Per-query min distance (`= -max_d IP`) against `docs`. + /// + /// # Panics + /// + /// If `scores.len() != self.num_vectors()` or the logical dims differ. + pub fn compute_max_sim(&mut self, docs: &PaneledQuantDocs, scores: &mut [f32]) { + self.compute(docs, scores, TileBudget::default()); + } + + #[allow(clippy::expect_used)] + fn compute(&mut self, docs: &PaneledQuantDocs, scores: &mut [f32], budget: TileBudget) { + let nq = self.query.nrows(); + assert_eq!(scores.len(), nq, "scores length must equal query count"); + assert_eq!(self.dim, docs.dim, "query dim != doc dim"); + + let k = self.query.padded_ncols(); + let padded = self.query.padded_nrows(); + assert_eq!(docs.codes.vector_dim(), k, "doc row stride must equal k"); + + self.arena.reset(); + let plan = Plan::::new( + k * size_of::(), + k * size_of::(), + size_of::(), + budget, + ); + let mut drain = MinMaxMax::new( + &self.meta, + &docs.meta, + &mut self.state[..padded], + self.dim as f32, + ); + let strip_len = plan.strip_len(); + let mut buf = + Poly::<[i32], _>::new_uninit_slice(strip_len, ScopedAllocator::new(&self.arena)) + .expect("strip fits the arena"); + let mut scratch = Strip::::from_uninit(&mut buf, strip_len); + drive::<_, _, _, _, Strip<'_, i32, A_PANEL, B_PANEL>, _>( + self.arch, + QueryWalk::new(self.query.as_view(), plan.a_panels), + DocWalk::::new(docs.codes.as_view(), plan.b_panels), + &I8Kernel, + &mut scratch, + &mut drain, + ); + + for (s, &raw) in scores.iter_mut().zip(self.state.iter()) { + *s = -raw; + } + } +} + +/// A prepared 4-bit MinMax document set. Codes are row-major with rows padded to an +/// even length, which the integer microkernel requires; metadata is kept alongside. +pub struct PaneledQuantDocs { + codes: Mat>, + meta: Vec, + dim: usize, +} + +impl PaneledQuantDocs { + #[allow(clippy::expect_used)] + pub fn build(docs: MatRef<'_, Standard>) -> Self { + let (nv, dim) = (docs.num_vectors(), docs.vector_dim()); + let repr = Standard::::new(nv, dim.next_multiple_of(2)).expect("codes fit in memory"); + let d_mat = quantize(docs); + + let mut codes = Mat::from_fn(repr, || 0u8); + let mut meta = Vec::with_capacity(nv); + for r in 0..nv { + let row = d_mat.get_row(r).expect("row < nv"); + let dst = codes.get_row_mut(r).expect("row < nv"); + for (j, d) in dst.iter_mut().take(dim).enumerate() { + *d = row.vector().get(j).expect("col < dim") as u8; + } + meta.push(row.meta()); + } + Self { codes, meta, dim } + } + + pub fn num_vectors(&self) -> usize { + self.codes.num_vectors() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::multi_vector::distance::{MaxSim, QueryMatRef}; + use diskann_vector::DistanceFunctionMut; + + fn rnd(seed: u64, idx: usize) -> f32 { + let x = seed + .wrapping_mul(6364136223846793005) + .wrapping_add(idx as u64) + .wrapping_mul(1442695040888963407); + ((x >> 33) as f32 / (1u64 << 31) as f32) - 1.0 + } + + #[allow(clippy::expect_used)] + fn reference(q: &[f32], nq: usize, d: &[f32], nd: usize, dim: usize) -> Vec { + let quantize = |data: &[f32], n: usize| -> Mat> { + let input = + MatRef::new(Standard::::new(n, dim).expect("n×dim"), data).expect("slice"); + super::quantize(input) + }; + let q_mat = quantize(q, nq); + let d_mat = quantize(d, nd); + let query: QueryMatRef<_> = q_mat.as_view().into(); + let mut out = vec![0.0f32; nq]; + MaxSim::new(&mut out).evaluate(query, d_mat.as_view()); + out + } + + /// (nq, nd, dim): every B-remainder class, an A-panel remainder (17), a + /// multi-B-tile case, and the odd-dim even-K contract. + const CASES: &[(usize, usize, usize)] = &[ + (1, 1, 64), + (5, 1, 128), + (16, 4, 64), + (16, 5, 128), + (16, 6, 64), + (16, 7, 256), + (17, 9, 64), + (32, 16, 256), + (64, 1250, 64), + (5, 3, 63), + (17, 9, 65), + (8, 33, 127), + ]; + + #[allow(clippy::expect_used)] + fn run(nq: usize, nd: usize, dim: usize, seed: u64, budget: TileBudget) { + let q: Vec = (0..nq * dim).map(|i| rnd(seed, i)).collect(); + let d: Vec = (0..nd * dim).map(|i| rnd(seed + 1, i)).collect(); + + let q_mat = MatRef::new(Standard::::new(nq, dim).expect("nq×dim"), q.as_slice()) + .expect("query slice"); + let d_mat = MatRef::new(Standard::::new(nd, dim).expect("nd×dim"), d.as_slice()) + .expect("doc slice"); + + let mut query = PaneledQuantQuery::build(q_mat).expect("V3 checked by caller"); + let docs = PaneledQuantDocs::build(d_mat); + let mut got = vec![0.0f32; nq]; + query.compute(&docs, &mut got, budget); + + let want = reference(&q, nq, &d, nd, dim); + for i in 0..nq { + assert!( + (got[i] - want[i]).abs() <= 1e-4 * want[i].abs().max(1.0), + "({nq},{nd},{dim}) row {i}: paneled-i8 {} != reference {}", + got[i], + want[i], + ); + } + } + + #[test] + fn paneled_i8_matches_minmax_reference() { + if V3::new_checked().is_none() { + return; + } + for &(nq, nd, dim) in CASES { + run(nq, nd, dim, 1, TileBudget::default()); + } + } + + /// A tiny budget clamps the planner to one panel per tile, forcing multiple A- and + /// B-tiles — the cross-tile metadata offsets the default budget never reaches. + #[test] + fn paneled_i8_multi_tile_tiny_budget() { + if V3::new_checked().is_none() { + return; + } + let budget = TileBudget { l2_a: 1, l1_b: 1 }; + for &(nq, nd, dim) in &[(48usize, 22usize, 64usize), (33, 37, 128), (35, 19, 65)] { + run(nq, nd, dim, 3, budget); + } + } + + /// Arena reuse across differently-sized doc sets stays correct. + #[test] + #[allow(clippy::expect_used)] + fn paneled_i8_arena_reuse() { + if V3::new_checked().is_none() { + return; + } + const NQ: usize = 17; + const DIM: usize = 128; + let q: Vec = (0..NQ * DIM).map(|i| rnd(5, i)).collect(); + let q_mat = MatRef::new(Standard::::new(NQ, DIM).expect("nq×dim"), q.as_slice()) + .expect("query slice"); + let mut query = PaneledQuantQuery::build(q_mat).expect("V3 checked by caller"); + + for (call, &nd) in [251usize, 3, 64, 1].iter().enumerate() { + let d: Vec = (0..nd * DIM).map(|i| rnd(6 + call as u64, i)).collect(); + let d_mat = MatRef::new(Standard::::new(nd, DIM).expect("nd×dim"), d.as_slice()) + .expect("doc slice"); + let docs = PaneledQuantDocs::build(d_mat); + let mut got = vec![0.0f32; NQ]; + query.compute_max_sim(&docs, &mut got); + + let want = reference(&q, NQ, &d, nd, DIM); + for i in 0..NQ { + assert!( + (got[i] - want[i]).abs() <= 1e-4 * want[i].abs().max(1.0), + "call {call} (nd={nd}) row {i}: paneled-i8 {} != reference {}", + got[i], + want[i], + ); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs new file mode 100644 index 000000000..f91412a6b --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs @@ -0,0 +1,245 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Paneled MaxSim: a [`TileWalk`] lends cache-sized views, each [`Paneled`] into +//! register-sized panels plus a typed tail. [`Accumulate`] folds one (A-panel, +//! B-panel) pair into an accumulator slot; [`Drain`] turns the finished accumulator +//! into output it owns. [`Scratch`] is the write-side mirror of [`Paneled`], so the +//! driver assumes no layout on either side. +//! +//! Panel widths are const parameters (`R` = A rows, `N` = B rows) carried by the panel +//! types; a kernel relays them into its leaf rather than declaring them. Literals in +//! [`views`]. Instantiated for f32 ([`float`]) and 4-bit MinMax ([`minmax`]). +//! +//! Sibling to [`tiler`](super::tiler), which keeps postprocess and reduce separate. + +use super::TileBudget; + +mod arena; +mod float; +mod leaves; +mod minmax; +mod strip; +mod views; + +pub(crate) use strip::{Block, Short, Strip, StripRef}; + +pub use float::{PaneledF32Docs, PaneledF32Query}; +pub use minmax::{PaneledQuantDocs, PaneledQuantQuery}; + +// ── Tile planning ──────────────────────────────────────────────── + +/// Panel counts per tile. `a_panels` A-panels sit resident in L2; as many B-panels as +/// co-fit L1 alongside one A-panel and the accumulator. +#[derive(Clone, Copy)] +struct Plan { + a_panels: usize, + b_panels: usize, +} + +impl Plan { + fn new(a_row_bytes: usize, b_row_bytes: usize, acc_bytes: usize, budget: TileBudget) -> Self { + let a_row_bytes = a_row_bytes.max(1); + let b_row_bytes = b_row_bytes.max(1); + let a_panels = (budget.l2_a / (a_row_bytes * R)).max(1); + let a_panel_bytes = R * a_row_bytes; + let per_b_row = b_row_bytes + R * acc_bytes; + let b_budget = budget.l1_b.saturating_sub(a_panel_bytes); + let b_panels = ((b_budget / per_b_row) / N).max(1); + Self { a_panels, b_panels } + } + + /// Accumulator elements: one A-panel × a whole B-tile. + fn strip_len(&self) -> usize { + R * self.b_panels * N + } +} + +// ── Accumulator ────────────────────────────────────────────────── + +/// Per-lifetime half of [`Scratch`] — same sealed-`Bounds` trick as [`TileAt`]. +pub(crate) trait ScratchAt<'a, B: sealed::Sealed = sealed::Bounds<&'a mut Self>> { + type Block; + /// Short trailing slot; a distinct type so it selects its own [`Accumulate`] impl. + type Short; + /// What [`Drain`] reads. + type Ref; +} + +/// [`Paneled`]'s write side: a buffer that carves itself into per-B-panel slots, so +/// the driver can't tell a contiguous strip from a padded or structure-of-arrays one. +/// Not [`Paneled`] itself — `Panel: Copy` and `panels(&self)` can't yield disjoint +/// `&mut`, and the carve needs a runtime `cols` (a scratch is sized to capacity). +/// Allocation stays on the concrete type. +pub(crate) trait Scratch: for<'a> ScratchAt<'a> { + /// Carve the live `cols` columns into one slot per B-panel plus the short trailer, + /// which must come from the same call to be provably disjoint. + fn split( + &mut self, + cols: usize, + ) -> ( + impl Iterator>::Block>, + Option<>::Short>, + ); + + fn as_ref(&self, cols: usize) -> >::Ref; +} + +// ── Data side ──────────────────────────────────────────────────── + +/// Misuse guard for the implicit-bounds parameter: private, so no downstream impl can +/// override the default with a type that drops the implied bound. +mod sealed { + pub trait Sealed {} + pub struct Bounds(#[allow(dead_code)] T); + impl Sealed for Bounds {} +} + +/// Per-lifetime half of [`TileWalk`]. The defaulted `B = Bounds<&'a Self>` carries the +/// `Self: 'a` implied bound through well-formedness — a plain GAT `where Self: 'a` +/// collapses to `'static` under the driver's `for<'a>` bound on stable. +pub(crate) trait TileAt<'a, B: sealed::Sealed = sealed::Bounds<&'a Self>> { + type View: Paneled; +} + +/// A **lending** walk: `next` reborrows `&mut self`, so a view may borrow a buffer the +/// walk reuses on the following call. `reset` rewinds — B is re-walked once per A-tile. +pub(crate) trait TileWalk: for<'a> TileAt<'a> { + fn next(&mut self) -> Option>::View>>; + fn reset(&mut self); +} + +/// A lent view plus where it starts. Lifetime-free — the borrow lives on `V`. +pub(crate) struct Tile { + pub(crate) view: V, + /// Position in the walk's own unit: A-panels for a query walk, B-rows for a doc + /// walk. Only the [`Drain`] turns it into an output index. + pub(crate) at: usize, +} + +/// A view that knows how it breaks into panels. `Tail` is distinct from `Panel` so the +/// short trailing panel selects its own [`Accumulate`] impl; a view that cannot tail +/// says [`NoTail`]. +pub(crate) trait Paneled { + type Panel: Copy; + type Tail: Copy; + + fn rows(&self) -> usize; + fn panels(&self) -> impl Iterator + '_; + fn tail(&self) -> Option; +} + +/// `Tail` for a view padded to whole panels. Uninhabited, so `tail()` provably +/// returns `None`. +#[derive(Clone, Copy)] +pub(crate) enum NoTail {} + +// ── Compute side ───────────────────────────────────────────────── + +/// One A-panel × one B-panel → an accumulator slot. Pinned on all three as type +/// parameters, so the walks' panel types select the impl. +pub(crate) trait Accumulate { + fn accumulate(&self, arch: Arch, a: A, b: B, out: O); +} + +/// [`NoTail`] is uninhabited, so this one impl discharges the driver's A-tail bounds +/// for every kernel. +impl Accumulate for K { + fn accumulate(&self, _: Arch, a: NoTail, _: B, _: O) { + match a {} + } +} + +/// Where a finished accumulator sits in the global problem. The driver counts panels +/// and never converts a count into a row — only the [`Drain`] owns `R`. B is a row +/// rather than a panel index because a tile is not panel-quantized. +#[derive(Clone, Copy)] +pub(crate) struct At { + pub a_panel: usize, + pub b_row: usize, +} + +/// Consume a finished accumulator. The drain owns its output, so dequant, reduction +/// and scatter all live behind this one call and may be fused. +/// +/// Implementations must initialize their output to the reduction's identity, and must +/// clamp their writes when the output is not padded to whole panels. +pub(crate) trait Drain { + fn drain(&mut self, arch: Arch, acc: >::Ref, at: At); +} + +// ── Driver ─────────────────────────────────────────────────────── + +type PanelOf<'a, W> = <>::View as Paneled>::Panel; +type TailOf<'a, W> = <>::View as Paneled>::Tail; +type BlockOf<'x, S> = >::Block; +type ShortOf<'x, S> = >::Short; + +/// One A-panel against a whole B-tile. Factored out so the driver's A-panel and +/// A-tail arms share the tail-dispatch. +#[inline(always)] +fn fill(arch: Arch, kernel: &K, a: A, b_view: &BV, scratch: &mut S, cols: usize) +where + Arch: Copy, + A: Copy, + BV: Paneled, + S: Scratch, + K: for<'x> Accumulate> + + for<'x> Accumulate>, +{ + let (blocks, acc_tail) = scratch.split(cols); + for (b, out) in b_view.panels().zip(blocks) { + kernel.accumulate(arch, a, b, out); + } + match (b_view.tail(), acc_tail) { + (Some(b), Some(out)) => kernel.accumulate(arch, a, b, out), + (None, None) => {} + _ => unreachable!("B view and accumulator disagree on tail"), + } +} + +/// Drive one A source against one B source. The walks carry the plan, `scratch` the +/// accumulator, `drain` the output — so this does no stride arithmetic and knows +/// nothing about where results go. B is re-walked once per A-tile. +/// +/// `S` is not inferable (the `for<'x>` bounds project through it); call sites +/// turbofish it. +pub(super) fn drive( + arch: Arch, + mut a_walk: AW, + mut b_walk: BW, + kernel: &K, + scratch: &mut S, + drain: &mut D, +) where + Arch: Copy, + AW: TileWalk, + BW: TileWalk, + S: Scratch, + K: for<'a, 'b, 'x> Accumulate, PanelOf<'b, BW>, BlockOf<'x, S>> + + for<'a, 'b, 'x> Accumulate, TailOf<'b, BW>, ShortOf<'x, S>> + + for<'a, 'b, 'x> Accumulate, PanelOf<'b, BW>, BlockOf<'x, S>> + + for<'a, 'b, 'x> Accumulate, TailOf<'b, BW>, ShortOf<'x, S>>, + D: Drain, +{ + while let Some(a_tile) = a_walk.next() { + b_walk.reset(); + while let Some(b_tile) = b_walk.next() { + let cols = b_tile.view.rows(); + let mut at = At { + a_panel: a_tile.at, + b_row: b_tile.at, + }; + + for a in a_tile.view.panels() { + fill(arch, kernel, a, &b_tile.view, scratch, cols); + drain.drain(arch, scratch.as_ref(cols), at); + at.a_panel += 1; + } + if let Some(a) = a_tile.view.tail() { + fill(arch, kernel, a, &b_tile.view, scratch, cols); + drain.drain(arch, scratch.as_ref(cols), at); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs new file mode 100644 index 000000000..6ccba69b1 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs @@ -0,0 +1,100 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! The default [`Scratch`]: one A-panel × one B-tile, `R` rows A-major, carved into +//! `N`-column [`Block`]s. Borrows its buffer, so the allocator stays at the call site. + +use core::mem::MaybeUninit; + +use super::{Scratch, ScratchAt}; +use crate::alloc::{AllocatorCore, Poly}; + +/// Marker for element types where all-zero is a valid value. +pub(crate) trait ZeroInit: Copy {} +impl ZeroInit for i32 {} +impl ZeroInit for f32 {} + +pub(crate) struct Strip<'a, T, const R: usize, const N: usize>(&'a mut [T]); + +/// One B-panel's slot: `R` rows by `N` columns (fewer inside a [`Short`]). +pub(crate) struct Block<'a, T, const R: usize, const N: usize>(&'a mut [T]); + +/// A finished [`Strip`], narrowed to its live columns. +pub(crate) struct StripRef<'a, T, const R: usize>(&'a [T]); + +/// A short trailing slot: same payload, different type, so the runtime-width path is +/// a separate [`Accumulate`](super::Accumulate) impl. +#[derive(Clone, Copy)] +pub(crate) struct Short

(pub(crate) P); + +impl<'a, T: ZeroInit, const R: usize, const N: usize> Strip<'a, T, R, N> { + /// Zeroing is what makes the `&mut [T]` sound; the kernels overwrite every live + /// column before it is read. + pub(crate) fn from_uninit( + poly: &'a mut Poly<[MaybeUninit], A>, + len: usize, + ) -> Self { + let ptr = poly.as_mut_ptr().cast::(); + // SAFETY: the poly owns `len` `T`-sized slots; `T: ZeroInit` ⇒ all-zero is a + // valid `T`, so zeroing initializes every element and the slice is sound. + Self(unsafe { + core::ptr::write_bytes(ptr, 0, len); + core::slice::from_raw_parts_mut(ptr, len) + }) + } +} + +impl Strip<'_, T, R, N> { + fn cols_capacity(&self) -> usize { + self.0.len() / R + } +} + +impl<'a, T, const R: usize, const N: usize> ScratchAt<'a> for Strip<'_, T, R, N> { + type Block = Block<'a, T, R, N>; + type Short = Short>; + type Ref = StripRef<'a, T, R>; +} + +impl Scratch for Strip<'_, T, R, N> { + fn split( + &mut self, + cols: usize, + ) -> ( + impl Iterator>, + Option>>, + ) { + debug_assert!( + cols <= self.cols_capacity(), + "strip must hold the whole B-tile" + ); + let (live, _) = self.0.split_at_mut(R * cols); + let (head, rest) = live.split_at_mut(R * (cols - cols % N)); + ( + head.chunks_mut(R * N).map(Block), + (!rest.is_empty()).then_some(Short(Block(rest))), + ) + } + + fn as_ref(&self, cols: usize) -> StripRef<'_, T, R> { + StripRef(&self.0[..R * cols]) + } +} + +impl Block<'_, T, R, N> { + pub(crate) fn cols(&self) -> usize { + self.0.len() / R + } + pub(crate) fn as_mut_ptr(&mut self) -> *mut T { + self.0.as_mut_ptr() + } +} + +impl StripRef<'_, T, R> { + pub(crate) fn cols(&self) -> usize { + self.0.len() / R + } + pub(crate) fn as_ptr(&self) -> *const T { + self.0.as_ptr() + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs new file mode 100644 index 000000000..a5f9e0aaa --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs @@ -0,0 +1,197 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! The two views and the walks that lend them. Both views are the real matrix types — +//! [`BlockTransposedRef`] for A, [`MatRef`] for B behind a [`Rows`] adapter (a +//! row-major matrix doesn't imply a panel height). Each sub-views itself, so a walk is +//! a cursor and nothing else. + +use super::{NoTail, Paneled, Tile, TileAt, TileWalk}; +use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; + +/// Rows per A-panel (the block-transposed group and the drain's state block). +pub(crate) const A_PANEL: usize = 16; +/// Rows per full B-panel (the kernel's micro-panel / max unroll). +pub(crate) const B_PANEL: usize = 4; + +// ── Panels ─────────────────────────────────────────────────────── + +/// One block-transposed A block: `R` rows × `k` `T`. +pub(crate) struct QPanel<'a, T, const R: usize> { + data: &'a [T], + k: usize, +} +/// One row-major B panel: exactly `N` rows × `k` `T`. `k` travels with the A panel, +/// and the row count is `N` by construction, so neither is stored. +pub(crate) struct DPanel<'a, T, const N: usize> { + data: &'a [T], +} +/// The short trailing B panel: `rows` in `1..N`. Distinct from [`DPanel`] so it +/// selects its own [`Accumulate`](super::Accumulate) impl, and so a full panel carries +/// no runtime row count. +pub(crate) struct DTail<'a, T> { + data: &'a [T], + rows: usize, +} + +impl Clone for QPanel<'_, T, R> { + fn clone(&self) -> Self { + *self + } +} +impl Copy for QPanel<'_, T, R> {} +impl Clone for DPanel<'_, T, N> { + fn clone(&self) -> Self { + *self + } +} +impl Copy for DPanel<'_, T, N> {} +impl Clone for DTail<'_, T> { + fn clone(&self) -> Self { + *self + } +} +impl Copy for DTail<'_, T> {} + +impl QPanel<'_, T, R> { + pub(crate) fn as_ptr(&self) -> *const T { + self.data.as_ptr() + } + pub(crate) fn k(&self) -> usize { + self.k + } +} +impl DPanel<'_, T, N> { + pub(crate) fn as_ptr(&self) -> *const T { + self.data.as_ptr() + } +} +impl DTail<'_, T> { + pub(crate) fn as_ptr(&self) -> *const T { + self.data.as_ptr() + } + pub(crate) fn rows(&self) -> usize { + self.rows + } +} + +// ── Views ──────────────────────────────────────────────────────── + +/// A block-transposed matrix' blocks *are* the A-panels, so the real type is the view. +/// The remainder block is zero-padded to a full `R` rows, hence [`NoTail`]; rows past +/// `nrows()` score against padding and are dropped by the caller. `P` only widens +/// `padded_ncols`, the contraction length the kernel sees. +impl<'a, T: Copy, const R: usize, const P: usize> Paneled for BlockTransposedRef<'a, T, R, P> { + type Panel = QPanel<'a, T, R>; + type Tail = NoTail; + + fn rows(&self) -> usize { + self.nrows() + } + fn panels(&self) -> impl Iterator> + '_ { + let (v, k) = (*self, self.padded_ncols()); + (0..self.num_blocks()).filter_map(move |b| { + Some(QPanel { + data: v.block_slice(b)?, + k, + }) + }) + } + fn tail(&self) -> Option { + None + } +} + +/// Cut a row-major matrix into `N`-row panels. All the geometry stays on the matrix. +pub(crate) struct Rows(pub(crate) V); + +impl<'a, const N: usize, T: Copy> Paneled for Rows>> { + type Panel = DPanel<'a, T, N>; + type Tail = DTail<'a, T>; + + fn rows(&self) -> usize { + self.0.num_vectors() + } + fn panels(&self) -> impl Iterator> + '_ { + let (data, k) = (self.0.as_slice(), self.0.vector_dim()); + (0..self.rows() / N).map(move |p| DPanel { + data: &data[p * N * k..(p + 1) * N * k], + }) + } + fn tail(&self) -> Option> { + let (data, k, n) = (self.0.as_slice(), self.0.vector_dim(), self.rows()); + let rem = n % N; + (rem > 0).then(|| DTail { + data: &data[(n - rem) * k..], + rows: rem, + }) + } +} + +// ── Walks ──────────────────────────────────────────────────────── + +/// A cursor over A-blocks, lending `tile_panels` of them at a time. +pub(crate) struct QueryWalk<'s, T: Copy, const R: usize, const P: usize = 1> { + src: BlockTransposedRef<'s, T, R, P>, + tile_panels: usize, + cur: usize, +} + +/// A cursor over B-rows, lending `tile_panels` `N`-row panels' worth at a time. +pub(crate) struct DocWalk<'s, T: Copy, const N: usize> { + src: MatRef<'s, Standard>, + tile_panels: usize, + cur: usize, +} + +impl<'s, T: Copy, const R: usize, const P: usize> QueryWalk<'s, T, R, P> { + pub(crate) fn new(src: BlockTransposedRef<'s, T, R, P>, tile_panels: usize) -> Self { + Self { + src, + tile_panels, + cur: 0, + } + } +} +impl<'s, T: Copy, const N: usize> DocWalk<'s, T, N> { + pub(crate) fn new(src: MatRef<'s, Standard>, tile_panels: usize) -> Self { + Self { + src, + tile_panels, + cur: 0, + } + } +} + +impl<'a, T: Copy, const R: usize, const P: usize> TileAt<'a> for QueryWalk<'_, T, R, P> { + type View = BlockTransposedRef<'a, T, R, P>; +} +impl TileWalk for QueryWalk<'_, T, R, P> { + fn next(&mut self) -> Option>> { + let view = self.src.block_range(self.cur, self.tile_panels)?; + let at = self.cur; + self.cur += view.num_blocks(); + Some(Tile { view, at }) + } + fn reset(&mut self) { + self.cur = 0; + } +} + +impl<'a, T: Copy, const N: usize> TileAt<'a> for DocWalk<'_, T, N> { + type View = Rows>>; +} +impl TileWalk for DocWalk<'_, T, N> { + fn next(&mut self) -> Option>>>> { + let view = self.src.row_range(self.cur, self.tile_panels * N)?; + let at = self.cur; + self.cur += view.num_vectors(); + Some(Tile { + view: Rows(view), + at, + }) + } + fn reset(&mut self) { + self.cur = 0; + } +} diff --git a/diskann-quantization/src/multi_vector/distance/mod.rs b/diskann-quantization/src/multi_vector/distance/mod.rs index 74fef8f5f..8ed4df3f6 100644 --- a/diskann-quantization/src/multi_vector/distance/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/mod.rs @@ -72,3 +72,9 @@ pub use kernels::{QuantTiledDocs, QuantTiledQuery}; /// postprocess reuse the same tiled pipeline. #[cfg(target_arch = "x86_64")] pub use kernels::{QuantTiledF16Docs, QuantTiledF16Query}; + +/// Standalone POC entries for the **paneled** rebuild (V3/AVX2 only) — views own +/// their panel decomposition (B through the real `MatRef`), the tail is a distinct +/// panel type, and one `Drain` seam replaces the tiler's postprocess + reduce pair. +#[cfg(target_arch = "x86_64")] +pub use kernels::{PaneledF32Docs, PaneledF32Query, PaneledQuantDocs, PaneledQuantQuery}; diff --git a/diskann-quantization/src/multi_vector/matrix.rs b/diskann-quantization/src/multi_vector/matrix.rs index 9a3c02ea1..944ba71a9 100644 --- a/diskann-quantization/src/multi_vector/matrix.rs +++ b/diskann-quantization/src/multi_vector/matrix.rs @@ -865,6 +865,23 @@ impl<'a, T: Copy> MatRef<'a, Standard> { unsafe { std::slice::from_raw_parts(self.ptr.as_ptr().cast::(), len) } } + /// A view over `count` consecutive rows starting at `start`. + /// + /// `count` is clipped to the rows available, so a walk can request a fixed tile + /// size and get a short final tile. Returns `None` once `start` reaches + /// [`num_vectors()`](Self::num_vectors), which ends such a walk. + #[allow(clippy::expect_used)] + pub fn row_range(&self, start: usize, count: usize) -> Option { + let count = count.min(self.num_vectors().checked_sub(start)?); + if count == 0 { + return None; + } + let k = self.vector_dim(); + let repr = + Standard::::new(count, k).expect("sub-view of a valid matrix cannot overflow"); + MatRef::new(repr, &self.as_slice()[start * k..(start + count) * k]).ok() + } + /// Return a [`MatrixView`] over the backing data. #[allow(clippy::expect_used)] #[inline] @@ -1418,6 +1435,37 @@ mod tests { // Standard // ////////////// + #[test] + fn standard_row_range_partitions_rows() { + let m = Mat::>::from_fn(Standard::new(7, 3).unwrap(), { + let mut i = 0; + move || { + i += 1; + i + } + }); + let v = m.as_view(); + + for tile in 1..10 { + let mut cur = 0; + while let Some(t) = v.row_range(cur, tile) { + assert!(t.num_vectors() <= tile); + assert_eq!(t.vector_dim(), 3); + assert_eq!( + t.as_slice(), + &v.as_slice()[cur * 3..(cur + t.num_vectors()) * 3] + ); + cur += t.num_vectors(); + } + assert_eq!(cur, 7, "tile size {tile} must cover every row"); + } + + assert!(v.row_range(7, 1).is_none(), "start at the end ends a walk"); + assert!(v.row_range(8, 1).is_none(), "start past the end"); + assert!(v.row_range(0, 0).is_none(), "zero-row view is not lent"); + assert_eq!(v.row_range(5, 99).unwrap().num_vectors(), 2, "count clips"); + } + #[test] fn standard_representation() { let repr = Standard::::new(4, 3).unwrap(); From 29809503b549f9bf86e6ee6ce288ba2437bcfcf2 Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Sun, 2 Aug 2026 00:54:52 +0530 Subject: [PATCH 4/5] Improve design as per new feedabcks --- .../distance/kernels/paneled/float.rs | 72 +++--- .../distance/kernels/paneled/leaves.rs | 118 +++++----- .../distance/kernels/paneled/minmax.rs | 78 ++++--- .../distance/kernels/paneled/mod.rs | 52 +++-- .../distance/kernels/paneled/strip.rs | 115 +++++++--- .../distance/kernels/paneled/views.rs | 210 +++++++++++++----- 6 files changed, 418 insertions(+), 227 deletions(-) diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs index c490b4dc5..92c05346c 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs @@ -9,31 +9,37 @@ use core::mem::size_of; use diskann_wide::arch::x86_64::V3; use super::arena::ResettableArena; -use super::views::{A_PANEL, B_PANEL, DPanel, DTail, DocWalk, QPanel, QueryWalk}; -use super::{ - Accumulate, At, Block, Drain, Plan, Short, Strip, StripRef, TileBudget, drive, leaves, -}; +use super::leaves::{A_PANEL, B_PANEL}; +use super::views::{DPanel, DocWalk, QPanel, QueryWalk}; +use super::{Accumulate, At, Block, Drain, Plan, Strip, StripRef, TileBudget, drive, leaves}; use crate::alloc::{Poly, ScopedAllocator}; +use crate::bits::{Dynamic, Static}; use crate::multi_vector::{BlockTransposed, Mat, MatRef, Standard}; // ── Kernel ─────────────────────────────────────────────────────── pub(crate) struct F32Kernel; -impl<'a, 'b, 'x, const R: usize, const N: usize> - Accumulate, DPanel<'b, f32, N>, Block<'x, f32, R, N>> for F32Kernel +impl<'a, 'b, 'x> + Accumulate< + V3, + QPanel<'a, f32, A_PANEL>, + DPanel<'b, f32, B_PANEL, Static>, + Block<'x, f32, A_PANEL, B_PANEL, Static>, + > for F32Kernel { + #[inline(always)] fn accumulate( &self, arch: V3, - a: QPanel<'a, f32, R>, - b: DPanel<'b, f32, N>, - mut out: Block<'x, f32, R, N>, + a: QPanel<'a, f32, A_PANEL>, + b: DPanel<'b, f32, B_PANEL, Static>, + mut out: Block<'x, f32, A_PANEL, B_PANEL, Static>, ) { - // SAFETY: `a` is an R×k block-transposed f32 block; `b` is N rows of k f32; - // `out` is N columns of R f32 at stride R. + // SAFETY: `a` is an A_PANEL×k block-transposed f32 block; `b` is B_PANEL rows + // of k f32; `out` is B_PANEL columns of A_PANEL f32 at stride A_PANEL. unsafe { - leaves::f32_store_microkernel::( + leaves::f32_store_microkernel::( arch, a.as_ptr(), b.as_ptr(), @@ -44,25 +50,32 @@ impl<'a, 'b, 'x, const R: usize, const N: usize> } } -impl<'a, 'b, 'x, const R: usize, const N: usize> - Accumulate, DTail<'b, f32>, Short>> for F32Kernel +impl<'a, 'b, 'x> + Accumulate< + V3, + QPanel<'a, f32, A_PANEL>, + DPanel<'b, f32, B_PANEL, Dynamic>, + Block<'x, f32, A_PANEL, B_PANEL, Dynamic>, + > for F32Kernel { + #[inline(always)] fn accumulate( &self, arch: V3, - a: QPanel<'a, f32, R>, - b: DTail<'b, f32>, - mut out: Short>, + a: QPanel<'a, f32, A_PANEL>, + b: DPanel<'b, f32, B_PANEL, Dynamic>, + mut out: Block<'x, f32, A_PANEL, B_PANEL, Dynamic>, ) { - debug_assert_eq!(out.0.cols(), b.rows()); - let (ap, bp, op, k) = (a.as_ptr(), b.as_ptr(), out.0.as_mut_ptr(), a.k()); - // SAFETY: as the full-width impl, with a runtime width in 1..N. + debug_assert_eq!(out.cols(), b.rows()); + debug_assert!(b.rows() < B_PANEL); + let (ap, bp, op, k) = (a.as_ptr(), b.as_ptr(), out.as_mut_ptr(), a.k()); + // SAFETY: as the full-width impl, with a runtime width in 1..B_PANEL. unsafe { match b.rows() { - 3 => leaves::f32_store_microkernel::<3, R>(arch, ap, bp, k, op), - 2 => leaves::f32_store_microkernel::<2, R>(arch, ap, bp, k, op), - 1 => leaves::f32_store_microkernel::<1, R>(arch, ap, bp, k, op), - other => unreachable!("tail width {other} out of 1..{N}"), + 3 => leaves::f32_store_microkernel::<3>(arch, ap, bp, k, op), + 2 => leaves::f32_store_microkernel::<2>(arch, ap, bp, k, op), + 1 => leaves::f32_store_microkernel::<1>(arch, ap, bp, k, op), + other => unreachable!("tail width {other} out of 1..{B_PANEL}"), } } } @@ -82,11 +95,12 @@ impl<'o> RawMax<'o> { } } -impl Drain> for RawMax<'_> { - fn drain(&mut self, arch: V3, acc: StripRef<'_, f32, R>, at: At) { - let out = &mut self.out[at.a_panel * R..][..R]; - // SAFETY: `out` is R f32; `acc` is `cols` columns of R f32. - unsafe { leaves::fold_strip::(arch, out.as_mut_ptr(), acc.as_ptr(), acc.cols()) } +impl Drain> for RawMax<'_> { + #[inline(always)] + fn drain(&mut self, arch: V3, acc: StripRef<'_, f32, A_PANEL>, at: At) { + let out = &mut self.out[at.a_panel * A_PANEL..][..A_PANEL]; + // SAFETY: `out` is A_PANEL f32; `acc` is `cols` columns of A_PANEL f32. + unsafe { leaves::fold_strip(arch, out.as_mut_ptr(), acc.as_ptr(), acc.cols()) } } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs index 993bd33fb..0ace42297 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs @@ -1,13 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. -//! SIMD leaves. The store-out micro-kernels are byte-identical to `tiler`'s, so an A/B -//! between the two experiments measures only the abstraction; [`score_fold_strip`] is -//! the one that differs — it fuses `tiler`'s separate dequant and max passes. -//! -//! Each leaf const-asserts the A-panel width it was previously assuming against its own -//! register shape, so a caller that disagrees fails to compile. This is the only guard -//! past the raw pointers, where the type-level widths stop. +//! V3 SIMD leaves, and the panel geometry they impose. The store-out micro-kernels are +//! byte-identical to `tiler`'s, so an A/B between the two experiments measures only the +//! abstraction; [`score_fold_strip`] is the one that differs — it fuses `tiler`'s +//! separate dequant and max passes. use diskann_wide::arch::x86_64::V3; use diskann_wide::{SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDMulAdd, SIMDReinterpret, SIMDVector}; @@ -19,41 +16,47 @@ diskann_wide::alias!(i32s = ::i32x8); diskann_wide::alias!(u32s = ::u32x8); diskann_wide::alias!(f32s = ::f32x8); -/// Integer store-out micro-kernel: `AROWS` A-rows × `UNROLL` B-rows. +/// Rows per A-panel: every leaf below emits exactly two 32-bit SIMD registers of rows. +/// Derived rather than written, so a wider ISA's leaves get their own width. +pub(super) const A_PANEL: usize = 2 * f32s::LANES; + +/// Max B-rows per kernel call. Not derived — a register-budget choice: `B_PANEL` × two +/// accumulator registers, plus two A registers, must fit the architectural file. +pub(super) const B_PANEL: usize = 4; + +/// Integer store-out micro-kernel: [`A_PANEL`] A-rows × `UNROLL` B-rows. /// /// # Safety /// -/// 1. `a_packed` points to an `AROWS × k` block-transposed `i16` block (`k` even). +/// 1. `a_packed` points to an `A_PANEL × k` block-transposed `i16` block (`k` even). /// 2. `b` points to `UNROLL` rows of `k` contiguous `u8` (`k` even). -/// 3. `partial` is valid for `UNROLL` columns of `AROWS` `i32` at stride `AROWS`. +/// 3. `partial` is valid for `UNROLL` columns of `A_PANEL` `i32` at stride `A_PANEL`. #[inline(always)] -pub(super) unsafe fn int_store_microkernel( +pub(super) unsafe fn int_store_microkernel( arch: V3, a_packed: *const i16, b: *const u8, k: usize, partial: *mut i32, ) { + // The i16 half-loads and the i32 stores must span the same rows; that relation + // holds across two different register types, so it is not self-evident. const { assert!( - AROWS == 2 * i32s::LANES, - "leaf emits exactly two i32 registers of rows" - ); - assert!( - AROWS == i16s::LANES, - "leaf loads AROWS i16 per A column-pair half" - ); + i16s::LANES == A_PANEL, + "leaf loads A_PANEL i16 per A column-pair half" + ) } let mut p0 = [i32s::default(arch); UNROLL]; let mut p1 = [i32s::default(arch); UNROLL]; let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); - let a_pair_stride = 2 * AROWS; - let a_half = AROWS; + let a_pair_stride = 2 * A_PANEL; + let a_half = A_PANEL; let pairs = k / 2; for p in 0..pairs { - // SAFETY: precondition 1 — the A block has `pairs` col-pairs of 2·AROWS i16. + // SAFETY: precondition 1 — the A block has `pairs` col-pairs of 2·A_PANEL i16. let (a0, a1) = unsafe { ( i16s::load_simd(arch, a_packed.add(a_pair_stride * p)), @@ -78,44 +81,38 @@ pub(super) unsafe fn int_store_microkernel( +pub(super) unsafe fn f32_store_microkernel( arch: V3, a_packed: *const f32, b: *const f32, k: usize, partial: *mut f32, ) { - const { - assert!( - AROWS == 2 * f32s::LANES, - "leaf emits exactly two f32 registers of rows" - ) - } let mut p0 = [f32s::default(arch); UNROLL]; let mut p1 = [f32s::default(arch); UNROLL]; let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); - let a_stride = AROWS; + let a_stride = A_PANEL; let a_half = f32s::LANES; for i in 0..k { - // SAFETY: precondition 1 — the A block has `k` columns of AROWS f32. + // SAFETY: precondition 1 — the A block has `k` columns of A_PANEL f32. let (a0, a1) = unsafe { ( f32s::load_simd(arch, a_packed.add(a_stride * i)), @@ -131,34 +128,28 @@ pub(super) unsafe fn f32_store_microkernel( - arch: V3, - state: *mut f32, - acc: *const f32, - cols: usize, -) { - const { assert!(R == 2 * f32s::LANES, "fold reads two f32 registers of rows") } +pub(super) unsafe fn fold_strip(arch: V3, state: *mut f32, acc: *const f32, cols: usize) { let lanes = f32s::LANES; - // SAFETY: `state` writable for R; `acc` valid for `cols` columns of R. + // SAFETY: `state` writable for A_PANEL; `acc` valid for `cols` columns of A_PANEL. unsafe { let mut m0 = f32s::load_simd(arch, state); let mut m1 = f32s::load_simd(arch, state.add(lanes)); for c in 0..cols { - let col = acc.add(c * R); + let col = acc.add(c * A_PANEL); m0 = m0.max_simd(f32s::load_simd(arch, col)); m1 = m1.max_simd(f32s::load_simd(arch, col.add(lanes))); } @@ -167,15 +158,15 @@ pub(super) unsafe fn fold_strip( } } -/// 4-bit MinMax dequant of an `R`×`cols` A-major `i32` strip, folded straight into the -/// running max — the score never reaches memory. +/// 4-bit MinMax dequant of an [`A_PANEL`]×`cols` A-major `i32` strip, folded straight +/// into the running max — the score never reaches memory. /// /// # Safety /// -/// `acc` valid for `cols` columns of `R` `i32` (stride `R`); `state` writable for `R` -/// `f32`; `q_meta.len() >= R`; `d_meta.len() >= cols`. +/// `acc` valid for `cols` columns of `A_PANEL` `i32` (stride `A_PANEL`); `state` +/// writable for `A_PANEL` `f32`; `q_meta.len() >= A_PANEL`; `d_meta.len() >= cols`. #[inline(always)] -pub(super) unsafe fn score_fold_strip( +pub(super) unsafe fn score_fold_strip( arch: V3, acc: *const i32, state: *mut f32, @@ -184,19 +175,18 @@ pub(super) unsafe fn score_fold_strip( d_meta: &[MinMaxCompensation], dim: f32, ) { - const { assert!(R == 2 * f32s::LANES, "fold reads two f32 registers of rows") } let lanes = f32s::LANES; - let mut qa = [0.0f32; R]; - let mut qb = [0.0f32; R]; - let mut qn = [0.0f32; R]; - for i in 0..R { + let mut qa = [0.0f32; A_PANEL]; + let mut qb = [0.0f32; A_PANEL]; + let mut qn = [0.0f32; A_PANEL]; + for i in 0..A_PANEL { let qm = q_meta[i]; qa[i] = qm.a; qb[i] = qm.b; qn[i] = qm.n; } - // SAFETY: each array holds exactly R = 2·LANES f32; `state` writable for R. + // SAFETY: each array holds exactly A_PANEL = 2·LANES f32; `state` writable for A_PANEL. let (qa0, qa1, qb0, qb1, qn0, qn1, mut m0, mut m1) = unsafe { ( f32s::load_simd(arch, qa.as_ptr()), @@ -214,8 +204,8 @@ pub(super) unsafe fn score_fold_strip( let a_c = f32s::splat(arch, dm.a); let b_c = f32s::splat(arch, dm.b); let c_c = f32s::splat(arch, dm.n + dm.b * dim); - let col = c * R; - // SAFETY: `col + 2·LANES <= cols*R`; `acc` valid for that many i32. + let col = c * A_PANEL; + // SAFETY: `col + 2·LANES <= cols*A_PANEL`; `acc` valid for that many i32. unsafe { let raw0 = i32s::load_simd(arch, acc.add(col)).simd_cast(); let raw1 = i32s::load_simd(arch, acc.add(col + lanes)).simd_cast(); @@ -226,7 +216,7 @@ pub(super) unsafe fn score_fold_strip( } } - // SAFETY: `state` writable for R f32. + // SAFETY: `state` writable for A_PANEL f32. unsafe { m0.store_simd(state); m1.store_simd(state.add(lanes)); diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs index 0b1d16ba2..a1da139e0 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs @@ -12,14 +12,14 @@ use diskann_utils::ReborrowMut; use diskann_wide::arch::x86_64::V3; use super::arena::ResettableArena; -use super::views::{A_PANEL, B_PANEL, DPanel, DTail, DocWalk, QPanel, QueryWalk}; -use super::{ - Accumulate, At, Block, Drain, Plan, Short, Strip, StripRef, TileBudget, drive, leaves, -}; +use super::leaves::{A_PANEL, B_PANEL}; +use super::views::{DPanel, DocWalk, QPanel, QueryWalk}; +use super::{Accumulate, At, Block, Drain, Plan, Strip, StripRef, TileBudget, drive, leaves}; use crate::CompressInto; use crate::algorithms::Transform; use crate::algorithms::transforms::NullTransform; use crate::alloc::{Poly, ScopedAllocator}; +use crate::bits::{Dynamic, Static}; use crate::minmax::{MinMaxCompensation, MinMaxMeta, MinMaxQuantizer}; use crate::multi_vector::{BlockTransposed, Defaulted, Mat, MatRef, Standard}; use crate::num::Positive; @@ -28,20 +28,26 @@ use crate::num::Positive; pub(crate) struct I8Kernel; -impl<'a, 'b, 'x, const R: usize, const N: usize> - Accumulate, DPanel<'b, u8, N>, Block<'x, i32, R, N>> for I8Kernel +impl<'a, 'b, 'x> + Accumulate< + V3, + QPanel<'a, i16, A_PANEL>, + DPanel<'b, u8, B_PANEL, Static>, + Block<'x, i32, A_PANEL, B_PANEL, Static>, + > for I8Kernel { + #[inline(always)] fn accumulate( &self, arch: V3, - a: QPanel<'a, i16, R>, - b: DPanel<'b, u8, N>, - mut out: Block<'x, i32, R, N>, + a: QPanel<'a, i16, A_PANEL>, + b: DPanel<'b, u8, B_PANEL, Static>, + mut out: Block<'x, i32, A_PANEL, B_PANEL, Static>, ) { - // SAFETY: `a` is an R×k block-transposed i16 block; `b` is N rows of k u8; - // `out` is N columns of R i32 at stride R (`k` even). + // SAFETY: `a` is an A_PANEL×k block-transposed i16 block; `b` is B_PANEL rows + // of k u8; `out` is B_PANEL columns of A_PANEL i32 at stride A_PANEL (`k` even). unsafe { - leaves::int_store_microkernel::( + leaves::int_store_microkernel::( arch, a.as_ptr(), b.as_ptr(), @@ -52,25 +58,32 @@ impl<'a, 'b, 'x, const R: usize, const N: usize> } } -impl<'a, 'b, 'x, const R: usize, const N: usize> - Accumulate, DTail<'b, u8>, Short>> for I8Kernel +impl<'a, 'b, 'x> + Accumulate< + V3, + QPanel<'a, i16, A_PANEL>, + DPanel<'b, u8, B_PANEL, Dynamic>, + Block<'x, i32, A_PANEL, B_PANEL, Dynamic>, + > for I8Kernel { + #[inline(always)] fn accumulate( &self, arch: V3, - a: QPanel<'a, i16, R>, - b: DTail<'b, u8>, - mut out: Short>, + a: QPanel<'a, i16, A_PANEL>, + b: DPanel<'b, u8, B_PANEL, Dynamic>, + mut out: Block<'x, i32, A_PANEL, B_PANEL, Dynamic>, ) { - debug_assert_eq!(out.0.cols(), b.rows()); - let (ap, bp, op, k) = (a.as_ptr(), b.as_ptr(), out.0.as_mut_ptr(), a.k()); - // SAFETY: as the full-width impl, with a runtime width in 1..N. + debug_assert_eq!(out.cols(), b.rows()); + debug_assert!(b.rows() < B_PANEL); + let (ap, bp, op, k) = (a.as_ptr(), b.as_ptr(), out.as_mut_ptr(), a.k()); + // SAFETY: as the full-width impl, with a runtime width in 1..B_PANEL. unsafe { match b.rows() { - 3 => leaves::int_store_microkernel::<3, R>(arch, ap, bp, k, op), - 2 => leaves::int_store_microkernel::<2, R>(arch, ap, bp, k, op), - 1 => leaves::int_store_microkernel::<1, R>(arch, ap, bp, k, op), - other => unreachable!("tail width {other} out of 1..{N}"), + 3 => leaves::int_store_microkernel::<3>(arch, ap, bp, k, op), + 2 => leaves::int_store_microkernel::<2>(arch, ap, bp, k, op), + 1 => leaves::int_store_microkernel::<1>(arch, ap, bp, k, op), + other => unreachable!("tail width {other} out of 1..{B_PANEL}"), } } } @@ -105,18 +118,19 @@ impl<'m, 'o> MinMaxMax<'m, 'o> { } } -impl Drain> for MinMaxMax<'_, '_> { - fn drain(&mut self, arch: V3, acc: StripRef<'_, i32, R>, at: At) { +impl Drain> for MinMaxMax<'_, '_> { + #[inline(always)] + fn drain(&mut self, arch: V3, acc: StripRef<'_, i32, A_PANEL>, at: At) { let cols = acc.cols(); - let lo = at.a_panel * R; - let q = &self.query_meta[lo..lo + R]; + let lo = at.a_panel * A_PANEL; + let q = &self.query_meta[lo..lo + A_PANEL]; let d = &self.doc_meta[at.b_row..at.b_row + cols]; let dim = self.dim; - let out = &mut self.out[lo..][..R]; - // SAFETY: `acc` is `cols` columns of R i32; `out` is R writable f32; - // `q.len() == R`; `d.len() == cols`. + let out = &mut self.out[lo..][..A_PANEL]; + // SAFETY: `acc` is `cols` columns of A_PANEL i32; `out` is A_PANEL writable + // f32; `q.len() == A_PANEL`; `d.len() == cols`. unsafe { - leaves::score_fold_strip::(arch, acc.as_ptr(), out.as_mut_ptr(), cols, q, d, dim); + leaves::score_fold_strip(arch, acc.as_ptr(), out.as_mut_ptr(), cols, q, d, dim); } } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs index f91412a6b..642348e27 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs @@ -7,9 +7,10 @@ //! into output it owns. [`Scratch`] is the write-side mirror of [`Paneled`], so the //! driver assumes no layout on either side. //! -//! Panel widths are const parameters (`R` = A rows, `N` = B rows) carried by the panel -//! types; a kernel relays them into its leaf rather than declaring them. Literals in -//! [`views`]. Instantiated for f32 ([`float`]) and 4-bit MinMax ([`minmax`]). +//! Panel widths are geometry, so they live with the leaves that impose them +//! ([`leaves::A_PANEL`], [`leaves::B_PANEL`]); the panel and accumulator types carry +//! them as const parameters (`R` = A rows, `N` = B rows) so the driver stays width- +//! agnostic. Instantiated for f32 ([`float`]) and 4-bit MinMax ([`minmax`]). //! //! Sibling to [`tiler`](super::tiler), which keeps postprocess and reduce separate. @@ -22,7 +23,7 @@ mod minmax; mod strip; mod views; -pub(crate) use strip::{Block, Short, Strip, StripRef}; +pub(crate) use strip::{Block, Strip, StripRef}; pub use float::{PaneledF32Docs, PaneledF32Query}; pub use minmax::{PaneledQuantDocs, PaneledQuantQuery}; @@ -62,6 +63,8 @@ pub(crate) trait ScratchAt<'a, B: sealed::Sealed = sealed::Bounds<&'a mut Self>> type Block; /// Short trailing slot; a distinct type so it selects its own [`Accumulate`] impl. type Short; + /// Named so `Block`/`Short` reach bounds without a nested projection. + type Blocks: TailIterator; /// What [`Drain`] reads. type Ref; } @@ -73,14 +76,8 @@ pub(crate) trait ScratchAt<'a, B: sealed::Sealed = sealed::Bounds<&'a mut Self>> /// Allocation stays on the concrete type. pub(crate) trait Scratch: for<'a> ScratchAt<'a> { /// Carve the live `cols` columns into one slot per B-panel plus the short trailer, - /// which must come from the same call to be provably disjoint. - fn split( - &mut self, - cols: usize, - ) -> ( - impl Iterator>::Block>, - Option<>::Short>, - ); + /// which comes off the same cursor and so is provably disjoint. + fn blocks(&mut self, cols: usize) -> >::Blocks; fn as_ref(&self, cols: usize) -> >::Ref; } @@ -117,16 +114,26 @@ pub(crate) struct Tile { pub(crate) at: usize, } +/// An iterator whose short trailing element has its own type. `tail` consumes the +/// exhausted iterator, so the trailer comes off the cursor the loop was already +/// advancing instead of being recomputed from the source. +pub(crate) trait TailIterator: ExactSizeIterator { + type Tail; + fn tail(self) -> Option; +} + /// A view that knows how it breaks into panels. `Tail` is distinct from `Panel` so the /// short trailing panel selects its own [`Accumulate`] impl; a view that cannot tail /// says [`NoTail`]. pub(crate) trait Paneled { type Panel: Copy; type Tail: Copy; + /// Named so it carries `ExactSizeIterator` and the tail type into bounds, and so a + /// k-fracturing driver could hold one across an outer loop. + type Panels: TailIterator; fn rows(&self) -> usize; - fn panels(&self) -> impl Iterator + '_; - fn tail(&self) -> Option; + fn panels(&self) -> Self::Panels; } /// `Tail` for a view padded to whole panels. Uninhabited, so `tail()` provably @@ -145,6 +152,7 @@ pub(crate) trait Accumulate { /// [`NoTail`] is uninhabited, so this one impl discharges the driver's A-tail bounds /// for every kernel. impl Accumulate for K { + #[inline(always)] fn accumulate(&self, _: Arch, a: NoTail, _: B, _: O) { match a {} } @@ -187,11 +195,16 @@ where K: for<'x> Accumulate> + for<'x> Accumulate>, { - let (blocks, acc_tail) = scratch.split(cols); - for (b, out) in b_view.panels().zip(blocks) { + let mut panels = b_view.panels(); + let mut blocks = scratch.blocks(cols); + // What [`TailIterator`]'s `ExactSizeIterator` bound is for: `zip` stops on the + // shorter side and silently drops the longer side's pending item, which would leave + // both cursors at zero and slip past the tail check below. + debug_assert_eq!(panels.len(), blocks.len(), "B view and accumulator desync"); + for (b, out) in panels.by_ref().zip(blocks.by_ref()) { kernel.accumulate(arch, a, b, out); } - match (b_view.tail(), acc_tail) { + match (panels.tail(), blocks.tail()) { (Some(b), Some(out)) => kernel.accumulate(arch, a, b, out), (None, None) => {} _ => unreachable!("B view and accumulator disagree on tail"), @@ -231,12 +244,13 @@ pub(super) fn drive( b_row: b_tile.at, }; - for a in a_tile.view.panels() { + let mut a_panels = a_tile.view.panels(); + for a in a_panels.by_ref() { fill(arch, kernel, a, &b_tile.view, scratch, cols); drain.drain(arch, scratch.as_ref(cols), at); at.a_panel += 1; } - if let Some(a) = a_tile.view.tail() { + if let Some(a) = a_panels.tail() { fill(arch, kernel, a, &b_tile.view, scratch, cols); drain.drain(arch, scratch.as_ref(cols), at); } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs index 6ccba69b1..899bc7ee3 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs @@ -4,28 +4,96 @@ //! The default [`Scratch`]: one A-panel × one B-tile, `R` rows A-major, carved into //! `N`-column [`Block`]s. Borrows its buffer, so the allocator stays at the call site. +use core::marker::PhantomData; use core::mem::MaybeUninit; -use super::{Scratch, ScratchAt}; +use super::{Scratch, ScratchAt, TailIterator}; use crate::alloc::{AllocatorCore, Poly}; +use crate::bits::{Dynamic, Length, Static}; /// Marker for element types where all-zero is a valid value. pub(crate) trait ZeroInit: Copy {} impl ZeroInit for i32 {} impl ZeroInit for f32 {} +/// Owns the accumulator's live region as a checked slice — the anchor the unchecked +/// [`Block`]s below are derived from. pub(crate) struct Strip<'a, T, const R: usize, const N: usize>(&'a mut [T]); -/// One B-panel's slot: `R` rows by `N` columns (fewer inside a [`Short`]). -pub(crate) struct Block<'a, T, const R: usize, const N: usize>(&'a mut [T]); +/// One B-panel's slot: `R` rows by `L` columns, where `L.value() <= N`. A full slot is +/// `Static` (a ZST, so the whole handle is one pointer); the trailing slot is +/// `Dynamic` in `1..N`, and its distinct type selects the short +/// [`Accumulate`](super::Accumulate) impl. +pub(crate) struct Block<'a, T, const R: usize, const N: usize, L: Length> { + ptr: *mut T, + cols: L, + _lifetime: PhantomData<&'a mut [T]>, +} /// A finished [`Strip`], narrowed to its live columns. pub(crate) struct StripRef<'a, T, const R: usize>(&'a [T]); -/// A short trailing slot: same payload, different type, so the runtime-width path is -/// a separate [`Accumulate`](super::Accumulate) impl. -#[derive(Clone, Copy)] -pub(crate) struct Short

(pub(crate) P); +/// Carves a [`Strip`] into per-B-panel slots. The trailer comes off the same cursor. +/// +/// # Safety invariants +/// +/// `ptr` is valid for writes of `R * (full * N + tail_cols)` `T` for `'a`, and only +/// ever advances — so every slot handed out is disjoint from every other. +pub(crate) struct BlockIter<'a, T, const R: usize, const N: usize> { + ptr: *mut T, + full: usize, + tail_cols: usize, + _lifetime: PhantomData<&'a mut [T]>, +} + +impl<'a, T, const R: usize, const N: usize, L: Length> Block<'a, T, R, N, L> { + /// # Safety + /// + /// `ptr` must be valid for writes of `R * cols.value()` `T` for `'a`, must not + /// alias any other live `Block`, and `cols.value()` must be at most `N`. + unsafe fn new(ptr: *mut T, cols: L) -> Self { + debug_assert!(cols.value() <= N, "block wider than its panel"); + Self { + ptr, + cols, + _lifetime: PhantomData, + } + } +} + +impl<'a, T, const R: usize, const N: usize> Iterator for BlockIter<'a, T, R, N> { + type Item = Block<'a, T, R, N, Static>; + + fn next(&mut self) -> Option { + if self.full == 0 { + return None; + } + let ptr = self.ptr; + // SAFETY: the invariant covers `full` more slots of `R * N`, so the bump stays + // inside the allocation and the yielded slot is disjoint from all later ones. + self.ptr = unsafe { self.ptr.add(R * N) }; + self.full -= 1; + // SAFETY: as above — `ptr` covers exactly `R * N` writable `T`. + Some(unsafe { Block::new(ptr, Static) }) + } + + fn size_hint(&self) -> (usize, Option) { + (self.full, Some(self.full)) + } +} + +impl ExactSizeIterator for BlockIter<'_, T, R, N> {} + +impl<'a, T, const R: usize, const N: usize> TailIterator for BlockIter<'a, T, R, N> { + type Tail = Block<'a, T, R, N, Dynamic>; + + fn tail(self) -> Option { + debug_assert_eq!(self.full, 0, "tail taken before the blocks are exhausted"); + // SAFETY: the blocks are exhausted, so the cursor sits on the trailer, which + // the invariant covers for `R * tail_cols` writable `T`. + (self.tail_cols > 0).then(|| unsafe { Block::new(self.ptr, Dynamic(self.tail_cols)) }) + } +} impl<'a, T: ZeroInit, const R: usize, const N: usize> Strip<'a, T, R, N> { /// Zeroing is what makes the `&mut [T]` sound; the kernels overwrite every live @@ -51,29 +119,26 @@ impl Strip<'_, T, R, N> { } impl<'a, T, const R: usize, const N: usize> ScratchAt<'a> for Strip<'_, T, R, N> { - type Block = Block<'a, T, R, N>; - type Short = Short>; + type Block = Block<'a, T, R, N, Static>; + type Short = Block<'a, T, R, N, Dynamic>; + type Blocks = BlockIter<'a, T, R, N>; type Ref = StripRef<'a, T, R>; } impl Scratch for Strip<'_, T, R, N> { - fn split( - &mut self, - cols: usize, - ) -> ( - impl Iterator>, - Option>>, - ) { + fn blocks(&mut self, cols: usize) -> BlockIter<'_, T, R, N> { debug_assert!( cols <= self.cols_capacity(), "strip must hold the whole B-tile" ); - let (live, _) = self.0.split_at_mut(R * cols); - let (head, rest) = live.split_at_mut(R * (cols - cols % N)); - ( - head.chunks_mut(R * N).map(Block), - (!rest.is_empty()).then_some(Short(Block(rest))), - ) + // The checked slice is what establishes `BlockIter`'s invariant: `R * cols` + // elements are live, and the `&mut self` borrow keeps them exclusive for `'_`. + BlockIter { + ptr: self.0.as_mut_ptr(), + full: cols / N, + tail_cols: cols % N, + _lifetime: PhantomData, + } } fn as_ref(&self, cols: usize) -> StripRef<'_, T, R> { @@ -81,12 +146,12 @@ impl Scratch for Strip<'_, T, R, N> { } } -impl Block<'_, T, R, N> { +impl Block<'_, T, R, N, L> { pub(crate) fn cols(&self) -> usize { - self.0.len() / R + self.cols.value() } pub(crate) fn as_mut_ptr(&mut self) -> *mut T { - self.0.as_mut_ptr() + self.ptr } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs index a5f9e0aaa..f40bae86f 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs @@ -6,32 +6,28 @@ //! row-major matrix doesn't imply a panel height). Each sub-views itself, so a walk is //! a cursor and nothing else. -use super::{NoTail, Paneled, Tile, TileAt, TileWalk}; -use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; +use core::marker::PhantomData; -/// Rows per A-panel (the block-transposed group and the drain's state block). -pub(crate) const A_PANEL: usize = 16; -/// Rows per full B-panel (the kernel's micro-panel / max unroll). -pub(crate) const B_PANEL: usize = 4; +use super::{NoTail, Paneled, TailIterator, Tile, TileAt, TileWalk}; +use crate::bits::{Dynamic, Length, Static}; +use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; // ── Panels ─────────────────────────────────────────────────────── /// One block-transposed A block: `R` rows × `k` `T`. pub(crate) struct QPanel<'a, T, const R: usize> { - data: &'a [T], + ptr: *const T, k: usize, + _lifetime: PhantomData<&'a [T]>, } -/// One row-major B panel: exactly `N` rows × `k` `T`. `k` travels with the A panel, -/// and the row count is `N` by construction, so neither is stored. -pub(crate) struct DPanel<'a, T, const N: usize> { - data: &'a [T], -} -/// The short trailing B panel: `rows` in `1..N`. Distinct from [`DPanel`] so it -/// selects its own [`Accumulate`](super::Accumulate) impl, and so a full panel carries -/// no runtime row count. -pub(crate) struct DTail<'a, T> { - data: &'a [T], - rows: usize, + +/// One row-major B panel: `L` rows × `k` `T`, where `L.value() <= N`. `k` travels with +/// the A panel, so it is not stored. A full panel is `Static` — a ZST, so the whole +/// handle is one pointer; the trailing panel is `Dynamic` in `1..N`. +pub(crate) struct DPanel<'a, T, const N: usize, L: Length> { + ptr: *const T, + rows: L, + _lifetime: PhantomData<&'a [T]>, } impl Clone for QPanel<'_, T, R> { @@ -40,38 +36,141 @@ impl Clone for QPanel<'_, T, R> { } } impl Copy for QPanel<'_, T, R> {} -impl Clone for DPanel<'_, T, N> { +impl Clone for DPanel<'_, T, N, L> { fn clone(&self) -> Self { *self } } -impl Copy for DPanel<'_, T, N> {} -impl Clone for DTail<'_, T> { - fn clone(&self) -> Self { - *self +impl Copy for DPanel<'_, T, N, L> {} + +impl<'a, T, const R: usize> QPanel<'a, T, R> { + /// # Safety + /// + /// `ptr` must be valid for reads of `R * k` `T` for `'a`. + unsafe fn new(ptr: *const T, k: usize) -> Self { + Self { + ptr, + k, + _lifetime: PhantomData, + } + } +} + +impl<'a, T, const N: usize, L: Length> DPanel<'a, T, N, L> { + /// # Safety + /// + /// `ptr` must be valid for reads of `rows.value() * k` `T` for `'a`, where `k` is + /// the contraction length carried by the A panel it is paired with, and + /// `rows.value()` must be at most `N`. + unsafe fn new(ptr: *const T, rows: L) -> Self { + debug_assert!(rows.value() <= N, "panel taller than its type claims"); + Self { + ptr, + rows, + _lifetime: PhantomData, + } } } -impl Copy for DTail<'_, T> {} impl QPanel<'_, T, R> { pub(crate) fn as_ptr(&self) -> *const T { - self.data.as_ptr() + self.ptr } pub(crate) fn k(&self) -> usize { self.k } } -impl DPanel<'_, T, N> { +impl DPanel<'_, T, N, L> { pub(crate) fn as_ptr(&self) -> *const T { - self.data.as_ptr() + self.ptr + } + pub(crate) fn rows(&self) -> usize { + self.rows.value() } } -impl DTail<'_, T> { - pub(crate) fn as_ptr(&self) -> *const T { - self.data.as_ptr() + +// ── Panel iterators ────────────────────────────────────────────── + +/// Walks a block-transposed matrix' blocks. The remainder block is zero-padded to a +/// full `R` rows, so there is no tail. +pub(crate) struct QPanelIter<'a, T: Copy, const R: usize, const P: usize> { + view: BlockTransposedRef<'a, T, R, P>, + k: usize, + cur: usize, + end: usize, +} + +impl<'a, T: Copy, const R: usize, const P: usize> Iterator for QPanelIter<'a, T, R, P> { + type Item = QPanel<'a, T, R>; + + fn next(&mut self) -> Option { + let data = self.view.block_slice(self.cur)?; + self.cur += 1; + // SAFETY: `block_slice` returns a checked `R * k` slice borrowed from the view. + Some(unsafe { QPanel::new(data.as_ptr(), self.k) }) } - pub(crate) fn rows(&self) -> usize { - self.rows + + fn size_hint(&self) -> (usize, Option) { + let n = self.end - self.cur; + (n, Some(n)) + } +} + +impl ExactSizeIterator for QPanelIter<'_, T, R, P> {} + +impl TailIterator for QPanelIter<'_, T, R, P> { + type Tail = NoTail; + fn tail(self) -> Option { + None + } +} + +/// Walks a row-major matrix' full `N`-row panels; [`TailIterator::tail`] hands back the +/// short trailer from the same cursor. +/// +/// # Safety invariants +/// +/// `ptr` is valid for reads of `(full * N + tail_rows) * k` `T` for `'a`, and only ever +/// advances. +pub(crate) struct DPanelIter<'a, T, const N: usize> { + ptr: *const T, + k: usize, + full: usize, + tail_rows: usize, + _lifetime: PhantomData<&'a [T]>, +} + +impl<'a, T, const N: usize> Iterator for DPanelIter<'a, T, N> { + type Item = DPanel<'a, T, N, Static>; + + fn next(&mut self) -> Option { + if self.full == 0 { + return None; + } + let ptr = self.ptr; + // SAFETY: the invariant covers `full` more panels of `N * k`, so the bump stays + // inside the allocation. + self.ptr = unsafe { self.ptr.add(N * self.k) }; + self.full -= 1; + // SAFETY: as above — `ptr` covers exactly `N * k` readable `T`. + Some(unsafe { DPanel::new(ptr, Static) }) + } + + fn size_hint(&self) -> (usize, Option) { + (self.full, Some(self.full)) + } +} + +impl ExactSizeIterator for DPanelIter<'_, T, N> {} + +impl<'a, T, const N: usize> TailIterator for DPanelIter<'a, T, N> { + type Tail = DPanel<'a, T, N, Dynamic>; + + fn tail(self) -> Option { + debug_assert_eq!(self.full, 0, "tail taken before the panels are exhausted"); + // SAFETY: the panels are exhausted, so the cursor sits on the trailer, which + // the invariant covers for `tail_rows * k` readable `T`. + (self.tail_rows > 0).then(|| unsafe { DPanel::new(self.ptr, Dynamic(self.tail_rows)) }) } } @@ -84,21 +183,18 @@ impl DTail<'_, T> { impl<'a, T: Copy, const R: usize, const P: usize> Paneled for BlockTransposedRef<'a, T, R, P> { type Panel = QPanel<'a, T, R>; type Tail = NoTail; + type Panels = QPanelIter<'a, T, R, P>; fn rows(&self) -> usize { self.nrows() } - fn panels(&self) -> impl Iterator> + '_ { - let (v, k) = (*self, self.padded_ncols()); - (0..self.num_blocks()).filter_map(move |b| { - Some(QPanel { - data: v.block_slice(b)?, - k, - }) - }) - } - fn tail(&self) -> Option { - None + fn panels(&self) -> QPanelIter<'a, T, R, P> { + QPanelIter { + view: *self, + k: self.padded_ncols(), + cur: 0, + end: self.num_blocks(), + } } } @@ -106,25 +202,23 @@ impl<'a, T: Copy, const R: usize, const P: usize> Paneled for BlockTransposedRef pub(crate) struct Rows(pub(crate) V); impl<'a, const N: usize, T: Copy> Paneled for Rows>> { - type Panel = DPanel<'a, T, N>; - type Tail = DTail<'a, T>; + type Panel = DPanel<'a, T, N, Static>; + type Tail = DPanel<'a, T, N, Dynamic>; + type Panels = DPanelIter<'a, T, N>; fn rows(&self) -> usize { self.0.num_vectors() } - fn panels(&self) -> impl Iterator> + '_ { - let (data, k) = (self.0.as_slice(), self.0.vector_dim()); - (0..self.rows() / N).map(move |p| DPanel { - data: &data[p * N * k..(p + 1) * N * k], - }) - } - fn tail(&self) -> Option> { - let (data, k, n) = (self.0.as_slice(), self.0.vector_dim(), self.rows()); - let rem = n % N; - (rem > 0).then(|| DTail { - data: &data[(n - rem) * k..], - rows: rem, - }) + fn panels(&self) -> DPanelIter<'a, T, N> { + let (n, k) = (self.rows(), self.0.vector_dim()); + // The checked slice is what establishes `DPanelIter`'s invariant. + DPanelIter { + ptr: self.0.as_slice().as_ptr(), + k, + full: n / N, + tail_rows: n % N, + _lifetime: PhantomData, + } } } From a550ecfe18a5b5327096e9400a682dcfb9f033da Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Fri, 7 Aug 2026 00:55:28 +0530 Subject: [PATCH 5/5] Improve Design --- .../distance/kernels/paneled/float.rs | 49 ++-- .../distance/kernels/paneled/leaves.rs | 155 ++++++++---- .../distance/kernels/paneled/minmax.rs | 62 ++--- .../distance/kernels/paneled/mod.rs | 222 +++++++++++------- .../distance/kernels/paneled/strip.rs | 148 ++++-------- .../distance/kernels/paneled/views.rs | 171 +++++++++----- 6 files changed, 433 insertions(+), 374 deletions(-) diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs index 92c05346c..5ecbf10be 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/float.rs @@ -11,7 +11,7 @@ use diskann_wide::arch::x86_64::V3; use super::arena::ResettableArena; use super::leaves::{A_PANEL, B_PANEL}; use super::views::{DPanel, DocWalk, QPanel, QueryWalk}; -use super::{Accumulate, At, Block, Drain, Plan, Strip, StripRef, TileBudget, drive, leaves}; +use super::{Accumulate, Block, Drain, Plan, Region, Strip, TileBudget, drive, leaves}; use crate::alloc::{Poly, ScopedAllocator}; use crate::bits::{Dynamic, Static}; use crate::multi_vector::{BlockTransposed, Mat, MatRef, Standard}; @@ -25,7 +25,7 @@ impl<'a, 'b, 'x> V3, QPanel<'a, f32, A_PANEL>, DPanel<'b, f32, B_PANEL, Static>, - Block<'x, f32, A_PANEL, B_PANEL, Static>, + Block<'x, f32, A_PANEL, B_PANEL>, > for F32Kernel { #[inline(always)] @@ -34,19 +34,9 @@ impl<'a, 'b, 'x> arch: V3, a: QPanel<'a, f32, A_PANEL>, b: DPanel<'b, f32, B_PANEL, Static>, - mut out: Block<'x, f32, A_PANEL, B_PANEL, Static>, + out: Block<'x, f32, A_PANEL, B_PANEL>, ) { - // SAFETY: `a` is an A_PANEL×k block-transposed f32 block; `b` is B_PANEL rows - // of k f32; `out` is B_PANEL columns of A_PANEL f32 at stride A_PANEL. - unsafe { - leaves::f32_store_microkernel::( - arch, - a.as_ptr(), - b.as_ptr(), - a.k(), - out.as_mut_ptr(), - ); - } + leaves::f32_store_microkernel::(arch, a, b, out); } } @@ -55,7 +45,7 @@ impl<'a, 'b, 'x> V3, QPanel<'a, f32, A_PANEL>, DPanel<'b, f32, B_PANEL, Dynamic>, - Block<'x, f32, A_PANEL, B_PANEL, Dynamic>, + Block<'x, f32, A_PANEL, B_PANEL>, > for F32Kernel { #[inline(always)] @@ -64,26 +54,21 @@ impl<'a, 'b, 'x> arch: V3, a: QPanel<'a, f32, A_PANEL>, b: DPanel<'b, f32, B_PANEL, Dynamic>, - mut out: Block<'x, f32, A_PANEL, B_PANEL, Dynamic>, + out: Block<'x, f32, A_PANEL, B_PANEL>, ) { - debug_assert_eq!(out.cols(), b.rows()); - debug_assert!(b.rows() < B_PANEL); - let (ap, bp, op, k) = (a.as_ptr(), b.as_ptr(), out.as_mut_ptr(), a.k()); - // SAFETY: as the full-width impl, with a runtime width in 1..B_PANEL. - unsafe { - match b.rows() { - 3 => leaves::f32_store_microkernel::<3>(arch, ap, bp, k, op), - 2 => leaves::f32_store_microkernel::<2>(arch, ap, bp, k, op), - 1 => leaves::f32_store_microkernel::<1>(arch, ap, bp, k, op), - other => unreachable!("tail width {other} out of 1..{B_PANEL}"), - } + // The leaf checks that the width it unrolls for is the width `b` actually has. + match b.rows() { + 3 => leaves::f32_store_microkernel::<3, _>(arch, a, b, out), + 2 => leaves::f32_store_microkernel::<2, _>(arch, a, b, out), + 1 => leaves::f32_store_microkernel::<1, _>(arch, a, b, out), + other => unreachable!("tail width {other} out of 1..{B_PANEL}"), } } } // ── Drain ──────────────────────────────────────────────────────── -/// Running max over an output it owns, padded to whole A-panels by the caller. +/// Running max, over an output the caller has padded to whole A-panels. pub(crate) struct RawMax<'o> { out: &'o mut [f32], } @@ -97,10 +82,10 @@ impl<'o> RawMax<'o> { impl Drain> for RawMax<'_> { #[inline(always)] - fn drain(&mut self, arch: V3, acc: StripRef<'_, f32, A_PANEL>, at: At) { - let out = &mut self.out[at.a_panel * A_PANEL..][..A_PANEL]; - // SAFETY: `out` is A_PANEL f32; `acc` is `cols` columns of A_PANEL f32. - unsafe { leaves::fold_strip(arch, out.as_mut_ptr(), acc.as_ptr(), acc.cols()) } + fn drain(&mut self, arch: V3, acc: &Strip<'_, f32, A_PANEL, B_PANEL>, region: Region) { + // This output is padded to whole panels, so the stride is ours to state. + let out = &mut self.out[region.a.start..][..A_PANEL]; + leaves::fold_strip(arch, out, acc, region.b.len()); } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs index 0ace42297..3f38477b4 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/leaves.rs @@ -1,14 +1,21 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. -//! V3 SIMD leaves, and the panel geometry they impose. The store-out micro-kernels are -//! byte-identical to `tiler`'s, so an A/B between the two experiments measures only the -//! abstraction; [`score_fold_strip`] is the one that differs — it fuses `tiler`'s -//! separate dequant and max passes. +//! V3 SIMD leaves, and the panel geometry they impose. The store-out micro-kernels do +//! the same math as `tiler`'s, so an A/B between the two experiments measures the +//! abstraction rather than a different kernel; [`score_fold_strip`] is the one with no +//! counterpart there — it fuses `tiler`'s separate dequant and max passes. +//! +//! Every leaf is safe to call: they take the panel and slot handles whole, so each one's +//! requirements are discharged here rather than promised at every call site. The +//! remaining `unsafe` is the SIMD loads and stores, each in bounds by the lines above it. use diskann_wide::arch::x86_64::V3; use diskann_wide::{SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDMulAdd, SIMDReinterpret, SIMDVector}; +use super::strip::{Block, Strip}; +use super::views::{DPanel, QPanel}; +use crate::bits::Length; use crate::minmax::MinMaxCompensation; diskann_wide::alias!(i16s = ::i16x16); @@ -16,28 +23,26 @@ diskann_wide::alias!(i32s = ::i32x8); diskann_wide::alias!(u32s = ::u32x8); diskann_wide::alias!(f32s = ::f32x8); -/// Rows per A-panel: every leaf below emits exactly two 32-bit SIMD registers of rows. -/// Derived rather than written, so a wider ISA's leaves get their own width. +/// Rows per A-panel: every leaf below emits exactly two 32-bit SIMD registers of rows, +/// so a wider ISA's leaves get their own width. pub(super) const A_PANEL: usize = 2 * f32s::LANES; -/// Max B-rows per kernel call. Not derived — a register-budget choice: `B_PANEL` × two -/// accumulator registers, plus two A registers, must fit the architectural file. +/// Max B-rows per kernel call — a register-budget choice: `B_PANEL` × two accumulator +/// registers, plus two A registers, must fit the architectural file. pub(super) const B_PANEL: usize = 4; -/// Integer store-out micro-kernel: [`A_PANEL`] A-rows × `UNROLL` B-rows. +/// Integer store-out micro-kernel: [`A_PANEL`] A-rows × `UNROLL` B-rows. `k` must be +/// even — the loads take column pairs. /// -/// # Safety +/// # Panics /// -/// 1. `a_packed` points to an `A_PANEL × k` block-transposed `i16` block (`k` even). -/// 2. `b` points to `UNROLL` rows of `k` contiguous `u8` (`k` even). -/// 3. `partial` is valid for `UNROLL` columns of `A_PANEL` `i32` at stride `A_PANEL`. +/// If the panels disagree on `k`, or `b` is not `UNROLL` rows — both driver bugs. #[inline(always)] -pub(super) unsafe fn int_store_microkernel( +pub(super) fn int_store_microkernel( arch: V3, - a_packed: *const i16, - b: *const u8, - k: usize, - partial: *mut i32, + a: QPanel<'_, i16, A_PANEL>, + b: DPanel<'_, u8, B_PANEL, L>, + mut out: Block<'_, i32, A_PANEL, B_PANEL>, ) { // The i16 half-loads and the i32 stores must span the same rows; that relation // holds across two different register types, so it is not self-evident. @@ -47,6 +52,19 @@ pub(super) unsafe fn int_store_microkernel( "leaf loads A_PANEL i16 per A column-pair half" ) } + // Bounds the store below inside the slot's `A_PANEL * B_PANEL`. + const { assert!(UNROLL <= B_PANEL, "unroll wider than the slot") } + + let k = a.k(); + assert_eq!( + k, + b.k(), + "panels paired across different contraction lengths" + ); + assert_eq!(b.rows(), UNROLL, "panel height must match the unroll"); + debug_assert_eq!(k % 2, 0, "the integer leaf loads column pairs"); + let (a_packed, b_ptr, partial) = (a.as_ptr(), b.as_ptr(), out.as_mut_ptr()); + let mut p0 = [i32s::default(arch); UNROLL]; let mut p1 = [i32s::default(arch); UNROLL]; let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); @@ -56,7 +74,8 @@ pub(super) unsafe fn int_store_microkernel( let pairs = k / 2; for p in 0..pairs { - // SAFETY: precondition 1 — the A block has `pairs` col-pairs of 2·A_PANEL i16. + // SAFETY: `a` is `A_PANEL * k` i16, so it holds `pairs` column-pairs of + // `2 * A_PANEL`, and `p < pairs`. let (a0, a1) = unsafe { ( i16s::load_simd(arch, a_packed.add(a_pair_stride * p)), @@ -65,12 +84,13 @@ pub(super) unsafe fn int_store_microkernel( }; for j in 0..UNROLL { - // SAFETY: precondition 2 — B row j is `offsets[j]` in, `2*p+1 < k`. + // SAFETY: `b` is `UNROLL * k` u8 (both asserted), row j starts at + // `offsets[j] = k * j`, and `2 * p + 1 < k`. let (d0, d1) = unsafe { let base = 2 * p + offsets[j]; ( - u32::from(b.add(base).read()), - u32::from(b.add(base + 1).read()), + u32::from(b_ptr.add(base).read()), + u32::from(b_ptr.add(base + 1).read()), ) }; let packed = d0 | (d1 << 16); @@ -81,7 +101,8 @@ pub(super) unsafe fn int_store_microkernel( } for j in 0..UNROLL { - // SAFETY: precondition 3 — column j occupies [j*A_PANEL, j*A_PANEL+A_PANEL) i32. + // SAFETY: the slot is `A_PANEL * B_PANEL` writable i32 and `UNROLL <= B_PANEL`, + // so column j occupies `[j * A_PANEL, (j + 1) * A_PANEL)` inside it. unsafe { p0[j].store_simd(partial.add(j * A_PANEL)); p1[j].store_simd(partial.add(j * A_PANEL + i32s::LANES)); @@ -91,19 +112,28 @@ pub(super) unsafe fn int_store_microkernel( /// f32 store-out micro-kernel: [`A_PANEL`] A-rows × `UNROLL` B-rows of inner product. /// -/// # Safety +/// # Panics /// -/// 1. `a_packed` points to an `A_PANEL × k` block-transposed `f32` block (`PACK = 1`). -/// 2. `b` points to `UNROLL` rows of `k` contiguous `f32`. -/// 3. `partial` is valid for `UNROLL` columns of `A_PANEL` `f32` at stride `A_PANEL`. +/// If the panels disagree on `k`, or `b` is not `UNROLL` rows — both driver bugs. #[inline(always)] -pub(super) unsafe fn f32_store_microkernel( +pub(super) fn f32_store_microkernel( arch: V3, - a_packed: *const f32, - b: *const f32, - k: usize, - partial: *mut f32, + a: QPanel<'_, f32, A_PANEL>, + b: DPanel<'_, f32, B_PANEL, L>, + mut out: Block<'_, f32, A_PANEL, B_PANEL>, ) { + // Bounds the store below inside the slot's `A_PANEL * B_PANEL`. + const { assert!(UNROLL <= B_PANEL, "unroll wider than the slot") } + + let k = a.k(); + assert_eq!( + k, + b.k(), + "panels paired across different contraction lengths" + ); + assert_eq!(b.rows(), UNROLL, "panel height must match the unroll"); + let (a_packed, b_ptr, partial) = (a.as_ptr(), b.as_ptr(), out.as_mut_ptr()); + let mut p0 = [f32s::default(arch); UNROLL]; let mut p1 = [f32s::default(arch); UNROLL]; let offsets: [usize; UNROLL] = core::array::from_fn(|j| k * j); @@ -112,7 +142,8 @@ pub(super) unsafe fn f32_store_microkernel( let a_half = f32s::LANES; for i in 0..k { - // SAFETY: precondition 1 — the A block has `k` columns of A_PANEL f32. + // SAFETY: `a` is `A_PANEL * k` f32 and `i < k`, so both halves of column `i` + // are in bounds. let (a0, a1) = unsafe { ( f32s::load_simd(arch, a_packed.add(a_stride * i)), @@ -120,15 +151,17 @@ pub(super) unsafe fn f32_store_microkernel( ) }; for j in 0..UNROLL { - // SAFETY: precondition 2 — B row j is `offsets[j]` in, `i < k`. - let bj = unsafe { f32s::splat(arch, b.add(i + offsets[j]).read_unaligned()) }; + // SAFETY: `b` is `UNROLL * k` f32 (both asserted), row j starts at + // `offsets[j] = k * j`, and `i < k`. + let bj = unsafe { f32s::splat(arch, b_ptr.add(i + offsets[j]).read_unaligned()) }; p0[j] = a0.mul_add_simd(bj, p0[j]); p1[j] = a1.mul_add_simd(bj, p1[j]); } } for j in 0..UNROLL { - // SAFETY: precondition 3 — column j occupies [j*A_PANEL, j*A_PANEL+A_PANEL) f32. + // SAFETY: the slot is `A_PANEL * B_PANEL` writable f32 and `UNROLL <= B_PANEL`, + // so column j occupies `[j * A_PANEL, (j + 1) * A_PANEL)` inside it. unsafe { p0[j].store_simd(partial.add(j * A_PANEL)); p1[j].store_simd(partial.add(j * A_PANEL + f32s::LANES)); @@ -138,13 +171,24 @@ pub(super) unsafe fn f32_store_microkernel( /// Fold an [`A_PANEL`]×`cols` A-major f32 strip into the running max. /// -/// # Safety +/// # Panics /// -/// `state` writable for `A_PANEL` `f32`; `acc` valid for `cols` columns of `A_PANEL` `f32`. +/// If `state` is not exactly one A-panel, or the strip holds fewer than `cols` columns. #[inline(always)] -pub(super) unsafe fn fold_strip(arch: V3, state: *mut f32, acc: *const f32, cols: usize) { +pub(super) fn fold_strip( + arch: V3, + state: &mut [f32], + strip: &Strip<'_, f32, A_PANEL, B_PANEL>, + cols: usize, +) { let lanes = f32s::LANES; - // SAFETY: `state` writable for A_PANEL; `acc` valid for `cols` columns of A_PANEL. + // The slicing is the bounds check: both loads below are inside these lengths. + let acc = strip.columns(cols); + assert_eq!(state.len(), A_PANEL, "the fold writes a whole A-panel"); + let (state, acc) = (state.as_mut_ptr(), acc.as_ptr()); + + // SAFETY: `state` is A_PANEL = 2·LANES writable f32, and `acc` is `cols * A_PANEL` + // readable f32, so column `c < cols` and both its halves are in bounds. unsafe { let mut m0 = f32s::load_simd(arch, state); let mut m1 = f32s::load_simd(arch, state.add(lanes)); @@ -161,32 +205,38 @@ pub(super) unsafe fn fold_strip(arch: V3, state: *mut f32, acc: *const f32, cols /// 4-bit MinMax dequant of an [`A_PANEL`]×`cols` A-major `i32` strip, folded straight /// into the running max — the score never reaches memory. /// -/// # Safety +/// # Panics /// -/// `acc` valid for `cols` columns of `A_PANEL` `i32` (stride `A_PANEL`); `state` -/// writable for `A_PANEL` `f32`; `q_meta.len() >= A_PANEL`; `d_meta.len() >= cols`. +/// If `state` or `q_meta` is not one A-panel, or `d_meta` or the strip holds fewer than +/// `cols` columns. #[inline(always)] -pub(super) unsafe fn score_fold_strip( +pub(super) fn score_fold_strip( arch: V3, - acc: *const i32, - state: *mut f32, + strip: &Strip<'_, i32, A_PANEL, B_PANEL>, + state: &mut [f32], cols: usize, q_meta: &[MinMaxCompensation], d_meta: &[MinMaxCompensation], dim: f32, ) { let lanes = f32s::LANES; + // The slicing is the bounds check: every load below is inside these lengths. + let acc = strip.columns(cols); + let d_meta = &d_meta[..cols]; + assert_eq!(state.len(), A_PANEL, "the fold writes a whole A-panel"); + assert_eq!(q_meta.len(), A_PANEL, "one compensation per A-panel row"); + let (state, acc) = (state.as_mut_ptr(), acc.as_ptr()); let mut qa = [0.0f32; A_PANEL]; let mut qb = [0.0f32; A_PANEL]; let mut qn = [0.0f32; A_PANEL]; - for i in 0..A_PANEL { - let qm = q_meta[i]; + for (i, qm) in q_meta.iter().enumerate() { qa[i] = qm.a; qb[i] = qm.b; qn[i] = qm.n; } - // SAFETY: each array holds exactly A_PANEL = 2·LANES f32; `state` writable for A_PANEL. + // SAFETY: each array holds exactly A_PANEL = 2·LANES f32, and `state` is A_PANEL + // writable f32, so both halves of each are in bounds. let (qa0, qa1, qb0, qb1, qn0, qn1, mut m0, mut m1) = unsafe { ( f32s::load_simd(arch, qa.as_ptr()), @@ -200,12 +250,13 @@ pub(super) unsafe fn score_fold_strip( ) }; - for (c, dm) in d_meta.iter().enumerate().take(cols) { + for (c, dm) in d_meta.iter().enumerate() { let a_c = f32s::splat(arch, dm.a); let b_c = f32s::splat(arch, dm.b); let c_c = f32s::splat(arch, dm.n + dm.b * dim); let col = c * A_PANEL; - // SAFETY: `col + 2·LANES <= cols*A_PANEL`; `acc` valid for that many i32. + // SAFETY: `acc` is `cols * A_PANEL` readable i32 and `c < cols`, so the column + // and both its halves are in bounds. unsafe { let raw0 = i32s::load_simd(arch, acc.add(col)).simd_cast(); let raw1 = i32s::load_simd(arch, acc.add(col + lanes)).simd_cast(); @@ -216,7 +267,7 @@ pub(super) unsafe fn score_fold_strip( } } - // SAFETY: `state` writable for A_PANEL f32. + // SAFETY: `state` is A_PANEL writable f32, as loaded above. unsafe { m0.store_simd(state); m1.store_simd(state.add(lanes)); diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs index a1da139e0..6c5fbf9ac 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/minmax.rs @@ -2,7 +2,7 @@ // Licensed under the MIT license. //! 4-bit MinMax instantiation. The interesting half is [`MinMaxMax`]: dequant needs -//! per-vector metadata indexed by [`At`] and the reduction needs the dequantized +//! per-vector metadata indexed by [`Region`] and the reduction needs the dequantized //! score, so both ride in one [`Drain`] and the score never reaches memory. use core::mem::size_of; @@ -14,7 +14,7 @@ use diskann_wide::arch::x86_64::V3; use super::arena::ResettableArena; use super::leaves::{A_PANEL, B_PANEL}; use super::views::{DPanel, DocWalk, QPanel, QueryWalk}; -use super::{Accumulate, At, Block, Drain, Plan, Strip, StripRef, TileBudget, drive, leaves}; +use super::{Accumulate, Block, Drain, Plan, Region, Strip, TileBudget, drive, leaves}; use crate::CompressInto; use crate::algorithms::Transform; use crate::algorithms::transforms::NullTransform; @@ -33,7 +33,7 @@ impl<'a, 'b, 'x> V3, QPanel<'a, i16, A_PANEL>, DPanel<'b, u8, B_PANEL, Static>, - Block<'x, i32, A_PANEL, B_PANEL, Static>, + Block<'x, i32, A_PANEL, B_PANEL>, > for I8Kernel { #[inline(always)] @@ -42,19 +42,9 @@ impl<'a, 'b, 'x> arch: V3, a: QPanel<'a, i16, A_PANEL>, b: DPanel<'b, u8, B_PANEL, Static>, - mut out: Block<'x, i32, A_PANEL, B_PANEL, Static>, + out: Block<'x, i32, A_PANEL, B_PANEL>, ) { - // SAFETY: `a` is an A_PANEL×k block-transposed i16 block; `b` is B_PANEL rows - // of k u8; `out` is B_PANEL columns of A_PANEL i32 at stride A_PANEL (`k` even). - unsafe { - leaves::int_store_microkernel::( - arch, - a.as_ptr(), - b.as_ptr(), - a.k(), - out.as_mut_ptr(), - ); - } + leaves::int_store_microkernel::(arch, a, b, out); } } @@ -63,7 +53,7 @@ impl<'a, 'b, 'x> V3, QPanel<'a, i16, A_PANEL>, DPanel<'b, u8, B_PANEL, Dynamic>, - Block<'x, i32, A_PANEL, B_PANEL, Dynamic>, + Block<'x, i32, A_PANEL, B_PANEL>, > for I8Kernel { #[inline(always)] @@ -72,19 +62,14 @@ impl<'a, 'b, 'x> arch: V3, a: QPanel<'a, i16, A_PANEL>, b: DPanel<'b, u8, B_PANEL, Dynamic>, - mut out: Block<'x, i32, A_PANEL, B_PANEL, Dynamic>, + out: Block<'x, i32, A_PANEL, B_PANEL>, ) { - debug_assert_eq!(out.cols(), b.rows()); - debug_assert!(b.rows() < B_PANEL); - let (ap, bp, op, k) = (a.as_ptr(), b.as_ptr(), out.as_mut_ptr(), a.k()); - // SAFETY: as the full-width impl, with a runtime width in 1..B_PANEL. - unsafe { - match b.rows() { - 3 => leaves::int_store_microkernel::<3>(arch, ap, bp, k, op), - 2 => leaves::int_store_microkernel::<2>(arch, ap, bp, k, op), - 1 => leaves::int_store_microkernel::<1>(arch, ap, bp, k, op), - other => unreachable!("tail width {other} out of 1..{B_PANEL}"), - } + // The leaf checks that the width it unrolls for is the width `b` actually has. + match b.rows() { + 3 => leaves::int_store_microkernel::<3, _>(arch, a, b, out), + 2 => leaves::int_store_microkernel::<2, _>(arch, a, b, out), + 1 => leaves::int_store_microkernel::<1, _>(arch, a, b, out), + other => unreachable!("tail width {other} out of 1..{B_PANEL}"), } } } @@ -93,7 +78,7 @@ impl<'a, 'b, 'x> /// Fused 4-bit MinMax dequant + running max: rewrites each raw integer dot into the /// MinMax inner product using per-vector `a`/`b`/`n` metadata, then folds it straight -/// into the output it owns. +/// into the output — the score never reaches memory. pub(crate) struct MinMaxMax<'m, 'o> { query_meta: &'m [MinMaxCompensation], doc_meta: &'m [MinMaxCompensation], @@ -120,18 +105,15 @@ impl<'m, 'o> MinMaxMax<'m, 'o> { impl Drain> for MinMaxMax<'_, '_> { #[inline(always)] - fn drain(&mut self, arch: V3, acc: StripRef<'_, i32, A_PANEL>, at: At) { - let cols = acc.cols(); - let lo = at.a_panel * A_PANEL; - let q = &self.query_meta[lo..lo + A_PANEL]; - let d = &self.doc_meta[at.b_row..at.b_row + cols]; + fn drain(&mut self, arch: V3, acc: &Strip<'_, i32, A_PANEL, B_PANEL>, region: Region) { + let (a, b) = (region.a, region.b); + // A-indexed buffers here are padded to whole panels, so the stride is ours to + // state — which is what makes the leaf's one-compensation-per-row check hold. + let q = &self.query_meta[a.start..][..A_PANEL]; + let d = &self.doc_meta[b.range()]; let dim = self.dim; - let out = &mut self.out[lo..][..A_PANEL]; - // SAFETY: `acc` is `cols` columns of A_PANEL i32; `out` is A_PANEL writable - // f32; `q.len() == A_PANEL`; `d.len() == cols`. - unsafe { - leaves::score_fold_strip(arch, acc.as_ptr(), out.as_mut_ptr(), cols, q, d, dim); - } + let out = &mut self.out[a.start..][..A_PANEL]; + leaves::score_fold_strip(arch, acc, out, b.len(), q, d, dim); } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs index 642348e27..de4b83ba0 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/mod.rs @@ -7,13 +7,14 @@ //! into output it owns. [`Scratch`] is the write-side mirror of [`Paneled`], so the //! driver assumes no layout on either side. //! -//! Panel widths are geometry, so they live with the leaves that impose them -//! ([`leaves::A_PANEL`], [`leaves::B_PANEL`]); the panel and accumulator types carry -//! them as const parameters (`R` = A rows, `N` = B rows) so the driver stays width- -//! agnostic. Instantiated for f32 ([`float`]) and 4-bit MinMax ([`minmax`]). +//! Panel widths are geometry, so they live with the leaves that impose them; the panel +//! and accumulator types carry them as const parameters (`R` = A rows, `N` = B rows) so +//! the driver stays width-agnostic. //! //! Sibling to [`tiler`](super::tiler), which keeps postprocess and reduce separate. +use core::ops::Range; + use super::TileBudget; mod arena; @@ -23,7 +24,7 @@ mod minmax; mod strip; mod views; -pub(crate) use strip::{Block, Strip, StripRef}; +pub(crate) use strip::{Block, Strip}; pub use float::{PaneledF32Docs, PaneledF32Query}; pub use minmax::{PaneledQuantDocs, PaneledQuantQuery}; @@ -58,68 +59,63 @@ impl Plan { // ── Accumulator ────────────────────────────────────────────────── -/// Per-lifetime half of [`Scratch`] — same sealed-`Bounds` trick as [`TileAt`]. -pub(crate) trait ScratchAt<'a, B: sealed::Sealed = sealed::Bounds<&'a mut Self>> { +/// Per-lifetime half of [`Scratch`] — same implied-bound trick as [`TileAt`]. +pub(crate) trait ScratchAt<'a, B = &'a mut Self> { + /// One B-panel's slot. type Block; - /// Short trailing slot; a distinct type so it selects its own [`Accumulate`] impl. - type Short; - /// Named so `Block`/`Short` reach bounds without a nested projection. - type Blocks: TailIterator; - /// What [`Drain`] reads. - type Ref; + /// Named so `Block` reaches bounds without a nested projection. + type Slots: Slots; +} + +/// Hands out one accumulator slot per B-panel, each disjoint from the last. +/// +/// Infallible: a scratch is sized for the widest tile it will ever see, so "ran out" is +/// not a state the driver can reach, and the fill loop carries no `Option`. Overdrawing +/// is a planning bug; implementations must fail loudly rather than alias. +pub(crate) trait Slots { + type Block; + fn next(&mut self) -> Self::Block; } /// [`Paneled`]'s write side: a buffer that carves itself into per-B-panel slots, so /// the driver can't tell a contiguous strip from a padded or structure-of-arrays one. /// Not [`Paneled`] itself — `Panel: Copy` and `panels(&self)` can't yield disjoint -/// `&mut`, and the carve needs a runtime `cols` (a scratch is sized to capacity). -/// Allocation stays on the concrete type. +/// `&mut`. Allocation stays on the concrete type. +/// +/// Slots are uniform: a short B-tail writes a prefix of a full-width slot, and the +/// [`Drain`] is told how much is live, so no second width has to be threaded through. pub(crate) trait Scratch: for<'a> ScratchAt<'a> { - /// Carve the live `cols` columns into one slot per B-panel plus the short trailer, - /// which comes off the same cursor and so is provably disjoint. - fn blocks(&mut self, cols: usize) -> >::Blocks; - - fn as_ref(&self, cols: usize) -> >::Ref; + fn slots(&mut self) -> >::Slots; } // ── Data side ──────────────────────────────────────────────────── -/// Misuse guard for the implicit-bounds parameter: private, so no downstream impl can -/// override the default with a type that drops the implied bound. -mod sealed { - pub trait Sealed {} - pub struct Bounds(#[allow(dead_code)] T); - impl Sealed for Bounds {} -} - -/// Per-lifetime half of [`TileWalk`]. The defaulted `B = Bounds<&'a Self>` carries the +/// Per-lifetime half of [`TileWalk`]. The defaulted `B = &'a Self` carries the /// `Self: 'a` implied bound through well-formedness — a plain GAT `where Self: 'a` /// collapses to `'static` under the driver's `for<'a>` bound on stable. -pub(crate) trait TileAt<'a, B: sealed::Sealed = sealed::Bounds<&'a Self>> { +pub(crate) trait TileAt<'a, B = &'a Self> { type View: Paneled; } /// A **lending** walk: `next` reborrows `&mut self`, so a view may borrow a buffer the /// walk reuses on the following call. `reset` rewinds — B is re-walked once per A-tile. pub(crate) trait TileWalk: for<'a> TileAt<'a> { - fn next(&mut self) -> Option>::View>>; + fn next(&mut self) -> Option<>::View>; fn reset(&mut self); } -/// A lent view plus where it starts. Lifetime-free — the borrow lives on `V`. -pub(crate) struct Tile { - pub(crate) view: V, - /// Position in the walk's own unit: A-panels for a query walk, B-rows for a doc - /// walk. Only the [`Drain`] turns it into an output index. - pub(crate) at: usize, -} - /// An iterator whose short trailing element has its own type. `tail` consumes the /// exhausted iterator, so the trailer comes off the cursor the loop was already /// advancing instead of being recomputed from the source. +/// +/// Each element carries its own [`Geo`], because the view is what knows how it is laid +/// out; the driver only forwards what it is handed. +/// +/// `ExactSizeIterator` binds implementors, not the driver: a view that cannot state its +/// panel count exactly does not know its own geometry. pub(crate) trait TailIterator: ExactSizeIterator { type Tail; - fn tail(self) -> Option; + fn tail(self) -> Option<(Self::Tail, Geo)>; } /// A view that knows how it breaks into panels. `Tail` is distinct from `Panel` so the @@ -130,9 +126,12 @@ pub(crate) trait Paneled { type Tail: Copy; /// Named so it carries `ExactSizeIterator` and the tail type into bounds, and so a /// k-fracturing driver could hold one across an outer loop. - type Panels: TailIterator; + type Panels: TailIterator; + + /// Where this view sits in the global problem. A view is a sub-view of something, + /// so it is the only thing that can say where it came from. + fn geo(&self) -> Geo; - fn rows(&self) -> usize; fn panels(&self) -> Self::Panels; } @@ -158,22 +157,48 @@ impl Accumulate for K { } } -/// Where a finished accumulator sits in the global problem. The driver counts panels -/// and never converts a count into a row — only the [`Drain`] owns `R`. B is a row -/// rather than a panel index because a tile is not panel-quantized. +/// A contiguous run of **real** vectors, `[start, end)`, in the global problem's +/// numbering. Purely logical: a view that pads reports what is real, and a drain that +/// pads owns that itself — so no consumer has to know anyone else's padding rule. +/// +/// This is the driver's whole vocabulary for position, so a walk must yield vectors +/// that are contiguous and monotone. A gathering or permuting walk cannot be described +/// this way; it needs an id source, which is a different trait rather than a wider +/// struct. #[derive(Clone, Copy)] -pub(crate) struct At { - pub a_panel: usize, - pub b_row: usize, +pub(crate) struct Geo { + pub(crate) start: usize, + pub(crate) end: usize, +} + +impl Geo { + /// Real vectors covered — the reduction's live width. + pub(crate) fn len(self) -> usize { + self.end - self.start + } + + /// For indexing a buffer held in the same numbering. + pub(crate) fn range(self) -> Range { + self.start..self.end + } +} + +/// Where a finished accumulator sits in the global problem. Both sides mean the same +/// thing — real vectors — whatever granularity the driver happens to drain them at. +#[derive(Clone, Copy)] +pub(crate) struct Region { + pub(crate) a: Geo, + pub(crate) b: Geo, } /// Consume a finished accumulator. The drain owns its output, so dequant, reduction /// and scatter all live behind this one call and may be fused. /// -/// Implementations must initialize their output to the reduction's identity, and must -/// clamp their writes when the output is not padded to whole panels. +/// `region` also carries the live extent: the accumulator is sized for the widest tile, +/// so a drain that folded its capacity instead of `region.b.len()` would fold the +/// *previous* tile's values back in. pub(crate) trait Drain { - fn drain(&mut self, arch: Arch, acc: >::Ref, at: At); + fn drain(&mut self, arch: Arch, acc: &S, region: Region); } // ── Driver ─────────────────────────────────────────────────────── @@ -181,33 +206,28 @@ pub(crate) trait Drain { type PanelOf<'a, W> = <>::View as Paneled>::Panel; type TailOf<'a, W> = <>::View as Paneled>::Tail; type BlockOf<'x, S> = >::Block; -type ShortOf<'x, S> = >::Short; /// One A-panel against a whole B-tile. Factored out so the driver's A-panel and /// A-tail arms share the tail-dispatch. #[inline(always)] -fn fill(arch: Arch, kernel: &K, a: A, b_view: &BV, scratch: &mut S, cols: usize) +fn fill(arch: Arch, kernel: &K, a: A, b_view: &BV, scratch: &mut S) where Arch: Copy, A: Copy, BV: Paneled, S: Scratch, K: for<'x> Accumulate> - + for<'x> Accumulate>, + + for<'x> Accumulate>, { let mut panels = b_view.panels(); - let mut blocks = scratch.blocks(cols); - // What [`TailIterator`]'s `ExactSizeIterator` bound is for: `zip` stops on the - // shorter side and silently drops the longer side's pending item, which would leave - // both cursors at zero and slip past the tail check below. - debug_assert_eq!(panels.len(), blocks.len(), "B view and accumulator desync"); - for (b, out) in panels.by_ref().zip(blocks.by_ref()) { - kernel.accumulate(arch, a, b, out); + let mut slots = scratch.slots(); + // The geos are the drain's business; this side only feeds the kernel. + for (b, _) in panels.by_ref() { + kernel.accumulate(arch, a, b, slots.next()); } - match (panels.tail(), blocks.tail()) { - (Some(b), Some(out)) => kernel.accumulate(arch, a, b, out), - (None, None) => {} - _ => unreachable!("B view and accumulator disagree on tail"), + // The tail draws from the same cursor as the full panels. + if let Some((b, _)) = panels.tail() { + kernel.accumulate(arch, a, b, slots.next()); } } @@ -230,30 +250,66 @@ pub(super) fn drive( BW: TileWalk, S: Scratch, K: for<'a, 'b, 'x> Accumulate, PanelOf<'b, BW>, BlockOf<'x, S>> - + for<'a, 'b, 'x> Accumulate, TailOf<'b, BW>, ShortOf<'x, S>> + + for<'a, 'b, 'x> Accumulate, TailOf<'b, BW>, BlockOf<'x, S>> + for<'a, 'b, 'x> Accumulate, PanelOf<'b, BW>, BlockOf<'x, S>> - + for<'a, 'b, 'x> Accumulate, TailOf<'b, BW>, ShortOf<'x, S>>, + + for<'a, 'b, 'x> Accumulate, TailOf<'b, BW>, BlockOf<'x, S>>, D: Drain, { - while let Some(a_tile) = a_walk.next() { + while let Some(a_view) = a_walk.next() { b_walk.reset(); - while let Some(b_tile) = b_walk.next() { - let cols = b_tile.view.rows(); - let mut at = At { - a_panel: a_tile.at, - b_row: b_tile.at, - }; - - let mut a_panels = a_tile.view.panels(); - for a in a_panels.by_ref() { - fill(arch, kernel, a, &b_tile.view, scratch, cols); - drain.drain(arch, scratch.as_ref(cols), at); - at.a_panel += 1; + while let Some(b_view) = b_walk.next() { + let b = b_view.geo(); + let mut a_panels = a_view.panels(); + for (panel, a) in a_panels.by_ref() { + fill(arch, kernel, panel, &b_view, scratch); + drain.drain(arch, scratch, Region { a, b }); } - if let Some(a) = a_panels.tail() { - fill(arch, kernel, a, &b_tile.view, scratch, cols); - drain.drain(arch, scratch.as_ref(cols), at); + if let Some((panel, a)) = a_panels.tail() { + fill(arch, kernel, panel, &b_view, scratch); + drain.drain(arch, scratch, Region { a, b }); } } } } + +#[cfg(test)] +mod tests { + use core::mem::size_of; + + use super::leaves::{A_PANEL, B_PANEL}; + use super::views::{DPanel, QPanel}; + use super::{Block, Strip}; + use crate::bits::{Dynamic, Static}; + + /// Handles stay thin on both sides: the geometry a leaf needs is const parameters, + /// so only what is genuinely runtime — the contraction length, and a tail's row + /// count — is ever stored. + /// + /// A guard, not a curiosity: each side lost this once already to a refactor that + /// swapped a pointer for a slice to simplify a cursor, and neither showed up in a + /// test. The choice is a type-system one — an A/B of two builds cannot resolve a + /// difference this small, so diff the emitted asm rather than re-benchmarking. + #[test] + fn handles_stay_thin() { + let word = size_of::<*const u8>(); + + // Both panels carry `k`: the price of each stating its own extent rather than + // borrowing the other's, which is what lets the leaves be safe. + assert_eq!(size_of::>(), 2 * word); + assert_eq!( + size_of::>>(), + 2 * word + ); + + // Only the trailing panel pays for a runtime row count: `Static` is a ZST. + assert_eq!( + size_of::>(), + 3 * word + ); + + // The slot's extent is in the type; the strip is the checked anchor it is + // carved from, so only the strip keeps a length. + assert_eq!(size_of::>(), word); + assert_eq!(size_of::>(), 2 * word); + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs index 899bc7ee3..faf8a44ce 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/strip.rs @@ -7,91 +7,52 @@ use core::marker::PhantomData; use core::mem::MaybeUninit; -use super::{Scratch, ScratchAt, TailIterator}; +use super::{Scratch, ScratchAt}; use crate::alloc::{AllocatorCore, Poly}; -use crate::bits::{Dynamic, Length, Static}; /// Marker for element types where all-zero is a valid value. pub(crate) trait ZeroInit: Copy {} impl ZeroInit for i32 {} impl ZeroInit for f32 {} -/// Owns the accumulator's live region as a checked slice — the anchor the unchecked -/// [`Block`]s below are derived from. +/// The accumulator buffer, sized to the widest B-tile the plan allows. pub(crate) struct Strip<'a, T, const R: usize, const N: usize>(&'a mut [T]); -/// One B-panel's slot: `R` rows by `L` columns, where `L.value() <= N`. A full slot is -/// `Static` (a ZST, so the whole handle is one pointer); the trailing slot is -/// `Dynamic` in `1..N`, and its distinct type selects the short -/// [`Accumulate`](super::Accumulate) impl. -pub(crate) struct Block<'a, T, const R: usize, const N: usize, L: Length> { - ptr: *mut T, - cols: L, - _lifetime: PhantomData<&'a mut [T]>, -} - -/// A finished [`Strip`], narrowed to its live columns. -pub(crate) struct StripRef<'a, T, const R: usize>(&'a [T]); - -/// Carves a [`Strip`] into per-B-panel slots. The trailer comes off the same cursor. +/// One B-panel's slot: `R` rows by `N` columns, A-major. /// -/// # Safety invariants +/// Slots are packed, so column `c` of the strip sits at `c * R` whichever slot it falls +/// in — the live columns stay one contiguous run, which is what lets a +/// [`Drain`](super::Drain) fold a whole tile in one pass. /// -/// `ptr` is valid for writes of `R * (full * N + tail_cols)` `T` for `'a`, and only -/// ever advances — so every slot handed out is disjoint from every other. -pub(crate) struct BlockIter<'a, T, const R: usize, const N: usize> { +/// Thin: `R * N` is already in the type, and the `split_at_mut` that carves the slot is +/// where the length still means something, so the handle keeps only the pointer. +pub(crate) struct Block<'a, T, const R: usize, const N: usize> { ptr: *mut T, - full: usize, - tail_cols: usize, _lifetime: PhantomData<&'a mut [T]>, } -impl<'a, T, const R: usize, const N: usize, L: Length> Block<'a, T, R, N, L> { - /// # Safety - /// - /// `ptr` must be valid for writes of `R * cols.value()` `T` for `'a`, must not - /// alias any other live `Block`, and `cols.value()` must be at most `N`. - unsafe fn new(ptr: *mut T, cols: L) -> Self { - debug_assert!(cols.value() <= N, "block wider than its panel"); - Self { - ptr, - cols, - _lifetime: PhantomData, - } - } -} +/// Splits slots off the front of a [`Strip`], so disjointness is the borrow checker's +/// job rather than an invariant to uphold by hand. +pub(crate) struct BlockSlots<'a, T, const R: usize, const N: usize>(&'a mut [T]); -impl<'a, T, const R: usize, const N: usize> Iterator for BlockIter<'a, T, R, N> { - type Item = Block<'a, T, R, N, Static>; +impl<'a, T, const R: usize, const N: usize> super::Slots for BlockSlots<'a, T, R, N> { + type Block = Block<'a, T, R, N>; - fn next(&mut self) -> Option { - if self.full == 0 { - return None; + /// # Panics + /// + /// If the strip holds fewer slots than the B-tile has panels — a planning bug, and + /// the reason this can be infallible everywhere else. + fn next(&mut self) -> Block<'a, T, R, N> { + // `take` re-lends the buffer for `'a` rather than the `&mut self` borrow, which + // is what lets a slot outlive the cursor call that produced it. + let (slot, rest) = core::mem::take(&mut self.0).split_at_mut(R * N); + self.0 = rest; + // `split_at_mut` is what proves the slots disjoint, so the pointer inherits that + // guarantee rather than resting on an invariant upheld by hand. + Block { + ptr: slot.as_mut_ptr(), + _lifetime: PhantomData, } - let ptr = self.ptr; - // SAFETY: the invariant covers `full` more slots of `R * N`, so the bump stays - // inside the allocation and the yielded slot is disjoint from all later ones. - self.ptr = unsafe { self.ptr.add(R * N) }; - self.full -= 1; - // SAFETY: as above — `ptr` covers exactly `R * N` writable `T`. - Some(unsafe { Block::new(ptr, Static) }) - } - - fn size_hint(&self) -> (usize, Option) { - (self.full, Some(self.full)) - } -} - -impl ExactSizeIterator for BlockIter<'_, T, R, N> {} - -impl<'a, T, const R: usize, const N: usize> TailIterator for BlockIter<'a, T, R, N> { - type Tail = Block<'a, T, R, N, Dynamic>; - - fn tail(self) -> Option { - debug_assert_eq!(self.full, 0, "tail taken before the blocks are exhausted"); - // SAFETY: the blocks are exhausted, so the cursor sits on the trailer, which - // the invariant covers for `R * tail_cols` writable `T`. - (self.tail_cols > 0).then(|| unsafe { Block::new(self.ptr, Dynamic(self.tail_cols)) }) } } @@ -112,54 +73,33 @@ impl<'a, T: ZeroInit, const R: usize, const N: usize> Strip<'a, T, R, N> { } } -impl Strip<'_, T, R, N> { - fn cols_capacity(&self) -> usize { - self.0.len() / R - } -} - impl<'a, T, const R: usize, const N: usize> ScratchAt<'a> for Strip<'_, T, R, N> { - type Block = Block<'a, T, R, N, Static>; - type Short = Block<'a, T, R, N, Dynamic>; - type Blocks = BlockIter<'a, T, R, N>; - type Ref = StripRef<'a, T, R>; + type Block = Block<'a, T, R, N>; + type Slots = BlockSlots<'a, T, R, N>; } impl Scratch for Strip<'_, T, R, N> { - fn blocks(&mut self, cols: usize) -> BlockIter<'_, T, R, N> { - debug_assert!( - cols <= self.cols_capacity(), - "strip must hold the whole B-tile" - ); - // The checked slice is what establishes `BlockIter`'s invariant: `R * cols` - // elements are live, and the `&mut self` borrow keeps them exclusive for `'_`. - BlockIter { - ptr: self.0.as_mut_ptr(), - full: cols / N, - tail_cols: cols % N, - _lifetime: PhantomData, - } - } - - fn as_ref(&self, cols: usize) -> StripRef<'_, T, R> { - StripRef(&self.0[..R * cols]) + fn slots(&mut self) -> BlockSlots<'_, T, R, N> { + BlockSlots(&mut *self.0) } } -impl Block<'_, T, R, N, L> { - pub(crate) fn cols(&self) -> usize { - self.cols.value() - } +impl Block<'_, T, R, N> { pub(crate) fn as_mut_ptr(&mut self) -> *mut T { self.ptr } } -impl StripRef<'_, T, R> { - pub(crate) fn cols(&self) -> usize { - self.0.len() / R - } - pub(crate) fn as_ptr(&self) -> *const T { - self.0.as_ptr() +impl Strip<'_, T, R, N> { + /// The live prefix, `cols` columns of `R`. The rest is capacity left over from a + /// wider tile, so the caller states what it wants and gets a bounds check rather + /// than a promise. + /// + /// # Panics + /// + /// If `cols * R` exceeds the strip — a planning bug, since the strip is sized to the + /// widest B-tile the plan allows. + pub(crate) fn columns(&self, cols: usize) -> &[T] { + &self.0[..cols * R] } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs b/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs index f40bae86f..0f39d0ca6 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/paneled/views.rs @@ -1,14 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. -//! The two views and the walks that lend them. Both views are the real matrix types — -//! [`BlockTransposedRef`] for A, [`MatRef`] for B behind a [`Rows`] adapter (a -//! row-major matrix doesn't imply a panel height). Each sub-views itself, so a walk is -//! a cursor and nothing else. +//! The two views and the walks that lend them. The views are the real matrix types +//! rather than wrappers around them — the B side needs one adapter, [`RowPanels`], only +//! because a row-major matrix doesn't imply a panel height. Each sub-views itself, so a +//! walk is a cursor and nothing else. use core::marker::PhantomData; -use super::{NoTail, Paneled, TailIterator, Tile, TileAt, TileWalk}; +use super::{Geo, NoTail, Paneled, TailIterator, TileAt, TileWalk}; use crate::bits::{Dynamic, Length, Static}; use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; @@ -21,11 +21,15 @@ pub(crate) struct QPanel<'a, T, const R: usize> { _lifetime: PhantomData<&'a [T]>, } -/// One row-major B panel: `L` rows × `k` `T`, where `L.value() <= N`. `k` travels with -/// the A panel, so it is not stored. A full panel is `Static` — a ZST, so the whole -/// handle is one pointer; the trailing panel is `Dynamic` in `1..N`. +/// One row-major B panel: `L` rows × `k` `T`, where `L.value() <= N`. A full panel is +/// `Static` — a ZST; only the trailing panel is `Dynamic`, in `1..N`. +/// +/// `k` is carried rather than borrowed from the A panel it pairs with: a handle whose +/// validity rests on a field in some *other* value cannot be checked where it is used, +/// which is what would force the leaves to be `unsafe`. pub(crate) struct DPanel<'a, T, const N: usize, L: Length> { ptr: *const T, + k: usize, rows: L, _lifetime: PhantomData<&'a [T]>, } @@ -59,13 +63,13 @@ impl<'a, T, const R: usize> QPanel<'a, T, R> { impl<'a, T, const N: usize, L: Length> DPanel<'a, T, N, L> { /// # Safety /// - /// `ptr` must be valid for reads of `rows.value() * k` `T` for `'a`, where `k` is - /// the contraction length carried by the A panel it is paired with, and + /// `ptr` must be valid for reads of `rows.value() * k` `T` for `'a`, and /// `rows.value()` must be at most `N`. - unsafe fn new(ptr: *const T, rows: L) -> Self { + unsafe fn new(ptr: *const T, k: usize, rows: L) -> Self { debug_assert!(rows.value() <= N, "panel taller than its type claims"); Self { ptr, + k, rows, _lifetime: PhantomData, } @@ -84,6 +88,9 @@ impl DPanel<'_, T, N, L> { pub(crate) fn as_ptr(&self) -> *const T { self.ptr } + pub(crate) fn k(&self) -> usize { + self.k + } pub(crate) fn rows(&self) -> usize { self.rows.value() } @@ -95,23 +102,32 @@ impl DPanel<'_, T, N, L> { /// full `R` rows, so there is no tail. pub(crate) struct QPanelIter<'a, T: Copy, const R: usize, const P: usize> { view: BlockTransposedRef<'a, T, R, P>, - k: usize, cur: usize, - end: usize, + base: Geo, } impl<'a, T: Copy, const R: usize, const P: usize> Iterator for QPanelIter<'a, T, R, P> { - type Item = QPanel<'a, T, R>; + type Item = (QPanel<'a, T, R>, Geo); fn next(&mut self) -> Option { let data = self.view.block_slice(self.cur)?; + // The trailing block is padded to a full `R`, so the real extent clamps to what + // the view says is real — the kernel still writes the whole panel. + let start = self.base.start + self.cur * R; + let geo = Geo { + start, + end: (start + R).min(self.base.end), + }; self.cur += 1; + let k = self.view.padded_ncols(); // SAFETY: `block_slice` returns a checked `R * k` slice borrowed from the view. - Some(unsafe { QPanel::new(data.as_ptr(), self.k) }) + Some((unsafe { QPanel::new(data.as_ptr(), k) }, geo)) } fn size_hint(&self) -> (usize, Option) { - let n = self.end - self.cur; + // `next` bails without bumping `cur` once `block_slice` says no, so `cur` stops + // at `num_blocks()` and this can't wrap. + let n = self.view.num_blocks() - self.cur; (n, Some(n)) } } @@ -120,7 +136,7 @@ impl ExactSizeIterator for QPanelIter<' impl TailIterator for QPanelIter<'_, T, R, P> { type Tail = NoTail; - fn tail(self) -> Option { + fn tail(self) -> Option<(NoTail, Geo)> { None } } @@ -128,36 +144,45 @@ impl TailIterator for QPanelIter<'_, T, /// Walks a row-major matrix' full `N`-row panels; [`TailIterator::tail`] hands back the /// short trailer from the same cursor. /// +/// `rest` is what has not been lent yet — the panel count and the trailer's height are +/// read off it. +/// /// # Safety invariants /// -/// `ptr` is valid for reads of `(full * N + tail_rows) * k` `T` for `'a`, and only ever -/// advances. +/// `ptr` is valid for reads of `rest.len() * k` `T` for `'a`. `ptr` and `rest.start` +/// advance together — `N` rows *is* `N * k` elements — so the two cannot drift apart. pub(crate) struct DPanelIter<'a, T, const N: usize> { ptr: *const T, k: usize, - full: usize, - tail_rows: usize, + rest: Geo, _lifetime: PhantomData<&'a [T]>, } impl<'a, T, const N: usize> Iterator for DPanelIter<'a, T, N> { - type Item = DPanel<'a, T, N, Static>; + type Item = (DPanel<'a, T, N, Static>, Geo); fn next(&mut self) -> Option { - if self.full == 0 { + if self.rest.len() < N { return None; } let ptr = self.ptr; - // SAFETY: the invariant covers `full` more panels of `N * k`, so the bump stays - // inside the allocation. + // SAFETY: `rest.len() >= N`, so the invariant covers another `N * k` and the + // bump stays inside the allocation. self.ptr = unsafe { self.ptr.add(N * self.k) }; - self.full -= 1; + // A full panel is exactly `N` real rows — the short trailer belongs to the tail, + // so nothing here clamps. + let geo = Geo { + start: self.rest.start, + end: self.rest.start + N, + }; + self.rest.start = geo.end; // SAFETY: as above — `ptr` covers exactly `N * k` readable `T`. - Some(unsafe { DPanel::new(ptr, Static) }) + Some((unsafe { DPanel::new(ptr, self.k, Static) }, geo)) } fn size_hint(&self) -> (usize, Option) { - (self.full, Some(self.full)) + let n = self.rest.len() / N; + (n, Some(n)) } } @@ -166,57 +191,80 @@ impl ExactSizeIterator for DPanelIter<'_, T, N> {} impl<'a, T, const N: usize> TailIterator for DPanelIter<'a, T, N> { type Tail = DPanel<'a, T, N, Dynamic>; - fn tail(self) -> Option { - debug_assert_eq!(self.full, 0, "tail taken before the panels are exhausted"); + fn tail(self) -> Option<(Self::Tail, Geo)> { + let rows = self.rest.len(); + debug_assert!(rows < N, "tail taken before the panels are exhausted"); + if rows == 0 { + return None; + } // SAFETY: the panels are exhausted, so the cursor sits on the trailer, which - // the invariant covers for `tail_rows * k` readable `T`. - (self.tail_rows > 0).then(|| unsafe { DPanel::new(self.ptr, Dynamic(self.tail_rows)) }) + // the invariant covers for `rows * k` readable `T`. + let panel = unsafe { DPanel::new(self.ptr, self.k, Dynamic(rows)) }; + // What is left over *is* the trailer, so the cursor's own extent is its geo. + Some((panel, self.rest)) } } // ── Views ──────────────────────────────────────────────────────── -/// A block-transposed matrix' blocks *are* the A-panels, so the real type is the view. -/// The remainder block is zero-padded to a full `R` rows, hence [`NoTail`]; rows past -/// `nrows()` score against padding and are dropped by the caller. `P` only widens -/// `padded_ncols`, the contraction length the kernel sees. -impl<'a, T: Copy, const R: usize, const P: usize> Paneled for BlockTransposedRef<'a, T, R, P> { +/// A block-transposed sub-view plus where it came from. `block_range` hands back a view +/// that has lost its origin, so the walk pairs it back up here — the same job +/// [`RowPanels`] does on the B side. +/// +/// The remainder block is zero-padded to a full `R` rows, hence [`NoTail`]; the padding +/// rows are excluded from the [`Geo`]. `P` only widens `padded_ncols`, the contraction +/// length the kernel sees. +pub(crate) struct BlockPanels<'a, T: Copy, const R: usize, const P: usize> { + view: BlockTransposedRef<'a, T, R, P>, + start: usize, +} + +impl<'a, T: Copy, const R: usize, const P: usize> Paneled for BlockPanels<'a, T, R, P> { type Panel = QPanel<'a, T, R>; type Tail = NoTail; type Panels = QPanelIter<'a, T, R, P>; - fn rows(&self) -> usize { - self.nrows() + fn geo(&self) -> Geo { + Geo { + start: self.start, + end: self.start + self.view.nrows(), + } } + fn panels(&self) -> QPanelIter<'a, T, R, P> { QPanelIter { - view: *self, - k: self.padded_ncols(), + view: self.view, cur: 0, - end: self.num_blocks(), + base: self.geo(), } } } -/// Cut a row-major matrix into `N`-row panels. All the geometry stays on the matrix. -pub(crate) struct Rows(pub(crate) V); +/// Cut a row-major matrix into `N`-row panels. All the geometry stays on the matrix; +/// only the origin, which sub-viewing drops, is carried alongside. +pub(crate) struct RowPanels { + view: V, + start: usize, +} -impl<'a, const N: usize, T: Copy> Paneled for Rows>> { +impl<'a, const N: usize, T: Copy> Paneled for RowPanels>> { type Panel = DPanel<'a, T, N, Static>; type Tail = DPanel<'a, T, N, Dynamic>; type Panels = DPanelIter<'a, T, N>; - fn rows(&self) -> usize { - self.0.num_vectors() + fn geo(&self) -> Geo { + Geo { + start: self.start, + end: self.start + self.view.num_vectors(), + } } + fn panels(&self) -> DPanelIter<'a, T, N> { - let (n, k) = (self.rows(), self.0.vector_dim()); // The checked slice is what establishes `DPanelIter`'s invariant. DPanelIter { - ptr: self.0.as_slice().as_ptr(), - k, - full: n / N, - tail_rows: n % N, + ptr: self.view.as_slice().as_ptr(), + k: self.view.vector_dim(), + rest: self.geo(), _lifetime: PhantomData, } } @@ -258,14 +306,14 @@ impl<'s, T: Copy, const N: usize> DocWalk<'s, T, N> { } impl<'a, T: Copy, const R: usize, const P: usize> TileAt<'a> for QueryWalk<'_, T, R, P> { - type View = BlockTransposedRef<'a, T, R, P>; + type View = BlockPanels<'a, T, R, P>; } impl TileWalk for QueryWalk<'_, T, R, P> { - fn next(&mut self) -> Option>> { + fn next(&mut self) -> Option> { let view = self.src.block_range(self.cur, self.tile_panels)?; - let at = self.cur; + let start = self.cur * R; self.cur += view.num_blocks(); - Some(Tile { view, at }) + Some(BlockPanels { view, start }) } fn reset(&mut self) { self.cur = 0; @@ -273,17 +321,14 @@ impl TileWalk for QueryWalk<'_, T, R, P } impl<'a, T: Copy, const N: usize> TileAt<'a> for DocWalk<'_, T, N> { - type View = Rows>>; + type View = RowPanels>>; } impl TileWalk for DocWalk<'_, T, N> { - fn next(&mut self) -> Option>>>> { + fn next(&mut self) -> Option>>> { let view = self.src.row_range(self.cur, self.tile_panels * N)?; - let at = self.cur; + let start = self.cur; self.cur += view.num_vectors(); - Some(Tile { - view: Rows(view), - at, - }) + Some(RowPanels { view, start }) } fn reset(&mut self) { self.cur = 0;