Merge feat/neon-int8: aarch64 int8 dot product, verified on a Pi 5
CI / test (push) Failing after 2s

SDOT and plain-NEON kernels for dot_i8, tested bit-exact against scalar
on real ARM. At equal recall the quantised index is 1.18x f32 on a Pi 5
and builds 2.3x faster. Also corrects an unmeasured claim that it was
slower than f32 on ARM.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-21 17:34:51 -07:00
co-authored by Claude Opus 5
6 changed files with 244 additions and 16 deletions
+31 -5
View File
@@ -108,11 +108,37 @@ 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 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.) 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 #### On ARM (Raspberry Pi 5, Cortex-A76)
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 `dot_i8` has two aarch64 kernels: `SDOT` for CPUs with the ARMv8.2
kernel would remove that caveat. On an x86-64 deployment, turning it on is a dot-product extension (Cortex-A76 and later, Neoverse-N1, all Apple Silicon)
win on every axis measured. and plain NEON (`vmull_s8` + `vpadalq_s16`) otherwise. Medians of three runs
at N = 100 000, ef = 64, recall@10 0.9940 in every int8 row against f32's
0.9945:
| int8 kernel | build | QPS | vs f32 |
|---|---:|---:|---:|
| *(f32 baseline)* | 33 413 ms | 6 164 | 1.00x |
| scalar (what v2.7.0 shipped) | 18 950 ms | ~6 190 | 1.00x |
| plain NEON | ~17 000 ms | 6 640 | 1.08x |
| **SDOT** | **14 464 ms** | **7 267** | **1.18x** |
These are Pi 5 numbers, not "ARM" numbers: a Pi has far less memory bandwidth
and cache than an Apple M-series or a flagship phone, so the ratios will move
on other hardware. The plain-NEON row is that code on an A76 with `SDOT`
disabled, not a measurement of a pre-A76 core.
**A correction.** Until this was measured, this section said aarch64 "falls
back to the scalar loop, where the original trade still applies" — that is,
that quantised search was ~13% slower than f32 on ARM. That was extrapolated
from x86 and it was wrong. On x86-64 the portable baseline is SSE2 while the
f32 kernels are hand-written AVX2, so scalar int8 lost; on aarch64 NEON *is*
the baseline, the compiler vectorises the scalar loop well, and scalar int8
already matched f32 for search while building 1.76x faster.
So on every configuration measured — x86-64 AVX2, and Pi 5 with each of the
three int8 kernels — the quantised index is at least as fast as f32 at equal
recall, builds faster, and holds a quarter of the vectors.
A measurement trap worth recording: the synthetic `clustered` generator in the A measurement trap worth recording: the synthetic `clustered` generator in the
`clawhdf5-ann` tests draws clusters far tighter than any real embedding, so `clawhdf5-ann` tests draws clusters far tighter than any real embedding, so
+25
View File
@@ -1,5 +1,30 @@
# Changelog # Changelog
## Unreleased
### Performance
- `clawhdf5-accel`: **`dot_i8` has aarch64 kernels** — `SDOT` for CPUs with
the ARMv8.2 dot-product extension (Cortex-A76 and later, Neoverse-N1, every
Apple Silicon generation) and plain NEON (`vmull_s8` + `vpadalq_s16`) for
the rest, selected at runtime. `SDOT` is issued through inline assembly,
because the `vdotq_s32` intrinsic is still behind the unstable
`stdarch_neon_dotprod` feature. On a Raspberry Pi 5 at N = 100 000 and
equal recall, the quantised index answers **1.18x the queries per second**
of f32 (7 267 vs 6 164) and builds **2.3x faster** (14 464 vs 33 413 ms).
Both kernels are tested bit-for-bit against scalar on real hardware, each
explicitly — dispatch only ever takes one path on a given CPU, so testing
through it alone would have left the plain-NEON fallback unexercised on any
machine with `SDOT`.
### Corrections
- The v2.7.0 entry for `dot_i8` said `quantized_index` stayed off by default
because "aarch64 falls back to the scalar loop", implying the ~13% search
penalty measured on x86 applied on ARM too. It did not. That figure came
from scalar int8 against hand-written AVX2 f32 kernels on x86, whose
portable baseline is SSE2; on aarch64 NEON is the baseline, and measured on
a Pi 5 the scalar int8 loop already matched f32 for search while building
1.76x faster. The claim was extrapolated rather than measured.
## v2.7.0 (2026-09-20) ## v2.7.0 (2026-09-20)
### Upgrade Notes ### Upgrade Notes
+6 -4
View File
@@ -44,10 +44,12 @@ 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 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 at 100K); because quantised distances are approximate and `ef` cannot
compensate, the query path then re-scores the candidate pool against the compensate, the query path then re-scores the candidate pool against the
exact embeddings, which holds recall at the f32 index's level. On AVX2 it is exact embeddings, which holds recall at the f32 index's level. It is also
also 1.63x the QPS and 1.8x the build speed (`clawhdf5_accel::dot_i8`); it faster at equal recall: 1.63x the QPS on x86-64 (AVX2) and 1.18x on a
stays off by default only because that kernel is AVX2-only and aarch64 falls Raspberry Pi 5 (`clawhdf5_accel::dot_i8`, NEON `SDOT` via inline asm since
back to scalar. `hybrid_search` keeps one incremental BM25 the intrinsic is unstable; plain NEON on pre-dotprod cores). The aarch64
code is `cfg`'d out on x86, so x86 CI never compiles or lints it — test it
on real ARM (`rpivision02`, 10.0.2.3, is a Pi 5). `hybrid_search` keeps one incremental BM25
index for the life of the store and never writes the store: Hebbian 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 activation boosts are persisted by the next checkpoint (or on drop), not per
query. Measure any search-path change with query. Measure any search-path change with
+3 -4
View File
@@ -442,10 +442,9 @@ 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 (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 approximate, so the query path re-scores the candidate pool against the exact
embeddings the store already holds, which keeps recall at the `f32` index's 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 level. It is also **faster**: 1.63x the queries per second at equal recall on
the build speed at equal recall — because the int8 kernel is SIMD too. It x86-64 (AVX2) and 1.18x on a Raspberry Pi 5 (NEON `SDOT`), with index builds
stays off by default only because that kernel is AVX2-only and aarch64 falls 1.8x and 2.3x faster respectively. See `BENCHMARKS.md`, "Quantising the index copy".
back to a scalar loop. See `BENCHMARKS.md`, "Quantising the index copy".
| `parallel` | no | Rayon parallel search | | `parallel` | no | Rayon parallel search |
| `fast-math` | no | BLAS matrix-vector multiply | | `fast-math` | no | BLAS matrix-vector multiply |
| `accelerate` | no | Apple Accelerate / AMX (macOS) | | `accelerate` | no | Apple Accelerate / AMX (macOS) |
+52 -3
View File
@@ -124,11 +124,24 @@ pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
/// Dot product of two `i8` slices, widened to `i32`. /// Dot product of two `i8` slices, widened to `i32`.
/// ///
/// The kernel behind int8-quantised vector search. Uses the AVX2 path /// The kernel behind int8-quantised vector search. On x86-64 it uses the AVX2
/// whenever AVX2 is present including on AVX-512 machines, where it is /// path whenever AVX2 is present (including on AVX-512 machines, where it is
/// what the f32 kernels use too on a default build. /// what the f32 kernels use too on a default build). On aarch64 it uses the
/// ARMv8.2 `SDOT` instruction when the CPU has the dot-product extension, and
/// plain NEON otherwise.
pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 { pub fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
match detect_backend() { match detect_backend() {
#[cfg(target_arch = "aarch64")]
Backend::Neon => {
if std::arch::is_aarch64_feature_detected!("dotprod") {
// SAFETY: the dotprod extension was just detected at runtime.
unsafe { neon::dot_i8_dotprod(a, b) }
} else {
// SAFETY: NEON is always available on aarch64.
unsafe { neon::dot_i8(a, b) }
}
}
#[cfg(target_arch = "x86_64")] #[cfg(target_arch = "x86_64")]
// SAFETY: both variants imply AVX2 was detected at runtime (the // SAFETY: both variants imply AVX2 was detected at runtime (the
// AVX-512 backend is only selected on CPUs that also have AVX2). // AVX-512 backend is only selected on CPUs that also have AVX2).
@@ -760,6 +773,42 @@ mod dot_i8_tests {
} }
} }
/// Dispatch only ever takes one path on a given CPU, so on a machine with
/// the dot-product extension the plain-NEON kernel would otherwise go
/// untested. Check each aarch64 kernel against scalar directly.
#[cfg(target_arch = "aarch64")]
#[test]
fn every_aarch64_kernel_matches_scalar_exactly() {
for len in [0, 1, 7, 15, 16, 17, 31, 32, 33, 63, 64, 100, 384, 385, 1536] {
let a = codes(len, 7 + len as u64);
let b = codes(len, 7000 + len as u64);
let want = scalar::dot_i8(&a, &b);
// SAFETY: NEON is always available on aarch64.
assert_eq!(unsafe { neon::dot_i8(&a, &b) }, want, "neon, len {len}");
if std::arch::is_aarch64_feature_detected!("dotprod") {
// SAFETY: the dotprod extension was just detected.
assert_eq!(
unsafe { neon::dot_i8_dotprod(&a, &b) },
want,
"dotprod, len {len}"
);
}
}
// The extremes, through both kernels.
let lo = vec![-128i8; 4096];
let hi = vec![127i8; 4096];
// SAFETY: NEON is always available on aarch64.
assert_eq!(unsafe { neon::dot_i8(&lo, &lo) }, 4096 * 128 * 128);
// SAFETY: NEON is always available on aarch64.
assert_eq!(unsafe { neon::dot_i8(&lo, &hi) }, -4096 * 128 * 127);
if std::arch::is_aarch64_feature_detected!("dotprod") {
// SAFETY: the dotprod extension was just detected.
assert_eq!(unsafe { neon::dot_i8_dotprod(&lo, &lo) }, 4096 * 128 * 128);
// SAFETY: the dotprod extension was just detected.
assert_eq!(unsafe { neon::dot_i8_dotprod(&lo, &hi) }, -4096 * 128 * 127);
}
}
#[test] #[test]
fn extremes_do_not_overflow() { fn extremes_do_not_overflow() {
// -128 * -128 is the largest product; a long run of it must still fit. // -128 * -128 is the largest product; a long run of it must still fit.
+127
View File
@@ -180,3 +180,130 @@ pub fn checksum_fletcher32(data: &[u8]) -> u32 {
(sum2 << 16) | sum1 (sum2 << 16) | sum1
} }
/// NEON dot product of two `i8` slices, widened to `i32`, for any aarch64 CPU.
///
/// `vmull_s8` multiplies eight lanes into `i16` — even `-128 * -128` is 16 384,
/// inside `i16` — and `vpadalq_s16` adds adjacent pairs of those into `i32`
/// accumulators, so nothing can overflow before the final horizontal sum.
///
/// CPUs with the ARMv8.2 dot-product extension should use
/// [`dot_i8_dotprod`], which does the multiply and the accumulate in one
/// instruction.
///
/// # Safety
/// Caller must ensure aarch64 target (NEON always available).
// SAFETY: NEON is always available on aarch64 targets; caller guarantees aarch64.
#[target_feature(enable = "neon")]
pub unsafe fn dot_i8(a: &[i8], b: &[i8]) -> i32 {
assert_eq!(a.len(), b.len());
let len = a.len();
let mut i = 0;
let mut acc0 = vdupq_n_s32(0);
let mut acc1 = vdupq_n_s32(0);
while i + 16 <= len {
// SAFETY: NEON is available per the # Safety contract, and both
// 16-byte loads start at an index checked against `len` above.
unsafe {
let va = vld1q_s8(a.as_ptr().add(i));
let vb = vld1q_s8(b.as_ptr().add(i));
acc0 = vpadalq_s16(acc0, vmull_s8(vget_low_s8(va), vget_low_s8(vb)));
acc1 = vpadalq_s16(acc1, vmull_high_s8(va, vb));
}
i += 16;
}
let mut sum = vaddvq_s32(vaddq_s32(acc0, acc1));
while i < len {
sum += i32::from(a[i]) * i32::from(b[i]);
i += 1;
}
sum
}
/// One `SDOT`: for each of the four `i32` lanes of `acc`, add the dot
/// product of the corresponding four `i8` pairs from `a` and `b`.
///
/// Written as inline assembly because the `vdotq_s32` intrinsic is still
/// behind the unstable `stdarch_neon_dotprod` feature; inline assembly is
/// stable on aarch64.
///
/// # Safety
/// Caller must ensure the CPU supports the `dotprod` extension.
#[inline]
#[target_feature(enable = "neon,dotprod")]
unsafe fn sdot(acc: int32x4_t, a: int8x16_t, b: int8x16_t) -> int32x4_t {
let mut acc = acc;
// SAFETY: `dotprod` is enabled for this function and the caller
// guarantees the CPU supports it. The instruction reads only its three
// vector registers and touches no memory.
unsafe {
std::arch::asm!(
"sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b",
acc = inout(vreg) acc,
a = in(vreg) a,
b = in(vreg) b,
options(pure, nomem, nostack),
);
}
acc
}
/// NEON dot product of two `i8` slices using the ARMv8.2 dot-product
/// extension (`SDOT`): sixteen multiply-accumulates per instruction, straight
/// into `i32` lanes.
///
/// Present on the cores this crate actually runs on — Cortex-A76 and later
/// (Raspberry Pi 5, current Android phones), Neoverse-N1 (Graviton2, Ampere
/// Altra), and every Apple Silicon generation.
///
/// # Safety
/// Caller must verify `is_aarch64_feature_detected!("dotprod")`.
// SAFETY: caller has verified the dotprod extension at runtime.
#[target_feature(enable = "neon,dotprod")]
pub unsafe fn dot_i8_dotprod(a: &[i8], b: &[i8]) -> i32 {
assert_eq!(a.len(), b.len());
let len = a.len();
let mut i = 0;
let mut acc0 = vdupq_n_s32(0);
let mut acc1 = vdupq_n_s32(0);
// Two independent accumulators so consecutive SDOTs are not serialised on
// one register.
while i + 32 <= len {
// SAFETY: dotprod is available per the # Safety contract, and every
// 16-byte load starts at an index checked against `len` above.
unsafe {
acc0 = sdot(
acc0,
vld1q_s8(a.as_ptr().add(i)),
vld1q_s8(b.as_ptr().add(i)),
);
acc1 = sdot(
acc1,
vld1q_s8(a.as_ptr().add(i + 16)),
vld1q_s8(b.as_ptr().add(i + 16)),
);
}
i += 32;
}
if i + 16 <= len {
// SAFETY: as above; the load is bounds-checked by this condition.
unsafe {
acc0 = sdot(
acc0,
vld1q_s8(a.as_ptr().add(i)),
vld1q_s8(b.as_ptr().add(i)),
);
}
i += 16;
}
let mut sum = vaddvq_s32(vaddq_s32(acc0, acc1));
while i < len {
sum += i32::from(a[i]) * i32::from(b[i]);
i += 1;
}
sum
}