quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat(accel): aarch64 int8 dot product — SDOT and plain NEON

Browse Source 
`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:
osobh
2026-09-21 17:34:51 -07:00
co-authored by Claude Opus 5
parent 7b16dc90d6
commit 56a8c2f3d0
2 changed files with 179 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-accel/src/neon.rs
+127
View File
@@ -180,3 +180,130 @@ pub fn checksum_fletcher32(data: &[u8]) -> u32 {
(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
}