onnxruntime
60b4dd6a - [MLAS] Add ARM64 NEON fp32 RoPE kernel (#29836)

Commit
31 days ago
[MLAS] Add ARM64 NEON fp32 RoPE kernel (#29836) ## Summary ARM64 has no NEON implementation for fp32 rotary embedding. `MlasRopeDispatchNeon` only ever assigns `HRope`, and only when FP16 vector acceleration is available; `SRope` is never assigned, so `MlasRotaryEmbedOneRow<float>` always falls back to the scalar reference implementation on ARM64 regardless of hardware. `test_rope.cpp`'s shared test registration guard reflects this too, excluding `MLAS_TARGET_ARM64` entirely. This adds a NEON fp32 kernel (`RopeKernel_Fp32`, non-interleaved and interleaved variants) to `rotary_embedding_kernel_neon.cpp`/`.h` and assigns it to `d.SRope` unconditionally in `MlasRopeDispatchNeon`, outside the `MlasFp16AccelerationSupported()` guard since fp32 doesn't depend on FP16 vector support. The interleaved path uses `vld2q_f32`/`vst2q_f32` to deinterleave/reinterleave the real/imag pairs directly, instead of the shuffle/permute sequence the AVX2 implementation uses for the same case. `test_rope.cpp`'s ARM64 guard now registers the existing 14 fp32 cases (dim 6/16/24/32/42/64/70 x interleaved). The fp16 half of that guard is left AMD64/RVV-only, since ARM64 fp16 already has its own dedicated coverage in `test_rope_neon_fp16.cpp` — this change is scoped to fp32 only. ## Testing `onnxruntime_mlas_test`, full suite, no regressions: - Apple M1 (macOS): 27975/27975 passed - Neoverse-N1 (Oracle Cloud A1, Ubuntu): 34769/34769 passed `RoPE_fp32/*`: 14/14 passed on both. ## Benchmark (Neoverse-N1, Oracle Cloud A1, median of 3 runs, `RoPE<float>`) | dim | interleaved | before (scalar) | after (NEON) | speedup | |------|-------------|-----------------|--------------|---------| | 128 | no | 514 ns | 24.4 ns | 21.1x | | 256 | no | 1028 ns | 43.1 ns | 23.9x | | 512 | no | 2067 ns | 80.6 ns | 25.6x | | 1024 | no | 4161 ns | 156 ns | 26.7x | | 128 | yes | 264 ns | 44.4 ns | 5.9x | | 256 | yes | 521 ns | 85.0 ns | 6.1x | | 512 | yes | 1035 ns | 165 ns | 6.3x | | 1024 | yes | 2064 ns | 329 ns | 6.3x | Numbers are from Neoverse-N1 only; it's a dedicated instance and gives a cleaner signal than a shared laptop.
Author
Parents
Loading