quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases 1 Wiki Activity
Files
00b0cb00351855fe0d82d430754f6aee4c497f11
clawhdf5/crates/clawhdf5-accel/src/lib.rs
T
osobhandClaude Opus 5 56a8c2f3d0 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]>
2026-09-21 17:34:51 -07:00

821 lines
28 KiB
Rust
Raw Blame History

//! SIMD-accelerated operations for clawhdf5.
//!
//! This crate provides runtime-dispatched SIMD acceleration for common
//! vector operations used in HDF5 processing: dot products, cosine similarity,
//! L2 distance, f16 conversion, and checksums.
//!
//! All public functions automatically select the best available SIMD backend
//! at runtime. Every operation has a portable scalar fallback.
pub mod scalar;
#[cfg(target_arch = "aarch64")]
pub mod neon;
#[cfg(target_arch = "x86_64")]
pub mod avx2;
#[cfg(all(target_arch = "x86_64", feature = "avx512"))]
pub mod avx512;
pub mod checksum;
pub mod convert;
// ---------------------------------------------------------------------------
// Cache-line size detection (TVL — Tensor Virtualization Layout)
// ---------------------------------------------------------------------------
/// Cache line size in bytes for the target architecture.
///
/// ARM64 (Apple M-series, Cortex-A76+) uses 128-byte cache lines.
/// x86_64 uses 64-byte cache lines. Other architectures default to 64.
#[cfg(target_arch = "aarch64")]
pub const CACHE_LINE_SIZE: usize = 128;
#[cfg(target_arch = "x86_64")]
pub const CACHE_LINE_SIZE: usize = 64;
#[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
pub const CACHE_LINE_SIZE: usize = 64;
/// Round `size` up to the next multiple of [`CACHE_LINE_SIZE`].
#[inline]
pub fn align_to_cache_line(size: usize) -> usize {
(size + CACHE_LINE_SIZE - 1) & !(CACHE_LINE_SIZE - 1)
}
/// Available SIMD backends.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Backend {
/// ARM NEON (always available on aarch64)
Neon,
/// x86_64 AVX2 + FMA
Avx2,
/// x86_64 AVX-512F
Avx512,
/// x86_64 SSE4.1
Sse4,
/// WebAssembly SIMD128
WasmSimd128,
/// Portable scalar fallback
Scalar,
}
/// The best available SIMD backend, detected once per process. Every kernel
/// dispatches through this, so it sits in the innermost loop of every search.
pub fn detect_backend() -> Backend {
static BACKEND: std::sync::OnceLock<Backend> = std::sync::OnceLock::new();
*BACKEND.get_or_init(detect_backend_uncached)
}
fn detect_backend_uncached() -> Backend {
#[cfg(target_arch = "aarch64")]
{
return Backend::Neon; // Always available on aarch64
}
#[cfg(target_arch = "x86_64")]
{
#[cfg(feature = "avx512")]
{
if is_x86_feature_detected!("avx512f") {
return Backend::Avx512;
}
}
if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
return Backend::Avx2;
}
if is_x86_feature_detected!("sse4.1") {
return Backend::Sse4;
}
}
#[cfg(target_arch = "wasm32")]
{
return Backend::WasmSimd128;
}
#[allow(unreachable_code)]
Backend::Scalar
}
// ---------------------------------------------------------------------------
// Public API — auto-dispatched
// ---------------------------------------------------------------------------
/// Compute the dot product of two f32 slices.
pub fn dot_product(a: &[f32], b: &[f32]) -> f32 {
match detect_backend() {
#[cfg(target_arch = "aarch64")]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Neon => unsafe { neon::dot_product(a, b) },
#[cfg(all(target_arch = "x86_64", feature = "avx512"))]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Avx512 => unsafe { avx512::dot_product(a, b) },
#[cfg(target_arch = "x86_64")]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Avx2 => unsafe { avx2::dot_product(a, b) },
_ => scalar::dot_product(a, b),
}
}
/// Dot product of two `i8` slices, widened to `i32`.
///
/// 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).
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()
}
/// Compute cosine similarity between two vectors (fused single-pass).
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
match detect_backend() {
#[cfg(target_arch = "aarch64")]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Neon => unsafe { neon::cosine_similarity(a, b) },
#[cfg(all(target_arch = "x86_64", feature = "avx512"))]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Avx512 => unsafe { avx512::cosine_similarity(a, b) },
#[cfg(target_arch = "x86_64")]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Avx2 => unsafe { avx2::cosine_similarity(a, b) },
_ => scalar::cosine_similarity(a, b),
}
}
/// Compute cosine similarity between a query and multiple vectors.
///
/// Results are stored as `(index, similarity)` pairs.
pub fn batch_cosine(query: &[f32], vectors: &[&[f32]], results: &mut [(usize, f32)]) {
assert!(results.len() >= vectors.len());
for (i, v) in vectors.iter().enumerate() {
results[i] = (i, cosine_similarity(query, v));
}
}
/// Compute cosine similarity with pre-normalized query vector.
///
/// `query_normed` must already be unit-length. `norms` contains the L2 norms
/// of each vector in `vectors`.
pub fn batch_cosine_prenorm(
query_normed: &[f32],
vectors: &[&[f32]],
norms: &[f32],
results: &mut [(usize, f32)],
) {
assert!(results.len() >= vectors.len());
assert!(norms.len() >= vectors.len());
for (i, v) in vectors.iter().enumerate() {
let dot = dot_product(query_normed, v);
let sim = if norms[i] == 0.0 { 0.0 } else { dot / norms[i] };
results[i] = (i, sim);
}
}
/// Compute L2 (Euclidean) distance between two vectors.
pub fn l2_distance(a: &[f32], b: &[f32]) -> f32 {
match detect_backend() {
#[cfg(target_arch = "aarch64")]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Neon => unsafe { neon::l2_distance(a, b) },
#[cfg(all(target_arch = "x86_64", feature = "avx512"))]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Avx512 => unsafe { avx512::l2_distance(a, b) },
#[cfg(target_arch = "x86_64")]
// SAFETY: detect_backend() verified the CPU feature for this variant is available at runtime.
Backend::Avx2 => unsafe { avx2::l2_distance(a, b) },
_ => scalar::l2_distance(a, b),
}
}
/// Compute L2 norms for a batch of vectors.
pub fn batch_norms(vectors: &[&[f32]], norms: &mut [f32]) {
assert!(norms.len() >= vectors.len());
for (i, v) in vectors.iter().enumerate() {
norms[i] = vector_norm(v);
}
}
/// Convert a batch of f16 values (as raw u16 bits) to f32.
pub fn f16_to_f32_batch(input: &[u16], output: &mut [f32]) {
convert::f16_to_f32_batch(input, output);
}
/// Compute Fletcher-32 checksum.
pub fn checksum_fletcher32(data: &[u8]) -> u32 {
checksum::checksum_fletcher32(data)
}
#[cfg(test)]
mod tests {
use super::*;
const EPSILON: f32 = 1e-5;
fn approx_eq(a: f32, b: f32, eps: f32) -> bool {
(a - b).abs() < eps
}
// -----------------------------------------------------------------------
// Backend detection
// -----------------------------------------------------------------------
#[test]
fn test_detect_backend_returns_valid() {
let backend = detect_backend();
match backend {
Backend::Neon
| Backend::Avx2
| Backend::Avx512
| Backend::Sse4
| Backend::WasmSimd128
| Backend::Scalar => {}
}
}
#[test]
fn test_detect_backend_consistent() {
let b1 = detect_backend();
let b2 = detect_backend();
assert_eq!(b1, b2);
}
// -----------------------------------------------------------------------
// Dot product
// -----------------------------------------------------------------------
#[test]
fn test_dot_product_known_values() {
let a = [1.0, 2.0, 3.0, 4.0];
let b = [5.0, 6.0, 7.0, 8.0];
// 1*5 + 2*6 + 3*7 + 4*8 = 5 + 12 + 21 + 32 = 70
let result = dot_product(&a, &b);
assert!(approx_eq(result, 70.0, EPSILON), "got {result}");
}
#[test]
fn test_dot_product_zero_vectors() {
let a = [0.0f32; 16];
let b = [1.0f32; 16];
assert!(approx_eq(dot_product(&a, &b), 0.0, EPSILON));
}
#[test]
fn test_dot_product_unit_vectors() {
let mut a = [0.0f32; 3];
let mut b = [0.0f32; 3];
a[0] = 1.0;
b[0] = 1.0;
assert!(approx_eq(dot_product(&a, &b), 1.0, EPSILON));
}
#[test]
fn test_dot_product_large_random() {
let n = 1024;
let a: Vec<f32> = (0..n).map(|i| (i as f32) * 0.01).collect();
let b: Vec<f32> = (0..n).map(|i| ((n - i) as f32) * 0.01).collect();
let scalar_result = scalar::dot_product(&a, &b);
let simd_result = dot_product(&a, &b);
assert!(
approx_eq(scalar_result, simd_result, 0.1),
"scalar={scalar_result} simd={simd_result}"
);
}
#[test]
fn test_dot_product_negative_values() {
let a = [-1.0, -2.0, -3.0];
let b = [1.0, 2.0, 3.0];
assert!(approx_eq(dot_product(&a, &b), -14.0, EPSILON));
}
#[test]
fn test_dot_product_single_element() {
assert!(approx_eq(dot_product(&[3.0], &[4.0]), 12.0, EPSILON));
}
#[test]
fn test_dot_product_empty() {
assert!(approx_eq(dot_product(&[], &[]), 0.0, EPSILON));
}
#[test]
fn test_dot_product_scalar_vs_dispatch() {
let a: Vec<f32> = (0..384).map(|i| (i as f32).sin()).collect();
let b: Vec<f32> = (0..384).map(|i| (i as f32).cos()).collect();
let s = scalar::dot_product(&a, &b);
let d = dot_product(&a, &b);
assert!(approx_eq(s, d, 0.01), "scalar={s} dispatched={d}");
}
// -----------------------------------------------------------------------
// Vector norm
// -----------------------------------------------------------------------
#[test]
fn test_vector_norm_unit() {
let v = [1.0, 0.0, 0.0];
assert!(approx_eq(vector_norm(&v), 1.0, EPSILON));
}
#[test]
fn test_vector_norm_345() {
let v = [3.0, 4.0];
assert!(approx_eq(vector_norm(&v), 5.0, EPSILON));
}
#[test]
fn test_vector_norm_zero() {
let v = [0.0f32; 10];
assert!(approx_eq(vector_norm(&v), 0.0, EPSILON));
}
// -----------------------------------------------------------------------
// Cosine similarity
// -----------------------------------------------------------------------
#[test]
fn test_cosine_identical_is_one() {
let v = [1.0, 2.0, 3.0, 4.0, 5.0];
assert!(approx_eq(cosine_similarity(&v, &v), 1.0, EPSILON));
}
#[test]
fn test_cosine_opposite_is_neg_one() {
let a = [1.0, 2.0, 3.0];
let b = [-1.0, -2.0, -3.0];
assert!(approx_eq(cosine_similarity(&a, &b), -1.0, EPSILON));
}
#[test]
fn test_cosine_orthogonal_is_zero() {
let a = [1.0, 0.0, 0.0, 0.0];
let b = [0.0, 1.0, 0.0, 0.0];
assert!(approx_eq(cosine_similarity(&a, &b), 0.0, EPSILON));
}
#[test]
fn test_cosine_zero_vector() {
let a = [0.0f32; 4];
let b = [1.0, 2.0, 3.0, 4.0];
assert!(approx_eq(cosine_similarity(&a, &b), 0.0, EPSILON));
}
#[test]
fn test_cosine_near_zero_norm_clamped() {
// denom = 1e-4 * 1e-4 = 1e-8, comfortably below f32::EPSILON
// (~1.19e-7) but not exactly 0.0 — must still clamp to 0.0 so
// callers computing `1.0 - cosine_similarity(...)` treat these
// as maximally dissimilar, matching the pre-SIMD scalar guard.
let a = [1e-4f32];
let b = [1e-4f32];
assert_eq!(cosine_similarity(&a, &b), 0.0);
assert_eq!(scalar::cosine_similarity(&a, &b), 0.0);
}
#[test]
fn test_cosine_scalar_vs_dispatch() {
let a: Vec<f32> = (0..384).map(|i| (i as f32).sin()).collect();
let b: Vec<f32> = (0..384).map(|i| (i as f32 * 0.7).cos()).collect();
let s = scalar::cosine_similarity(&a, &b);
let d = cosine_similarity(&a, &b);
assert!(approx_eq(s, d, 1e-4), "scalar={s} dispatched={d}");
}
// -----------------------------------------------------------------------
// Batch cosine
// -----------------------------------------------------------------------
#[test]
fn test_batch_cosine_ranking_order() {
let query = [1.0, 0.0, 0.0];
let v0: Vec<f32> = vec![0.0, 1.0, 0.0]; // orthogonal = 0
let v1: Vec<f32> = vec![1.0, 0.0, 0.0]; // identical = 1
let v2: Vec<f32> = vec![0.5, 0.5, 0.0]; // in between
let vectors: Vec<&[f32]> = vec![&v0, &v1, &v2];
let mut results = vec![(0usize, 0.0f32); 3];
batch_cosine(&query, &vectors, &mut results);
// v1 should have highest similarity
assert!(results[1].1 > results[2].1);
assert!(results[2].1 > results[0].1);
}
#[test]
fn test_batch_cosine_scalar_vs_dispatch() {
let query: Vec<f32> = (0..32).map(|i| (i as f32).sin()).collect();
let v0: Vec<f32> = (0..32).map(|i| (i as f32).cos()).collect();
let v1: Vec<f32> = (0..32).map(|i| (i as f32 * 2.0).sin()).collect();
let vectors: Vec<&[f32]> = vec![&v0, &v1];
let mut scalar_results = vec![(0usize, 0.0f32); 2];
scalar::batch_cosine(&query, &vectors, &mut scalar_results);
let mut simd_results = vec![(0usize, 0.0f32); 2];
batch_cosine(&query, &vectors, &mut simd_results);
for i in 0..2 {
assert!(
approx_eq(scalar_results[i].1, simd_results[i].1, 1e-4),
"mismatch at {i}: scalar={} simd={}",
scalar_results[i].1,
simd_results[i].1
);
}
}
// -----------------------------------------------------------------------
// Batch cosine prenorm
// -----------------------------------------------------------------------
#[test]
fn test_batch_cosine_prenorm() {
let query = [1.0, 0.0, 0.0]; // already unit-length
let v0: Vec<f32> = vec![3.0, 4.0, 0.0];
let v1: Vec<f32> = vec![0.0, 0.0, 5.0];
let vectors: Vec<&[f32]> = vec![&v0, &v1];
let norms = [5.0, 5.0];
let mut results = vec![(0usize, 0.0f32); 2];
batch_cosine_prenorm(&query, &vectors, &norms, &mut results);
// dot(query, v0) = 3.0, sim = 3.0/5.0 = 0.6
assert!(approx_eq(results[0].1, 0.6, EPSILON));
// dot(query, v1) = 0.0, sim = 0.0
assert!(approx_eq(results[1].1, 0.0, EPSILON));
}
// -----------------------------------------------------------------------
// L2 distance
// -----------------------------------------------------------------------
#[test]
fn test_l2_distance_same_is_zero() {
let v = [1.0, 2.0, 3.0, 4.0];
assert!(approx_eq(l2_distance(&v, &v), 0.0, EPSILON));
}
#[test]
fn test_l2_distance_known_triangle() {
let a = [0.0, 0.0];
let b = [3.0, 4.0];
assert!(approx_eq(l2_distance(&a, &b), 5.0, EPSILON));
}
#[test]
fn test_l2_distance_unit_axes() {
let a = [1.0, 0.0, 0.0];
let b = [0.0, 1.0, 0.0];
assert!(approx_eq(l2_distance(&a, &b), 2.0f32.sqrt(), EPSILON));
}
#[test]
fn test_l2_distance_scalar_vs_dispatch() {
let a: Vec<f32> = (0..384).map(|i| (i as f32).sin()).collect();
let b: Vec<f32> = (0..384).map(|i| (i as f32).cos()).collect();
let s = scalar::l2_distance(&a, &b);
let d = l2_distance(&a, &b);
assert!(approx_eq(s, d, 0.01), "scalar={s} dispatched={d}");
}
// -----------------------------------------------------------------------
// Batch norms
// -----------------------------------------------------------------------
#[test]
fn test_batch_norms() {
let v0: Vec<f32> = vec![3.0, 4.0];
let v1: Vec<f32> = vec![0.0, 0.0];
let v2: Vec<f32> = vec![1.0, 0.0, 0.0];
let vectors: Vec<&[f32]> = vec![&v0, &v1, &v2];
let mut norms = vec![0.0f32; 3];
batch_norms(&vectors, &mut norms);
assert!(approx_eq(norms[0], 5.0, EPSILON));
assert!(approx_eq(norms[1], 0.0, EPSILON));
assert!(approx_eq(norms[2], 1.0, EPSILON));
}
// -----------------------------------------------------------------------
// f16 conversion
// -----------------------------------------------------------------------
#[test]
fn test_f16_to_f32_known_values() {
// f16 representation of 1.0 = 0x3C00
let input = [0x3C00u16, 0x4000, 0x0000]; // 1.0, 2.0, 0.0
let mut output = [0.0f32; 3];
f16_to_f32_batch(&input, &mut output);
assert!(approx_eq(output[0], 1.0, EPSILON), "got {}", output[0]);
assert!(approx_eq(output[1], 2.0, EPSILON), "got {}", output[1]);
assert!(approx_eq(output[2], 0.0, EPSILON), "got {}", output[2]);
}
#[test]
fn test_f16_to_f32_negative() {
// f16 -1.0 = 0xBC00
let input = [0xBC00u16];
let mut output = [0.0f32; 1];
f16_to_f32_batch(&input, &mut output);
assert!(approx_eq(output[0], -1.0, EPSILON), "got {}", output[0]);
}
#[test]
fn test_f16_to_f32_batch_larger() {
// Test with a larger batch to exercise SIMD paths
let input: Vec<u16> = (0..32).map(|_| 0x3C00u16).collect(); // all 1.0
let mut output = vec![0.0f32; 32];
f16_to_f32_batch(&input, &mut output);
for (i, &v) in output.iter().enumerate() {
assert!(approx_eq(v, 1.0, EPSILON), "mismatch at {i}: {v}");
}
}
#[test]
fn test_f16_to_f32_round_trip_accuracy() {
// Test several known f16 bit patterns
let cases: Vec<(u16, f32)> = vec![
(0x3C00, 1.0),
(0x4000, 2.0),
(0x3800, 0.5),
(0x4200, 3.0),
(0x4400, 4.0),
(0x0000, 0.0),
(0x8000, -0.0),
];
let input: Vec<u16> = cases.iter().map(|(bits, _)| *bits).collect();
let mut output = vec![0.0f32; cases.len()];
f16_to_f32_batch(&input, &mut output);
for (i, (_, expected)) in cases.iter().enumerate() {
assert!(
approx_eq(output[i], *expected, EPSILON),
"f16 0x{:04X}: expected {expected}, got {}",
input[i],
output[i]
);
}
}
// -----------------------------------------------------------------------
// Fletcher-32 checksum
// -----------------------------------------------------------------------
#[test]
fn test_fletcher32_empty() {
let result = checksum_fletcher32(&[]);
// Both sums remain 0xFFFF
assert_eq!(result, 0xFFFF_FFFF);
}
#[test]
fn test_fletcher32_known() {
let data = [0x00u8, 0x01, 0x00, 0x02];
let result = checksum_fletcher32(&data);
let scalar = scalar::checksum_fletcher32(&data);
assert_eq!(result, scalar);
}
#[test]
fn test_fletcher32_scalar_vs_dispatch() {
let data: Vec<u8> = (0..256).map(|i| i as u8).collect();
let s = scalar::checksum_fletcher32(&data);
let d = checksum_fletcher32(&data);
assert_eq!(s, d);
}
// -----------------------------------------------------------------------
// Performance sanity check
// -----------------------------------------------------------------------
#[test]
fn test_dot_product_384_dim_perf() {
use std::time::Instant;
let a: Vec<f32> = (0..384).map(|i| (i as f32).sin()).collect();
let b: Vec<f32> = (0..384).map(|i| (i as f32).cos()).collect();
// Warm up
for _ in 0..100 {
let _ = dot_product(&a, &b);
}
let start = Instant::now();
let iterations = 10_000;
let mut sum = 0.0f32;
for _ in 0..iterations {
sum += dot_product(&a, &b);
}
let elapsed = start.elapsed();
let per_call = elapsed / iterations;
// Prevent optimization
assert!(sum.abs() >= 0.0);
// In release mode, 384-dim dot product should be < 1µs.
// In debug mode, allow up to 20µs (no optimizations).
let limit_ns = if cfg!(debug_assertions) {
20_000
} else {
1_000
};
assert!(
per_call.as_nanos() < limit_ns,
"dot product too slow: {per_call:?} per call (limit {limit_ns}ns)"
);
}
// -----------------------------------------------------------------------
// Cache-line alignment (TVL)
// -----------------------------------------------------------------------
#[test]
fn test_cache_line_size_is_power_of_two() {
assert!(CACHE_LINE_SIZE.is_power_of_two());
}
#[test]
fn test_cache_line_size_platform() {
#[cfg(target_arch = "aarch64")]
assert_eq!(CACHE_LINE_SIZE, 128);
#[cfg(target_arch = "x86_64")]
assert_eq!(CACHE_LINE_SIZE, 64);
}
#[test]
fn test_align_to_cache_line() {
assert_eq!(align_to_cache_line(0), 0);
assert_eq!(align_to_cache_line(1), CACHE_LINE_SIZE);
assert_eq!(align_to_cache_line(CACHE_LINE_SIZE), CACHE_LINE_SIZE);
assert_eq!(
align_to_cache_line(CACHE_LINE_SIZE + 1),
CACHE_LINE_SIZE * 2
);
assert_eq!(
align_to_cache_line(CACHE_LINE_SIZE * 3),
CACHE_LINE_SIZE * 3
);
}
#[test]
fn test_align_to_cache_line_64_and_128() {
// Both 64 and 128 alignment scenarios
let val = align_to_cache_line(100);
assert_eq!(val % CACHE_LINE_SIZE, 0);
assert!(val >= 100);
assert!(val < 100 + CACHE_LINE_SIZE);
}
// -----------------------------------------------------------------------
// Edge cases / additional coverage
// -----------------------------------------------------------------------
#[test]
fn test_dot_product_non_aligned_length() {
// Test with lengths that don't align to SIMD widths (not multiple of 4, 8, 16)
for len in [1, 3, 5, 7, 9, 13, 17, 31, 33] {
let a: Vec<f32> = (0..len).map(|i| i as f32).collect();
let b: Vec<f32> = (0..len).map(|i| (i as f32) * 0.5).collect();
let s = scalar::dot_product(&a, &b);
let d = dot_product(&a, &b);
assert!(
approx_eq(s, d, 0.01),
"len={len}: scalar={s} dispatched={d}"
);
}
}
#[test]
fn test_cosine_non_aligned_length() {
for len in [1, 3, 5, 7, 9, 13, 17, 31, 33] {
let a: Vec<f32> = (0..len).map(|i| i as f32 + 1.0).collect();
let b: Vec<f32> = (0..len).map(|i| (i as f32 + 1.0) * 2.0).collect();
let s = scalar::cosine_similarity(&a, &b);
let d = cosine_similarity(&a, &b);
assert!(
approx_eq(s, d, 1e-4),
"len={len}: scalar={s} dispatched={d}"
);
}
}
#[test]
fn test_l2_distance_non_aligned_length() {
for len in [1, 3, 5, 7, 9, 13, 17, 31, 33] {
let a: Vec<f32> = (0..len).map(|i| i as f32).collect();
let b: Vec<f32> = (0..len).map(|i| (i as f32) + 1.0).collect();
let s = scalar::l2_distance(&a, &b);
let d = l2_distance(&a, &b);
assert!(
approx_eq(s, d, 0.01),
"len={len}: scalar={s} dispatched={d}"
);
}
}
}
#[cfg(test)]
mod dot_i8_tests {
use super::*;
fn codes(n: usize, seed: u64) -> Vec<i8> {
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}");
}
}
/// 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.
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);
}
}
Reference in New Issue View Git Blame Copy Permalink