Skip to content

Commit dfad43b

Browse files
committed
fix(cpu/matmul): add NEON simd_dot_f32 for aarch64 f16 gemv path
The f16 gemv_bt_via_f32 path extended cfg gates to include aarch64 but simd_dot_f32 only had an x86_64 implementation, causing a compile error on Apple Silicon and other aarch64 targets. Add a NEON implementation using dual vfmaq_f32 accumulators to hide FMA latency, matching the pattern used by the AVX2 path.
1 parent 2b0f433 commit dfad43b

1 file changed

Lines changed: 45 additions & 0 deletions

File tree

src/runtime/cpu/kernels/matmul.rs

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,51 @@ unsafe fn simd_dot_f32_avx512(a: *const f32, b: *const f32, len: usize) -> f32 {
6767
sum
6868
}
6969

70+
#[cfg(all(feature = "f16", target_arch = "aarch64"))]
71+
#[inline]
72+
unsafe fn simd_dot_f32(
73+
a: *const f32,
74+
b: *const f32,
75+
len: usize,
76+
_level: super::simd::SimdLevel,
77+
) -> f32 {
78+
simd_dot_f32_neon(a, b, len)
79+
}
80+
81+
#[cfg(all(feature = "f16", target_arch = "aarch64"))]
82+
#[target_feature(enable = "neon")]
83+
unsafe fn simd_dot_f32_neon(a: *const f32, b: *const f32, len: usize) -> f32 {
84+
use std::arch::aarch64::*;
85+
let mut offset = 0;
86+
let mut acc0 = vdupq_n_f32(0.0);
87+
let mut acc1 = vdupq_n_f32(0.0);
88+
// Process 8 floats per iteration with dual accumulators
89+
while offset + 8 <= len {
90+
let av0 = vld1q_f32(a.add(offset));
91+
let bv0 = vld1q_f32(b.add(offset));
92+
acc0 = vfmaq_f32(acc0, av0, bv0);
93+
let av1 = vld1q_f32(a.add(offset + 4));
94+
let bv1 = vld1q_f32(b.add(offset + 4));
95+
acc1 = vfmaq_f32(acc1, av1, bv1);
96+
offset += 8;
97+
}
98+
acc0 = vaddq_f32(acc0, acc1);
99+
// Handle remaining 4-float chunk
100+
while offset + 4 <= len {
101+
let av = vld1q_f32(a.add(offset));
102+
let bv = vld1q_f32(b.add(offset));
103+
acc0 = vfmaq_f32(acc0, av, bv);
104+
offset += 4;
105+
}
106+
let mut sum = vaddvq_f32(acc0);
107+
// Scalar tail
108+
while offset < len {
109+
sum += *a.add(offset) * *b.add(offset);
110+
offset += 1;
111+
}
112+
sum
113+
}
114+
70115
#[cfg(all(feature = "f16", target_arch = "x86_64"))]
71116
#[target_feature(enable = "avx2", enable = "fma")]
72117
unsafe fn simd_dot_f32_avx2(a: *const f32, b: *const f32, len: usize) -> f32 {

0 commit comments

Comments
 (0)