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) | 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`].