From 97e65f2adf72d09dc2a0212ac2c60a5fdf049a17 Mon Sep 17 00:00:00 2001 From: osobh Date: Sun, 20 Sep 2026 17:30:56 -0700 Subject: [PATCH 1/2] feat(accel): runtime-dispatched int8 dot product MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The int8-quantised HNSW index compared vectors with a scalar loop that the compiler vectorised for the x86-64 baseline (SSE2), while the f32 path it was measured against goes through `clawhdf5-accel` and runs AVX2. So the ~13% throughput cost recorded for `quantized_index` was a missing kernel rather than a property of int8. `clawhdf5_accel::dot_i8` adds a scalar fallback and an AVX2 path: sign-extend each 16-byte half to i16, then `madd_epi16`, which multiplies and sums adjacent pairs straight into i32 lanes. It is dispatched through the same detected backend as the f32 kernels, and the index now calls it. Integer arithmetic, so the SIMD path must agree with scalar bit for bit — tested at lengths that are and are not multiples of the block, and at the -128 extreme for overflow. Co-Authored-By: Claude Opus 5 (1M context) --- crates/clawhdf5-accel/src/avx2.rs | 49 +++++++++++++++++++++++++ crates/clawhdf5-accel/src/lib.rs | 56 +++++++++++++++++++++++++++++ crates/clawhdf5-accel/src/scalar.rs | 30 ++++++++++++++++ crates/clawhdf5-ann/src/hnsw.rs | 29 +++------------ 4 files changed, 140 insertions(+), 24 deletions(-) diff --git a/crates/clawhdf5-accel/src/avx2.rs b/crates/clawhdf5-accel/src/avx2.rs index dcdbf7e..f2da7e3 100644 --- a/crates/clawhdf5-accel/src/avx2.rs +++ b/crates/clawhdf5-accel/src/avx2.rs @@ -25,6 +25,55 @@ unsafe fn hsum_256(v: __m256) -> f32 { _mm_cvtss_f32(result) } +/// AVX2 dot product of two `i8` slices, widened to `i32`. +/// +/// Each 16-byte half is sign-extended to sixteen `i16` lanes and multiplied +/// pairwise with `madd_epi16`, which sums adjacent products straight into +/// eight `i32` lanes — the widening that an autovectorised scalar loop does +/// in several shuffles is one instruction here. A pair sum is at most +/// `2 * 127 * 127`, far inside `i32`. +/// +/// # Safety +/// Caller must verify is_x86_feature_detected!("avx2"). +// SAFETY: Caller must have verified AVX2 via is_x86_feature_detected!. +#[target_feature(enable = "avx2")] +pub unsafe fn dot_i8(a: &[i8], b: &[i8]) -> i32 { + // SAFETY: Caller guarantees AVX2 is available per the # Safety contract; + // every load reads 32 bytes at an index checked against `len` first. + unsafe { + assert_eq!(a.len(), b.len()); + let len = a.len(); + let mut i = 0; + let mut acc0 = _mm256_setzero_si256(); + let mut acc1 = _mm256_setzero_si256(); + + while i + 32 <= len { + let va = _mm256_loadu_si256(a.as_ptr().add(i).cast()); + let vb = _mm256_loadu_si256(b.as_ptr().add(i).cast()); + let a_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(va)); + let b_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(vb)); + let a_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(va, 1)); + let b_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(vb, 1)); + acc0 = _mm256_add_epi32(acc0, _mm256_madd_epi16(a_lo, b_lo)); + acc1 = _mm256_add_epi32(acc1, _mm256_madd_epi16(a_hi, b_hi)); + i += 32; + } + + // Horizontal sum of the eight i32 lanes. + let v = _mm256_add_epi32(acc0, acc1); + let s128 = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1)); + let s64 = _mm_add_epi32(s128, _mm_unpackhi_epi64(s128, s128)); + let s32 = _mm_add_epi32(s64, _mm_shuffle_epi32(s64, 0b01)); + let mut sum = _mm_cvtsi128_si32(s32); + + while i < len { + sum += i32::from(a[i]) * i32::from(b[i]); + i += 1; + } + sum + } +} + /// AVX2 dot product for f32 slices. /// /// # Safety diff --git a/crates/clawhdf5-accel/src/lib.rs b/crates/clawhdf5-accel/src/lib.rs index 95987ef..6fcd9b3 100644 --- a/crates/clawhdf5-accel/src/lib.rs +++ b/crates/clawhdf5-accel/src/lib.rs @@ -122,6 +122,23 @@ pub fn dot_product(a: &[f32], b: &[f32]) -> f32 { } } +/// Dot product of two `i8` slices, widened to `i32`. +/// +/// The kernel behind int8-quantised vector search. Uses the AVX2 path +/// whenever AVX2 is present — including on AVX-512 machines, where it is +/// what the f32 kernels use too on a default build. +pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 { + match detect_backend() { + #[cfg(target_arch = "x86_64")] + // SAFETY: both variants imply AVX2 was detected at runtime (the + // AVX-512 backend is only selected on CPUs that also have AVX2). + Backend::Avx2 | Backend::Avx512 if is_x86_feature_detected!("avx2") => unsafe { + avx2::dot_i8(a, b) + }, + _ => scalar::dot_i8(a, b), + } +} + /// Compute the L2 norm (magnitude) of a vector. pub fn vector_norm(v: &[f32]) -> f32 { dot_product(v, v).sqrt() @@ -713,3 +730,42 @@ mod tests { } } } + +#[cfg(test)] +mod dot_i8_tests { + use super::*; + + fn codes(n: usize, seed: u64) -> Vec { + let mut state = seed; + (0..n) + .map(|_| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + // Full range, including the extremes. + ((state >> 56) as u8) as i8 + }) + .collect() + } + + #[test] + fn dispatched_kernel_matches_scalar_exactly() { + // Integer arithmetic: the SIMD path must agree bit for bit, at every + // length — including ones that are not multiples of the 32-byte block, + // which exercise the tail. + for len in [0, 1, 7, 31, 32, 33, 63, 64, 100, 384, 385, 1536] { + let a = codes(len, 1 + len as u64); + let b = codes(len, 1000 + len as u64); + assert_eq!(dot_i8(&a, &b), scalar::dot_i8(&a, &b), "len {len}"); + } + } + + #[test] + fn extremes_do_not_overflow() { + // -128 * -128 is the largest product; a long run of it must still fit. + let a = vec![-128i8; 4096]; + assert_eq!(dot_i8(&a, &a), 4096 * 128 * 128); + let b = vec![127i8; 4096]; + assert_eq!(dot_i8(&a, &b), -4096 * 128 * 127); + } +} diff --git a/crates/clawhdf5-accel/src/scalar.rs b/crates/clawhdf5-accel/src/scalar.rs index 3c9a507..e8dd5b4 100644 --- a/crates/clawhdf5-accel/src/scalar.rs +++ b/crates/clawhdf5-accel/src/scalar.rs @@ -140,3 +140,33 @@ fn f16_to_f32_soft(h: u16) -> f32 { f32::from_bits(f32_bits) } + +/// Dot product of two `i8` slices, widened to `i32`. +/// +/// `dim` terms of at most `127 * 127` fit an `i32` for any realistic +/// dimension (over 130 000 terms before overflow is possible). +pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 { + assert_eq!(a.len(), b.len()); + // Four independent accumulators over 32-lane blocks: the widening product + // has to sit in a fixed-length chunk for the vectoriser to see it, and the + // separate accumulators keep it off one dependency chain. + const LANE: usize = 8; + let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>(); + let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>(); + let mut acc = [0i32; 4]; + for (x, y) in a_blocks.iter().zip(b_blocks) { + for (lane, slot) in acc.iter_mut().enumerate() { + let mut sum = 0i32; + for k in 0..LANE { + sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]); + } + *slot += sum; + } + } + let tail: i32 = a_tail + .iter() + .zip(b_tail) + .map(|(&x, &y)| i32::from(x) * i32::from(y)) + .sum(); + acc[0] + acc[1] + acc[2] + acc[3] + tail +} diff --git a/crates/clawhdf5-ann/src/hnsw.rs b/crates/clawhdf5-ann/src/hnsw.rs index a4104e4..a526bee 100644 --- a/crates/clawhdf5-ann/src/hnsw.rs +++ b/crates/clawhdf5-ann/src/hnsw.rs @@ -358,31 +358,12 @@ enum Query { Int8(Vec, f32), } -/// Sum of products, widened so it cannot overflow: `dim` terms of at most -/// `127 * 127`, so `i32` suffices for any realistic dimension. +/// Sum of products, widened so it cannot overflow. Runtime-dispatched to the +/// same SIMD backend as the f32 kernels, so the two storages are compared on +/// equal terms. +#[inline] fn dot_i8(a: &[i8], b: &[i8]) -> i32 { - // Four independent accumulators over 32-lane blocks: the widening product - // has to sit in a fixed-length chunk for the vectoriser to see it, and the - // separate accumulators keep it off one dependency chain. - const LANE: usize = 8; - let (a_blocks, a_tail) = a.as_chunks::<{ LANE * 4 }>(); - let (b_blocks, b_tail) = b.as_chunks::<{ LANE * 4 }>(); - let mut acc = [0i32; 4]; - for (x, y) in a_blocks.iter().zip(b_blocks) { - for (lane, slot) in acc.iter_mut().enumerate() { - let mut sum = 0i32; - for k in 0..LANE { - sum += i32::from(x[lane * LANE + k]) * i32::from(y[lane * LANE + k]); - } - *slot += sum; - } - } - let tail: i32 = a_tail - .iter() - .zip(b_tail) - .map(|(&x, &y)| i32::from(x) * i32::from(y)) - .sum(); - acc[0] + acc[1] + acc[2] + acc[3] + tail + clawhdf5_accel::dot_i8(a, b) } /// Magic for [`HnswIndex::graph_to_bytes`]. From dea02f5214663be1910da360956c6179a5960ce0 Mon Sep 17 00:00:00 2001 From: osobh Date: Sun, 20 Sep 2026 17:30:56 -0700 Subject: [PATCH 2/2] docs: the int8 index is faster than f32 on AVX2, not slower Measured with `clawhdf5_accel::dot_i8` in place, medians of three alternating runs at N = 100 000 x 384, same binary: build f32 3197 ms int8 1778 ms int8+re-score 1826 ms ef = 64 f32 13 399 QPS @ 0.9945 int8+re-score 21 848 QPS @ 0.9940 So at equal recall the quantised index is 1.63x the queries per second and 1.8x the build speed, holding a quarter of the vectors. The earlier "~13% of QPS" figure compared a scalar int8 loop against hand-written AVX2 f32 kernels and was measuring the missing kernel; it is kept in BENCHMARKS.md with that explanation rather than quietly replaced. Still off by default, now for portability rather than performance: the kernel is AVX2-only and aarch64 falls back to scalar, where the original trade applies. A NEON kernel would settle it. Co-Authored-By: Claude Opus 5 (1M context) --- BENCHMARKS.md | 30 +++++++++++++++++++++++++----- CHANGELOG.md | 14 ++++++++++++++ CLAUDE.md | 6 ++++-- README.md | 7 +++++-- 4 files changed, 48 insertions(+), 9 deletions(-) diff --git a/BENCHMARKS.md b/BENCHMARKS.md index bfea271..7668fb8 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -88,11 +88,31 @@ back**, because the loss is in the distances rather than in the graph Re-scoring closes the gap: the store already holds the exact embeddings, so the query path re-scores the candidate pool against them before fusion. That -is done automatically whenever the index is quantised. What it costs is -throughput — about 13% of QPS and 16% of build time at 100 000 x 384. So the -setting trades ~13% of query speed for ~36% of the process's memory at equal -recall. It is **off by default**: the right side of that trade depends on -whether the deployment is short of memory or short of CPU. +is done automatically whenever the index is quantised. + +**On AVX2 this costs nothing — it pays.** The first measurement of this put +the cost at ~13% of QPS and ~16% of build time, but that compared a scalar +int8 loop against `clawhdf5-accel`'s hand-written AVX2 kernels for `f32`: +the gap was a missing kernel, not a property of int8. With +`clawhdf5_accel::dot_i8` (AVX2: sign-extend to `i16`, then `madd_epi16`), +medians of three alternating runs at N = 100 000, same binary: + +| | f32 | int8 | int8 + re-score | +|---|---:|---:|---:| +| build | 3197 ms | **1778 ms** | 1826 ms | +| QPS at ef = 64 | 13 399 | 29 195 | **21 848** | +| recall@10 at ef = 64 | 0.9945 | 0.9625 | **0.9940** | + +So at equal recall the quantised index answers **1.63x as many queries per +second**, builds **1.8x faster**, and holds a quarter of the vectors. (Compare +only at equal `ef`: with re-scoring the harness raises `ef` to at least the +candidate pool, so the `ef = 16` and `ef = 32` rows are not like-for-like.) + +It is still **off by default**, for portability rather than performance: the +int8 kernel is AVX2-only, and on aarch64 — including `clawhdf5-android` — it +falls back to the scalar loop, where the original trade still applies. A NEON +kernel would remove that caveat. On an x86-64 deployment, turning it on is a +win on every axis measured. A measurement trap worth recording: the synthetic `clustered` generator in the `clawhdf5-ann` tests draws clusters far tighter than any real embedding, so diff --git a/CHANGELOG.md b/CHANGELOG.md index e4ea2a1..da4972d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,20 @@ ## Unreleased +### Performance +- `clawhdf5-accel`: **`dot_i8`, a runtime-dispatched int8 dot product** (AVX2: + sign-extend each half to `i16`, then `madd_epi16`; scalar fallback + elsewhere). The quantised HNSW index used a scalar loop while the `f32` path + it was measured against ran AVX2, so the ~13% throughput cost recorded for + `MemoryConfig::quantized_index` was a missing kernel rather than a property + of int8. With the kernel, at N = 100 000 x 384 and equal recall, the + quantised index answers **1.63x as many queries per second** (21 848 vs + 13 399 at ef=64, recall 0.9940 vs 0.9945) and builds **1.8x faster** (1778 + vs 3197 ms) — on top of holding a quarter of the vectors. Medians of three + alternating runs. It remains off by default only because the kernel is + AVX2-only and aarch64 falls back to the scalar loop. Integer arithmetic, so + the SIMD path is tested to agree with scalar bit for bit. + ### Correctness - `clawhdf5-format`: **datasets indexed by an Extensible Array returned wrong data beyond their first few dozen chunks.** One unlimited dimension gives a diff --git a/CLAUDE.md b/CLAUDE.md index 60bd9a0..b295dd3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -44,8 +44,10 @@ Cargo workspace with 16 crates under `crates/` (plus `libaec-sys`, an internal F which roughly halves a loaded store's memory (2.72x -> 1.74x the raw vectors at 100K); because quantised distances are approximate and `ef` cannot compensate, the query path then re-scores the candidate pool against the - exact embeddings, which holds recall at the f32 index's level and costs - ~13% of QPS. `hybrid_search` keeps one incremental BM25 + exact embeddings, which holds recall at the f32 index's level. On AVX2 it is + also 1.63x the QPS and 1.8x the build speed (`clawhdf5_accel::dot_i8`); it + stays off by default only because that kernel is AVX2-only and aarch64 falls + back to scalar. `hybrid_search` keeps one incremental BM25 index for the life of the store and never writes the store: Hebbian activation boosts are persisted by the next checkpoint (or on drop), not per query. Measure any search-path change with diff --git a/README.md b/README.md index 53e8b85..97a294d 100644 --- a/README.md +++ b/README.md @@ -437,8 +437,11 @@ ClawhDF5's agent memory design draws from 15+ recent papers: copy of the embeddings as `i8`, roughly halving a loaded store's memory (2.72x -> 1.74x the raw vectors at 100k x 384). Quantised distances are approximate, so the query path re-scores the candidate pool against the exact -embeddings the store already holds — recall matches the `f32` index, at about -13% fewer queries per second. See `BENCHMARKS.md`, "Quantising the index copy". +embeddings the store already holds, which keeps recall at the `f32` index's +level. On AVX2 it is also **faster** — 1.63x the queries per second and 1.8x +the build speed at equal recall — because the int8 kernel is SIMD too. It +stays off by default only because that kernel is AVX2-only and aarch64 falls +back to a scalar loop. See `BENCHMARKS.md`, "Quantising the index copy". | `parallel` | no | Rayon parallel search | | `fast-math` | no | BLAS matrix-vector multiply | | `accelerate` | no | Apple Accelerate / AMX (macOS) |