Merge feat/neon-int8: aarch64 int8 dot product, verified on a Pi 5
CI / test (push) Failing after 2s
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:
+31
-5
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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) |
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user