linalg/arm64: NEON rms_norm_f32 kernel (stacked on #2311)

[czoli1976]

Stacked on #2311 — needs the Ops::rms_norm_f32 slot + RmsNorm::eval fast path that PR adds. Single commit on top; review only the top commit. Rebases trivially to a standalone PR if #2311 merges first.

What

NEON (aarch64, 128-bit, 4-lane) implementation of tract_linalg::ops().rms_norm_f32, mirroring the AVX-512 kernel from the parent PR. 16 f32 lanes per inner iteration (4 v-registers × 4 lanes):

• Pass 1 — sum of squares via 4 parallel fmla chains (v0..v3) → 3-way fadd tree → vaddvq_f32 horizontal reduce.
• Pass 2 — broadcast inv_std into v0, fmul/st1 each 4-v-register chunk in place.
• Scalar tail for (len % 16 != 0).

Plugs into arm64::plug(). NEON is mandatory on aarch64 so no runtime feature detection — #[target_feature(enable = "neon")] is just for the inline-asm + intrinsic context. Generic scalar fallback in tract_linalg::generic covers non-aarch64.

Validation (M1, real hardware)

Surface	Result
4 new unit tests (trivial / 1024+7 tail / n=8 all-tail / empty)	4/4 pass
Full cargo test -p tract-linalg --lib	3718/0 pass, 0 regressions
cargo fmt --check, cargo clippy	clean (no new warnings; the 6 pre-existing tract-linalg warnings on main are unchanged)

Synthetic stress test (local, not in this commit)

Local-only stress test file covering every size 1..32, hidden ∈ {768..8192} × 5 tail residues, pathological distributions (all-zero / all-equal / mixed-sign large-magnitude / subnormal-mostly), 8 epsilons (0 → 100), 500 random sizes in [1, 8192], and huge rows up to 32768. 10/10 pass vs scalar reference. Happy to include the file in this PR if reviewers want it as test/.

E2E correctness — 64×RmsNorm@4096 ONNX chain (this PR vs main)

max_abs_diff   = 0.000e+00
mean_abs_diff  = 0.000e+00
bit-exact equal? True

Bit-exact across 64 chained layers. No drift, no FMA-reorder loss.

Performance (M1 P-core)

Kernel-level microbench (linalg/benches/rms_norm.rs, 30 samples × 3s)

Row N	composed (inline 2-loop, autovec)	NEON	gap
1024	1.04 µs / 0.98 Gelem/s	0.140 µs / 7.31 Gelem/s	7.43×
2048	2.08 µs / 0.98 Gelem/s	0.279 µs / 7.34 Gelem/s	7.47×
4096	4.17 µs / 0.98 Gelem/s	0.549 µs / 7.45 Gelem/s	7.59×

E2E on a 64×RmsNorm@4096 ONNX (real tract pipeline, bench mode)

Form	main (composed 4-call)	this PR	speedup
Opset-23 native RMSNormalization × 64	0.299 ms/i	0.097 ms/i	3.08×
Opset-17 composed pattern × 64 (declutter hoists)	0.265 ms/i	0.077 ms/i	3.44×

E2E gap is smaller than kernel-level because per-call dispatch overhead is ~constant; with the fast NEON kernel it becomes a larger fraction. For a real LLM where RmsNorm is ~5% of inference time, this is ~1–4% E2E.

Caveat / context

The fast path only fires on graphs where RmsNorm is hoisted as a single op — i.e., native opset-23 RMSNormalization, SimplifiedLayerNormalization contrib (#2288), or the "clean" composed pattern that detect_rms_norm matches. Many HuggingFace optimum-onnxruntime LLM exports use a Cast → Pow → ReduceMean → Add → Sqrt → Div → Mul → Cast → Mul variant that the current declutter doesn't recognize, so on those models this kernel doesn't fire and there's no gain. Worth a follow-up to broaden the declutter pattern; not in scope here.

Risk

• NEON is always present on aarch64 — no gating required, no host detection issue.
• Pure addition: no existing kernel modified, no Ops field semantics changed.
• Non-AVX-512 x86 / non-aarch64 keeps the generic scalar fallback from linalg/x86_64 + core/nn: fused AVX-512 RmsNorm kernel #2311 (unchanged).

🤖 Generated with Claude Code