From a1db87c5492cba3f00e633f7e260abf0b25117f5 Mon Sep 17 00:00:00 2001 From: Junkui Chen Date: Tue, 13 Oct 2026 19:08:26 +0800 Subject: [PATCH 1/2] Refactor tiled MinMax MaxSim and integrate benchmarks Prepare MinMax queries and synchronized panels using the existing matrix-kernel drivers. Add Scalar, V3, V4, and Neon execution paths and reuse the V3 maddubs wrapper. Integrate MinMax8/MinMax4 with the shared multi-vector benchmark, including input fixtures, checked allocation, regression checks, and CLI coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- diskann-benchmark/README.md | 14 +- diskann-benchmark/example/multi-vector.json | 30 + .../multi-vector-minmax8-neon.json | 47 + .../multi-vector-minmax8-v4.json | 47 + .../multi-vector-minmax8.json | 47 + diskann-benchmark/src/inputs/multi_vector.rs | 31 +- diskann-benchmark/src/main.rs | 178 ++++ diskann-benchmark/src/multi_vector/driver.rs | 88 +- diskann-benchmark/src/multi_vector/kernels.rs | 111 ++- diskann-benchmark/src/multi_vector/mod.rs | 163 +++- diskann-benchmark/src/utils/mod.rs | 2 +- .../maxsim/minmax8_x_minmax4.rs | 824 ++++++++++++++++++ .../maxsim/minmax8_x_minmax4/decode.rs | 101 +++ .../maxsim/minmax8_x_minmax4/layout.rs | 211 +++++ .../maxsim/minmax8_x_minmax4/panels.rs | 191 ++++ .../src/matrix_kernels/maxsim/mod.rs | 1 + diskann-quantization/src/minmax/mod.rs | 2 +- .../src/minmax/multi/factory.rs | 307 +++++++ .../src/minmax/multi/kernel.rs | 54 ++ diskann-quantization/src/minmax/multi/mod.rs | 5 + .../src/minmax/multi/query.rs | 60 ++ 21 files changed, 2473 insertions(+), 41 deletions(-) create mode 100644 diskann-benchmark/perf_test_inputs/multi-vector-minmax8-neon.json create mode 100644 diskann-benchmark/perf_test_inputs/multi-vector-minmax8-v4.json create mode 100644 diskann-benchmark/perf_test_inputs/multi-vector-minmax8.json create mode 100644 diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs create mode 100644 diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs create mode 100644 diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs create mode 100644 diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/panels.rs create mode 100644 diskann-quantization/src/minmax/multi/factory.rs create mode 100644 diskann-quantization/src/minmax/multi/kernel.rs create mode 100644 diskann-quantization/src/minmax/multi/query.rs diff --git a/diskann-benchmark/README.md b/diskann-benchmark/README.md index b48fdf18a0..0e01068d43 100644 --- a/diskann-benchmark/README.md +++ b/diskann-benchmark/README.md @@ -260,6 +260,19 @@ is running. This incremental saving allows benchmarks to be interrupted without In addition to the machine-readable JSON output files, a (hopefully) helpful summary of the results will be printed to `stdout`. +With `--features multi-vector`, select MinMax8 queries/MinMax4 documents using +`type: "multi-vector-op"` and `element_type: "minmax8"`; +see `example\multi-vector.json` and `perf_test_inputs\multi-vector-minmax8*.json`. + +```powershell +cargo run --release --package diskann-benchmark --features multi-vector -- run --input-file diskann-benchmark\perf_test_inputs\multi-vector-minmax8.json --output-file minmax-results.json +``` + +Quantization and query preparation are outside the timed region. Each run repeatedly scores +the same document matrix and reports nanoseconds per query/document vector pair (`ns/IP`), +including per-call scratch work. These results use a different workload and statistic from +benchmarks that sweep many distinct documents and report microseconds per document. + ### Streaming Runs Running the benchmark on a streaming workload is similar to other registered benchmarks, relying on the file formats and streaming runbooks of `big-ann-benchmarks` @@ -504,4 +517,3 @@ for a successful or unsuccessful match with the enclosed value. Doing these two error reporting in the event of a dispatch fail much easier for the user to understand and fix. Refer to implementations within the benchmarking framework for what some of this may look like. - diff --git a/diskann-benchmark/example/multi-vector.json b/diskann-benchmark/example/multi-vector.json index 2dd53693d0..1b48bb9011 100644 --- a/diskann-benchmark/example/multi-vector.json +++ b/diskann-benchmark/example/multi-vector.json @@ -64,6 +64,36 @@ { "num_query_vectors": 64, "num_doc_vectors": 32, "dim": 264, "loops_per_measurement": 2, "num_measurements": 1 } ] } + }, + { + "type": "multi-vector-op", + "content": { + "element_type": "minmax8", + "isa": "auto", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 7, "num_doc_vectors": 13, "dim": 63, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 64, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 65, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 2, "num_measurements": 1 } + ] + } + }, + { + "type": "multi-vector-op", + "content": { + "element_type": "minmax8", + "isa": "scalar", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 7, "num_doc_vectors": 13, "dim": 63, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 8, "num_doc_vectors": 32, "dim": 64, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 65, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 2, "num_measurements": 1 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 2, "num_measurements": 1 } + ] + } } ] } diff --git a/diskann-benchmark/perf_test_inputs/multi-vector-minmax8-neon.json b/diskann-benchmark/perf_test_inputs/multi-vector-minmax8-neon.json new file mode 100644 index 0000000000..d4255be252 --- /dev/null +++ b/diskann-benchmark/perf_test_inputs/multi-vector-minmax8-neon.json @@ -0,0 +1,47 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-op", + "content": { + "element_type": "minmax8", + "isa": "neon", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 } + ] + } + }, + { + "type": "multi-vector-op", + "content": { + "element_type": "minmax8", + "isa": "scalar", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 } + ] + } + } + ] +} diff --git a/diskann-benchmark/perf_test_inputs/multi-vector-minmax8-v4.json b/diskann-benchmark/perf_test_inputs/multi-vector-minmax8-v4.json new file mode 100644 index 0000000000..0a97c74af8 --- /dev/null +++ b/diskann-benchmark/perf_test_inputs/multi-vector-minmax8-v4.json @@ -0,0 +1,47 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-op", + "content": { + "element_type": "minmax8", + "isa": "x86-64-v4", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 } + ] + } + }, + { + "type": "multi-vector-op", + "content": { + "element_type": "minmax8", + "isa": "scalar", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 } + ] + } + } + ] +} diff --git a/diskann-benchmark/perf_test_inputs/multi-vector-minmax8.json b/diskann-benchmark/perf_test_inputs/multi-vector-minmax8.json new file mode 100644 index 0000000000..9abba77ac8 --- /dev/null +++ b/diskann-benchmark/perf_test_inputs/multi-vector-minmax8.json @@ -0,0 +1,47 @@ +{ + "search_directories": [], + "jobs": [ + { + "type": "multi-vector-op", + "content": { + "element_type": "minmax8", + "isa": "auto", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 } + ] + } + }, + { + "type": "multi-vector-op", + "content": { + "element_type": "minmax8", + "isa": "scalar", + "runs": [ + { "num_query_vectors": 16, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 63, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 64, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 65, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 127, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 8, "num_doc_vectors": 64, "dim": 129, "loops_per_measurement": 200, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 128, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 256, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 64, "dim": 768, "loops_per_measurement": 50, "num_measurements": 50 }, + { "num_query_vectors": 17, "num_doc_vectors": 65, "dim": 129, "loops_per_measurement": 100, "num_measurements": 50 }, + { "num_query_vectors": 32, "num_doc_vectors": 1250, "dim": 128, "loops_per_measurement": 10, "num_measurements": 50 } + ] + } + } + ] +} diff --git a/diskann-benchmark/src/inputs/multi_vector.rs b/diskann-benchmark/src/inputs/multi_vector.rs index 6f1ec9dddd..a241d7fab3 100644 --- a/diskann-benchmark/src/inputs/multi_vector.rs +++ b/diskann-benchmark/src/inputs/multi_vector.rs @@ -5,7 +5,7 @@ use std::num::NonZeroUsize; -use diskann_benchmark_runner::{utils::datatype::DataType, Checker, Input}; +use diskann_benchmark_runner::{Checker, Input}; use diskann_quantization::multi_vector::MaxSimIsa; use serde::{Deserialize, Serialize}; @@ -69,14 +69,36 @@ pub(crate) struct Run { pub(crate) num_measurements: NonZeroUsize, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum ElementType { + Float32, + Float16, + Int8, + /// MinMax8 queries and MinMax4 documents, quantized from f32 fixtures. + MinMax8, +} + +impl std::fmt::Display for ElementType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Self::Float32 => "float32", + Self::Float16 => "float16", + Self::Int8 => "int8", + Self::MinMax8 => "minmax8", + }) + } +} + /////////////////////// // Multi-Vector Op // /////////////////////// /// A complete multi-vector benchmark job. #[derive(Debug, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub(crate) struct MultiVectorOp { - pub(crate) element_type: DataType, + pub(crate) element_type: ElementType, pub(crate) isa: BenchIsa, pub(crate) runs: Vec, } @@ -95,6 +117,9 @@ impl Input for MultiVectorOp { } fn from_raw(raw: Self::Raw, _checker: &mut Checker) -> anyhow::Result { + if raw.element_type == ElementType::MinMax8 { + anyhow::ensure!(!raw.runs.is_empty(), "MinMax8 requires at least one run"); + } Ok(raw) } @@ -126,7 +151,7 @@ impl Input for MultiVectorOp { ]; Self { - element_type: DataType::Float32, + element_type: ElementType::Float32, isa: BenchIsa::Auto, runs, } diff --git a/diskann-benchmark/src/main.rs b/diskann-benchmark/src/main.rs index d25a56fc92..ab736a2e89 100644 --- a/diskann-benchmark/src/main.rs +++ b/diskann-benchmark/src/main.rs @@ -1006,6 +1006,8 @@ mod tests { #[test] fn multi_vector_integration() { + use crate::inputs::multi_vector::MultiVectorOp; + let path = example_directory().join("multi-vector.json"); let tempdir = tempfile::tempdir().unwrap(); let output_path = tempdir.path().join("output.json"); @@ -1014,6 +1016,10 @@ mod tests { let modified_input_path = tempdir.path().join("input.json"); let mut raw = value_from_file(&path); + for job in raw["jobs"].as_array().unwrap() { + let operation: MultiVectorOp = serde_json::from_value(job["content"].clone()).unwrap(); + assert_eq!(serde_json::to_value(operation).unwrap(), job["content"]); + } prefix_search_directories(&mut raw, &root_directory()); save_to_file(&modified_input_path, &raw); @@ -1040,6 +1046,26 @@ mod tests { // Check that the results file is generated. assert!(output_path.exists()); + + let input = value_from_file(input_path); + let output = value_from_file(output_path); + let jobs = input["jobs"].as_array().unwrap(); + let results = output.as_array().unwrap(); + assert_eq!(results.len(), jobs.len()); + for (result, job) in std::iter::zip(results, jobs) { + assert_eq!(result["input"], *job); + let runs = job["content"]["runs"].as_array().unwrap(); + let measurements = result["results"].as_array().unwrap(); + assert_eq!(measurements.len(), runs.len()); + for (measurement, run) in std::iter::zip(measurements, runs) { + assert_eq!(measurement["run"], *run); + assert_eq!( + measurement["latencies"].as_array().unwrap().len() as u64, + run["num_measurements"].as_u64().unwrap(), + ); + assert!(measurement["percentiles"].is_object()); + } + } } #[cfg(not(feature = "multi-vector"))] @@ -1086,6 +1112,158 @@ mod tests { ); } + #[test] + fn minmax8_isa_performance_inputs_have_matching_shapes() { + use crate::inputs::multi_vector::{BenchIsa, ElementType, MultiVectorOp}; + + let directory = project_directory().join("perf_test_inputs"); + let portable = value_from_file(&directory.join("multi-vector-minmax8.json")); + let baseline: MultiVectorOp = + serde_json::from_value(portable["jobs"][0]["content"].clone()).unwrap(); + assert_eq!(baseline.runs.len(), 12); + let primary = &baseline.runs[0]; + assert_eq!( + ( + primary.num_query_vectors.get(), + primary.num_doc_vectors.get(), + primary.dim.get(), + primary.loops_per_measurement.get(), + primary.num_measurements.get(), + ), + (16, 64, 256, 100, 50), + ); + for (file, isa) in [ + ("multi-vector-minmax8.json", BenchIsa::Auto), + ("multi-vector-minmax8-neon.json", BenchIsa::Neon), + ("multi-vector-minmax8-v4.json", BenchIsa::X86_64_V4), + ] { + let input = value_from_file(&directory.join(file)); + let jobs = input["jobs"].as_array().unwrap(); + assert_eq!(jobs.len(), 2, "{file}"); + for (job, expected_isa) in std::iter::zip(jobs, [isa, BenchIsa::Scalar]) { + assert_eq!(job["type"], MultiVectorOp::tag(), "{file}"); + let operation: MultiVectorOp = + serde_json::from_value(job["content"].clone()).unwrap(); + assert_eq!(operation.element_type, ElementType::MinMax8, "{file}"); + assert_eq!(operation.isa, expected_isa, "{file}"); + assert_eq!(operation.runs, baseline.runs, "{file}"); + assert_eq!(serde_json::to_value(operation).unwrap(), job["content"]); + } + } + } + + #[test] + fn multi_vector_rejects_invalid_element_selection() { + for (element_type, format, expected_error) in [ + ("float32", Some("minmax8"), "unknown field `format`"), + ("minmax4", None, "unknown variant `minmax4`"), + ] { + let tempdir = tempfile::tempdir().unwrap(); + let input_path = tempdir.path().join("input.json"); + let output_path = tempdir.path().join("output.json"); + let mut raw = value_from_file( + &project_directory() + .join("perf_test_inputs") + .join("multi-vector-minmax8.json"), + ); + raw["jobs"].as_array_mut().unwrap().truncate(1); + raw["jobs"][0]["content"]["element_type"] = serde_json::json!(element_type); + if let Some(format) = format { + raw["jobs"][0]["content"]["format"] = serde_json::json!(format); + } + save_to_file(&input_path, &raw); + let command = Commands::Run { + input_file: input_path, + output_file: output_path.clone(), + dry_run: false, + allow_debug: true, + }; + let error = Cli::from_commands(command, true) + .run(&mut Memory::new()) + .unwrap_err(); + assert!(format!("{error:#}").contains(expected_error), "{error:#}"); + assert!(!output_path.exists()); + } + } + + #[test] + #[cfg(feature = "multi-vector")] + fn minmax8_rejects_empty_runs() { + let tempdir = tempfile::tempdir().unwrap(); + let input_path = tempdir.path().join("input.json"); + let output_path = tempdir.path().join("output.json"); + let mut raw = value_from_file( + &project_directory() + .join("perf_test_inputs") + .join("multi-vector-minmax8.json"), + ); + raw["jobs"].as_array_mut().unwrap().truncate(1); + raw["jobs"][0]["content"]["runs"] = serde_json::json!([]); + save_to_file(&input_path, &raw); + let command = Commands::Run { + input_file: input_path, + output_file: output_path.clone(), + dry_run: false, + allow_debug: true, + }; + let error = Cli::from_commands(command, true) + .run(&mut Memory::new()) + .unwrap_err(); + assert!(format!("{error:#}").contains("MinMax8 requires at least one run")); + assert!(!output_path.exists()); + } + + #[test] + #[cfg(feature = "multi-vector")] + fn minmax8_rejects_unsupported_isa() { + use crate::inputs::multi_vector::BenchIsa; + use diskann_quantization::multi_vector::MaxSimIsa; + + for isa in [ + BenchIsa::Reference, + BenchIsa::X86_64_V3, + BenchIsa::X86_64_V4, + BenchIsa::Neon, + ] { + if !matches!(isa, BenchIsa::Reference) && MaxSimIsa::from(isa).is_available() { + continue; + } + let tempdir = tempfile::tempdir().unwrap(); + let input_path = tempdir.path().join("input.json"); + let output_path = tempdir.path().join("output.json"); + let path = match isa { + BenchIsa::Neon => project_directory() + .join("perf_test_inputs") + .join("multi-vector-minmax8-neon.json"), + BenchIsa::X86_64_V4 => project_directory() + .join("perf_test_inputs") + .join("multi-vector-minmax8-v4.json"), + _ => project_directory() + .join("perf_test_inputs") + .join("multi-vector-minmax8.json"), + }; + let mut raw = value_from_file(&path); + raw["jobs"].as_array_mut().unwrap().truncate(1); + raw["jobs"][0]["content"]["isa"] = serde_json::to_value(isa).unwrap(); + save_to_file(&input_path, &raw); + let command = Commands::Run { + input_file: input_path, + output_file: output_path.clone(), + dry_run: false, + allow_debug: true, + }; + let mut output = Memory::new(); + assert!(Cli::from_commands(command, true).run(&mut output).is_err()); + let output = String::from_utf8(output.into_inner()).unwrap(); + if matches!(isa, BenchIsa::Reference) { + assert!(output.contains("MinMax8 has no reference kernel; use scalar")); + } else { + assert!(output.contains("ISA unavailable on this CPU")); + } + assert!(!output_path.exists()); + } + } + #[test] fn quiet_suppresses_check_target_warning() { let cli = Cli::from_commands(Commands::Skeleton, true); diff --git a/diskann-benchmark/src/multi_vector/driver.rs b/diskann-benchmark/src/multi_vector/driver.rs index 5cec88bc02..256f778232 100644 --- a/diskann-benchmark/src/multi_vector/driver.rs +++ b/diskann-benchmark/src/multi_vector/driver.rs @@ -6,6 +6,7 @@ //! Shared benchmark infrastructure for multi-vector kernels: timing harness, //! data fixtures, result types. None of the contents are kernel-aware. +use anyhow::Context; use diskann_benchmark_runner::{ utils::{ fmt::Table, @@ -15,8 +16,16 @@ use diskann_benchmark_runner::{ Checker, Input, }; use diskann_quantization::multi_vector::{ - Mat, MatRef, MaxSimElement, MaxSimKernel, Overflow, Standard, + Defaulted, Mat, MatRef, MaxSimElement, MaxSimKernel, Overflow, Standard, }; +use diskann_quantization::{ + algorithms::{transforms::NullTransform, Transform}, + bits::{Representation, Unsigned}, + minmax::{MinMaxMeta, MinMaxQuantizer}, + num::Positive, + CompressInto, +}; +use diskann_utils::ReborrowMut; use rand::{ distr::{Distribution, StandardUniform}, rngs::StdRng, @@ -79,15 +88,56 @@ where StandardUniform: Distribution, { pub(super) fn new(run: &Run) -> Result { + let query_repr = Standard::new(run.num_query_vectors.get(), run.dim.get())?; + let doc_repr = Standard::new(run.num_doc_vectors.get(), run.dim.get())?; let mut rng = StdRng::seed_from_u64(0x12345); - let queries = Mat::from_fn( - Standard::new(run.num_query_vectors.get(), run.dim.get())?, - || StandardUniform.sample(&mut rng), - ); - let docs = Mat::from_fn( - Standard::new(run.num_doc_vectors.get(), run.dim.get())?, - || StandardUniform.sample(&mut rng), + let queries = Mat::from_fn(query_repr, || StandardUniform.sample(&mut rng)); + let docs = Mat::from_fn(doc_repr, || StandardUniform.sample(&mut rng)); + Ok(Self { queries, docs }) + } +} + +// MinMaxMeta::new does not check canonical allocation extents. +fn checked_minmax_meta( + rows: usize, + dim: usize, +) -> anyhow::Result> +where + Unsigned: Representation, +{ + anyhow::ensure!( + u32::try_from(dim).is_ok(), + "MinMax dimension exceeds u32::MAX" + ); + dim.checked_mul(BITS) + .and_then(|bits| bits.checked_add(7)) + .context("MinMax row size overflow")?; + let repr = MinMaxMeta::new(rows, dim); + let bytes = rows + .checked_mul(repr.ncols()) + .context("MinMax matrix size overflow")?; + std::alloc::Layout::array::(bytes).context("MinMax allocation exceeds the byte limit")?; + Ok(repr) +} + +pub(super) struct MinMax8Data { + pub(super) queries: Mat>, + pub(super) docs: Mat>, +} + +impl MinMax8Data { + pub(super) fn new(run: &Run) -> anyhow::Result { + let query_repr = checked_minmax_meta::<8>(run.num_query_vectors.get(), run.dim.get())?; + let doc_repr = checked_minmax_meta::<4>(run.num_doc_vectors.get(), run.dim.get())?; + let data = Data::::new(run)?; + let quantizer = MinMaxQuantizer::new( + Transform::Null(NullTransform::new(run.dim)), + Positive::new(1.0).unwrap(), ); + let mut queries = Mat::new(query_repr, Defaulted)?; + let mut docs = Mat::new(doc_repr, Defaulted)?; + quantizer.compress_into(data.queries.as_view(), queries.reborrow_mut())?; + quantizer.compress_into(data.docs.as_view(), docs.reborrow_mut())?; Ok(Self { queries, docs }) } } @@ -100,27 +150,35 @@ pub(super) fn run_with_kernel( run: &Run, doc: MatRef<'_, Standard>, kernel: &dyn MaxSimKernel, -) -> RunResult { +) -> anyhow::Result { let mut scores = vec![T::Score::default(); run.num_query_vectors.get()]; + run_with_compute(run, || { + kernel.compute_max_sim(doc, &mut scores)?; + std::hint::black_box(&mut scores); + Ok(()) + }) +} + +pub(super) fn run_with_compute( + run: &Run, + mut compute: impl FnMut() -> anyhow::Result<()>, +) -> anyhow::Result { 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() { - kernel - .compute_max_sim(doc, &mut scores) - .expect("scores.len() == kernel.nrows() by construction"); - std::hint::black_box(&mut scores); + compute()?; } latencies.push(start.elapsed().into()); } let percentiles = percentiles::compute_percentiles(&mut latencies).unwrap(); - RunResult { + Ok(RunResult { run: run.clone(), latencies, percentiles, - } + }) } ////////////////////// diff --git a/diskann-benchmark/src/multi_vector/kernels.rs b/diskann-benchmark/src/multi_vector/kernels.rs index 7e169cbd9f..ac24b802e8 100644 --- a/diskann-benchmark/src/multi_vector/kernels.rs +++ b/diskann-benchmark/src/multi_vector/kernels.rs @@ -8,23 +8,29 @@ //! A single generic [`Kernel`] carrier covers every element type accepted //! by [`MaxSimElement`]; `try_match` also rejects ISAs unavailable on the //! host so unsupported jobs fail at job-selection rather than mid-run. +//! MinMax8/MinMax4 uses a separate carrier with the same input and result format. use std::io::Write; use std::marker::PhantomData; use diskann_benchmark_runner::{ benchmark::{MatchContext, PassFail, Regression, Score}, - utils::{datatype::AsDataType, num::relative_change}, + utils::{ + datatype::{AsDataType, DataType}, + num::relative_change, + }, Benchmark, Checkpoint, Output, Registry, }; +use diskann_quantization::minmax::build_minmax_max_sim; use diskann_quantization::multi_vector::{build_max_sim, BoxErase, MaxSimElement, MaxSimIsa}; use rand::distr::{Distribution, StandardUniform}; use super::driver::{ - run_with_kernel, CheckResult, Comparison, Data, MultiVectorTolerance, RunResult, + run_with_compute, run_with_kernel, CheckResult, Comparison, Data, MinMax8Data, + MultiVectorTolerance, RunResult, }; -use crate::inputs::multi_vector::MultiVectorOp; -use crate::utils::DisplayWrapper; +use crate::inputs::multi_vector::{ElementType, MultiVectorOp}; +use crate::utils::{DisplayWrapper, DATA_TYPE_MISMATCH}; // ───────────────────────────────────────────────────────────────────────── // Kernel — generic carrier registered once per element type. @@ -48,8 +54,16 @@ where type Output = Vec; fn try_match(&self, from: &MultiVectorOp, context: &MatchContext) -> Score { + let data_type = match from.element_type { + ElementType::Float32 => DataType::Float32, + ElementType::Float16 => DataType::Float16, + ElementType::Int8 => DataType::Int8, + ElementType::MinMax8 => { + return context.fail(DATA_TYPE_MISMATCH, &"MinMax8 requires the quantized kernel"); + } + }; let mut score = context.success(0); - crate::utils::match_data_type::(&mut score, from.element_type); + crate::utils::match_data_type::(&mut score, data_type); let isa: MaxSimIsa = from.isa.into(); if !isa.is_available() { score.fail(1, &format_args!("ISA unavailable on this CPU: {}", isa)); @@ -68,7 +82,7 @@ where for run in input.runs.iter() { let data = Data::::new(run)?; let kernel = build_max_sim::(input.isa.into(), data.queries.as_view(), BoxErase)?; - results.push(run_with_kernel(run, data.docs.as_view(), &*kernel)); + results.push(run_with_kernel(run, data.docs.as_view(), &*kernel)?); } writeln!(output, "\n\n{}", DisplayWrapper(&*results))?; Ok(results) @@ -95,6 +109,16 @@ where before: &Vec, after: &Vec, ) -> anyhow::Result> { + CheckResult::compare(tolerance, before, after) + } +} + +impl CheckResult { + fn compare( + tolerance: &MultiVectorTolerance, + before: &[RunResult], + after: &[RunResult], + ) -> anyhow::Result> { anyhow::ensure!( before.len() == after.len(), "before has {} runs but after has {}", @@ -103,7 +127,7 @@ where ); let mut passed = true; - let checks: Vec = std::iter::zip(before.iter(), after.iter()) + let checks = std::iter::zip(before, after) .enumerate() .map(|(i, (b, a))| { anyhow::ensure!(b.run == a.run, "run {i} mismatched"); @@ -118,7 +142,6 @@ where before_min, after_min, }; - match relative_change(before_min, after_min) { Ok(change) => { if change > tolerance.min_time_regression.get() { @@ -127,19 +150,84 @@ where } Err(_) => passed = false, }; - Ok(comparison) }) .collect::>>()?; Ok(if passed { - PassFail::Pass(CheckResult { checks }) + PassFail::Pass(Self { checks }) } else { - PassFail::Fail(CheckResult { checks }) + PassFail::Fail(Self { checks }) }) } } +#[derive(Debug)] +pub(super) struct MinMax8Kernel; + +impl Benchmark for MinMax8Kernel { + type Input = MultiVectorOp; + type Output = Vec; + + fn try_match(&self, from: &MultiVectorOp, context: &MatchContext) -> Score { + let mut score = context.success(0); + if from.element_type != ElementType::MinMax8 { + score.fail(DATA_TYPE_MISMATCH, &"expected minmax8 element type"); + } + let isa: MaxSimIsa = from.isa.into(); + if isa == MaxSimIsa::Reference { + score.fail(1, &"MinMax8 has no reference kernel; use scalar"); + } else if !isa.is_available() { + score.fail(1, &format_args!("ISA unavailable on this CPU: {}", isa)); + } + score + } + + fn run( + &self, + input: &MultiVectorOp, + _: 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 { + let data = MinMax8Data::new(run)?; + let kernel = build_minmax_max_sim(input.isa.into(), data.queries.as_view(), BoxErase)?; + let mut scores = vec![0.0; run.num_query_vectors.get()]; + kernel.compute_max_sim(data.docs.as_view(), &mut scores)?; + std::hint::black_box(&mut scores); + results.push(run_with_compute(run, || { + kernel.compute_max_sim(data.docs.as_view(), &mut scores)?; + std::hint::black_box(&mut scores); + Ok(()) + })?); + } + writeln!(output, "\n\n{}", DisplayWrapper(&*results))?; + Ok(results) + } + + fn description(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "- Query: MinMax8; Document: MinMax4") + } +} + +impl Regression for MinMax8Kernel { + type Tolerances = MultiVectorTolerance; + type Pass = CheckResult; + type Fail = CheckResult; + + fn check( + &self, + tolerance: &MultiVectorTolerance, + _input: &MultiVectorOp, + before: &Vec, + after: &Vec, + ) -> anyhow::Result> { + CheckResult::compare(tolerance, before, after) + } +} + // ───────────────────────────────────────────────────────────────────────── // Registration. // ───────────────────────────────────────────────────────────────────────── @@ -148,5 +236,6 @@ pub(super) fn register(registry: &mut Registry) -> anyhow::Result<()> { registry.register_regression("multi-vector-op-f32", Kernel::::new())?; registry.register_regression("multi-vector-op-f16", Kernel::::new())?; registry.register_regression("multi-vector-op-i8", Kernel::::new())?; + registry.register_regression("multi-vector-op-minmax8", MinMax8Kernel)?; Ok(()) } diff --git a/diskann-benchmark/src/multi_vector/mod.rs b/diskann-benchmark/src/multi_vector/mod.rs index b184682f50..c8db980c0c 100644 --- a/diskann-benchmark/src/multi_vector/mod.rs +++ b/diskann-benchmark/src/multi_vector/mod.rs @@ -7,6 +7,8 @@ //! //! One `Benchmark` is registered per element type supported by //! [`MaxSimElement`]; the JSON `isa` field picks the kernel at run time. +//! The shared `multi-vector-op` input also accepts `element_type: "minmax8"` for +//! MinMax8 queries and MinMax4 documents. //! //! # Why two ISA enums? //! @@ -47,16 +49,21 @@ mod tests { use std::num::NonZeroUsize; use diskann_benchmark_runner::{ - benchmark::{PassFail, Regression}, - utils::{ - datatype::DataType, num::NonNegativeFinite, percentiles::compute_percentiles, - MicroSeconds, - }, + benchmark::{MatchContext, PassFail, Regression, TestScore}, + utils::{num::NonNegativeFinite, percentiles::compute_percentiles, MicroSeconds}, }; - use super::driver::{CheckResult, Comparison, MultiVectorTolerance, RunResult}; - use super::kernels::Kernel; - use crate::inputs::multi_vector::{BenchIsa, MultiVectorOp, Run}; + use diskann_quantization::{ + minmax::build_minmax_max_sim, + multi_vector::{BoxErase, MaxSim, MaxSimIsa, QueryMatRef}, + }; + use diskann_vector::DistanceFunctionMut; + + use super::driver::{ + run_with_compute, CheckResult, Comparison, MinMax8Data, MultiVectorTolerance, RunResult, + }; + use super::kernels::{Kernel, MinMax8Kernel}; + use crate::inputs::multi_vector::{BenchIsa, ElementType, MultiVectorOp, Run}; fn tiny_run() -> Run { Run { @@ -70,12 +77,41 @@ mod tests { fn tiny_op() -> MultiVectorOp { MultiVectorOp { - element_type: DataType::Float32, + element_type: ElementType::Float32, isa: BenchIsa::Auto, runs: vec![tiny_run()], } } + #[test] + fn multi_vector_element_type_selects_only_matching_kernel() { + for (element_type, selected) in [ + (ElementType::Float32, 0), + (ElementType::Float16, 1), + (ElementType::Int8, 2), + (ElementType::MinMax8, 3), + ] { + let input = MultiVectorOp { + element_type, + isa: BenchIsa::Scalar, + runs: vec![tiny_run()], + }; + let scores = [ + MatchContext::test(&Kernel::::new(), &input), + MatchContext::test(&Kernel::::new(), &input), + MatchContext::test(&Kernel::::new(), &input), + MatchContext::test(&MinMax8Kernel, &input), + ]; + for (index, score) in scores.into_iter().enumerate() { + assert_eq!( + matches!(score, TestScore::Success(_)), + index == selected, + "{input:?}: {score:?}" + ); + } + } + } + fn tiny_result(minimum: u64) -> RunResult { let mut latencies = vec![MicroSeconds::new(minimum)]; let percentiles = compute_percentiles(&mut latencies).unwrap(); @@ -207,4 +243,113 @@ mod tests { assert!(matches!(result, PassFail::Fail(_))); } + + #[test] + fn minmax8_fixtures_reject_unrepresentable_shapes() { + let row_bytes = diskann_quantization::minmax::Data::<8>::canonical_bytes(1); + let too_many_rows = isize::MAX as usize / row_bytes + 1; + for (queries, docs, dim) in [ + (usize::MAX, 1, 1), + (1, usize::MAX, 1), + (too_many_rows, 1, 1), + (1, too_many_rows, 1), + (1, 1, usize::MAX), + ( + 1, + isize::MAX as usize / (64 * std::mem::size_of::()) + 1, + 64, + ), + ] { + let run = Run { + num_query_vectors: NonZeroUsize::new(queries).unwrap(), + num_doc_vectors: NonZeroUsize::new(docs).unwrap(), + dim: NonZeroUsize::new(dim).unwrap(), + ..tiny_run() + }; + assert!(MinMax8Data::new(&run).is_err(), "{run:?}"); + } + } + + #[test] + fn minmax8_fixtures_match_reference() { + for (queries, docs, dim) in [ + (16, 64, 256), + (7, 5, 63), + (16, 16, 128), + (17, 13, 129), + (32, 65, 768), + ] { + let run = Run { + num_query_vectors: NonZeroUsize::new(queries).unwrap(), + num_doc_vectors: NonZeroUsize::new(docs).unwrap(), + dim: NonZeroUsize::new(dim).unwrap(), + ..tiny_run() + }; + let data = MinMax8Data::new(&run).unwrap(); + let mut expected = vec![0.0; queries]; + MaxSim::new(&mut expected).evaluate( + QueryMatRef::from(data.queries.as_view()), + data.docs.as_view(), + ); + assert!(expected.iter().all(|score| score.is_finite())); + for isa in [ + MaxSimIsa::Scalar, + MaxSimIsa::X86_64_V3, + MaxSimIsa::X86_64_V4, + MaxSimIsa::Neon, + MaxSimIsa::Auto, + ] { + if !isa.is_available() { + continue; + } + let kernel = build_minmax_max_sim(isa, data.queries.as_view(), BoxErase).unwrap(); + let mut actual = vec![f32::NAN; queries]; + kernel + .compute_max_sim(data.docs.as_view(), &mut actual) + .unwrap(); + assert_eq!(actual, expected, "{isa:?}: {run:?}"); + } + } + } + + #[test] + fn minmax8_uses_shared_regression_check() { + let input = MultiVectorOp { + element_type: ElementType::MinMax8, + isa: BenchIsa::Scalar, + runs: vec![tiny_run()], + }; + for (after, should_pass) in [(105, true), (106, false), (0, true)] { + let result = MinMax8Kernel + .check( + &tolerance(0.05), + &input, + &vec![tiny_result(100)], + &vec![tiny_result(after)], + ) + .unwrap(); + assert_eq!(matches!(result, PassFail::Pass(_)), should_pass); + } + } + + #[test] + fn timing_counts_computations_and_propagates_errors() { + let run = Run { + loops_per_measurement: NonZeroUsize::new(3).unwrap(), + num_measurements: NonZeroUsize::new(2).unwrap(), + ..tiny_run() + }; + let mut calls = 0; + let result = run_with_compute(&run, || { + calls += 1; + Ok(()) + }) + .unwrap(); + assert_eq!(calls, 6); + assert_eq!(result.latencies.len(), 2); + assert_eq!(result.run, run); + + let error = run_with_compute(&run, || anyhow::bail!("compute failed")).unwrap_err(); + assert_eq!(error.to_string(), "compute failed"); + } } diff --git a/diskann-benchmark/src/utils/mod.rs b/diskann-benchmark/src/utils/mod.rs index cd8e6510cf..52ce89d989 100644 --- a/diskann-benchmark/src/utils/mod.rs +++ b/diskann-benchmark/src/utils/mod.rs @@ -15,7 +15,7 @@ pub(crate) mod recall; pub(crate) mod streaming; pub(crate) mod tokio; -const DATA_TYPE_MISMATCH: u32 = 1000; +pub(crate) const DATA_TYPE_MISMATCH: u32 = 1000; pub(crate) fn match_data_type(score: &mut Score, data_type: DataType) where diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs new file mode 100644 index 0000000000..2150d4df90 --- /dev/null +++ b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs @@ -0,0 +1,824 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! B-first MinMax MaxSim. Preparation and synchronized panel iterators +//! establish the micro-kernel's inputs; integer contraction never reads canonical data. + +mod decode; +pub(crate) mod layout; +pub(crate) mod panels; + +use super::packed_f32_x_unpacked_f32::Params; +use crate::matrix_kernels::{ + Cache, driver, + num::{Bytes, value_or_one}, + util::LoadStore, +}; +use crate::minmax::MinMaxCompensation; +use decode::Decoder; +use diskann_wide::{ + Architecture, + arch::{Scalar, Target}, +}; +use layout::{APanel, AView}; +use panels::{BPanel, BTile, RawRows, Scratch}; + +pub(crate) struct Driver<'a, A, const MR: usize, const NR: usize, const PACK: usize> { + arch: A, + a: AView<'a, MR, PACK>, + b: RawRows<'a>, + c: &'a mut [f32], + scratch: Scratch, + b_rows: std::num::NonZeroUsize, +} + +impl<'a, A, const MR: usize, const NR: usize, const PACK: usize> Driver<'a, A, MR, NR, PACK> { + pub(crate) fn new( + arch: A, + a: AView<'a, MR, PACK>, + b: RawRows<'a>, + c: &'a mut [f32], + cache: Cache, + ) -> Self { + const { + assert!(NR > 0 && NR <= 8); + } + assert_eq!(a.rows(), c.len(), "output row mismatch"); + assert_eq!(a.layout().dim(), b.dim(), "document dimension mismatch"); + assert!( + a.layout().dim() <= layout::MAX_DIM, + "integer accumulator dimension limit" + ); + let k = a.layout().k().value().get(); + let params = Params::new( + cache, + Bytes::new(MR * k + std::mem::size_of::<[crate::minmax::MinMaxCompensation; MR]>()), + Bytes::new(k + std::mem::size_of::()), + NR, + ); + Self { + arch, + a, + b, + c, + scratch: Scratch::default(), + b_rows: value_or_one(b.rows().min(params.b_cols_in_l1.get())), + } + } +} + +impl driver::Drive + for Driver<'_, A, MR, NR, PACK> +where + A: Decoder + ExtraWide + LoadStore, +{ + fn drive(&mut self) { + self.arch.run_inline(&mut *self); + } +} + +impl Target + for &mut Driver<'_, A, MR, NR, PACK> +where + A: Decoder + ExtraWide + LoadStore, +{ + #[inline(always)] + fn run(self, arch: A) { + self.c.fill(f32::MAX); + let a = self.a; + self.b.visit_tiles( + self.b_rows, + #[inline(always)] + |raw| { + let b = self.scratch.prepare(arch, raw, a.layout()); + a.visit_panels( + #[inline(always)] + |a, start| { + let output = &mut self.c[start..start + a.valid_rows()]; + let mut panel = PanelKernel::<_, MR, NR, PACK>::new(arch, a, b, output); + driver::PanelKernel::panel_kernel(&mut panel); + arch.store(panel.scores, output); + }, + ); + }, + ); + } +} + +struct PanelKernel<'a, A, const MR: usize, const NR: usize, const PACK: usize> { + arch: A, + a: APanel<'a, MR, PACK>, + b: BTile<'a>, + compensation: PanelCompensation, + scores: [f32; MR], +} + +impl<'a, A, const MR: usize, const NR: usize, const PACK: usize> PanelKernel<'a, A, MR, NR, PACK> { + #[inline(always)] + fn new(arch: A, a: APanel<'a, MR, PACK>, b: BTile<'a>, output: &[f32]) -> Self + where + A: Architecture + LoadStore, + { + assert!( + std::ptr::eq(a.layout(), b.layout()), + "A and B must share the same layout" + ); + assert_eq!(output.len(), a.valid_rows()); + Self { + arch, + a, + b, + compensation: PanelCompensation::new(a.compensation(), a.valid_rows()), + scores: arch.load(output), + } + } + + #[inline(always)] + fn micro_kernel(&mut self, b: BPanel<'_, N>) + where + A: Architecture + ExtraWide, + { + let mut micro = MicroKernel { + arch: self.arch, + a: self.a, + b, + compensation: &self.compensation, + scores: &mut self.scores, + }; + driver::MicroKernel::micro_kernel(&mut micro); + } +} + +impl driver::PanelKernel + for PanelKernel<'_, A, MR, NR, PACK> +where + A: Architecture + ExtraWide, +{ + #[inline(always)] + fn panel_kernel(&mut self) { + let b = self.b; + let remainder = b.visit_panels::(|panel| self.micro_kernel(panel)); + if let Some(remainder) = remainder { + macro_rules! tail { + ($($n:literal),+) => {$( + if $n < NR { + if let Some(panel) = remainder.panel::<$n>() { self.micro_kernel(panel); } + } + )+}; + } + tail!(1, 2, 3, 4, 5, 6, 7); + } + } +} + +struct MicroKernel<'a, A, const MR: usize, const N: usize, const PACK: usize> { + arch: A, + a: APanel<'a, MR, PACK>, + b: BPanel<'a, N>, + compensation: &'a PanelCompensation, + scores: &'a mut [f32; MR], +} + +impl driver::MicroKernel + for MicroKernel<'_, A, MR, N, PACK> +where + A: Architecture + ExtraWide, +{ + #[inline(always)] + fn micro_kernel(&mut self) { + self.arch.run_inline(&mut *self); + } +} + +impl Target + for &mut MicroKernel<'_, A, MR, N, PACK> +where + A: Architecture + ExtraWide, +{ + #[inline(always)] + fn run(self, arch: A) { + let acc = contract(arch, self.a, self.b); + arch.reduce( + acc, + self.compensation, + self.b.compensation(), + self.a.layout().dim(), + self.scores, + ); + } +} + +/// Panel-local SIMD-friendly projection; owned query metadata remains array-of-structs. +struct PanelCompensation { + scale: [f32; MR], + bias: [f32; MR], + sum: [f32; MR], + valid_rows: usize, +} + +impl PanelCompensation { + #[inline(always)] + fn new(query: &[MinMaxCompensation; MR], valid_rows: usize) -> Self { + assert!(valid_rows <= MR); + Self { + scale: core::array::from_fn(|i| query[i].a), + bias: core::array::from_fn(|i| query[i].b), + sum: core::array::from_fn(|i| query[i].n), + valid_rows, + } + } +} + +#[inline(always)] +fn contract( + arch: A, + a: APanel<'_, MR, PACK>, + b: BPanel<'_, N>, +) -> [A::Accumulator; N] +where + A: ExtraWide, +{ + let mut acc = [arch.zero(); N]; + for group in 0..a.layout().k().value().get() / PACK { + // SAFETY: Shared layout pads K to 64, a multiple of PACK; group is in bounds. + let query = arch.load(unsafe { a.group(group) }, a.valid_rows()); + for (doc, acc) in acc.iter_mut().enumerate() { + // SAFETY: doc < N and B uses the same padded K as A. + let values = unsafe { b.group::(doc, group) }; + *acc = arch.dot(query, arch.splat(values), *acc); + } + } + acc +} + +/// Register representations remain private to each architecture backend. +trait ExtraWide: Architecture { + type Query: Copy; + type Splat: Copy; + type Accumulator: Copy; + fn load(self, values: &[[u8; PACK]; MR], valid_rows: usize) -> Self::Query; + fn splat(self, values: [u8; PACK]) -> Self::Splat; + fn zero(self) -> Self::Accumulator; + fn dot(self, query: Self::Query, doc: Self::Splat, acc: Self::Accumulator) + -> Self::Accumulator; + fn reduce( + self, + acc: [Self::Accumulator; N], + query: &PanelCompensation, + docs: &[MinMaxCompensation; N], + dim: usize, + scores: &mut [f32; MR], + ); +} + +impl ExtraWide<8, 4> for Scalar { + type Query = [[u16; 8]; 4]; + type Splat = [u16; 4]; + type Accumulator = [u32; 8]; + #[inline(always)] + fn load(self, values: &[[u8; 4]; 8], _: usize) -> Self::Query { + core::array::from_fn(|d| core::array::from_fn(|row| u16::from(values[row][d]))) + } + #[inline(always)] + fn splat(self, values: [u8; 4]) -> Self::Splat { + values.map(u16::from) + } + #[inline(always)] + fn zero(self) -> Self::Accumulator { + [0; 8] + } + #[inline(always)] + fn dot(self, a: Self::Query, b: Self::Splat, acc: Self::Accumulator) -> Self::Accumulator { + core::array::from_fn(|i| { + let dot = a[0][i] * b[0] + a[1][i] * b[1] + a[2][i] * b[2] + a[3][i] * b[3]; + acc[i].wrapping_add(u32::from(dot)) + }) + } + #[inline(always)] + fn reduce( + self, + acc: [Self::Accumulator; N], + query: &PanelCompensation<8>, + docs: &[MinMaxCompensation; N], + dim: usize, + scores: &mut [f32; 8], + ) { + for (acc, doc) in acc.iter().zip(docs) { + for (i, (&raw, best)) in acc + .iter() + .zip(scores.iter_mut()) + .take(query.valid_rows) + .enumerate() + { + let similarity = query.scale[i] * doc.a * raw as f32 + + query.sum[i] * doc.b + + doc.n * query.bias[i] + + query.bias[i] * doc.b * dim as f32; + *best = best.min(-similarity); + } + } + } +} + +#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] +#[inline(always)] +fn compensate( + arch: F::Arch, + acc: [T; N], + query: &PanelCompensation, + docs: &[MinMaxCompensation; N], + dim: usize, + scores: &mut [f32; MR], + convert: impl Fn(T, usize) -> F, +) where + F: diskann_wide::SIMDVector + + diskann_wide::SIMDMinMax + + std::ops::Add + + std::ops::Sub + + std::ops::Mul, +{ + const { + assert!(MR.is_multiple_of(F::LANES)); + } + for (part, output) in scores + .chunks_exact_mut(F::LANES) + .take(query.valid_rows.div_ceil(F::LANES)) + .enumerate() + { + let start = part * F::LANES; + // SAFETY: Each metadata and output slice contains exactly LANES elements. + let (scale, bias, sum, mut best) = unsafe { + ( + F::load_simd(arch, query.scale[start..].as_ptr()), + F::load_simd(arch, query.bias[start..].as_ptr()), + F::load_simd(arch, query.sum[start..].as_ptr()), + F::load_simd(arch, output.as_ptr()), + ) + }; + for (&raw, doc) in acc.iter().zip(docs) { + let mut similarity = (scale * F::splat(arch, doc.a)) * convert(raw, part); + similarity = similarity + sum * F::splat(arch, doc.b); + similarity = similarity + F::splat(arch, doc.n) * bias; + similarity = similarity + (bias * F::splat(arch, doc.b)) * F::splat(arch, dim as f32); + best = best.min_simd_standard(F::default(arch) - similarity); + } + // SAFETY: The output chunk contains LANES writable elements. + unsafe { best.store_simd(output.as_mut_ptr()) }; + } +} + +#[cfg(target_arch = "x86_64")] +mod x86_64 { + use super::*; + use diskann_wide::{ + SIMDDotProduct, SIMDReinterpret, SIMDVector, + arch::x86_64::{V3, V4}, + }; + + impl ExtraWide<16, 4> for V3 { + type Query = ( + ::u8x32, + Option<::u8x32>, + ); + type Splat = ::i8x32; + type Accumulator = [::i32x8; 2]; + #[inline(always)] + fn load(self, values: &[[u8; 4]; 16], rows: usize) -> Self::Query { + let ptr = values.as_flattened().as_ptr(); + // SAFETY: Each half of the patch spans 32 bytes. + unsafe { + ( + SIMDVector::load_simd(self, ptr), + (rows > 8).then(|| SIMDVector::load_simd(self, ptr.add(32))), + ) + } + } + #[inline(always)] + fn splat(self, value: [u8; 4]) -> Self::Splat { + diskann_wide::alias!(words = ::u32x8); + words::splat(self, u32::from_le_bytes(value)).reinterpret_simd() + } + #[inline(always)] + fn zero(self) -> Self::Accumulator { + [SIMDVector::default(self); 2] + } + #[inline(always)] + fn dot( + self, + a: Self::Query, + b: Self::Splat, + mut acc: Self::Accumulator, + ) -> Self::Accumulator { + diskann_wide::alias!(pairs = ::i16x16); + let dot = |a: ::u8x32, acc: ::i32x8| { + // B is a nibble, so each pair sum is at most 7650 and cannot saturate. + let products = self.mm256_maddubs_epi16(a, b); + acc.dot_simd(products, pairs::splat(self, 1)) + }; + acc[0] = dot(a.0, acc[0]); + if let Some(hi) = a.1 { + acc[1] = dot(hi, acc[1]); + } + acc + } + #[inline(always)] + fn reduce( + self, + acc: [Self::Accumulator; N], + query: &PanelCompensation<16>, + docs: &[MinMaxCompensation; N], + dim: usize, + scores: &mut [f32; 16], + ) { + diskann_wide::alias!(floats = ::f32x8); + compensate( + self, + acc, + query, + docs, + dim, + scores, + |acc: Self::Accumulator, part| { + floats::from_array(self, acc[part].to_array().map(|x| x as u32 as f32)) + }, + ); + } + } + + impl ExtraWide<16, 8> for V4 { + type Query = ( + ::u8x64, + Option<::u8x64>, + ); + type Splat = ::i8x64; + type Accumulator = [::i32x16; 2]; + #[inline(always)] + fn load(self, values: &[[u8; 8]; 16], rows: usize) -> Self::Query { + let ptr = values.as_flattened().as_ptr(); + // SAFETY: Each half of the patch spans 64 bytes. + unsafe { + ( + SIMDVector::load_simd(self, ptr), + (rows > 8).then(|| SIMDVector::load_simd(self, ptr.add(64))), + ) + } + } + #[inline(always)] + fn splat(self, value: [u8; 8]) -> Self::Splat { + diskann_wide::alias!(words = ::u64x8); + words::splat(self, u64::from_le_bytes(value)).reinterpret_simd() + } + #[inline(always)] + fn zero(self) -> Self::Accumulator { + [SIMDVector::default(self); 2] + } + #[inline(always)] + fn dot( + self, + a: Self::Query, + b: Self::Splat, + mut acc: Self::Accumulator, + ) -> Self::Accumulator { + acc[0] = acc[0].dot_simd(a.0, b); + if let Some(hi) = a.1 { + acc[1] = acc[1].dot_simd(hi, b); + } + acc + } + #[inline(always)] + fn reduce( + self, + acc: [Self::Accumulator; N], + query: &PanelCompensation<16>, + docs: &[MinMaxCompensation; N], + dim: usize, + scores: &mut [f32; 16], + ) { + diskann_wide::alias!(floats = ::f32x8); + let convert = |acc: Self::Accumulator, part: usize| { + // PACK8 uses two adjacent VNNI accumulator lanes per query row. + let lanes = acc[part].to_array(); + floats::from_array( + self, + core::array::from_fn(|i| { + lanes[2 * i].wrapping_add(lanes[2 * i + 1]) as u32 as f32 + }), + ) + }; + compensate(self, acc, query, docs, dim, scores, convert); + } + } +} + +#[cfg(target_arch = "aarch64")] +mod aarch64 { + use super::*; + use diskann_wide::{SIMDDotProduct, SIMDVector, arch::aarch64::Neon}; + + impl ExtraWide<8, 4> for Neon { + type Query = ( + ::u8x16, + Option<::u8x16>, + ); + type Splat = ::u8x16; + type Accumulator = [::u32x4; 2]; + #[inline(always)] + fn load(self, values: &[[u8; 4]; 8], rows: usize) -> Self::Query { + let ptr = values.as_flattened().as_ptr(); + // SAFETY: Each half of the patch spans 16 bytes. + unsafe { + ( + SIMDVector::load_simd(self, ptr), + (rows > 4).then(|| SIMDVector::load_simd(self, ptr.add(16))), + ) + } + } + #[inline(always)] + fn splat(self, value: [u8; 4]) -> Self::Splat { + Self::Splat::from_array(self, core::array::from_fn(|i| value[i % 4])) + } + #[inline(always)] + fn zero(self) -> Self::Accumulator { + [SIMDVector::default(self); 2] + } + #[inline(always)] + fn dot( + self, + a: Self::Query, + b: Self::Splat, + mut acc: Self::Accumulator, + ) -> Self::Accumulator { + acc[0] = acc[0].dot_simd(a.0, b); + if let Some(hi) = a.1 { + acc[1] = acc[1].dot_simd(hi, b); + } + acc + } + #[inline(always)] + fn reduce( + self, + acc: [Self::Accumulator; N], + query: &PanelCompensation<8>, + docs: &[MinMaxCompensation; N], + dim: usize, + scores: &mut [f32; 8], + ) { + diskann_wide::alias!(floats = ::f32x4); + compensate( + self, + acc, + query, + docs, + dim, + scores, + |acc: Self::Accumulator, part| { + floats::from_array(self, acc[part].to_array().map(|x| x as f32)) + }, + ); + } + } +} + +#[cfg(test)] +mod tests { + #[cfg(target_arch = "x86_64")] + #[test] + fn unsigned_accumulators_are_preserved_during_conversion_and_pair_folding() { + use diskann_wide::{ + SIMDVector, + arch::x86_64::{V3, V4}, + }; + let query = PanelCompensation::new( + &[MinMaxCompensation { + a: 1.0, + ..Default::default() + }; 16], + 16, + ); + let docs = [MinMaxCompensation { + a: 1.0, + ..Default::default() + }; 1]; + if let Some(arch) = V3::new_checked() { + arch.run_inline( + #[inline] + || { + let raw = 0xe000_0000u32; + let acc = [[::i32x8::splat(arch, raw as i32); 2]; 1]; + let mut scores = [f32::MAX; 16]; + arch.reduce(acc, &query, &docs, 4, &mut scores); + assert_eq!(scores, [-(raw as f32); 16]); + }, + ); + } + if let Some(arch) = V4::new_checked_miri() { + arch.run_inline( + #[inline] + || { + let codes = [0, 1, 2, 3, 12, 13, 14, 15]; + assert_eq!( + arch.splat(codes).to_array(), + core::array::from_fn(|i| codes[i % 8] as i8) + ); + let values: [[i32; 16]; 2] = core::array::from_fn(|part| { + core::array::from_fn(|lane| 0x7000_0000 + (part * 16 + lane) as i32 * 1024) + }); + let acc = + [values.map(|part| ::i32x16::from_array(arch, part))]; + let mut scores = [f32::MAX; 16]; + arch.reduce(acc, &query, &docs, 8, &mut scores); + let expected = core::array::from_fn(|row| { + let part = &values[row / 8]; + let lane = 2 * (row % 8); + -(part[lane].wrapping_add(part[lane + 1]) as u32 as f32) + }); + assert_eq!(scores, expected); + }, + ); + } + } + + use super::*; + use crate::{ + bits::{Representation, Unsigned}, + minmax::{MinMaxCompensation, MinMaxMeta}, + multi_vector::{BlockTransposed, Defaulted, Mat, MaxSim, QueryMatRef}, + }; + use diskann_utils::ReborrowMut; + use diskann_vector::DistanceFunctionMut; + use diskann_wide::arch::Scalar; + const NR: usize = 5; + const PACK: usize = 4; + + fn matrix(rows: usize, dim: usize) -> Mat> + where + Unsigned: Representation, + { + let mut result = Mat::new(MinMaxMeta::new(rows, dim), Defaulted).unwrap(); + for (i, mut row) in result.reborrow_mut().rows_mut().enumerate() { + let mut sum = 0u32; + for d in 0..dim { + let code = ((i * 17 + d * 13) % (1 << BITS)) as u32; + row.vector_mut().set(d, i64::from(code)).unwrap(); + sum += code; + } + let scale = (i + 1) as f32 / 10.0; + row.set_meta(MinMaxCompensation { + dim: dim as u32, + a: scale, + b: -(i as f32) / 7.0, + n: scale * sum as f32, + ..Default::default() + }); + } + result + } + + fn check_driver( + arch: A, + dimensions: &[usize], + panels: usize, + ) where + A: Decoder + ExtraWide + LoadStore, + { + for &dim in dimensions { + let query = matrix::<8>(panels * MR + 1, dim); + let docs = matrix::<4>(panels * NR + 1, dim); + let layout = layout::Layout::new(dim).unwrap(); + let mut values = + BlockTransposed::::try_new(query.num_vectors(), layout.padded_dim()) + .unwrap(); + let mut metadata = + vec![[MinMaxCompensation::default(); MR]; query.num_vectors().div_ceil(MR)]; + for (i, row) in query.as_view().rows().enumerate() { + let mut packed_row = values.get_row_mut(i).unwrap(); + for (d, &value) in row.vector().as_slice().iter().enumerate() { + packed_row.set(layout::Layout::position(d), value); + } + metadata[i / MR][i % MR] = row.meta(); + } + let a = AView::new(values.as_view(), &metadata, &layout).unwrap(); + let mut expected = vec![0.0; query.num_vectors()]; + MaxSim::new(&mut expected).evaluate(QueryMatRef::from(query.as_view()), docs.as_view()); + let k = a.layout().k().value().get(); + let a_bytes = MR * k + std::mem::size_of::<[MinMaxCompensation; MR]>(); + let b_bytes = k + std::mem::size_of::(); + for b_rows in [NR, 2 * NR] { + let cache = Cache::new(value_or_one(a_bytes + b_rows * b_bytes), value_or_one(1)); + let mut scores = vec![123.0; query.num_vectors() + 2]; + let mut driver = Driver::<_, MR, NR, PACK>::new( + arch, + a, + RawRows::new(docs.as_view()), + &mut scores[1..query.num_vectors() + 1], + cache, + ); + for _ in 0..2 { + driver::Drive::drive(&mut driver); + } + assert_eq!(&scores[1..query.num_vectors() + 1], expected); + assert_eq!((scores[0], scores[query.num_vectors() + 1]), (123.0, 123.0)); + } + } + } + + fn check_driver_architectures(dimensions: &[usize], panels: usize) { + check_driver::<_, 8, 5, 4>(Scalar::new(), dimensions, panels); + #[cfg(target_arch = "x86_64")] + { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check_driver::<_, 16, 5, 4>(arch, dimensions, panels); + } + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() { + check_driver::<_, 16, 8, 8>(arch, dimensions, panels); + } + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check_driver::<_, 8, 8, 4>(arch, dimensions, panels); + } + } + + #[test] + fn query_panels_and_document_tiles_keep_metadata_aligned() { + check_driver_architectures(&[1, 63, 65, 129], 3); + } + + #[test] + fn small_odd_dimension_panels_match_pairwise() { + check_driver_architectures(&[3], 1); + } + + fn check_inactive_half, const MR: usize, const PACK: usize>(arch: A) { + arch.run_inline( + #[inline] + || { + let query = PanelCompensation::new( + &[MinMaxCompensation { + a: 1.0, + b: 1.0, + n: 1.0, + ..Default::default() + }; MR], + MR / 2, + ); + let docs = [MinMaxCompensation { + a: 1.0, + b: 1.0, + n: 1.0, + ..Default::default() + }]; + let mut scores = [123.0; MR]; + arch.reduce([arch.zero()], &query, &docs, 1, &mut scores); + assert!(scores[..MR / 2].iter().all(|&score| score == -3.0)); + assert!(scores[MR / 2..].iter().all(|&score| score == 123.0)); + }, + ); + } + + #[test] + fn compensation_skips_the_inactive_query_half() { + check_inactive_half::<_, 8, 4>(Scalar::new()); + #[cfg(target_arch = "x86_64")] + { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check_inactive_half::<_, 16, 4>(arch); + } + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() { + check_inactive_half::<_, 16, 8>(arch); + } + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check_inactive_half::<_, 8, 4>(arch); + } + } + + #[test] + fn reused_b_scratch_fully_overwrites_codes_padding_and_metadata() { + let docs = matrix::<4>(NR + 1, 65); + let layout = layout::Layout::new(65).unwrap(); + let mut scratch = Scratch::default(); + // First populate a full panel, then decode a smaller tail into the same buffers. + RawRows::new(docs.as_view()).visit_tiles(value_or_one(NR), |raw| { + let tile = scratch.prepare(Scalar::new(), raw, &layout); + let mut visited = 0; + let tail = tile.visit_panels::(|b| { + for row in 0..NR { + assert!(b.compensation()[row].a.is_finite()); + visited += 1; + } + }); + if let Some(tail) = tail { + let panel = tail.panel::<1>().unwrap(); + assert_eq!(panel.compensation()[0], docs.get_row(NR).unwrap().meta()); + // SAFETY: The layout contains 128 columns, so the last pack is complete. + assert_eq!(unsafe { panel.group::(0, 128 / PACK - 1) }, [0; PACK]); + visited += 1; + } + assert_eq!(visited, raw.rows()); + }); + } +} diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs new file mode 100644 index 0000000000..a628fb681f --- /dev/null +++ b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs @@ -0,0 +1,101 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Nibble expansion, independent of query packing and metadata. + +pub(super) const BYTES: usize = 32; + +use diskann_wide::{Architecture, arch::Scalar}; + +pub(super) trait Decoder: Architecture { + fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]); + + #[expect( + clippy::expect_used, + reason = "a full decoder block contains exactly BYTES bytes" + )] + #[inline(always)] + fn unpack(self, packed: &[u8], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { + assert!(packed.len() <= BYTES, "packed block exceeds decoder width"); + if packed.len() == BYTES { + self.block(packed.try_into().expect("full decoder block"), low, high); + } else { + let mut tail = [0; BYTES]; + tail[..packed.len()].copy_from_slice(packed); + self.block(&tail, low, high); + } + } +} + +impl Decoder for Scalar { + #[inline(always)] + fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { + for (i, &byte) in packed.iter().enumerate() { + low[i] = byte & 15; + high[i] = byte >> 4; + } + } +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use super::*; + use diskann_wide::{ + SIMDVector, + arch::x86_64::{V3, V4}, + }; + + impl Decoder for V3 { + #[inline(always)] + fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { + diskann_wide::alias!(words = ::u32x8); + diskann_wide::alias!(bytes = ::u8x32); + // Native u32 shifts avoid AVX2's emulated per-byte shifts. + let mask = words::splat(self, 0x0f0f_0f0f); + // SAFETY: packed spans 32 bytes; SIMD loads do not require alignment. + let codes = unsafe { words::load_simd(self, packed.as_ptr().cast()) }; + let lows = bytes::from_underlying(self, (codes & mask).to_underlying()); + let highs = bytes::from_underlying(self, ((codes >> 4) & mask).to_underlying()); + // SAFETY: Both output channels span 32 writable bytes. + unsafe { + lows.store_simd(low.as_mut_ptr()); + highs.store_simd(high.as_mut_ptr()); + } + } + } + + impl Decoder for V4 { + #[inline(always)] + fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { + self.retarget().block(packed, low, high); + } + #[inline(always)] + fn unpack(self, packed: &[u8], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { + assert!(packed.len() <= BYTES, "packed block exceeds decoder width"); + diskann_wide::alias!(bytes = ::u8x32); + // SAFETY: Only the available bytes are read; the other lanes are zero. + let codes = unsafe { bytes::load_simd_first(self, packed.as_ptr(), packed.len()) }; + self.retarget().block(&codes.to_array(), low, high); + } + } +} + +#[cfg(target_arch = "aarch64")] +impl Decoder for diskann_wide::arch::aarch64::Neon { + #[inline(always)] + fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { + use diskann_wide::SIMDVector; + diskann_wide::alias!(bytes = ::u8x16); + let mask = bytes::splat(self, 15); + for offset in [0, 16] { + // SAFETY: Each load and store stays within its 32-byte channel. + unsafe { + let codes = bytes::load_simd(self, packed.as_ptr().add(offset)); + (codes & mask).store_simd(low.as_mut_ptr().add(offset)); + ((codes >> 4) & mask).store_simd(high.as_mut_ptr().add(offset)); + } + } + } +} diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs new file mode 100644 index 0000000000..f4bf60993f --- /dev/null +++ b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs @@ -0,0 +1,211 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Shared order and borrowed panels for MinMax query packing and document decoding. +//! Each 64-dimensional block stores even coordinates followed by odd coordinates. +//! Missing coordinates are zero-padded before applying the block-transposed PACK layout. + +use super::decode::{self, Decoder}; +use crate::{ + matrix_kernels::{ + blocks::packed, + num::{DimK, value_or_one}, + ptr::Slice, + }, + minmax::MinMaxCompensation, + multi_vector::{BlockTransposedRef, Overflow}, +}; + +const BLOCK: usize = 2 * decode::BYTES; +pub(crate) const MAX_DIM: usize = u32::MAX as usize / (255 * 15); + +#[derive(Debug)] +pub(crate) struct Layout { + dim: usize, + padded: usize, +} + +impl Layout { + pub(crate) fn new(dim: usize) -> Result { + let padded = dim + .checked_next_multiple_of(BLOCK) + .ok_or_else(|| Overflow::for_type::(1, dim))?; + Ok(Self { dim, padded }) + } + pub(crate) fn dim(&self) -> usize { + self.dim + } + pub(crate) fn padded_dim(&self) -> usize { + self.padded + } + pub(super) fn k(&self) -> DimK { + DimK::new(value_or_one(self.padded)) + } + pub(crate) fn position(d: usize) -> usize { + d / BLOCK * BLOCK + d % 2 * decode::BYTES + d % BLOCK / 2 + } + + #[expect( + clippy::expect_used, + reason = "padded blocks contain two complete decoder channels" + )] + #[inline(always)] + pub(super) fn decode_row(&self, arch: impl Decoder, packed: &[u8], output: &mut [u8]) { + assert_eq!(packed.len(), self.dim.div_ceil(2)); + assert_eq!(output.len(), self.padded); + let (blocks, tail) = packed.as_chunks::<{ decode::BYTES }>(); + let (decoded, decoded_tail) = output.split_at_mut(blocks.len() * BLOCK); + for (source, block) in blocks.iter().zip(decoded.as_chunks_mut::().0) { + let (low, high) = block.split_at_mut(decode::BYTES); + arch.block( + source, + low.try_into().expect("low channel"), + high.try_into().expect("high channel"), + ); + } + if !tail.is_empty() { + let (low, high) = decoded_tail.split_at_mut(decode::BYTES); + arch.unpack( + tail, + low.try_into().expect("low channel"), + high.try_into().expect("high channel"), + ); + } + if !self.dim.is_multiple_of(2) { + output[Self::position(self.dim)] = 0; + } + } +} + +#[derive(Clone, Copy)] +pub(crate) struct AView<'a, const MR: usize, const PACK: usize> { + values: packed::View<'a, u8, MR, PACK>, + compensation: &'a [[MinMaxCompensation; MR]], + rows: usize, + layout: &'a Layout, +} + +impl<'a, const MR: usize, const PACK: usize> AView<'a, MR, PACK> { + pub(crate) fn new( + values: BlockTransposedRef<'a, u8, MR, PACK>, + compensation: &'a [[MinMaxCompensation; MR]], + layout: &'a Layout, + ) -> Option { + const { assert!(BLOCK.is_multiple_of(PACK)) }; + assert_eq!(values.ncols(), layout.padded_dim(), "query layout mismatch"); + assert_eq!( + values.num_blocks(), + compensation.len(), + "query metadata mismatch" + ); + Some(Self { + values: packed::View::from_block_transposed(values)?, + compensation, + rows: values.nrows(), + layout, + }) + } + + pub(super) fn rows(self) -> usize { + self.rows + } + pub(super) fn layout(self) -> &'a Layout { + self.layout + } + #[inline(always)] + pub(super) fn visit_panels(self, mut visit: impl FnMut(APanel<'_, MR, PACK>, usize)) { + // SAFETY: Each packed panel has matching metadata and zero row/column padding. + unsafe { + self.values.visit_panels(self.layout.k(), |values, block| { + visit( + APanel { + values, + compensation: &self.compensation[block], + valid_rows: (self.rows - block * MR).min(MR), + layout: self.layout, + }, + block * MR, + ); + }); + } + } +} + +#[derive(Clone, Copy)] +pub(super) struct APanel<'a, const MR: usize, const PACK: usize> { + values: packed::Panel<'a, u8, MR, PACK>, + compensation: &'a [MinMaxCompensation; MR], + valid_rows: usize, + layout: &'a Layout, +} + +impl APanel<'_, MR, PACK> { + pub(super) fn layout(&self) -> &Layout { + self.layout + } + + pub(super) fn valid_rows(&self) -> usize { + self.valid_rows + } + pub(super) fn compensation(&self) -> &[MinMaxCompensation; MR] { + self.compensation + } + + /// # Safety + /// + /// group < K / PACK. Query padding makes every group a complete MR x PACK patch. + pub(super) unsafe fn group(&self, group: usize) -> &[[u8; PACK]; MR] { + // SAFETY: The caller provides an in-bounds group; each spans MR * PACK elements. + let patch: Slice<'_, u8> = unsafe { + self.values + .as_ptr() + .add(self.values.pack_stride() * group) + .truncate(self.values.pack_stride()) + }; + // SAFETY: This span contains exactly MR * PACK consecutive bytes. + unsafe { &*patch.as_ptr().cast::<[[u8; PACK]; MR]>() } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn check_decode(arch: impl Decoder) { + for dim in 0..=129 { + let layout = Layout::new(dim).unwrap(); + let packed = vec![0xa5; dim.div_ceil(2)]; + let mut output = vec![0xfe; layout.padded + 2]; + layout.decode_row(arch, &packed, &mut output[1..layout.padded + 1]); + let mut expected = vec![0; layout.padded]; + for block in (0..dim).step_by(BLOCK) { + let count = (dim - block).min(BLOCK); + expected[block..block + count.div_ceil(2)].fill(5); + expected[block + decode::BYTES..block + decode::BYTES + count / 2].fill(10); + } + assert_eq!(&output[1..layout.padded + 1], expected, "dim={dim}"); + assert_eq!((output[0], output[layout.padded + 1]), (0xfe, 0xfe)); + } + assert!(Layout::new(usize::MAX).is_err()); + } + + #[test] + fn all_partial_blocks_ignore_unused_nibbles_and_overwrite_padding() { + check_decode(diskann_wide::arch::Scalar::new()); + #[cfg(target_arch = "x86_64")] + { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check_decode(arch); + } + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() { + check_decode(arch); + } + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check_decode(arch); + } + } +} diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/panels.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/panels.rs new file mode 100644 index 0000000000..8b08a13d10 --- /dev/null +++ b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/panels.rs @@ -0,0 +1,191 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Canonical B tiles and prepared panels whose values and compensation cannot drift apart. + +use super::{decode::Decoder, layout::Layout}; +use crate::{ + matrix_kernels::{blocks::unpacked, num::Elements, ptr::Slice}, + minmax::{MinMaxCompensation, MinMaxMeta}, + multi_vector::MatRef, +}; +use std::num::NonZeroUsize; + +#[derive(Clone, Copy)] +pub(crate) struct RawRows<'a> { + doc: MatRef<'a, MinMaxMeta<4>>, + start: usize, + rows: usize, +} + +impl<'a> RawRows<'a> { + pub(crate) fn new(doc: MatRef<'a, MinMaxMeta<4>>) -> Self { + Self { + doc, + start: 0, + rows: doc.num_vectors(), + } + } + pub(super) fn dim(self) -> usize { + self.doc.repr().intrinsic_dim() + } + pub(super) fn rows(self) -> usize { + self.rows + } + #[inline(always)] + pub(super) fn visit_tiles(self, rows: NonZeroUsize, mut visit: impl FnMut(RawRows<'_>)) { + for offset in (0..self.rows).step_by(rows.get()) { + visit(Self { + doc: self.doc, + start: self.start + offset, + rows: (self.rows - offset).min(rows.get()), + }); + } + } +} + +#[derive(Default)] +pub(super) struct Scratch { + values: Vec, + compensation: Vec, +} + +impl Scratch { + #[expect( + clippy::expect_used, + reason = "tile extents are nonzero and rows retain canonical provenance" + )] + #[inline(always)] + pub(super) fn prepare<'a>( + &'a mut self, + arch: impl Decoder, + raw: RawRows<'_>, + layout: &'a Layout, + ) -> BTile<'a> { + assert_eq!(raw.dim(), layout.dim(), "document dimension mismatch"); + let k = layout.k(); + let rows = NonZeroUsize::new(raw.rows).expect("nonempty B tile"); + let len = raw + .rows + .checked_mul(k.value().get()) + .expect("B scratch extent overflow"); + if self.values.len() < len { + self.values.resize(len, 0); + } + if self.compensation.len() < raw.rows { + self.compensation + .resize(raw.rows, MinMaxCompensation::default()); + } + let values = &mut self.values[..len]; + let compensation = &mut self.compensation[..raw.rows]; + for (i, (output, metadata)) in values + .chunks_exact_mut(k.value().get()) + .zip(compensation.iter_mut()) + .enumerate() + { + let source = raw + .doc + .get_row(raw.start + i) + .expect("canonical row in bounds"); + *metadata = source.meta(); + layout.decode_row(arch, source.vector().as_slice(), output); + } + BTile { + // SAFETY: Preparation completely writes rows * K values and one metadata per row. + values: unsafe { unpacked::View::new(Slice::new(values), rows, k) }, + compensation, + layout, + } + } +} + +#[derive(Clone, Copy)] +pub(super) struct BTile<'a> { + values: unpacked::View<'a, u8>, + compensation: &'a [MinMaxCompensation], + layout: &'a Layout, +} + +impl BTile<'_> { + pub(super) fn layout(&self) -> &Layout { + self.layout + } + #[inline(always)] + pub(super) fn visit_panels( + &self, + mut visit: impl FnMut(BPanel<'_, N>), + ) -> Option> { + const { assert!(N > 0) }; + let metadata = self.compensation.as_chunks::().0; + // SAFETY: Scratch::prepare binds the decoded stride and value/metadata row counts. + let remainder = unsafe { + self.values.visit_panels( + self.layout.k(), + |values: unpacked::Panel<'_, u8, N>, start| { + visit(BPanel { + values, + compensation: &metadata[start / N], + layout: self.layout, + }); + }, + ) + }; + remainder.map(|values| Remainder { + compensation: &self.compensation[values.start()..], + values, + layout: self.layout, + }) + } +} + +pub(super) struct Remainder<'a, const CAP: usize> { + values: unpacked::Remainder<'a, u8, CAP>, + compensation: &'a [MinMaxCompensation], + layout: &'a Layout, +} + +impl Remainder<'_, CAP> { + #[expect( + clippy::expect_used, + reason = "the metadata remainder has the same extent as values" + )] + pub(super) fn panel(&self) -> Option> { + let values = self.values.try_as_panel::()?; + Some(BPanel { + values, + compensation: self.compensation.try_into().expect("exact metadata panel"), + layout: self.layout, + }) + } +} + +#[derive(Clone, Copy)] +pub(super) struct BPanel<'a, const N: usize> { + values: unpacked::Panel<'a, u8, N>, + compensation: &'a [MinMaxCompensation; N], + layout: &'a Layout, +} + +impl BPanel<'_, N> { + pub(super) fn compensation(&self) -> &[MinMaxCompensation; N] { + self.compensation + } + + /// # Safety + /// + /// row < N and group < K / PACK. Layout preparation completely writes K bytes per row. + pub(super) unsafe fn group(&self, row: usize, group: usize) -> [u8; PACK] { + // SAFETY: The caller selects a complete group inside one decoded row. + unsafe { + self.values + .as_ptr() + .add(self.values.stride(self.layout.k()) * row + Elements::new(group * PACK)) + .truncate(Elements::new(PACK)) + .as_ptr() + .cast::<[u8; PACK]>() + .read_unaligned() + } + } +} diff --git a/diskann-quantization/src/matrix_kernels/maxsim/mod.rs b/diskann-quantization/src/matrix_kernels/maxsim/mod.rs index 412f9d98c5..759ec44a41 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/mod.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/mod.rs @@ -16,6 +16,7 @@ //! C[i] = max(j in 0..N, dot(A[i, :], B[:, j])) //! ``` +pub(crate) mod minmax8_x_minmax4; pub(crate) mod packed_f32_x_unpacked_f16; pub(crate) mod packed_f32_x_unpacked_f32; pub(crate) mod packed_i8_x_unpacked_i8; diff --git a/diskann-quantization/src/minmax/mod.rs b/diskann-quantization/src/minmax/mod.rs index 41c29c3d03..d717868152 100644 --- a/diskann-quantization/src/minmax/mod.rs +++ b/diskann-quantization/src/minmax/mod.rs @@ -133,7 +133,7 @@ mod vectors; // Exports // ///////////// -pub use multi::{MinMaxKernel, MinMaxMeta}; +pub use multi::{MinMaxErase, MinMaxKernel, MinMaxMaxSimKernel, MinMaxMeta, build_minmax_max_sim}; pub use quantizer::{L2Loss, MinMaxQuantizer}; pub use recompress::{RecompressError, Recompressor}; pub use vectors::{ diff --git a/diskann-quantization/src/minmax/multi/factory.rs b/diskann-quantization/src/minmax/multi/factory.rs new file mode 100644 index 0000000000..47bf68212e --- /dev/null +++ b/diskann-quantization/src/minmax/multi/factory.rs @@ -0,0 +1,307 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::kernel::{MinMaxErase, MinMaxMaxSimKernel}; +use super::query::PackedQuery; +use crate::{ + matrix_kernels::{ + self as mk, + maxsim::minmax8_x_minmax4::{Driver, layout::MAX_DIM, panels::RawRows}, + }, + minmax::MinMaxMeta, + multi_vector::{MatRef, MaxSimError, MaxSimIsa, NotSupported, Overflow}, +}; +#[cfg(target_arch = "aarch64")] +use diskann_wide::arch::aarch64::Neon; +#[cfg(target_arch = "x86_64")] +use diskann_wide::arch::x86_64::{V3, V4}; +use diskann_wide::{Architecture, arch::Scalar}; + +#[derive(Debug)] +struct Prepared { + arch: A, + query: PackedQuery, +} + +impl MinMaxMaxSimKernel<8, 4> + for Prepared +where + A: Architecture, + for<'a> Driver<'a, A, MR, NR, PACK>: mk::Drive, +{ + fn nrows(&self) -> usize { + self.query.rows() + } + fn compute_max_sim( + &self, + doc: MatRef<'_, MinMaxMeta<4>>, + scores: &mut [f32], + ) -> Result<(), MaxSimError> { + if scores.len() != self.nrows() { + return Err(MaxSimError::InvalidBufferLength(scores.len(), self.nrows())); + } + if doc.repr().intrinsic_dim() != self.query.dim() { + return Err(MaxSimError::UnequalDim( + doc.repr().intrinsic_dim(), + self.query.dim(), + )); + } + if doc.num_vectors() == 0 { + scores.fill(f32::MAX); + return Ok(()); + } + if self.query.dim() == 0 { + scores.fill(0.0); + return Ok(()); + } + let Some(a) = self.query.view() else { + return Ok(()); + }; + let mut driver = Driver::<_, MR, NR, PACK>::new( + self.arch, + a, + RawRows::new(doc), + scores, + mk::Cache::detect(), + ); + mk::Drive::drive(&mut driver); + Ok(()) + } +} + +struct BuildAndErase(E); + +macro_rules! builder { + ($arch:ty, $mr:literal, $nr:literal, $pack:literal) => { + impl> + diskann_wide::arch::Target1< + $arch, + Result, + MatRef<'_, MinMaxMeta<8>>, + > for BuildAndErase + { + fn run( + self, + arch: $arch, + query: MatRef<'_, MinMaxMeta<8>>, + ) -> Result { + Ok(self.0.erase(Prepared::<_, $mr, $nr, $pack> { + arch, + query: PackedQuery::new(query)?, + })) + } + } + }; +} + +builder!(Scalar, 8, 5, 4); +#[cfg(target_arch = "x86_64")] +builder!(V3, 16, 5, 4); +#[cfg(target_arch = "x86_64")] +builder!(V4, 16, 8, 8); +#[cfg(target_arch = "aarch64")] +builder!(Neon, 8, 8, 4); + +/// Prepare a MinMax8-query x MinMax4-document MaxSim kernel. +/// +/// Auto selects the available architecture once. Canonical documents remain borrowed. +/// B-first traversal prepares a document tile once and reuses it across query panels. +/// Values and compensation are always exposed through synchronized panel iterators. +/// +/// # Errors +/// +/// Returns NotSupported for unavailable ISA or Reference requests, dimensions beyond the +/// u32 accumulator limit (1,122,867), or unrepresentable prepared-query extents. +pub fn build_minmax_max_sim>( + isa: MaxSimIsa, + query: MatRef<'_, MinMaxMeta<8>>, + erase: E, +) -> Result { + if query.repr().intrinsic_dim() > MAX_DIM { + return Err(NotSupported { + isa, + reason: "MinMax8 dimension exceeds the u32 accumulator limit", + }); + } + let result = match isa { + MaxSimIsa::Auto => diskann_wide::arch::dispatch1_no_features(BuildAndErase(erase), query), + MaxSimIsa::Scalar => Scalar::new().run1(BuildAndErase(erase), query), + #[cfg(target_arch = "x86_64")] + MaxSimIsa::X86_64_V3 => V3::new_checked() + .ok_or(NotSupported { + isa, + reason: "AVX2/FMA unavailable", + })? + .run1(BuildAndErase(erase), query), + #[cfg(target_arch = "x86_64")] + MaxSimIsa::X86_64_V4 => V4::new_checked() + .ok_or(NotSupported { + isa, + reason: "AVX-512 unavailable", + })? + .run1(BuildAndErase(erase), query), + #[cfg(target_arch = "aarch64")] + MaxSimIsa::Neon => Neon::new_checked() + .ok_or(NotSupported { + isa, + reason: "NEON/dotprod unavailable", + })? + .run1(BuildAndErase(erase), query), + #[cfg(not(target_arch = "x86_64"))] + MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => { + return Err(NotSupported { + isa, + reason: "x86-64 ISA unavailable on this target", + }); + } + #[cfg(not(target_arch = "aarch64"))] + MaxSimIsa::Neon => { + return Err(NotSupported { + isa, + reason: "NEON unavailable on this target", + }); + } + MaxSimIsa::Reference => { + return Err(NotSupported { + isa, + reason: "reference kernel unsupported", + }); + } + }; + result.map_err(|_| NotSupported { + isa, + reason: "prepared MinMax query exceeds the allocation limit", + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + CompressInto, + algorithms::{Transform, transforms::NullTransform}, + bits::{Representation, Unsigned}, + minmax::MinMaxQuantizer, + multi_vector::{BoxErase, Defaulted, Mat, MaxSim, QueryMatRef, Standard}, + num::Positive, + }; + use diskann_utils::ReborrowMut; + use diskann_vector::DistanceFunctionMut; + use std::num::NonZeroUsize; + + fn data(rows: usize, dim: usize, seed: usize) -> Mat> + where + Unsigned: Representation, + { + let mut result = Mat::new(MinMaxMeta::new(rows, dim), Defaulted).unwrap(); + if dim != 0 { + let mut i = seed; + let source = Mat::from_fn(Standard::new(rows, dim).unwrap(), || { + i += 1; + ((i * 17 % 257) as f32 - 128.0) / 37.0 + }); + let quantizer = MinMaxQuantizer::new( + Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())), + Positive::new(1.0).unwrap(), + ); + quantizer + .compress_into(source.as_view(), result.reborrow_mut()) + .unwrap(); + } + result + } + + #[test] + fn tiled_scalar_matches_pairwise() { + for dim in [0, 1, 3, 4, 5, 63, 64, 65, 127, 128, 129, 250, 256, 257] { + for q in (0..=17).chain([32, 33]) { + let query = data::<8>(q, dim, 3); + for isa in [ + MaxSimIsa::Scalar, + MaxSimIsa::Auto, + MaxSimIsa::X86_64_V3, + MaxSimIsa::X86_64_V4, + MaxSimIsa::Neon, + ] + .into_iter() + .filter(|isa| isa.is_available()) + { + let kernel = build_minmax_max_sim(isa, query.as_view(), BoxErase).unwrap(); + for d in [0, 1, 2, 3, 4, 5, 6, 7, 13, 129] { + let doc = data::<4>(d, dim, 11); + let mut expected = vec![0.0; q]; + MaxSim::new(&mut expected) + .evaluate(QueryMatRef::from(query.as_view()), doc.as_view()); + let mut scores = vec![123.0; q + 2]; + for _ in 0..2 { + kernel + .compute_max_sim(doc.as_view(), &mut scores[1..q + 1]) + .unwrap(); + assert_eq!(&scores[1..q + 1], expected, "{q}x{d}x{dim}"); + assert_eq!((scores[0], scores[q + 1]), (123.0, 123.0)); + } + } + } + } + } + } + + #[test] + fn shape_errors_and_unsupported_requests() { + let query = data::<8>(9, 65, 3); + let kernel = build_minmax_max_sim(MaxSimIsa::Scalar, query.as_view(), BoxErase).unwrap(); + let doc = data::<4>(7, 65, 11); + let mut short = [123.0; 8]; + assert!(kernel.compute_max_sim(doc.as_view(), &mut short).is_err()); + assert_eq!(short, [123.0; 8]); + for rows in [0, 7] { + let wrong = data::<4>(rows, 64, 11); + let mut scores = [123.0; 9]; + assert!( + kernel + .compute_max_sim(wrong.as_view(), &mut scores) + .is_err() + ); + assert_eq!(scores, [123.0; 9]); + } + for isa in [ + MaxSimIsa::X86_64_V3, + MaxSimIsa::X86_64_V4, + MaxSimIsa::Neon, + MaxSimIsa::Reference, + ] { + assert_eq!( + build_minmax_max_sim(isa, query.as_view(), BoxErase).is_err(), + isa == MaxSimIsa::Reference || !isa.is_available() + ); + } + let excessive = Mat::new(MinMaxMeta::<8>::new(0, MAX_DIM + 1), Defaulted).unwrap(); + assert!(build_minmax_max_sim(MaxSimIsa::Scalar, excessive.as_view(), BoxErase).is_err()); + } + + #[test] + fn prepared_query_ownership_nan_and_concurrent_calls() { + let query = data::<8>(9, 65, 3); + let mut doc = data::<4>(13, 65, 11); + let kernel = build_minmax_max_sim(MaxSimIsa::Scalar, query.as_view(), BoxErase).unwrap(); + for mut row in doc.reborrow_mut().rows_mut() { + let mut meta = row.meta(); + meta.a = f32::NAN; + row.set_meta(meta); + } + drop(query); + std::thread::scope(|scope| { + for _ in 0..3 { + let kernel = &kernel; + let doc = &doc; + scope.spawn(move || { + let mut scores = [0.0; 9]; + kernel.compute_max_sim(doc.as_view(), &mut scores).unwrap(); + assert_eq!(scores, [f32::MAX; 9]); + }); + } + }); + } +} diff --git a/diskann-quantization/src/minmax/multi/kernel.rs b/diskann-quantization/src/minmax/multi/kernel.rs new file mode 100644 index 0000000000..0281e0ba0c --- /dev/null +++ b/diskann-quantization/src/minmax/multi/kernel.rs @@ -0,0 +1,54 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use crate::{ + bits::{Representation, Unsigned}, + minmax::MinMaxMeta, + multi_vector::{BoxErase, MatRef, MaxSimError}, +}; + +/// Prepared MinMax MaxSim computation with independent per-call scratch. +pub trait MinMaxMaxSimKernel: + Send + Sync + std::fmt::Debug +where + Unsigned: Representation + Representation, +{ + /// Number of query vectors whose distances are written by this kernel. + fn nrows(&self) -> usize; + + /// Overwrite one distance per query row. Empty documents produce f32::MAX; + /// nonempty zero-dimensional documents produce zero. NaN distances are ignored. + /// + /// # Errors + /// + /// Shape errors leave scores unchanged. + fn compute_max_sim( + &self, + doc: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError>; +} + +/// Choose how the factory packages a concrete prepared kernel. +pub trait MinMaxErase +where + Unsigned: Representation + Representation, +{ + type Output; + fn erase + 'static>( + self, + kernel: K, + ) -> Self::Output; +} + +impl MinMaxErase for BoxErase +where + Unsigned: Representation + Representation, +{ + type Output = Box>; + fn erase + 'static>(self, kernel: K) -> Self::Output { + Box::new(kernel) + } +} diff --git a/diskann-quantization/src/minmax/multi/mod.rs b/diskann-quantization/src/minmax/multi/mod.rs index d6ad4c8369..6524eae50c 100644 --- a/diskann-quantization/src/minmax/multi/mod.rs +++ b/diskann-quantization/src/minmax/multi/mod.rs @@ -5,8 +5,13 @@ //! Multi-vector support for MinMax quantized vectors. +mod factory; +mod kernel; mod max_sim; mod meta; +mod query; +pub use factory::build_minmax_max_sim; +pub use kernel::{MinMaxErase, MinMaxMaxSimKernel}; pub use max_sim::MinMaxKernel; pub use meta::MinMaxMeta; diff --git a/diskann-quantization/src/minmax/multi/query.rs b/diskann-quantization/src/minmax/multi/query.rs new file mode 100644 index 0000000000..120b214448 --- /dev/null +++ b/diskann-quantization/src/minmax/multi/query.rs @@ -0,0 +1,60 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use crate::{ + matrix_kernels::maxsim::minmax8_x_minmax4::layout::{AView, Layout}, + minmax::{MinMaxCompensation, MinMaxMeta}, + multi_vector::{BlockTransposed, MatRef, Overflow}, +}; + +#[derive(Debug)] +pub(super) struct PackedQuery { + values: BlockTransposed, + compensation: Vec<[MinMaxCompensation; MR]>, + layout: Layout, +} + +impl PackedQuery { + #[expect( + clippy::expect_used, + reason = "enumerated query rows match the packed allocation" + )] + pub(super) fn new(query: MatRef<'_, MinMaxMeta<8>>) -> Result { + const { assert!(MR > 0 && PACK > 0 && MR.is_multiple_of(PACK)) }; + let rows = query.num_vectors(); + let layout = Layout::new(query.repr().intrinsic_dim())?; + Overflow::check_byte_budget::<[MinMaxCompensation; MR]>( + rows.div_ceil(MR), + rows, + layout.dim(), + )?; + let mut values = BlockTransposed::try_new(rows, layout.padded_dim())?; + let mut compensation = vec![[MinMaxCompensation::default(); MR]; rows.div_ceil(MR)]; + for (row, source) in query.rows().enumerate() { + let mut output = values.get_row_mut(row).expect("query row in bounds"); + for (d, &value) in source.vector().as_slice().iter().enumerate() { + output.set(Layout::position(d), value); + } + compensation[row / MR][row % MR] = source.meta(); + } + Ok(Self { + values, + compensation, + layout, + }) + } + + pub(super) fn rows(&self) -> usize { + self.values.nrows() + } + + pub(super) fn dim(&self) -> usize { + self.layout.dim() + } + + pub(super) fn view(&self) -> Option> { + AView::new(self.values.as_view(), &self.compensation, &self.layout) + } +} From a79c823af6b29221500b74f18d05f73e4d148021 Mon Sep 17 00:00:00 2001 From: Junkui Chen Date: Wed, 14 Oct 2026 18:17:07 +0800 Subject: [PATCH 2/2] Address MinMax MaxSim review feedback Separate layout and operand preparation, validate canonical extents in MinMaxMeta, and use shared span interfaces with backend-specific panel registration. Keep SIMD queries fixed-size and move prepared byte accounting into operand helpers. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- diskann-benchmark/src/multi_vector/driver.rs | 29 +- .../maxsim/minmax8_x_minmax4.rs | 367 ++++++++----- .../maxsim/minmax8_x_minmax4/a.rs | 182 +++++++ .../maxsim/minmax8_x_minmax4/b.rs | 512 ++++++++++++++++++ .../maxsim/minmax8_x_minmax4/decode.rs | 101 ---- .../maxsim/minmax8_x_minmax4/layout.rs | 189 +------ .../maxsim/minmax8_x_minmax4/panels.rs | 191 ------- diskann-quantization/src/minmax/mod.rs | 5 +- .../src/minmax/multi/factory.rs | 4 +- diskann-quantization/src/minmax/multi/meta.rs | 107 +++- diskann-quantization/src/minmax/multi/mod.rs | 2 +- .../src/minmax/multi/query.rs | 78 ++- 12 files changed, 1133 insertions(+), 634 deletions(-) create mode 100644 diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/a.rs create mode 100644 diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/b.rs delete mode 100644 diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs delete mode 100644 diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/panels.rs diff --git a/diskann-benchmark/src/multi_vector/driver.rs b/diskann-benchmark/src/multi_vector/driver.rs index 256f778232..339c7005a7 100644 --- a/diskann-benchmark/src/multi_vector/driver.rs +++ b/diskann-benchmark/src/multi_vector/driver.rs @@ -6,7 +6,6 @@ //! Shared benchmark infrastructure for multi-vector kernels: timing harness, //! data fixtures, result types. None of the contents are kernel-aware. -use anyhow::Context; use diskann_benchmark_runner::{ utils::{ fmt::Table, @@ -20,7 +19,6 @@ use diskann_quantization::multi_vector::{ }; use diskann_quantization::{ algorithms::{transforms::NullTransform, Transform}, - bits::{Representation, Unsigned}, minmax::{MinMaxMeta, MinMaxQuantizer}, num::Positive, CompressInto, @@ -97,29 +95,6 @@ where } } -// MinMaxMeta::new does not check canonical allocation extents. -fn checked_minmax_meta( - rows: usize, - dim: usize, -) -> anyhow::Result> -where - Unsigned: Representation, -{ - anyhow::ensure!( - u32::try_from(dim).is_ok(), - "MinMax dimension exceeds u32::MAX" - ); - dim.checked_mul(BITS) - .and_then(|bits| bits.checked_add(7)) - .context("MinMax row size overflow")?; - let repr = MinMaxMeta::new(rows, dim); - let bytes = rows - .checked_mul(repr.ncols()) - .context("MinMax matrix size overflow")?; - std::alloc::Layout::array::(bytes).context("MinMax allocation exceeds the byte limit")?; - Ok(repr) -} - pub(super) struct MinMax8Data { pub(super) queries: Mat>, pub(super) docs: Mat>, @@ -127,8 +102,8 @@ pub(super) struct MinMax8Data { impl MinMax8Data { pub(super) fn new(run: &Run) -> anyhow::Result { - let query_repr = checked_minmax_meta::<8>(run.num_query_vectors.get(), run.dim.get())?; - let doc_repr = checked_minmax_meta::<4>(run.num_doc_vectors.get(), run.dim.get())?; + let query_repr = MinMaxMeta::<8>::try_new(run.num_query_vectors.get(), run.dim.get())?; + let doc_repr = MinMaxMeta::<4>::try_new(run.num_doc_vectors.get(), run.dim.get())?; let data = Data::::new(run)?; let quantizer = MinMaxQuantizer::new( Transform::Null(NullTransform::new(run.dim)), diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs index 2150d4df90..64ce26c828 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4.rs @@ -6,24 +6,38 @@ //! B-first MinMax MaxSim. Preparation and synchronized panel iterators //! establish the micro-kernel's inputs; integer contraction never reads canonical data. -mod decode; +pub(crate) mod a; +pub(crate) mod b; pub(crate) mod layout; -pub(crate) mod panels; use super::packed_f32_x_unpacked_f32::Params; use crate::matrix_kernels::{ - Cache, driver, - num::{Bytes, value_or_one}, + Cache, bounds, driver, + num::{Elements, value_or_one}, + ptr::Slice, util::LoadStore, }; use crate::minmax::MinMaxCompensation; -use decode::Decoder; +use a::{APanel, AView}; +use b::{BPanel, BTile, RawRows, Scratch, UnpackNibbles}; use diskann_wide::{ Architecture, arch::{Scalar, Target}, }; -use layout::{APanel, AView}; -use panels::{BPanel, BTile, RawRows, Scratch}; + +diskann_wide::alias!(f32x4 = f32x4); +diskann_wide::alias!(f32x8 = f32x8); +diskann_wide::alias!(i8x32 = i8x32); +diskann_wide::alias!(i8x64 = i8x64); +diskann_wide::alias!(i16x16 = i16x16); +diskann_wide::alias!(i32x8 = i32x8); +diskann_wide::alias!(i32x16 = i32x16); +diskann_wide::alias!(u8x16 = u8x16); +diskann_wide::alias!(u8x32 = u8x32); +diskann_wide::alias!(u8x64 = u8x64); +diskann_wide::alias!(u32x4 = u32x4); +diskann_wide::alias!(u32x8 = u32x8); +diskann_wide::alias!(u64x8 = u64x8); pub(crate) struct Driver<'a, A, const MR: usize, const NR: usize, const PACK: usize> { arch: A, @@ -43,7 +57,7 @@ impl<'a, A, const MR: usize, const NR: usize, const PACK: usize> Driver<'a, A, M cache: Cache, ) -> Self { const { - assert!(NR > 0 && NR <= 8); + assert!(NR > 0); } assert_eq!(a.rows(), c.len(), "output row mismatch"); assert_eq!(a.layout().dim(), b.dim(), "document dimension mismatch"); @@ -51,13 +65,7 @@ impl<'a, A, const MR: usize, const NR: usize, const PACK: usize> Driver<'a, A, M a.layout().dim() <= layout::MAX_DIM, "integer accumulator dimension limit" ); - let k = a.layout().k().value().get(); - let params = Params::new( - cache, - Bytes::new(MR * k + std::mem::size_of::<[crate::minmax::MinMaxCompensation; MR]>()), - Bytes::new(k + std::mem::size_of::()), - NR, - ); + let params = Params::new(cache, a.panel_bytes(), Scratch::row_bytes(a.layout()), NR); Self { arch, a, @@ -72,7 +80,8 @@ impl<'a, A, const MR: usize, const NR: usize, const PACK: usize> Driver<'a, A, M impl driver::Drive for Driver<'_, A, MR, NR, PACK> where - A: Decoder + ExtraWide + LoadStore, + A: UnpackNibbles + ExtraWide + LoadStore, + for<'b> PanelKernel<'b, A, MR, NR, PACK>: driver::PanelKernel, { fn drive(&mut self) { self.arch.run_inline(&mut *self); @@ -82,7 +91,8 @@ where impl Target for &mut Driver<'_, A, MR, NR, PACK> where - A: Decoder + ExtraWide + LoadStore, + A: UnpackNibbles + ExtraWide + LoadStore, + for<'b> PanelKernel<'b, A, MR, NR, PACK>: driver::PanelKernel, { #[inline(always)] fn run(self, arch: A) { @@ -121,9 +131,10 @@ impl<'a, A, const MR: usize, const NR: usize, const PACK: usize> PanelKernel<'a, where A: Architecture + LoadStore, { - assert!( - std::ptr::eq(a.layout(), b.layout()), - "A and B must share the same layout" + assert_eq!( + a.layout(), + b.layout(), + "A and B must use matching MinMax layouts" ); assert_eq!(output.len(), a.valid_rows()); Self { @@ -140,37 +151,35 @@ impl<'a, A, const MR: usize, const NR: usize, const PACK: usize> PanelKernel<'a, where A: Architecture + ExtraWide, { - let mut micro = MicroKernel { + let mut kernel = MicroKernel { arch: self.arch, a: self.a, b, compensation: &self.compensation, scores: &mut self.scores, }; - driver::MicroKernel::micro_kernel(&mut micro); + driver::MicroKernel::micro_kernel(&mut kernel); } } -impl driver::PanelKernel - for PanelKernel<'_, A, MR, NR, PACK> -where - A: Architecture + ExtraWide, -{ - #[inline(always)] - fn panel_kernel(&mut self) { - let b = self.b; - let remainder = b.visit_panels::(|panel| self.micro_kernel(panel)); - if let Some(remainder) = remainder { - macro_rules! tail { - ($($n:literal),+) => {$( - if $n < NR { - if let Some(panel) = remainder.panel::<$n>() { self.micro_kernel(panel); } - } - )+}; +macro_rules! panel_kernel { + ($arch:ty, $mr:literal, $nr:literal, $pack:literal, [$($n:literal),+ $(,)?]) => { + impl driver::PanelKernel for PanelKernel<'_, $arch, $mr, $nr, $pack> { + #[inline(always)] + fn panel_kernel(&mut self) { + let b = self.b; + let remainder = b.visit_panels::<$nr>(|panel| self.micro_kernel(panel)); + if let Some(remainder) = remainder { + $( + const { assert!($n > 0 && $n < $nr) }; + if let Some(panel) = remainder.panel::<$n>() { + self.micro_kernel(panel); + } + )+ + } } - tail!(1, 2, 3, 4, 5, 6, 7); } - } + }; } struct MicroKernel<'a, A, const MR: usize, const N: usize, const PACK: usize> { @@ -242,24 +251,48 @@ where { let mut acc = [arch.zero(); N]; for group in 0..a.layout().k().value().get() / PACK { - // SAFETY: Shared layout pads K to 64, a multiple of PACK; group is in bounds. - let query = arch.load(unsafe { a.group(group) }, a.valid_rows()); + // SAFETY: Preparation pads K to a multiple of PACK; each A group spans MR * PACK. + let query = unsafe { arch.load(a.group(group)) }; for (doc, acc) in acc.iter_mut().enumerate() { - // SAFETY: doc < N and B uses the same padded K as A. - let values = unsafe { b.group::(doc, group) }; - *acc = arch.dot(query, arch.splat(values), *acc); + // SAFETY: doc < N and B uses the same padded K; each B group spans PACK. + let values = unsafe { arch.splat(b.group::(doc, group)) }; + *acc = arch.dot(query, values, *acc); } } acc } +/// Read one complete prepared B group. +/// +/// # Safety +/// +/// `values.len()` must be exactly `PACK`. +#[inline(always)] +unsafe fn group(values: Slice<'_, u8>) -> [u8; PACK] { + bounds::check_eq!(values.len(), PACK); + core::array::from_fn(|i| { + // SAFETY: i < PACK, so the offset selects one readable element. + unsafe { *values.add(Elements::new(i)).as_unit().as_ref() } + }) +} + /// Register representations remain private to each architecture backend. trait ExtraWide: Architecture { type Query: Copy; type Splat: Copy; type Accumulator: Copy; - fn load(self, values: &[[u8; PACK]; MR], valid_rows: usize) -> Self::Query; - fn splat(self, values: [u8; PACK]) -> Self::Splat; + /// Load a complete prepared A patch, including zero-padded query rows. + /// + /// # Safety + /// + /// `slice.len()` must be exactly `MR * PACK`. + unsafe fn load(self, slice: Slice<'_, u8>) -> Self::Query; + /// Broadcast one group of unpacked 4-bit B codes. + /// + /// # Safety + /// + /// `slice.len()` must be exactly `PACK`. + unsafe fn splat(self, slice: Slice<'_, u8>) -> Self::Splat; fn zero(self) -> Self::Accumulator; fn dot(self, query: Self::Query, doc: Self::Splat, acc: Self::Accumulator) -> Self::Accumulator; @@ -273,17 +306,22 @@ trait ExtraWide: Architecture { ); } +panel_kernel!(Scalar, 8, 5, 4, [1, 2, 3, 4]); + impl ExtraWide<8, 4> for Scalar { type Query = [[u16; 8]; 4]; type Splat = [u16; 4]; type Accumulator = [u32; 8]; #[inline(always)] - fn load(self, values: &[[u8; 4]; 8], _: usize) -> Self::Query { - core::array::from_fn(|d| core::array::from_fn(|row| u16::from(values[row][d]))) + unsafe fn load(self, slice: Slice<'_, u8>) -> Self::Query { + // SAFETY: By precondition, the patch contains MR * PACK = 32 bytes. + let values = unsafe { slice.as_std_slice(32) }; + core::array::from_fn(|d| core::array::from_fn(|row| u16::from(values[row * 4 + d]))) } #[inline(always)] - fn splat(self, values: [u8; 4]) -> Self::Splat { - values.map(u16::from) + unsafe fn splat(self, slice: Slice<'_, u8>) -> Self::Splat { + // SAFETY: By precondition, the span contains one complete PACK4 group. + unsafe { group::<4>(slice) }.map(u16::from) } #[inline(always)] fn zero(self) -> Self::Accumulator { @@ -377,28 +415,29 @@ mod x86_64 { arch::x86_64::{V3, V4}, }; + panel_kernel!(V3, 16, 5, 4, [1, 2, 3, 4]); + panel_kernel!(V4, 16, 8, 8, [1, 2, 3, 4, 5, 6, 7]); + impl ExtraWide<16, 4> for V3 { - type Query = ( - ::u8x32, - Option<::u8x32>, - ); - type Splat = ::i8x32; - type Accumulator = [::i32x8; 2]; + type Query = [u8x32; 2]; + type Splat = i8x32; + type Accumulator = [i32x8; 2]; #[inline(always)] - fn load(self, values: &[[u8; 4]; 16], rows: usize) -> Self::Query { - let ptr = values.as_flattened().as_ptr(); - // SAFETY: Each half of the patch spans 32 bytes. + unsafe fn load(self, slice: Slice<'_, u8>) -> Self::Query { + bounds::check_eq!(slice.len(), 64); + // SAFETY: The 64-byte patch contains two complete 32-byte halves. unsafe { - ( - SIMDVector::load_simd(self, ptr), - (rows > 8).then(|| SIMDVector::load_simd(self, ptr.add(32))), - ) + [ + SIMDVector::load_simd(self, slice.as_ptr()), + SIMDVector::load_simd(self, slice.add(Elements::new(32)).as_ptr()), + ] } } #[inline(always)] - fn splat(self, value: [u8; 4]) -> Self::Splat { - diskann_wide::alias!(words = ::u32x8); - words::splat(self, u32::from_le_bytes(value)).reinterpret_simd() + unsafe fn splat(self, slice: Slice<'_, u8>) -> Self::Splat { + // SAFETY: By precondition, the span contains one complete PACK4 group. + let value = unsafe { group(slice) }; + u32x8::::splat(self, u32::from_le_bytes(value)).reinterpret_simd() } #[inline(always)] fn zero(self) -> Self::Accumulator { @@ -411,16 +450,13 @@ mod x86_64 { b: Self::Splat, mut acc: Self::Accumulator, ) -> Self::Accumulator { - diskann_wide::alias!(pairs = ::i16x16); - let dot = |a: ::u8x32, acc: ::i32x8| { + let dot = |a: u8x32, acc: i32x8| { // B is a nibble, so each pair sum is at most 7650 and cannot saturate. let products = self.mm256_maddubs_epi16(a, b); - acc.dot_simd(products, pairs::splat(self, 1)) + acc.dot_simd(products, i16x16::::splat(self, 1)) }; - acc[0] = dot(a.0, acc[0]); - if let Some(hi) = a.1 { - acc[1] = dot(hi, acc[1]); - } + acc[0] = dot(a[0], acc[0]); + acc[1] = dot(a[1], acc[1]); acc } #[inline(always)] @@ -432,7 +468,6 @@ mod x86_64 { dim: usize, scores: &mut [f32; 16], ) { - diskann_wide::alias!(floats = ::f32x8); compensate( self, acc, @@ -441,34 +476,32 @@ mod x86_64 { dim, scores, |acc: Self::Accumulator, part| { - floats::from_array(self, acc[part].to_array().map(|x| x as u32 as f32)) + f32x8::::from_array(self, acc[part].to_array().map(|x| x as u32 as f32)) }, ); } } impl ExtraWide<16, 8> for V4 { - type Query = ( - ::u8x64, - Option<::u8x64>, - ); - type Splat = ::i8x64; - type Accumulator = [::i32x16; 2]; + type Query = [u8x64; 2]; + type Splat = i8x64; + type Accumulator = [i32x16; 2]; #[inline(always)] - fn load(self, values: &[[u8; 8]; 16], rows: usize) -> Self::Query { - let ptr = values.as_flattened().as_ptr(); - // SAFETY: Each half of the patch spans 64 bytes. + unsafe fn load(self, slice: Slice<'_, u8>) -> Self::Query { + bounds::check_eq!(slice.len(), 128); + // SAFETY: The 128-byte patch contains two complete 64-byte halves. unsafe { - ( - SIMDVector::load_simd(self, ptr), - (rows > 8).then(|| SIMDVector::load_simd(self, ptr.add(64))), - ) + [ + SIMDVector::load_simd(self, slice.as_ptr()), + SIMDVector::load_simd(self, slice.add(Elements::new(64)).as_ptr()), + ] } } #[inline(always)] - fn splat(self, value: [u8; 8]) -> Self::Splat { - diskann_wide::alias!(words = ::u64x8); - words::splat(self, u64::from_le_bytes(value)).reinterpret_simd() + unsafe fn splat(self, slice: Slice<'_, u8>) -> Self::Splat { + // SAFETY: By precondition, the span contains one complete PACK8 group. + let value = unsafe { group(slice) }; + u64x8::::splat(self, u64::from_le_bytes(value)).reinterpret_simd() } #[inline(always)] fn zero(self) -> Self::Accumulator { @@ -481,10 +514,8 @@ mod x86_64 { b: Self::Splat, mut acc: Self::Accumulator, ) -> Self::Accumulator { - acc[0] = acc[0].dot_simd(a.0, b); - if let Some(hi) = a.1 { - acc[1] = acc[1].dot_simd(hi, b); - } + acc[0] = acc[0].dot_simd(a[0], b); + acc[1] = acc[1].dot_simd(a[1], b); acc } #[inline(always)] @@ -496,11 +527,10 @@ mod x86_64 { dim: usize, scores: &mut [f32; 16], ) { - diskann_wide::alias!(floats = ::f32x8); let convert = |acc: Self::Accumulator, part: usize| { // PACK8 uses two adjacent VNNI accumulator lanes per query row. let lanes = acc[part].to_array(); - floats::from_array( + f32x8::::from_array( self, core::array::from_fn(|i| { lanes[2 * i].wrapping_add(lanes[2 * i + 1]) as u32 as f32 @@ -517,26 +547,27 @@ mod aarch64 { use super::*; use diskann_wide::{SIMDDotProduct, SIMDVector, arch::aarch64::Neon}; + panel_kernel!(Neon, 8, 8, 4, [1, 2, 3, 4, 5, 6, 7]); + impl ExtraWide<8, 4> for Neon { - type Query = ( - ::u8x16, - Option<::u8x16>, - ); - type Splat = ::u8x16; - type Accumulator = [::u32x4; 2]; + type Query = [u8x16; 2]; + type Splat = u8x16; + type Accumulator = [u32x4; 2]; #[inline(always)] - fn load(self, values: &[[u8; 4]; 8], rows: usize) -> Self::Query { - let ptr = values.as_flattened().as_ptr(); - // SAFETY: Each half of the patch spans 16 bytes. + unsafe fn load(self, slice: Slice<'_, u8>) -> Self::Query { + bounds::check_eq!(slice.len(), 32); + // SAFETY: The 32-byte patch contains two complete 16-byte halves. unsafe { - ( - SIMDVector::load_simd(self, ptr), - (rows > 4).then(|| SIMDVector::load_simd(self, ptr.add(16))), - ) + [ + SIMDVector::load_simd(self, slice.as_ptr()), + SIMDVector::load_simd(self, slice.add(Elements::new(16)).as_ptr()), + ] } } #[inline(always)] - fn splat(self, value: [u8; 4]) -> Self::Splat { + unsafe fn splat(self, slice: Slice<'_, u8>) -> Self::Splat { + // SAFETY: By precondition, the span contains one complete PACK4 group. + let value = unsafe { group::<4>(slice) }; Self::Splat::from_array(self, core::array::from_fn(|i| value[i % 4])) } #[inline(always)] @@ -550,10 +581,8 @@ mod aarch64 { b: Self::Splat, mut acc: Self::Accumulator, ) -> Self::Accumulator { - acc[0] = acc[0].dot_simd(a.0, b); - if let Some(hi) = a.1 { - acc[1] = acc[1].dot_simd(hi, b); - } + acc[0] = acc[0].dot_simd(a[0], b); + acc[1] = acc[1].dot_simd(a[1], b); acc } #[inline(always)] @@ -565,7 +594,6 @@ mod aarch64 { dim: usize, scores: &mut [f32; 8], ) { - diskann_wide::alias!(floats = ::f32x4); compensate( self, acc, @@ -574,7 +602,7 @@ mod aarch64 { dim, scores, |acc: Self::Accumulator, part| { - floats::from_array(self, acc[part].to_array().map(|x| x as f32)) + f32x4::::from_array(self, acc[part].to_array().map(|x| x as f32)) }, ); } @@ -606,7 +634,7 @@ mod tests { #[inline] || { let raw = 0xe000_0000u32; - let acc = [[::i32x8::splat(arch, raw as i32); 2]; 1]; + let acc = [[i32x8::::splat(arch, raw as i32); 2]; 1]; let mut scores = [f32::MAX; 16]; arch.reduce(acc, &query, &docs, 4, &mut scores); assert_eq!(scores, [-(raw as f32); 16]); @@ -618,15 +646,16 @@ mod tests { #[inline] || { let codes = [0, 1, 2, 3, 12, 13, 14, 15]; + // SAFETY: codes contains one complete PACK8 group. + let splat = unsafe { arch.splat(Slice::new(&codes)) }; assert_eq!( - arch.splat(codes).to_array(), + splat.to_array(), core::array::from_fn(|i| codes[i % 8] as i8) ); let values: [[i32; 16]; 2] = core::array::from_fn(|part| { core::array::from_fn(|lane| 0x7000_0000 + (part * 16 + lane) as i32 * 1024) }); - let acc = - [values.map(|part| ::i32x16::from_array(arch, part))]; + let acc = [values.map(|part| i32x16::::from_array(arch, part))]; let mut scores = [f32::MAX; 16]; arch.reduce(acc, &query, &docs, 8, &mut scores); let expected = core::array::from_fn(|row| { @@ -681,7 +710,8 @@ mod tests { dimensions: &[usize], panels: usize, ) where - A: Decoder + ExtraWide + LoadStore, + A: UnpackNibbles + ExtraWide + LoadStore, + for<'a> PanelKernel<'a, A, MR, NR, PACK>: driver::PanelKernel, { for &dim in dimensions { let query = matrix::<8>(panels * MR + 1, dim); @@ -705,6 +735,8 @@ mod tests { let k = a.layout().k().value().get(); let a_bytes = MR * k + std::mem::size_of::<[MinMaxCompensation; MR]>(); let b_bytes = k + std::mem::size_of::(); + assert_eq!(a.panel_bytes().value(), a_bytes); + assert_eq!(Scratch::row_bytes(a.layout()).value(), b_bytes); for b_rows in [NR, 2 * NR] { let cache = Cache::new(value_or_one(a_bytes + b_rows * b_bytes), value_or_one(1)); let mut scores = vec![123.0; query.num_vectors() + 2]; @@ -715,6 +747,7 @@ mod tests { &mut scores[1..query.num_vectors() + 1], cache, ); + assert_eq!(driver.b_rows.get(), b_rows.min(docs.num_vectors())); for _ in 0..2 { driver::Drive::drive(&mut driver); } @@ -751,6 +784,13 @@ mod tests { check_driver_architectures(&[3], 1); } + panel_kernel!(Scalar, 8, 9, 4, [1, 2, 3, 4, 5, 6, 7, 8]); + + #[test] + fn registered_panels_are_not_limited_to_eight_columns() { + check_driver::<_, 8, 9, 4>(Scalar::new(), &[3], 1); + } + fn check_inactive_half, const MR: usize, const PACK: usize>(arch: A) { arch.run_inline( #[inline] @@ -796,12 +836,88 @@ mod tests { } } + #[test] + fn compatible_layouts_and_document_group_bounds() { + use crate::matrix_kernels::test_util::panic_message_for; + use layout::Layout; + + let a_layout = Layout::new(65).unwrap(); + let b_layout = Layout::new(65).unwrap(); + assert!(!std::ptr::eq(&a_layout, &b_layout)); + let values = BlockTransposed::::new(1, 128); + let metadata = [[MinMaxCompensation::default(); 8]]; + let a = AView::new(values.as_view(), &metadata, &a_layout).unwrap(); + let docs = matrix::<4>(2, 65); + let mut scratch = Scratch::default(); + let b = scratch.prepare(Scalar::new(), RawRows::new(docs.as_view()), &b_layout); + a.visit_panels(|a, _| { + let mut panel = PanelKernel::<_, 8, 5, 4>::new(Scalar::new(), a, b, &[f32::MAX]); + driver::PanelKernel::panel_kernel(&mut panel); + assert_eq!(panel.scores[0], 0.0); + }); + assert!( + b.visit_panels::<2>(|panel| { + for row in 0..2 { + let source = docs.get_row(row).unwrap(); + for group in 0..32 { + let expected: [u8; 4] = core::array::from_fn(|i| { + let column = group * 4 + i; + let dim = column / 64 * 64 + + 2 * (column % 32) + + usize::from(column % 64 >= 32); + if dim < 65 { + source.vector().get(dim).unwrap() as u8 + } else { + 0 + } + }); + // SAFETY: Both rows contain 32 complete PACK4 groups. + let values = unsafe { panel.group::<4>(row, group).as_std_slice(4) }; + assert_eq!(values, expected.as_slice()); + } + } + let message = panic_message_for(|| { + // SAFETY: cfg(test) checks each logical index before accessing the pointer. + unsafe { panel.group::<4>(0, 32) }; + }); + assert!( + message.contains("document group out of bounds"), + "{message}" + ); + let message = panic_message_for(|| { + // SAFETY: cfg(test) checks the row before accessing the pointer. + unsafe { panel.group::<4>(2, 0) }; + }); + assert!(message.contains("document row out of bounds"), "{message}"); + }) + .is_none() + ); + + // The same padded K does not make different original dimensions compatible. + let other_layout = Layout::new(66).unwrap(); + let other_docs = matrix::<4>(2, 66); + let other_b = scratch.prepare( + Scalar::new(), + RawRows::new(other_docs.as_view()), + &other_layout, + ); + a.visit_panels(|a, _| { + let message = panic_message_for(|| { + PanelKernel::<_, 8, 5, 4>::new(Scalar::new(), a, other_b, &[f32::MAX]); + }); + assert!( + message.contains("A and B must use matching MinMax layouts"), + "{message}" + ); + }); + } + #[test] fn reused_b_scratch_fully_overwrites_codes_padding_and_metadata() { let docs = matrix::<4>(NR + 1, 65); let layout = layout::Layout::new(65).unwrap(); let mut scratch = Scratch::default(); - // First populate a full panel, then decode a smaller tail into the same buffers. + // First populate a full panel, then unpack a smaller tail into the same buffers. RawRows::new(docs.as_view()).visit_tiles(value_or_one(NR), |raw| { let tile = scratch.prepare(Scalar::new(), raw, &layout); let mut visited = 0; @@ -815,7 +931,8 @@ mod tests { let panel = tail.panel::<1>().unwrap(); assert_eq!(panel.compensation()[0], docs.get_row(NR).unwrap().meta()); // SAFETY: The layout contains 128 columns, so the last pack is complete. - assert_eq!(unsafe { panel.group::(0, 128 / PACK - 1) }, [0; PACK]); + let values = unsafe { panel.group::(0, 128 / PACK - 1).as_std_slice(PACK) }; + assert_eq!(values, &[0; PACK]); visited += 1; } assert_eq!(visited, raw.rows()); diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/a.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/a.rs new file mode 100644 index 0000000000..34e44a62a3 --- /dev/null +++ b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/a.rs @@ -0,0 +1,182 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Borrowed packed A panels with synchronized compensation and valid-row extents. + +use super::layout::{BLOCK, Layout}; +use crate::{ + matrix_kernels::{ + blocks::packed, + bounds::{self, Bound}, + num::Bytes, + ptr::Slice, + }, + minmax::MinMaxCompensation, + multi_vector::BlockTransposedRef, +}; + +#[derive(Clone, Copy)] +pub(crate) struct AView<'a, const MR: usize, const PACK: usize> { + values: packed::View<'a, u8, MR, PACK>, + compensation: &'a [[MinMaxCompensation; MR]], + rows: usize, + layout: &'a Layout, +} + +impl<'a, const MR: usize, const PACK: usize> AView<'a, MR, PACK> { + /// Bind already-prepared values to row-matched compensation. + /// + /// The owner must establish the even/odd ordering and zero row/column padding. + /// This checks extents, not the contents of the prepared buffer. + pub(crate) fn new( + values: BlockTransposedRef<'a, u8, MR, PACK>, + compensation: &'a [[MinMaxCompensation; MR]], + layout: &'a Layout, + ) -> Option { + const { assert!(BLOCK.is_multiple_of(PACK)) }; + assert_eq!(values.ncols(), layout.padded_dim(), "query layout mismatch"); + assert_eq!( + values.num_blocks(), + compensation.len(), + "query metadata mismatch" + ); + Some(Self { + values: packed::View::from_block_transposed(values)?, + compensation, + rows: values.nrows(), + layout, + }) + } + + pub(super) fn rows(self) -> usize { + self.rows + } + pub(super) fn layout(self) -> &'a Layout { + self.layout + } + + /// Working-set bytes per panel, including row/column padding and compensation. + pub(super) fn panel_bytes(self) -> Bytes { + let values = self.values.block_stride(self.layout.k()).bytes(); + Bytes::new(values.value() + std::mem::size_of::<[MinMaxCompensation; MR]>()) + } + + #[inline(always)] + pub(super) fn visit_panels(self, mut visit: impl FnMut(APanel<'_, MR, PACK>, usize)) { + // SAFETY: Each packed panel has matching metadata and zero row/column padding. + unsafe { + self.values.visit_panels(self.layout.k(), |values, block| { + visit( + APanel { + values, + compensation: &self.compensation[block], + valid_rows: (self.rows - block * MR).min(MR), + layout: self.layout, + }, + block * MR, + ); + }); + } + } +} + +#[derive(Clone, Copy)] +pub(super) struct APanel<'a, const MR: usize, const PACK: usize> { + values: packed::Panel<'a, u8, MR, PACK>, + compensation: &'a [MinMaxCompensation; MR], + valid_rows: usize, + layout: &'a Layout, +} + +impl APanel<'_, MR, PACK> { + pub(super) fn layout(&self) -> &Layout { + self.layout + } + + pub(super) fn valid_rows(&self) -> usize { + self.valid_rows + } + pub(super) fn compensation(&self) -> &[MinMaxCompensation; MR] { + self.compensation + } + + /// # Safety + /// + /// group < K / PACK. Query padding makes every group a complete MR x PACK patch. + pub(super) unsafe fn group(&self, group: usize) -> Slice<'_, u8> { + bounds::check_lt!( + Bound::new(group), + self.values.packs(self.layout.k()), + "query group out of bounds" + ); + // SAFETY: The caller provides an in-bounds group; each spans MR * PACK elements. + unsafe { + self.values + .as_ptr() + .add(self.values.pack_stride() * group) + .truncate(self.values.pack_stride()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::multi_vector::BlockTransposed; + + fn check_panel_bytes() { + for dim in [1, 63, 64, 65, 129] { + let layout = Layout::new(dim).unwrap(); + for rows in [1, MR, MR + 1] { + let values = BlockTransposed::::new(rows, layout.padded_dim()); + let metadata = vec![[MinMaxCompensation::default(); MR]; rows.div_ceil(MR)]; + let view = AView::new(values.as_view(), &metadata, &layout).unwrap(); + let bytes = std::mem::size_of_val(values.as_slice()) + + std::mem::size_of_val(metadata.as_slice()); + assert_eq!(view.panel_bytes().value() * metadata.len(), bytes); + } + } + } + + #[test] + fn panel_bytes_include_padding_and_compensation() { + check_panel_bytes::<8, 4>(); + check_panel_bytes::<16, 4>(); + check_panel_bytes::<16, 8>(); + } + + #[test] + fn query_groups_keep_geometry_and_reject_invalid_indices() { + use crate::matrix_kernels::test_util::panic_message_for; + + let layout = Layout::new(128).unwrap(); + let mut values = BlockTransposed::::new(8, 128); + for row in 0..8 { + for col in 0..128 { + values + .get_row_mut(row) + .unwrap() + .set(col, (row * 17 + col) as u8); + } + } + let metadata = [[MinMaxCompensation::default(); 8]]; + let view = AView::new(values.as_view(), &metadata, &layout).unwrap(); + view.visit_panels(|panel, _| { + for group in 0..32 { + // SAFETY: All 32 complete groups contain MR * PACK = 32 bytes. + let values = unsafe { panel.group(group).as_std_slice(32) }; + let expected: [[u8; 4]; 8] = core::array::from_fn(|row| { + core::array::from_fn(|i| (row * 17 + group * 4 + i) as u8) + }); + assert_eq!(values, expected.as_flattened()); + } + let message = panic_message_for(|| { + // SAFETY: Under cfg(test), the logical group check precedes pointer arithmetic. + unsafe { panel.group(32) }; + }); + assert!(message.contains("query group out of bounds"), "{message}"); + }); + } +} diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/b.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/b.rs new file mode 100644 index 0000000000..f2cf88da2f --- /dev/null +++ b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/b.rs @@ -0,0 +1,512 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! B-side unpacking, scratch preparation, and synchronized value/compensation panels. + +use super::layout::{BLOCK, Layout, PACKED_BLOCK}; +use crate::{ + matrix_kernels::{ + blocks::unpacked, + bounds::{self, Bound}, + num::{Bytes, Elements}, + ptr::Slice, + }, + minmax::{MinMaxCompensation, MinMaxMeta}, + multi_vector::MatRef, +}; +use diskann_wide::{Architecture, arch::Scalar}; +use std::num::NonZeroUsize; + +#[inline(always)] +fn unpack_row(layout: &Layout, arch: impl UnpackNibbles, packed: &[u8], output: &mut [u8]) { + assert_eq!(packed.len(), layout.dim().div_ceil(2)); + assert_eq!(output.len(), layout.padded_dim()); + let (blocks, tail) = packed.as_chunks::(); + let (unpacked, unpacked_tail) = output.split_at_mut(blocks.len() * BLOCK); + for (source, block) in blocks.iter().zip(unpacked.as_chunks_mut::().0) { + let (low, high) = block.split_at_mut(PACKED_BLOCK); + arch.unpack(source, low, high); + } + if !tail.is_empty() { + let (low, high) = unpacked_tail.split_at_mut(PACKED_BLOCK); + arch.unpack(tail, low, high); + } + if !layout.dim().is_multiple_of(2) { + output[Layout::position(layout.dim())] = 0; + } +} + +#[derive(Clone, Copy)] +pub(crate) struct RawRows<'a> { + doc: MatRef<'a, MinMaxMeta<4>>, + start: usize, + rows: usize, +} + +impl<'a> RawRows<'a> { + pub(crate) fn new(doc: MatRef<'a, MinMaxMeta<4>>) -> Self { + Self { + doc, + start: 0, + rows: doc.num_vectors(), + } + } + pub(super) fn dim(self) -> usize { + self.doc.repr().intrinsic_dim() + } + pub(super) fn rows(self) -> usize { + self.rows + } + #[inline(always)] + pub(super) fn visit_tiles(self, rows: NonZeroUsize, mut visit: impl FnMut(RawRows<'_>)) { + for offset in (0..self.rows).step_by(rows.get()) { + visit(Self { + doc: self.doc, + start: self.start + offset, + rows: (self.rows - offset).min(rows.get()), + }); + } + } +} + +#[derive(Default)] +pub(super) struct Scratch { + values: Vec, + compensation: Vec, +} + +impl Scratch { + /// Working-set bytes per prepared row, including K padding and compensation. + pub(super) fn row_bytes(layout: &Layout) -> Bytes { + Bytes::new(layout.padded_dim() + std::mem::size_of::()) + } + + #[expect( + clippy::expect_used, + reason = "tile extents are nonzero and rows retain canonical provenance" + )] + #[inline(always)] + pub(super) fn prepare<'a>( + &'a mut self, + arch: impl UnpackNibbles, + raw: RawRows<'_>, + layout: &'a Layout, + ) -> BTile<'a> { + assert_eq!(raw.dim(), layout.dim(), "document dimension mismatch"); + let k = layout.k(); + let rows = NonZeroUsize::new(raw.rows).expect("nonempty B tile"); + let len = raw + .rows + .checked_mul(k.value().get()) + .expect("B scratch extent overflow"); + if self.values.len() < len { + self.values.resize(len, 0); + } + if self.compensation.len() < raw.rows { + self.compensation + .resize(raw.rows, MinMaxCompensation::default()); + } + let values = &mut self.values[..len]; + let compensation = &mut self.compensation[..raw.rows]; + for (i, (output, metadata)) in values + .chunks_exact_mut(k.value().get()) + .zip(compensation.iter_mut()) + .enumerate() + { + let source = raw + .doc + .get_row(raw.start + i) + .expect("canonical row in bounds"); + *metadata = source.meta(); + unpack_row(layout, arch, source.vector().as_slice(), output); + } + BTile { + // SAFETY: Preparation completely writes rows * K values and one metadata per row. + values: unsafe { unpacked::View::new(Slice::new(values), rows, k) }, + compensation, + layout, + } + } +} + +#[derive(Clone, Copy)] +pub(super) struct BTile<'a> { + values: unpacked::View<'a, u8>, + compensation: &'a [MinMaxCompensation], + layout: &'a Layout, +} + +impl BTile<'_> { + pub(super) fn layout(&self) -> &Layout { + self.layout + } + #[inline(always)] + pub(super) fn visit_panels( + &self, + mut visit: impl FnMut(BPanel<'_, N>), + ) -> Option> { + const { assert!(N > 0) }; + let metadata = self.compensation.as_chunks::().0; + // SAFETY: Scratch::prepare binds the unpacked stride and value/metadata row counts. + let remainder = unsafe { + self.values.visit_panels( + self.layout.k(), + |values: unpacked::Panel<'_, u8, N>, start| { + visit(BPanel { + values, + compensation: &metadata[start / N], + layout: self.layout, + }); + }, + ) + }; + remainder.map(|values| Remainder { + compensation: &self.compensation[values.start()..], + values, + layout: self.layout, + }) + } +} + +pub(super) struct Remainder<'a, const CAP: usize> { + values: unpacked::Remainder<'a, u8, CAP>, + compensation: &'a [MinMaxCompensation], + layout: &'a Layout, +} + +impl Remainder<'_, CAP> { + #[expect( + clippy::expect_used, + reason = "the metadata remainder has the same extent as values" + )] + pub(super) fn panel(&self) -> Option> { + let values = self.values.try_as_panel::()?; + Some(BPanel { + values, + compensation: self.compensation.try_into().expect("exact metadata panel"), + layout: self.layout, + }) + } +} + +#[derive(Clone, Copy)] +pub(super) struct BPanel<'a, const N: usize> { + values: unpacked::Panel<'a, u8, N>, + compensation: &'a [MinMaxCompensation; N], + layout: &'a Layout, +} + +impl BPanel<'_, N> { + pub(super) fn compensation(&self) -> &[MinMaxCompensation; N] { + self.compensation + } + + /// # Safety + /// + /// row < N and group < K / PACK. Layout preparation completely writes K bytes per row. + pub(super) unsafe fn group( + &self, + row: usize, + group: usize, + ) -> Slice<'_, u8> { + const { assert!(PACK > 0) }; + bounds::check_lt!(Bound::new(row), N, "document row out of bounds"); + bounds::check_lt!( + Bound::new(group), + self.layout.k().value().get() / PACK, + "document group out of bounds" + ); + // SAFETY: The caller selects a complete group inside one unpacked row. + unsafe { + self.values + .as_ptr() + .add(self.values.stride(self.layout.k()) * row + Elements::new(group * PACK)) + .truncate(Elements::new(PACK)) + } + } +} + +pub(super) trait UnpackNibbles: Architecture { + /// Expand low/high nibbles into equally sized channels and zero-fill unused output. + /// Channel lengths must be at least `packed.len()`; no SIMD-width alignment is required. + fn unpack(self, packed: &[u8], low: &mut [u8], high: &mut [u8]); +} + +#[inline(always)] +fn assert_unpack_lengths(packed: &[u8], low: &[u8], high: &[u8]) { + assert_eq!(low.len(), high.len(), "unpack channel length mismatch"); + assert!(packed.len() <= low.len(), "unpack output is too short"); +} + +impl UnpackNibbles for Scalar { + #[inline(always)] + fn unpack(self, packed: &[u8], low: &mut [u8], high: &mut [u8]) { + assert_unpack_lengths(packed, low, high); + for (i, (low, high)) in low.iter_mut().zip(high).enumerate() { + let byte = packed.get(i).copied().unwrap_or(0); + *low = byte & 15; + *high = byte >> 4; + } + } +} + +#[cfg(target_arch = "x86_64")] +mod x86 { + use super::*; + use diskann_wide::{ + SIMDVector, + arch::x86_64::{V3, V4}, + }; + + #[inline(always)] + fn unpack( + arch: V3, + packed: &[u8], + low: &mut [u8], + high: &mut [u8], + load: impl Fn(&[u8]) -> ::u32x8, + ) { + assert_unpack_lengths(packed, low, high); + diskann_wide::alias!(words = ::u32x8); + diskann_wide::alias!(bytes = ::u8x32); + let mask = words::splat(arch, 0x0f0f_0f0f); + for ((source, low), high) in packed + .chunks(32) + .zip(low.chunks_mut(32)) + .zip(high.chunks_mut(32)) + { + let codes = load(source); + let lows = bytes::from_underlying(arch, (codes & mask).to_underlying()); + let highs = bytes::from_underlying(arch, ((codes >> 4) & mask).to_underlying()); + // SAFETY: The channels have equal lengths. Complete chunks admit full stores; + // partial chunks are written only up to their actual length. + unsafe { + if low.len() == 32 { + lows.store_simd(low.as_mut_ptr()); + highs.store_simd(high.as_mut_ptr()); + } else { + lows.store_simd_first(low.as_mut_ptr(), low.len()); + highs.store_simd_first(high.as_mut_ptr(), high.len()); + } + } + } + let written = packed.len().next_multiple_of(32).min(low.len()); + if written < low.len() { + low[written..].fill(0); + high[written..].fill(0); + } + } + + impl UnpackNibbles for V3 { + #[inline(always)] + fn unpack(self, packed: &[u8], low: &mut [u8], high: &mut [u8]) { + diskann_wide::alias!(words = ::u32x8); + unpack( + self, + packed, + low, + high, + #[inline(always)] + |source| { + if source.len() == 32 { + // SAFETY: The complete input chunk contains 32 readable bytes. + unsafe { words::load_simd(self, source.as_ptr().cast()) } + } else { + let mut tail = [0; 32]; + tail[..source.len()].copy_from_slice(source); + // SAFETY: The local buffer contains 32 initialized bytes. + unsafe { words::load_simd(self, tail.as_ptr().cast()) } + } + }, + ); + } + } + + impl UnpackNibbles for V4 { + #[inline(always)] + fn unpack(self, packed: &[u8], low: &mut [u8], high: &mut [u8]) { + diskann_wide::alias!(bytes = ::u8x32); + diskann_wide::alias!(words = ::u32x8); + let arch = self.retarget(); + unpack( + arch, + packed, + low, + high, + #[inline(always)] + |source| { + if source.len() == 32 { + // SAFETY: The complete input chunk contains 32 readable bytes. + unsafe { words::load_simd(arch, source.as_ptr().cast()) } + } else { + // SAFETY: Read only the available bytes; the other lanes are zero. + let codes = + unsafe { bytes::load_simd_first(self, source.as_ptr(), source.len()) }; + // SAFETY: The SIMD array contains 32 initialized bytes. + unsafe { words::load_simd(arch, codes.to_array().as_ptr().cast()) } + } + }, + ); + } + } +} + +#[cfg(target_arch = "aarch64")] +impl UnpackNibbles for diskann_wide::arch::aarch64::Neon { + #[inline(always)] + fn unpack(self, packed: &[u8], low: &mut [u8], high: &mut [u8]) { + assert_unpack_lengths(packed, low, high); + use diskann_wide::SIMDVector; + diskann_wide::alias!(bytes = ::u8x16); + let mask = bytes::splat(self, 15); + for ((source, low), high) in packed + .chunks(16) + .zip(low.chunks_mut(16)) + .zip(high.chunks_mut(16)) + { + // SAFETY: Access only the available input/output bytes of each chunk. + unsafe { + let codes = bytes::load_simd_first(self, source.as_ptr(), source.len()); + if low.len() == 16 { + (codes & mask).store_simd(low.as_mut_ptr()); + ((codes >> 4) & mask).store_simd(high.as_mut_ptr()); + } else { + (codes & mask).store_simd_first(low.as_mut_ptr(), low.len()); + ((codes >> 4) & mask).store_simd_first(high.as_mut_ptr(), high.len()); + } + } + } + + let written = packed.len().next_multiple_of(16).min(low.len()); + if written < low.len() { + low[written..].fill(0); + high[written..].fill(0); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn row_bytes_include_unpacked_padding_and_compensation() { + use crate::multi_vector::{Defaulted, Mat}; + + for dim in [1, 63, 64, 65, 250, 256] { + let layout = Layout::new(dim).unwrap(); + for rows in [1, 3] { + let docs = Mat::new(MinMaxMeta::<4>::new(rows, dim), Defaulted).unwrap(); + let mut scratch = Scratch::default(); + scratch.prepare(Scalar::new(), RawRows::new(docs.as_view()), &layout); + let bytes = std::mem::size_of_val(scratch.values.as_slice()) + + std::mem::size_of_val(scratch.compensation.as_slice()); + assert_eq!(Scratch::row_bytes(&layout).value() * rows, bytes); + assert!(Scratch::row_bytes(&layout).value() > docs.repr().ncols()); + } + } + } + + fn check(arch: impl UnpackNibbles) { + for len in [0, 1, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129] { + let input: Vec<_> = (0..len).map(|i| (i * 37 + 0xf3) as u8).collect(); + for padding in [0, 1, 31] { + let extent = len + padding; + let mut low = vec![0xa5; extent + 2]; + let mut high = low.clone(); + arch.unpack(&input, &mut low[1..extent + 1], &mut high[1..extent + 1]); + for i in 0..extent { + let byte = input.get(i).copied().unwrap_or(0); + assert_eq!(low[i + 1], byte & 15, "len={len}, padding={padding}, i={i}"); + assert_eq!( + high[i + 1], + byte >> 4, + "len={len}, padding={padding}, i={i}" + ); + } + assert_eq!( + (low[0], low[extent + 1], high[0], high[extent + 1]), + (0xa5, 0xa5, 0xa5, 0xa5) + ); + } + } + } + + #[test] + fn unpacking_is_independent_of_layout_and_register_width() { + check(Scalar::new()); + #[cfg(target_arch = "x86_64")] + { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check(arch); + } + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() { + check(arch); + } + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check(arch); + } + } + + #[test] + #[should_panic(expected = "unpack channel length mismatch")] + fn mismatched_channels_are_rejected() { + Scalar::new().unpack(&[0x12], &mut [0; 1], &mut [0; 2]); + } + + #[test] + #[should_panic(expected = "unpack output is too short")] + fn short_channels_are_rejected() { + Scalar::new().unpack(&[0x12, 0x34], &mut [0; 1], &mut [0; 1]); + } + + fn check_unpack(arch: impl UnpackNibbles) { + for dim in 0..=129 { + let layout = Layout::new(dim).unwrap(); + let mut packed: Vec<_> = (0..dim.div_ceil(2)) + .map(|i| (i * 37 + 0xa5) as u8) + .collect(); + if dim % 2 == 1 { + *packed.last_mut().unwrap() |= 0xf0; + } + let padded = layout.padded_dim(); + let mut output = vec![0xfe; padded + 2]; + unpack_row(&layout, arch, &packed, &mut output[1..padded + 1]); + let mut expected = vec![0; padded]; + for block in (0..dim).step_by(BLOCK) { + let count = (dim - block).min(BLOCK); + for i in 0..count.div_ceil(2) { + expected[block + i] = packed[block / 2 + i] & 15; + } + for i in 0..count / 2 { + expected[block + PACKED_BLOCK + i] = packed[block / 2 + i] >> 4; + } + } + assert_eq!(&output[1..padded + 1], expected, "dim={dim}"); + assert_eq!((output[0], output[padded + 1]), (0xfe, 0xfe)); + } + assert!(Layout::new(usize::MAX).is_err()); + } + + #[test] + fn all_partial_blocks_ignore_unused_nibbles_and_overwrite_padding() { + check_unpack(Scalar::new()); + #[cfg(target_arch = "x86_64")] + { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check_unpack(arch); + } + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() { + check_unpack(arch); + } + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check_unpack(arch); + } + } +} diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs deleted file mode 100644 index a628fb681f..0000000000 --- a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/decode.rs +++ /dev/null @@ -1,101 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Nibble expansion, independent of query packing and metadata. - -pub(super) const BYTES: usize = 32; - -use diskann_wide::{Architecture, arch::Scalar}; - -pub(super) trait Decoder: Architecture { - fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]); - - #[expect( - clippy::expect_used, - reason = "a full decoder block contains exactly BYTES bytes" - )] - #[inline(always)] - fn unpack(self, packed: &[u8], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { - assert!(packed.len() <= BYTES, "packed block exceeds decoder width"); - if packed.len() == BYTES { - self.block(packed.try_into().expect("full decoder block"), low, high); - } else { - let mut tail = [0; BYTES]; - tail[..packed.len()].copy_from_slice(packed); - self.block(&tail, low, high); - } - } -} - -impl Decoder for Scalar { - #[inline(always)] - fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { - for (i, &byte) in packed.iter().enumerate() { - low[i] = byte & 15; - high[i] = byte >> 4; - } - } -} - -#[cfg(target_arch = "x86_64")] -mod x86 { - use super::*; - use diskann_wide::{ - SIMDVector, - arch::x86_64::{V3, V4}, - }; - - impl Decoder for V3 { - #[inline(always)] - fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { - diskann_wide::alias!(words = ::u32x8); - diskann_wide::alias!(bytes = ::u8x32); - // Native u32 shifts avoid AVX2's emulated per-byte shifts. - let mask = words::splat(self, 0x0f0f_0f0f); - // SAFETY: packed spans 32 bytes; SIMD loads do not require alignment. - let codes = unsafe { words::load_simd(self, packed.as_ptr().cast()) }; - let lows = bytes::from_underlying(self, (codes & mask).to_underlying()); - let highs = bytes::from_underlying(self, ((codes >> 4) & mask).to_underlying()); - // SAFETY: Both output channels span 32 writable bytes. - unsafe { - lows.store_simd(low.as_mut_ptr()); - highs.store_simd(high.as_mut_ptr()); - } - } - } - - impl Decoder for V4 { - #[inline(always)] - fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { - self.retarget().block(packed, low, high); - } - #[inline(always)] - fn unpack(self, packed: &[u8], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { - assert!(packed.len() <= BYTES, "packed block exceeds decoder width"); - diskann_wide::alias!(bytes = ::u8x32); - // SAFETY: Only the available bytes are read; the other lanes are zero. - let codes = unsafe { bytes::load_simd_first(self, packed.as_ptr(), packed.len()) }; - self.retarget().block(&codes.to_array(), low, high); - } - } -} - -#[cfg(target_arch = "aarch64")] -impl Decoder for diskann_wide::arch::aarch64::Neon { - #[inline(always)] - fn block(self, packed: &[u8; BYTES], low: &mut [u8; BYTES], high: &mut [u8; BYTES]) { - use diskann_wide::SIMDVector; - diskann_wide::alias!(bytes = ::u8x16); - let mask = bytes::splat(self, 15); - for offset in [0, 16] { - // SAFETY: Each load and store stays within its 32-byte channel. - unsafe { - let codes = bytes::load_simd(self, packed.as_ptr().add(offset)); - (codes & mask).store_simd(low.as_mut_ptr().add(offset)); - ((codes >> 4) & mask).store_simd(high.as_mut_ptr().add(offset)); - } - } - } -} diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs index f4bf60993f..05839358c4 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/layout.rs @@ -3,25 +3,28 @@ * Licensed under the MIT license. */ -//! Shared order and borrowed panels for MinMax query packing and document decoding. +//! Shared layout contract for MinMax A and B. //! Each 64-dimensional block stores even coordinates followed by odd coordinates. //! Missing coordinates are zero-padded before applying the block-transposed PACK layout. +//! This is a prepared-data format, not an ISA register width. -use super::decode::{self, Decoder}; use crate::{ - matrix_kernels::{ - blocks::packed, - num::{DimK, value_or_one}, - ptr::Slice, - }, - minmax::MinMaxCompensation, - multi_vector::{BlockTransposedRef, Overflow}, + matrix_kernels::num::{DimK, value_or_one}, + multi_vector::Overflow, }; -const BLOCK: usize = 2 * decode::BYTES; +// V4 benchmarks on Xeon 8370C (MR16/NR8/PACK8, 22 shapes with 16/64 query and +// document rows, selected dimensions 63-257) compared 32/64/128-dimensional blocks. +// Versus 64, geometric-mean latency was ~1.7% lower at 32 and ~5.8% higher at 128. +// Keep the AVX2-friendly shared format rather than specialize V4 for a small gain; +// this is a measured tradeoff on that CPU, not a universally optimal block size. +pub(super) const BLOCK: usize = 64; +pub(super) const PACKED_BLOCK: usize = BLOCK / 2; pub(crate) const MAX_DIM: usize = u32::MAX as usize / (255 * 15); -#[derive(Debug)] +/// Original dimension and physical K for the fixed prepared-data format. +/// Equal layouts are interchangeable; they do not identify a particular allocation. +#[derive(Debug, PartialEq, Eq)] pub(crate) struct Layout { dim: usize, padded: usize, @@ -44,168 +47,6 @@ impl Layout { DimK::new(value_or_one(self.padded)) } pub(crate) fn position(d: usize) -> usize { - d / BLOCK * BLOCK + d % 2 * decode::BYTES + d % BLOCK / 2 - } - - #[expect( - clippy::expect_used, - reason = "padded blocks contain two complete decoder channels" - )] - #[inline(always)] - pub(super) fn decode_row(&self, arch: impl Decoder, packed: &[u8], output: &mut [u8]) { - assert_eq!(packed.len(), self.dim.div_ceil(2)); - assert_eq!(output.len(), self.padded); - let (blocks, tail) = packed.as_chunks::<{ decode::BYTES }>(); - let (decoded, decoded_tail) = output.split_at_mut(blocks.len() * BLOCK); - for (source, block) in blocks.iter().zip(decoded.as_chunks_mut::().0) { - let (low, high) = block.split_at_mut(decode::BYTES); - arch.block( - source, - low.try_into().expect("low channel"), - high.try_into().expect("high channel"), - ); - } - if !tail.is_empty() { - let (low, high) = decoded_tail.split_at_mut(decode::BYTES); - arch.unpack( - tail, - low.try_into().expect("low channel"), - high.try_into().expect("high channel"), - ); - } - if !self.dim.is_multiple_of(2) { - output[Self::position(self.dim)] = 0; - } - } -} - -#[derive(Clone, Copy)] -pub(crate) struct AView<'a, const MR: usize, const PACK: usize> { - values: packed::View<'a, u8, MR, PACK>, - compensation: &'a [[MinMaxCompensation; MR]], - rows: usize, - layout: &'a Layout, -} - -impl<'a, const MR: usize, const PACK: usize> AView<'a, MR, PACK> { - pub(crate) fn new( - values: BlockTransposedRef<'a, u8, MR, PACK>, - compensation: &'a [[MinMaxCompensation; MR]], - layout: &'a Layout, - ) -> Option { - const { assert!(BLOCK.is_multiple_of(PACK)) }; - assert_eq!(values.ncols(), layout.padded_dim(), "query layout mismatch"); - assert_eq!( - values.num_blocks(), - compensation.len(), - "query metadata mismatch" - ); - Some(Self { - values: packed::View::from_block_transposed(values)?, - compensation, - rows: values.nrows(), - layout, - }) - } - - pub(super) fn rows(self) -> usize { - self.rows - } - pub(super) fn layout(self) -> &'a Layout { - self.layout - } - #[inline(always)] - pub(super) fn visit_panels(self, mut visit: impl FnMut(APanel<'_, MR, PACK>, usize)) { - // SAFETY: Each packed panel has matching metadata and zero row/column padding. - unsafe { - self.values.visit_panels(self.layout.k(), |values, block| { - visit( - APanel { - values, - compensation: &self.compensation[block], - valid_rows: (self.rows - block * MR).min(MR), - layout: self.layout, - }, - block * MR, - ); - }); - } - } -} - -#[derive(Clone, Copy)] -pub(super) struct APanel<'a, const MR: usize, const PACK: usize> { - values: packed::Panel<'a, u8, MR, PACK>, - compensation: &'a [MinMaxCompensation; MR], - valid_rows: usize, - layout: &'a Layout, -} - -impl APanel<'_, MR, PACK> { - pub(super) fn layout(&self) -> &Layout { - self.layout - } - - pub(super) fn valid_rows(&self) -> usize { - self.valid_rows - } - pub(super) fn compensation(&self) -> &[MinMaxCompensation; MR] { - self.compensation - } - - /// # Safety - /// - /// group < K / PACK. Query padding makes every group a complete MR x PACK patch. - pub(super) unsafe fn group(&self, group: usize) -> &[[u8; PACK]; MR] { - // SAFETY: The caller provides an in-bounds group; each spans MR * PACK elements. - let patch: Slice<'_, u8> = unsafe { - self.values - .as_ptr() - .add(self.values.pack_stride() * group) - .truncate(self.values.pack_stride()) - }; - // SAFETY: This span contains exactly MR * PACK consecutive bytes. - unsafe { &*patch.as_ptr().cast::<[[u8; PACK]; MR]>() } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn check_decode(arch: impl Decoder) { - for dim in 0..=129 { - let layout = Layout::new(dim).unwrap(); - let packed = vec![0xa5; dim.div_ceil(2)]; - let mut output = vec![0xfe; layout.padded + 2]; - layout.decode_row(arch, &packed, &mut output[1..layout.padded + 1]); - let mut expected = vec![0; layout.padded]; - for block in (0..dim).step_by(BLOCK) { - let count = (dim - block).min(BLOCK); - expected[block..block + count.div_ceil(2)].fill(5); - expected[block + decode::BYTES..block + decode::BYTES + count / 2].fill(10); - } - assert_eq!(&output[1..layout.padded + 1], expected, "dim={dim}"); - assert_eq!((output[0], output[layout.padded + 1]), (0xfe, 0xfe)); - } - assert!(Layout::new(usize::MAX).is_err()); - } - - #[test] - fn all_partial_blocks_ignore_unused_nibbles_and_overwrite_padding() { - check_decode(diskann_wide::arch::Scalar::new()); - #[cfg(target_arch = "x86_64")] - { - if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - check_decode(arch); - } - if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked_miri() { - check_decode(arch); - } - } - #[cfg(target_arch = "aarch64")] - if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - check_decode(arch); - } + d / BLOCK * BLOCK + d % 2 * PACKED_BLOCK + d % BLOCK / 2 } } diff --git a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/panels.rs b/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/panels.rs deleted file mode 100644 index 8b08a13d10..0000000000 --- a/diskann-quantization/src/matrix_kernels/maxsim/minmax8_x_minmax4/panels.rs +++ /dev/null @@ -1,191 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Canonical B tiles and prepared panels whose values and compensation cannot drift apart. - -use super::{decode::Decoder, layout::Layout}; -use crate::{ - matrix_kernels::{blocks::unpacked, num::Elements, ptr::Slice}, - minmax::{MinMaxCompensation, MinMaxMeta}, - multi_vector::MatRef, -}; -use std::num::NonZeroUsize; - -#[derive(Clone, Copy)] -pub(crate) struct RawRows<'a> { - doc: MatRef<'a, MinMaxMeta<4>>, - start: usize, - rows: usize, -} - -impl<'a> RawRows<'a> { - pub(crate) fn new(doc: MatRef<'a, MinMaxMeta<4>>) -> Self { - Self { - doc, - start: 0, - rows: doc.num_vectors(), - } - } - pub(super) fn dim(self) -> usize { - self.doc.repr().intrinsic_dim() - } - pub(super) fn rows(self) -> usize { - self.rows - } - #[inline(always)] - pub(super) fn visit_tiles(self, rows: NonZeroUsize, mut visit: impl FnMut(RawRows<'_>)) { - for offset in (0..self.rows).step_by(rows.get()) { - visit(Self { - doc: self.doc, - start: self.start + offset, - rows: (self.rows - offset).min(rows.get()), - }); - } - } -} - -#[derive(Default)] -pub(super) struct Scratch { - values: Vec, - compensation: Vec, -} - -impl Scratch { - #[expect( - clippy::expect_used, - reason = "tile extents are nonzero and rows retain canonical provenance" - )] - #[inline(always)] - pub(super) fn prepare<'a>( - &'a mut self, - arch: impl Decoder, - raw: RawRows<'_>, - layout: &'a Layout, - ) -> BTile<'a> { - assert_eq!(raw.dim(), layout.dim(), "document dimension mismatch"); - let k = layout.k(); - let rows = NonZeroUsize::new(raw.rows).expect("nonempty B tile"); - let len = raw - .rows - .checked_mul(k.value().get()) - .expect("B scratch extent overflow"); - if self.values.len() < len { - self.values.resize(len, 0); - } - if self.compensation.len() < raw.rows { - self.compensation - .resize(raw.rows, MinMaxCompensation::default()); - } - let values = &mut self.values[..len]; - let compensation = &mut self.compensation[..raw.rows]; - for (i, (output, metadata)) in values - .chunks_exact_mut(k.value().get()) - .zip(compensation.iter_mut()) - .enumerate() - { - let source = raw - .doc - .get_row(raw.start + i) - .expect("canonical row in bounds"); - *metadata = source.meta(); - layout.decode_row(arch, source.vector().as_slice(), output); - } - BTile { - // SAFETY: Preparation completely writes rows * K values and one metadata per row. - values: unsafe { unpacked::View::new(Slice::new(values), rows, k) }, - compensation, - layout, - } - } -} - -#[derive(Clone, Copy)] -pub(super) struct BTile<'a> { - values: unpacked::View<'a, u8>, - compensation: &'a [MinMaxCompensation], - layout: &'a Layout, -} - -impl BTile<'_> { - pub(super) fn layout(&self) -> &Layout { - self.layout - } - #[inline(always)] - pub(super) fn visit_panels( - &self, - mut visit: impl FnMut(BPanel<'_, N>), - ) -> Option> { - const { assert!(N > 0) }; - let metadata = self.compensation.as_chunks::().0; - // SAFETY: Scratch::prepare binds the decoded stride and value/metadata row counts. - let remainder = unsafe { - self.values.visit_panels( - self.layout.k(), - |values: unpacked::Panel<'_, u8, N>, start| { - visit(BPanel { - values, - compensation: &metadata[start / N], - layout: self.layout, - }); - }, - ) - }; - remainder.map(|values| Remainder { - compensation: &self.compensation[values.start()..], - values, - layout: self.layout, - }) - } -} - -pub(super) struct Remainder<'a, const CAP: usize> { - values: unpacked::Remainder<'a, u8, CAP>, - compensation: &'a [MinMaxCompensation], - layout: &'a Layout, -} - -impl Remainder<'_, CAP> { - #[expect( - clippy::expect_used, - reason = "the metadata remainder has the same extent as values" - )] - pub(super) fn panel(&self) -> Option> { - let values = self.values.try_as_panel::()?; - Some(BPanel { - values, - compensation: self.compensation.try_into().expect("exact metadata panel"), - layout: self.layout, - }) - } -} - -#[derive(Clone, Copy)] -pub(super) struct BPanel<'a, const N: usize> { - values: unpacked::Panel<'a, u8, N>, - compensation: &'a [MinMaxCompensation; N], - layout: &'a Layout, -} - -impl BPanel<'_, N> { - pub(super) fn compensation(&self) -> &[MinMaxCompensation; N] { - self.compensation - } - - /// # Safety - /// - /// row < N and group < K / PACK. Layout preparation completely writes K bytes per row. - pub(super) unsafe fn group(&self, row: usize, group: usize) -> [u8; PACK] { - // SAFETY: The caller selects a complete group inside one decoded row. - unsafe { - self.values - .as_ptr() - .add(self.values.stride(self.layout.k()) * row + Elements::new(group * PACK)) - .truncate(Elements::new(PACK)) - .as_ptr() - .cast::<[u8; PACK]>() - .read_unaligned() - } - } -} diff --git a/diskann-quantization/src/minmax/mod.rs b/diskann-quantization/src/minmax/mod.rs index d717868152..516e1bcc3d 100644 --- a/diskann-quantization/src/minmax/mod.rs +++ b/diskann-quantization/src/minmax/mod.rs @@ -133,7 +133,10 @@ mod vectors; // Exports // ///////////// -pub use multi::{MinMaxErase, MinMaxKernel, MinMaxMaxSimKernel, MinMaxMeta, build_minmax_max_sim}; +pub use multi::{ + MinMaxErase, MinMaxKernel, MinMaxMaxSimKernel, MinMaxMeta, MinMaxMetaError, + build_minmax_max_sim, +}; pub use quantizer::{L2Loss, MinMaxQuantizer}; pub use recompress::{RecompressError, Recompressor}; pub use vectors::{ diff --git a/diskann-quantization/src/minmax/multi/factory.rs b/diskann-quantization/src/minmax/multi/factory.rs index 47bf68212e..609b087dcb 100644 --- a/diskann-quantization/src/minmax/multi/factory.rs +++ b/diskann-quantization/src/minmax/multi/factory.rs @@ -8,7 +8,7 @@ use super::query::PackedQuery; use crate::{ matrix_kernels::{ self as mk, - maxsim::minmax8_x_minmax4::{Driver, layout::MAX_DIM, panels::RawRows}, + maxsim::minmax8_x_minmax4::{Driver, b::RawRows, layout::MAX_DIM}, }, minmax::MinMaxMeta, multi_vector::{MatRef, MaxSimError, MaxSimIsa, NotSupported, Overflow}, @@ -216,7 +216,7 @@ mod tests { #[test] fn tiled_scalar_matches_pairwise() { for dim in [0, 1, 3, 4, 5, 63, 64, 65, 127, 128, 129, 250, 256, 257] { - for q in (0..=17).chain([32, 33]) { + for q in (0..=17).chain([24, 25, 32, 33]) { let query = data::<8>(q, dim, 3); for isa in [ MaxSimIsa::Scalar, diff --git a/diskann-quantization/src/minmax/multi/meta.rs b/diskann-quantization/src/minmax/multi/meta.rs index f006d4afb4..a3d717ca88 100644 --- a/diskann-quantization/src/minmax/multi/meta.rs +++ b/diskann-quantization/src/minmax/multi/meta.rs @@ -11,11 +11,11 @@ use super::super::MinMaxQuantizer; use super::super::vectors::DataMutRef; use crate::CompressInto; use crate::bits::{Representation, Unsigned}; -use crate::minmax::{self, Data}; +use crate::minmax::{self, Data, MinMaxCompensation}; use crate::multi_vector::matrix::{ Defaulted, NewMut, NewOwned, NewRef, Repr, ReprMut, ReprOwned, SliceError, }; -use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, Standard}; +use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, Overflow, Standard}; use crate::scalar::InputContainsNaN; use crate::utils; @@ -26,12 +26,22 @@ use crate::utils; /// Metadata for MinMax quantized multi-vectors. /// /// Stores the intrinsic dimension (output dimension after transform) which is -/// needed to interpret each row of the quantized data. The row stride in bytes -/// is computed using [`minmax::Data::canonical_bytes`]. +/// needed to interpret each row of the quantized data. Construction validates the +/// canonical row stride and total allocation extent, including empty matrices. #[derive(Debug, Clone, Copy)] pub struct MinMaxMeta { nrows: usize, intrinsic_dim: usize, + row_bytes: usize, +} + +/// Failure to represent a canonical MinMax matrix's dimension or byte extent. +#[derive(Debug, thiserror::Error)] +pub enum MinMaxMetaError { + #[error("intrinsic dimension {0} exceeds the u32 MinMax metadata limit")] + DimensionTooLarge(usize), + #[error(transparent)] + Overflow(#[from] Overflow), } impl MinMaxMeta @@ -39,11 +49,42 @@ where Unsigned: Representation, { /// Creates new MinMax metadata with the given number of rows and intrinsic dimension. + /// + /// # Panics + /// + /// Panics if the dimension does not fit the canonical metadata or the row/buffer + /// size is not representable. Use [`Self::try_new`] to handle invalid input. + #[track_caller] + #[expect( + clippy::expect_used, + reason = "try_new is the fallible construction API" + )] pub fn new(nrows: usize, intrinsic_dim: usize) -> Self { - Self { + Self::try_new(nrows, intrinsic_dim).expect("invalid canonical MinMax matrix dimensions") + } + + /// Validate canonical geometry without allocating storage. + /// + /// # Errors + /// + /// Returns [`MinMaxMetaError`] if the dimension exceeds `u32::MAX`, size arithmetic + /// overflows, or a row or the complete buffer exceeds `isize::MAX` bytes. + pub fn try_new(nrows: usize, intrinsic_dim: usize) -> Result { + u32::try_from(intrinsic_dim) + .map_err(|_| MinMaxMetaError::DimensionTooLarge(intrinsic_dim))?; + let overflow = || Overflow::for_type::(nrows, intrinsic_dim); + intrinsic_dim.checked_mul(NBITS).ok_or_else(overflow)?; + let row_bytes = Data::::slice_bytes(intrinsic_dim) + .checked_add(std::mem::size_of::()) + .ok_or_else(overflow)?; + let bytes = nrows.checked_mul(row_bytes).ok_or_else(overflow)?; + Overflow::check_byte_budget::(row_bytes, 1, intrinsic_dim)?; + Overflow::check_byte_budget::(bytes, nrows, intrinsic_dim)?; + Ok(Self { nrows, intrinsic_dim, - } + row_bytes, + }) } /// Returns the `intrinsic_dim` @@ -54,11 +95,11 @@ where /// Returns the number of bytes from a canonical repr of /// a minmax quantized vector, see [minmax::Data::canonical_bytes] pub fn ncols(&self) -> usize { - Data::::canonical_bytes(self.intrinsic_dim) + self.row_bytes } fn bytes(&self) -> usize { - std::mem::size_of::() * self.nrows() * self.ncols() + self.nrows * self.row_bytes } } @@ -82,9 +123,7 @@ where /// - [`crate::minmax::MinMaxCompensation`] does not require alignment /// since it is always accessed by copying bytes first before casting. fn layout(&self) -> Result { - Ok(std::alloc::Layout::array::( - self.nrows() * self.ncols(), - )?) + Ok(std::alloc::Layout::array::(self.bytes())?) } /// Returns an immutable reference to the i-th row. @@ -374,6 +413,52 @@ mod tests { mod construction { use super::*; + fn checked_dimensions() + where + Unsigned: Representation, + { + for dim in [0, 1, 7, 8, 9, 63, 64, 65] { + for rows in [0, 1, 3] { + let meta = MinMaxMeta::::try_new(rows, dim).unwrap(); + assert_eq!(meta.ncols(), Data::::canonical_bytes(dim)); + assert_eq!(meta.bytes(), rows * meta.ncols()); + assert_eq!(meta.layout().unwrap().size(), meta.bytes()); + let bytes = vec![0; meta.bytes()]; + let view = MatRef::new(meta, &bytes).unwrap(); + assert_eq!(view.rows().count(), rows); + } + } + let row_bytes = Data::::canonical_bytes(1); + let last_rows = isize::MAX as usize / row_bytes; + assert!(MinMaxMeta::::try_new(last_rows, 1).is_ok()); + assert!(matches!( + MinMaxMeta::::try_new(last_rows + 1, 1), + Err(MinMaxMetaError::Overflow(_)) + )); + assert!(matches!( + MinMaxMeta::::try_new(usize::MAX, 1), + Err(MinMaxMetaError::Overflow(_)) + )); + if let Some(dim) = (u32::MAX as usize).checked_add(1) { + for rows in [0, 1] { + assert!(matches!(MinMaxMeta::::try_new(rows, dim), + Err(MinMaxMetaError::DimensionTooLarge(value)) if value == dim)); + } + } + assert!(MinMaxMeta::::try_new(usize::MAX, usize::MAX).is_err()); + } + + expand_to_bitrates!( + canonical_extents_are_validated_before_allocation, + checked_dimensions + ); + + #[test] + #[should_panic(expected = "invalid canonical MinMax matrix dimensions")] + fn infallible_constructor_also_rejects_overflow() { + MinMaxMeta::<8>::new(usize::MAX, 1); + } + /// Tests NewOwned (Mat::new with Defaulted) for various dimensions and sizes. fn test_new_owned() where diff --git a/diskann-quantization/src/minmax/multi/mod.rs b/diskann-quantization/src/minmax/multi/mod.rs index 6524eae50c..88670370e7 100644 --- a/diskann-quantization/src/minmax/multi/mod.rs +++ b/diskann-quantization/src/minmax/multi/mod.rs @@ -14,4 +14,4 @@ mod query; pub use factory::build_minmax_max_sim; pub use kernel::{MinMaxErase, MinMaxMaxSimKernel}; pub use max_sim::MinMaxKernel; -pub use meta::MinMaxMeta; +pub use meta::{MinMaxMeta, MinMaxMetaError}; diff --git a/diskann-quantization/src/minmax/multi/query.rs b/diskann-quantization/src/minmax/multi/query.rs index 120b214448..edd72bfa4e 100644 --- a/diskann-quantization/src/minmax/multi/query.rs +++ b/diskann-quantization/src/minmax/multi/query.rs @@ -4,7 +4,7 @@ */ use crate::{ - matrix_kernels::maxsim::minmax8_x_minmax4::layout::{AView, Layout}, + matrix_kernels::maxsim::minmax8_x_minmax4::{a::AView, layout::Layout}, minmax::{MinMaxCompensation, MinMaxMeta}, multi_vector::{BlockTransposed, MatRef, Overflow}, }; @@ -58,3 +58,79 @@ impl PackedQuery { AView::new(self.values.as_view(), &self.compensation, &self.layout) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::multi_vector::{Defaulted, Mat}; + use diskann_utils::ReborrowMut; + + fn check_packing() { + for rows in [0, 1, MR - 1, MR, MR + 1] { + for dim in [0usize, 1, 31, 32, 33, 63, 64, 65, 127, 128, 129] { + let mut query = Mat::new(MinMaxMeta::<8>::new(rows, dim), Defaulted).unwrap(); + for (row, mut source) in query.reborrow_mut().rows_mut().enumerate() { + for d in 0..dim { + source + .vector_mut() + .set(d, ((row * 37 + d * 19) % 255 + 1) as i64) + .unwrap(); + } + source.set_meta(MinMaxCompensation { + dim: dim as u32, + a: row as f32 + 1.0, + b: -(row as f32) - 2.0, + n: row as f32 + 3.0, + ..Default::default() + }); + } + let packed = PackedQuery::::new(query.as_view()).unwrap(); + assert_eq!(packed.view().is_none(), rows == 0 || dim == 0); + let padded = dim.next_multiple_of(64); + // Construct the order independently of Layout::position. + let order: Vec<_> = (0..padded) + .step_by(64) + .flat_map(|base| { + (0..2).flat_map(move |parity| (base + parity..base + 64).step_by(2)) + }) + .collect(); + for (index, group) in packed.values.as_slice().chunks_exact(MR * PACK).enumerate() { + let panel = index / (padded / PACK); + let k = index % (padded / PACK) * PACK; + for row in 0..MR { + for lane in 0..PACK { + let original_row = panel * MR + row; + let d = order[k + lane]; + let expected = if original_row < rows && d < dim { + ((original_row * 37 + d * 19) % 255 + 1) as u8 + } else { + 0 + }; + assert_eq!( + group[row * PACK + lane], + expected, + "rows={rows}, dim={dim}, MR={MR}, PACK={PACK}, panel={panel}, row={row}, k={k}, lane={lane}" + ); + } + } + } + for (panel, metadata) in packed.compensation.iter().enumerate() { + for (row, actual) in metadata.iter().enumerate() { + let expected = query + .get_row(panel * MR + row) + .map(|row| row.meta()) + .unwrap_or_default(); + assert_eq!(*actual, expected); + } + } + } + } + } + + #[test] + fn prepared_queries_preserve_order_padding_and_metadata() { + check_packing::<8, 4>(); + check_packing::<16, 4>(); + check_packing::<16, 8>(); + } +}