[MLAS] Add ARM64 NEON fp32 RoPE kernel

[kjg0724]

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.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

This PR adds an ARM64 NEON-optimized fp32 Rotary Positional Embedding (RoPE) kernel to MLAS and wires it into the ARM64 NEON dispatch path so MlasRotaryEmbedOneRow<float> no longer falls back to the scalar reference implementation on ARM64.

Changes:

• Implement fp32 NEON RoPE kernels (non-interleaved and interleaved) and expose them via rope_neon::RopeKernel_Fp32.
• Assign the fp32 kernel to MlasRopeDispatchNeon.SRope unconditionally (independent of fp16 acceleration support).
• Enable the existing fp32 RoPE short-execute tests on ARM64 while keeping fp16 registration limited to AMD64/RVV (ARM64 fp16 covered elsewhere).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/test/mlas/unittest/test_rope.cpp	Registers fp32 RoPE short-execute tests on ARM64; keeps fp16 registration on AMD64/RVV only.
onnxruntime/core/mlas/lib/rotary_embedding_kernel_neon.h	Declares the new fp32 NEON RoPE kernel entry point.
onnxruntime/core/mlas/lib/rotary_embedding_kernel_neon.cpp	Implements fp32 NEON RoPE (interleaved + non-interleaved) and hooks it into MlasRopeDispatchNeon.SRope.

[hariharans29]

Review: PR #29836 — [MLAS] Add ARM64 NEON fp32 RoPE kernel (head f0158d9)

Author: @kjg0724. 1 commit. CI: 82 / 86 checks OK. 3 files changed. Copilot AI reviewed with 1 Low-severity comment (comment-indentation nit).

Verdict: approve. Small, well-motivated PR that closes a real gap — MlasRotaryEmbedOneRow<float> was silently falling back to the scalar reference on ARM64 because MlasRopeDispatchNeon.SRope was never assigned. Benchmarks show 5.9x–26.7x speedup on Neoverse-N1 depending on interleaving mode. Correctness verifiable by direct comparison to the scalar reference math.

---

What it does

1. Adds rope_neon::RopeKernel_Fp32 in rotary_embedding_kernel_neon.cpp, with two template specializations (interleaved=false, interleaved=true) plus a public dispatcher that asserts dim % 2 == 0.
2. Assigns d.SRope = rope_neon::RopeKernel_Fp32 unconditionally inside MlasRopeDispatchNeon — outside the MLAS_F16VEC_INTRINSICS_SUPPORTED && MlasFp16AccelerationSupported() guard, because fp32 NEON doesn't depend on FP16 vector acceleration.
3. Declares the entry point in rotary_embedding_kernel_neon.h.
4. Extends the guard in test_rope.cpp so the existing 14 fp32 short-execute tests (dim 6/16/24/32/42/64/70 × {interleaved, non-interleaved}) run on ARM64 too. fp16 test registration is kept off ARM64 because test_rope_neon_fp16.cpp already covers that path with a dedicated suite.

---

Kernel correctness — both specializations verified

Non-interleaved path (half_dim = dim/2; input laid out [real_0..real_{half_dim-1}, imag_0..imag_{half_dim-1}]):

float32x4_t real_out = vfmsq_f32(vmulq_f32(real, cos_val), imag, sin_val);  // real*cos - imag*sin
float32x4_t imag_out = vfmaq_f32(vmulq_f32(real, sin_val), imag, cos_val);  // real*sin + imag*cos

• vfmsq_f32(a, b, c) = a − b·c → real·cos − imag·sin ✓
• vfmaq_f32(a, b, c) = a + b·c → real·sin + imag·cos ✓

Matches the scalar reference exactly. Loop bounds: main body i + 3 < half_dim reads/writes input/output[i..i+3] and input/output[j..j+3] (j = half_dim + i); tail for (; i < half_dim; i++, j++) handles the remainder. All in-bounds for even dim.

Interleaved path (input laid out [real_0, imag_0, real_1, imag_1, ...]):

float32x4x2_t v = vld2q_f32(input + i);   // deinterleave load
float32x4_t real = v.val[0];              // 4 real values
float32x4_t imag = v.val[1];              // 4 imag values
// ...same FMA math as above...
vst2q_f32(output + i, out);               // interleaved store

Using vld2q_f32/vst2q_f32 is the idiomatic NEON approach here, and materially cleaner than the shuffle/permute sequence the AVX2 code uses for the same case (the PR description flags this deliberately). The sin/cos lookups use i/2 because the cache is half-length (one entry per real/imag pair). Loop bounds: i + 7 < dim in the vectorized block, i + 1 < dim in the scalar tail — all in-bounds.

Dispatcher assertion: assert(dim % 2 == 0) in RopeKernel_Fp32 documents the precondition. In release builds this compiles out under NDEBUG, but since RoPE's dim is always even by construction (real/imag pairs) and callers reach here through the typed MlasRotaryEmbedOneRow<float> path, this is sufficient.

---

Dispatch wiring — correct

const MLAS_ROPE_DISPATCH MlasRopeDispatchNeon = []() {
    MLAS_ROPE_DISPATCH d;
    d.SRope = rope_neon::RopeKernel_Fp32;    // <-- new, unconditional
#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64)
    if (MlasFp16AccelerationSupported()) {
        d.HRope = rope_neon::RopeKernel_Fp16;
    }
#endif
    ...
}();

The compile-time inclusion of the whole file is already ARM64-gated by the MLAS build system, so d.SRope assignment is effectively ARM64-only. Placing it outside the FP16 guard is exactly right: fp32 NEON is present on every ARMv8 core, no FEAT_FP16 dependency.

---

Test coverage

// Before:
#if defined(MLAS_TARGET_AMD64) || (defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV))

// After:
#if defined(MLAS_TARGET_AMD64) || (defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV)) || defined(MLAS_TARGET_ARM64)
static size_t RoPERegisterAllShortExecuteTests() {
  size_t tests_registered = RoPEShortExecuteTest<float>::RegisterShortExecuteTests();
// fp16 RoPE on ARM64 already has dedicated coverage in test_rope_neon_fp16.cpp.
#if defined(MLAS_TARGET_AMD64) || (defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV))
  tests_registered += RoPEShortExecuteTest<MLFloat16>::RegisterShortExecuteTests();
#endif
  return tests_registered;
}

fp32 tests register on all three platforms; fp16 tests register on everything except ARM64. The scoping is correct — the ARM64 fp16 RoPE path is already exercised by the dedicated NEON-fp16 test file. Author's local testing shows 34769/34769 pass on Neoverse-N1 and 27975/27975 on Apple M1, so the newly-enabled fp32 cases hold up.

Only extra suggestion: the 14 registered cases cover dim ∈ {6, 16, 24, 32, 42, 64, 70} × 2 interleaving modes. That's a good mix but doesn't include any dim that hits only the scalar tail (dim=4 for non-interleaved, dim=2 for interleaved). If the goal is to explicitly exercise the tail loop, adding dim=2 or 4 to the case list would do it — but the existing cases (particularly dim=6 non-interleaved, which does 0 vector iterations and 3 scalar tail iterations) already touch the tail code paths transitively. Not a blocker.

---

Copilot review disposition

Copilot #1 (Low): "comment not indented consistently with surrounding block". Suggests indenting the // fp16 RoPE on ARM64 already has dedicated coverage in test_rope_neon_fp16.cpp. comment by 2 spaces.

Defensible either way. C/C++ convention leaves preprocessor directives at column 1, and placing a comment immediately preceding a #if at column 1 to document the guard is a common idiom. clang-format's default behavior with IndentPPDirectives: None won't touch it. If the author's local .clang-format is stricter and reformats it, they can accept the suggestion; otherwise it's fine to leave as-is. Not a correctness issue.

---

Benchmarks (from the PR description, Neoverse-N1, median of 3 runs)

dim	interleaved	scalar	new NEON	speedup
128	no	514 ns	24.4 ns	21.1×
256	no	1028 ns	43.1 ns	23.9×
512	no	2067 ns	80.6 ns	25.6×
1024	no	4161 ns	156 ns	26.7×
128	yes	264 ns	44.4 ns	5.9×
256	yes	521 ns	85.0 ns	6.1×
512	yes	1035 ns	165 ns	6.3×
1024	yes	2064 ns	329 ns	6.3×

The non-interleaved speedup being ~4× larger than the interleaved case is expected: interleaved requires the deinterleave/reinterleave via vld2/vst2, which on Neoverse-N1 execute as multiple micro-ops vs. the single-cycle vld1/vst1 used by the non-interleaved path. Both wins are still substantial and well worth the change.

---

Minor observations (none blocking)

1. vmulq_f32 + vfmsq_f32 / vfmaq_f32 pattern: one non-FMA multiply plus one FMA per output pair. An alternative would be to zero-initialize and use two pure FMAs, but that trades a multiply for a zero-materialize and doesn't help on ARMv8. The current idiom is what most NEON RoPE implementations use; the compiler often folds this into a pair of FMLAs anyway.
2. #include <cassert> is added for the single assert(dim % 2 == 0). Correct include for the header used. In release (NDEBUG), the assert is compiled out — fine, since the callers ensure the invariant.
3. Anonymous namespace for RopeKernel_Fp32_Impl<...> with the public entry point exposed via namespace rope_neon — clean internal/external linkage separation, matches the existing fp16 kernel structure in the same file.
4. No changes to MlasRopeDispatchNeon's fp16 assignment: correct — this PR is deliberately scoped to fp32 only.

---

Bottom line

Clean, focused perf improvement. Real correctness gap (scalar fallback on all ARM64) closed. Benchmarks are convincing, tests are enabled where they should be, and the scoping is right. The single Copilot comment is a style-level indentation nit that doesn't affect correctness. Approve.

[hariharans29]

Please address the copilot comments