1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
//! Hand-vectorized f32 primitives, for the kernels where the compiler will not do it for us.
//!
//! Two things LLVM cannot do to safe Rust float code, both of which cost ~2× each:
//!
//! 1. **Reassociate.** `a.iter().zip(b).map(|(x, y)| x * y).sum()` is a scalar loop — summing in a
//! different order is a different answer, and the compiler is not allowed to change the answer.
//! Independent accumulators grant that permission explicitly.
//! 2. **Contract.** Rust does not enable FP contraction, so `acc += x * y` emits a separate `fmul`
//! and `fadd` rather than one `fmla`. There is no way to ask for the fused form from safe Rust;
//! `vfmaq_f32` is.
//!
//! So the safe-Rust dot in this crate measured **17 GF/s against a ~130 GF/s single-core NEON
//! roofline** — 13%. That is not a memory wall (a pure streaming read of the same buffer is
//! *slower*); it is the multiply and the add not being the same instruction.
//!
//! Determinism is a hard requirement, not a nicety: these kernels are the oracle the GPU is gated
//! against, and a gate that moves is not a gate. Every routine here fixes its summation order (a
//! static number of accumulators, combined in a fixed tree), so it is bit-reproducible run to run.
//! It is NOT bit-identical to the scalar spelling — the order genuinely differs — which is why the
//! callers' parity tests carry a tolerance rather than an equality.
/// `a · b`. Panics in debug if the lengths differ.
#[inline]
// The `return` is load-bearing: it is the aarch64 arm of a cfg-gated pair, and clippy only sees one
// arm at a time.
#[allow(clippy::needless_return)]
pub fn dot(a: &[f32], b: &[f32]) -> f32 {
debug_assert_eq!(a.len(), b.len());
#[cfg(target_arch = "aarch64")]
{
// SAFETY: NEON is baseline on aarch64 (no runtime detection needed), and the helper reads
// only within `a`/`b`, whose lengths are equal.
return unsafe { dot_neon(a, b) };
}
#[cfg(not(target_arch = "aarch64"))]
dot_portable(a, b)
}
/// `dst += W · x`, `W` row-major with `dst.len()` rows of `x.len()` columns.
///
/// The shape an RNN step makes: no reuse to exploit (one vector against the whole matrix), so this
/// just streams `W` once and leans on [`dot`] for throughput.
#[inline]
pub fn gemv_acc(dst: &mut [f32], w: &[f32], x: &[f32]) {
let k = x.len();
debug_assert_eq!(w.len(), dst.len() * k);
for (i, d) in dst.iter_mut().enumerate() {
*d += dot(&w[i * k..(i + 1) * k], x);
}
}
/// Four independent 4-lane FMA chains: SIMD width from the lanes, latency hiding from the four
/// separate chains (an FMA is ~4 cycles deep but issues every cycle, so one chain idles 3/4 of the
/// pipeline).
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn dot_neon(a: &[f32], b: &[f32]) -> f32 {
use std::arch::aarch64::*;
unsafe {
let n = a.len();
let (pa, pb) = (a.as_ptr(), b.as_ptr());
let mut acc = [vdupq_n_f32(0.0); 4];
let blocks = n / 16;
for i in 0..blocks {
let base = i * 16;
for (l, acc) in acc.iter_mut().enumerate() {
let va = vld1q_f32(pa.add(base + l * 4));
let vb = vld1q_f32(pb.add(base + l * 4));
*acc = vfmaq_f32(*acc, va, vb);
}
}
let mut sum = vaddvq_f32(vaddq_f32(
vaddq_f32(acc[0], acc[1]),
vaddq_f32(acc[2], acc[3]),
));
for i in blocks * 16..n {
sum += *pa.add(i) * *pb.add(i);
}
sum
}
}
/// The same fixed summation order without intrinsics — the reference the NEON path is tested
/// against, and what runs off aarch64.
#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
pub(crate) fn dot_portable(a: &[f32], b: &[f32]) -> f32 {
let mut acc = [[0f32; 4]; 4];
let mut ca = a.chunks_exact(16);
let mut cb = b.chunks_exact(16);
for (x, y) in ca.by_ref().zip(cb.by_ref()) {
for (c, acc) in acc.iter_mut().enumerate() {
for l in 0..4 {
acc[l] += x[c * 4 + l] * y[c * 4 + l];
}
}
}
let mut lanes = [0f32; 4];
for (l, lane) in lanes.iter_mut().enumerate() {
*lane = (acc[0][l] + acc[1][l]) + (acc[2][l] + acc[3][l]);
}
let mut s = (lanes[0] + lanes[1]) + (lanes[2] + lanes[3]);
for (x, y) in ca.remainder().iter().zip(cb.remainder()) {
s += x * y;
}
s
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn neon_and_portable_agree() {
for n in [0usize, 1, 15, 16, 17, 64, 256, 259] {
let a: Vec<f32> = (0..n).map(|i| ((i % 37) as f32 - 18.0) * 0.031).collect();
let b: Vec<f32> = (0..n).map(|i| ((i % 41) as f32 - 20.0) * 0.027).collect();
let (fast, refr) = (dot(&a, &b), dot_portable(&a, &b));
assert!(
(fast - refr).abs() <= 1e-4 * refr.abs().max(1.0),
"n={n}: {fast} vs {refr}"
);
}
}
#[test]
fn dot_is_deterministic() {
let a: Vec<f32> = (0..1024)
.map(|i| ((i % 37) as f32 - 18.0) * 0.031)
.collect();
let b: Vec<f32> = (0..1024)
.map(|i| ((i % 41) as f32 - 20.0) * 0.027)
.collect();
assert_eq!(dot(&a, &b).to_bits(), dot(&a, &b).to_bits());
}
}