feat(accel): aarch64 int8 dot product — SDOT and plain NEON
`dot_i8` had an AVX2 kernel and a scalar fallback, so on aarch64 the
quantised HNSW index ran the scalar loop. It now dispatches to one of
two NEON kernels:
- `dot_i8_dotprod`: the ARMv8.2 dot-product instruction, `SDOT`, which
multiplies and accumulates sixteen i8 pairs into four i32 lanes per
instruction. Present on Cortex-A76 and later (Raspberry Pi 5, current
Android phones), Neoverse-N1 (Graviton2, Ampere Altra) and every Apple
Silicon generation. Issued as inline assembly because the `vdotq_s32`
intrinsic is still behind the unstable `stdarch_neon_dotprod` feature;
inline asm is stable on aarch64.
- `dot_i8`: plain NEON for cores without the extension — `vmull_s8`
widens to i16 (even -128 * -128 fits) and `vpadalq_s16` folds adjacent
pairs into i32 accumulators, so nothing overflows.
Selected at runtime with `is_aarch64_feature_detected!("dotprod")`.
Verified on a Raspberry Pi 5 (Cortex-A76, `asimddp` present), not just
compiled — the aarch64 code is cfg'd out on x86, so x86 CI never builds
or lints it:
- both kernels bit-exact against scalar at every length, tails and
extremes included. Each is tested directly rather than through
dispatch, because dispatch only takes one path on a given CPU: on the
Pi, testing through it alone would never have run the plain-NEON
fallback at all.
- mutation-checked: dropping the SDOT kernel's second accumulator fails
at length 32, and using the low half twice in the NEON kernel fails at
length 16 — the first lengths that exercise each.
- the ANN suite passes, including int8 recall against ground truth.
- clippy clean with -D warnings on aarch64.
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
This commit is contained in:
@@ -124,11 +124,24 @@ 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.
|
||||
/// The kernel behind int8-quantised vector search. On x86-64 it 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). 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 {
|
||||
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")]
|
||||
// SAFETY: both variants imply AVX2 was detected at runtime (the
|
||||
// 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]
|
||||
fn extremes_do_not_overflow() {
|
||||
// -128 * -128 is the largest product; a long run of it must still fit.
|
||||
|
||||
Reference in New Issue
Block a user