[MLAS] Add Apple Accelerate (vDSP) LayerNorm/RMSNorm kernel (depends on #32001) - #32066
Open
Justin Chu (justinchuby) wants to merge 21 commits into
Open
[MLAS] Add Apple Accelerate (vDSP) LayerNorm/RMSNorm kernel (depends on #32001)#32066Justin Chu (justinchuby) wants to merge 21 commits into
Justin Chu (justinchuby) wants to merge 21 commits into
Conversation
Add a new opt-in CMake option for Apple Accelerate/BNNS/vDSP support in MLAS. This PR adds build-system scaffolding only — no kernels, no behaviour change. - Option: onnxruntime_USE_APPLE_ACCELERATE, default OFF - FATAL_ERROR on non-Apple platforms when enabled - Links the system Accelerate framework via find_library (macOS/iOS/universal2) - Defines MLAS_USE_APPLE_ACCELERATE=1 compile definition when enabled - No effect whatsoever when disabled (default) Follow-up PRs will add kernels (Accelerate cblas, BNNS, vDSP) separately, each with portable fallback, numeric parity tests, and Apple-hardware benchmarks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…gling define
S1: Replace FATAL_ERROR with warning + disable on non-Apple platforms,
matching the idiom used by onnxruntime_USE_SVE and onnxruntime_USE_KLEIDIAI.
S2: Add --use_apple_accelerate to build.py/build_args.py so the option is
reachable through the standard build tooling.
S3: Remove MLAS_USE_APPLE_ACCELERATE=1 compile definition that nothing
consumes; the first kernel PR will introduce it alongside its reader.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Gate onnxruntime_USE_APPLE_ACCELERATE on APPLE + onnxruntime_target_platform == "arm64" (Apple Silicon). On Apple + non-arm64, warn-and-disable matching the existing idiom for SVE/KleidiAI. Update comments and help text to state macOS arm64 scope; remove universal2/iOS claims. No behaviour change when the option is OFF — all side-effectful statements remain inside the guarded block. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Fixes the Python format lint check. The argument declaration exceeded the line-length limit and ruff splits it across lines. Verified the parsed AST is unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ntract, CLI validation B1: Detect macOS arm64 robustly — check CMAKE_OSX_ARCHITECTURES for arm64/arm64e, fall back to CMAKE_SYSTEM_PROCESSOR when it is unset or empty. Fixes silent disable on plain cmake configure on Apple Silicon. B2: Rewrite PR body to match reality (warn-and-disable, macOS arm64 only, MLAS_USE_APPLE_ACCELERATE reinstated). N1: Gate on CMAKE_SYSTEM_NAME=Darwin instead of APPLE (excludes iOS, tvOS, visionOS). N2: Reintroduce MLAS_USE_APPLE_ACCELERATE=1 compile definition as observable contract for follow-up kernel PRs. N3: build.py raises BuildError on non-macOS — loud failure for explicit CLI opt-in. CMake side stays tolerant (warn-and-disable). N4: Move --use_apple_accelerate to CPU EP argument group in build_args.py. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…talyst The previous validation only checked is_macOS(), allowing explicit --use_apple_accelerate on Intel Macs, cross-compiled x86_64 builds, and non-macOS Apple targets (iOS, tvOS, visionOS, Mac Catalyst) to pass Python validation and then silently downgrade in CMake — exactly the silent-disable the explicit gate was meant to prevent. Now: - build_args.py rejects --use_apple_accelerate unless osx_arch is arm64 or arm64e, and rejects --ios, --tvos, --visionos, and --macos Catalyst explicitly. - build.py keeps the same checks as defence-in-depth (BuildError). - All Apple-target attributes are accessed via getattr() to avoid AttributeError on non-macOS hosts where add_apple_args is not called. - 9 new test cases in test_build_args.py covering each rejection path and the accepting arm64/arm64e paths. Body clarifications (for manual update): - PRIVATE compile definition: MLAS_USE_APPLE_ACCELERATE=1 is defined with target_compile_definitions(onnxruntime_mlas PRIVATE ...), so it is visible only within onnxruntime_mlas translation units, not leaked to downstream consumers. - Multiarch fallback: when CMAKE_OSX_ARCHITECTURES lists multiple architectures (universal2), the detection logic sets the arch to empty, which does not match arm64, so the option is warned-and-disabled. It does NOT attempt to build for the arm64 slice only. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Remove the dead-code copy of Apple Accelerate target validation from build.py (BuildError path was never reached because build_args.py parser.error() exits first). Keep the single source of truth in build_args.py where all other argument validation lives. Also improve the non-macOS error message to distinguish 'wrong OS' from 'wrong arch' — the old wording said 'only supported on macOS arm64' even when the host was macOS but wrong architecture. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add validation rules to build_args.py rejecting Apple Accelerate when targeting Android, WebAssembly (including --build_wasm_static_lib), or RISC-V rv64 — all cross-targets incompatible with macOS arm64. Update tests to assert specific diagnostic messages rather than bare SystemExit, preventing vacuous-test false positives. Test count: 13→17. Document CMake limitation: Catalyst (macabi) cannot be reliably detected at configure time without an external toolchain file setting PLATFORM_NAME. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- PLC0415: Hoist 'import io' and 'from contextlib import redirect_stderr' to module level. No circular-import or lazy-load reason to keep them inside the helper. - SIM105: Replace try/except SystemExit: pass with contextlib.suppress(SystemExit). stderr capture and the assertIn on the diagnostic fragment are preserved so the test still validates the specific error message, not just any argparse failure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…enchmark
onnxruntime_mlas_test and onnxruntime_mlas_benchmark are separate targets
that never received MLAS_USE_APPLE_ACCELERATE=1: it was applied only via a
direct target_compile_definitions(onnxruntime_mlas PRIVATE ...) call after
the shared target-definition foreach had already run, and PRIVATE compile
definitions of a linked static library do not propagate to a consumer's own
translation units (unlike PRIVATE link libraries of a STATIC library, which
CMake does forward to the final link line). Any follow-up kernel PR gated on
this define would have its tests/benchmarks silently compiled out.
Fix: express the definition through the shared mlas_private_compile_definitions
list, matching the MLAS_USE_SVE / MLAS_USE_ARM_NEON_NCHWC idiom, and move the
whole Apple Accelerate block before the
'foreach(mlas_target ${ONNXRUNTIME_MLAS_LIBS}) ... target_compile_definitions'
loop so onnxruntime_mlas picks it up the same way. onnxruntime_mlas_test and
onnxruntime_mlas_benchmark in cmake/onnxruntime_unittests.cmake already apply
${mlas_private_compile_definitions}, so they now see the define without any
change on their side.
Add a small configure-time regression canary at both consuming sites
(onnxruntime_mlas_test, onnxruntime_mlas_benchmark) that FATAL_ERRORs if
onnxruntime_USE_APPLE_ACCELERATE is ON but the define is missing from the
shared list, so this exact class of regression fails the build loudly instead
of silently compiling out gated code.
Verified: cmake if/foreach/endif/endforeach structure parses and executes to
end-of-file via a stubbed-command harness; the canary condition was exercised
standalone for true-positive (fires) and true-negative (silent) cases;
tools/ci_build/test_build_args.py (17 tests) still passes unchanged.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… arm64 ARM64 has no SIMD-vectorized MlasLayerNormF32 kernel today: GetMlasPlatform().LayerNormF32Kernel is only registered for AVX2 (x86) and RVV (RISC-V); see platform.cpp. Every ARM64 build falls back to the single-element-at-a-time Welford/sum-of-squares scalar loop in onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc's ComputeJob. This is the per-token hot path for every transformer decode step (RMSNorm/ LayerNorm called ~2x per layer in Llama/Phi/Gemma/Mistral), so it is a genuinely unaccelerated gap on Apple Silicon -- unlike Erf/Sigmoid/Tanh, which already have 4-wide NEON polynomial kernels compiled unconditionally on ARM64 (confirmed by reading erf.cpp/logistic.cpp/tanh.cpp), which is why PR microsoft#32036's vForce Tanh kernel lost to the existing NEON code on real hardware. This kernel's baseline comparison is different in kind: vDSP's purpose-built reduction/elementwise primitives against a plain scalar loop, not against an already-vectorized competitor. Adds MlasLayerNormKernelAppleAccelerate, registered in platform.cpp's MLAS_TARGET_ARM64 block, gated by the same MLAS_USE_APPLE_ACCELERATE && __APPLE__ && MLAS_TARGET_ARM64 triple-#if used by the sibling Tanh kernel. Depends on microsoft#32001's onnxruntime_USE_APPLE_ACCELERATE CMake option; no CMake changes needed here since that option already links Accelerate and defines MLAS_USE_APPLE_ACCELERATE=1 for all MLAS sources. Algorithm: centered two-pass (mean first, then sum of centered squares), matching the AVX2 LayerNorm kernel's corrected approach rather than the numerically-riskier uncentered "E[x^2] - mean^2" formula (which suffers catastrophic cancellation in fp32 for large-base/small-spread inputs, and which RVV's kernel still uses). RMSNorm uses vDSP_measqv directly (single-pass mean-of-squares, no centering needed). Follows the build-break precedent from PR microsoft#32036's history: does not include <Accelerate/Accelerate.h> (vecLib's cblas.h/cblas_new.h forward- declare BLAS enums without an inline definition on recent macOS SDKs, which is valid C but rejected by ISO C++). Instead forward-declares only the seven vDSP entry points actually used (vDSP_meanv/measqv/svesq/vsadd/vsmul/vmul/vma), extern "C", matching Apple's stable public ABI. Output may safely alias Input: Input is only ever read until the final vDSP_vma/vDSP_vmul call, whose operands are a scratch buffer/Scale/Bias, never Input directly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…kernel Covers, all gated identically to the kernel (MLAS_USE_APPLE_ACCELERATE && __APPLE__ && MLAS_TARGET_ARM64) so this file is an empty no-op on every other configuration: - ForcedReachability / LargeNormSizeHeapFallback: direct calls to MlasLayerNormKernelAppleAccelerate across representative NormSize values (including real transformer hidden dims 768-4096, and sizes above the 8192-float on-stack scratch buffer to exercise the heap fallback path), both Simplified/full LayerNorm, with/without bias, checked against an independent fp64-accumulated scalar reference. - PublicDispatchMatchesDirectKernelCall: proves the public MlasLayerNormF32 dispatch (the entry point every real ORT caller uses) actually resolves to this kernel rather than silently keeping the scalar fallback. - InPlaceAliasing: Input == Output, matching real ORT LayerNorm/ SkipLayerNorm call sites. - NoWriteDetectionPoisonBuffer: pre-fills Output with a poison sentinel (1e30f) before calling the kernel, so a silent no-write bug fails loudly instead of coincidentally matching a leftover 0.0f. - ZeroVarianceConstantRow, NanPropagation, InfPropagation, DenormalInputs, ZeroNormSizeNoWrite, NullMeanAndInvStdDevOutputs: edge cases specific to a whole-row reduction (a single NaN/Inf anywhere in the row contaminates every output element through the shared mean/sum-of-squares, unlike an elementwise function). The fp64 reference deliberately uses the same centered two-pass formula as the kernel (unlike the AVX2 kernel's fp64 reference, which intentionally uses an uncentered formula as an algorithmically-independent oracle) -- see the file-level comment for why this is still a meaningful check for a reduction whose fp64 accumulation cannot hit the fp32 cancellation failure mode that motivated the centered-formula choice in the first place. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Compares the public MlasLayerNormF32 dispatch (which resolves to MlasLayerNormKernelAppleAccelerate whenever this option is enabled) against ScalarLayerNormBaseline, a standalone reproduction of the exact scalar fallback in onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc's ComputeJob (Welford's algorithm for full LayerNorm, single-pass sum-of-squares for RMSNorm) -- the REAL alternative this kernel replaces on ARM64, not a hypothetical baseline. This differs from bench_transcendental.cpp's Tanh A/B, which compares against an ALREADY NEON-vectorized polynomial kernel (and found vForce consistently slower there, PR microsoft#32036) -- LayerNorm/RMSNorm has no SIMD kernel on ARM64 today, so this benchmark answers the actual open question for this kernel rather than assuming an answer from a different kernel's result. Sizes cover small-N crossover probes (1, 8, 64) plus representative transformer hidden dims (256, 768, 1024, 1600, 2048, 3072, 4096, 8192), for both full LayerNorm (with bias) and RMSNorm. Verified: onnxruntime_mlas_benchmark builds cleanly on Linux x86_64 with this file included (--build_micro_benchmarks); a --benchmark_filter smoke run confirms the harness executes without crashing and reports "no_kernel_registered" as expected on a non-Apple build (no LayerNorm kernel dispatch on Linux). No Apple Silicon hardware was available locally to collect real apple_accelerate_vdsp numbers -- see the PR description for CI-collected results. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Adds an apple_accelerate job to mac.yml (native arm64 Debug + Release, mirroring the job PR microsoft#32036 added for its Tanh kernel -- that PR was closed without merging, so this CI lane does not otherwise exist upstream yet) and a use_apple_accelerate input to macos-ci-build-and-test-workflow.yml that threads --use_apple_accelerate into the build. Without this, the new kernel and its tests/benchmark would never build or execute anywhere in CI: onnxruntime_USE_APPLE_ACCELERATE is opt-in (PR microsoft#32001) and no existing native arm64 lane enables it. Two new steps, both list-then-verify fail-loud (a --gtest_filter or --benchmark_filter that matches nothing still exits 0, which would make either step a silent no-op that appears to pass while providing zero coverage): - Running MLAS Apple Accelerate targeted tests: explicitly invokes onnxruntime_mlas_test --gtest_filter="LayerNormAppleAccelerate.*", since onnxruntime_mlas_test is never registered with ctest and so the existing "build.py --test" step never runs it. Runs on every native-arm64 entry regardless of build_config (Debug + Release). - Running MLAS LayerNorm/RMSNorm benchmark A/B: builds onnxruntime_mlas_benchmark (via --build_micro_benchmarks, reconfigured in-place to avoid the wheel-repackaging cost of build.py --build) and runs the BM_LayerNormDispatch/BM_RMSNormDispatch vs. BM_LayerNormPortableScalarBaseline/BM_RMSNormPortableScalarBaseline A/B, asserting the dispatch benchmarks actually report the "apple_accelerate_vdsp" label rather than silently measuring some other path. Gated to build_config == 'Release' (a Debug/-O0 run is not real performance evidence). Validated: both YAML files parse with python's yaml.safe_load; every "run:" block's bash was extracted and syntax-checked with `bash -n` (GitHub Actions ${{ }} expressions replaced with a placeholder token for this check only). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Real Apple Silicon CI (apple_accelerate / Debug job) caught this: the test originally asserted that a single +Inf anywhere in the row contaminates every output element with NaN, for both simplified (RMSNorm) and non-simplified (full LayerNorm) modes. That assumption is only correct for the non-simplified path. For RMSNorm (simplified=true) there is no mean subtraction, so a lone +Inf only makes mean-of-squares equal to Inf (finite-sum-plus-Inf is still Inf). 1/sqrt(Inf + eps) is exactly 0.0, so every *non*-Inf element computes finite * 0.0 == 0.0 (there is no Inf-Inf cancellation on this path); only the single element that was actually set to +Inf computes Inf * 0.0 == NaN. Verified independently against plain IEEE-754/Python semantics before fixing (1.0/sqrt(inf) == 0.0, finite*0.0 == 0.0, inf*0.0 == nan) rather than just changing the assertion to whatever the kernel produced. Full LayerNorm (simplified=false) is unaffected by this fix -- the mean becomes Inf there too, so (finite - Inf) = -Inf feeds into every centered sum-of-squares term including a NaN one (Inf - Inf for the Inf element itself), poisoning the shared variance/inv-std and hence every output element, exactly as originally asserted. This is a test-only fix (kernel behavior is correct IEEE-754 arithmetic, not a bug); no changes to layernorm.cpp/platform.cpp/mlasi.h. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
CI's Python format / lintrunner check (CLANGFORMAT) flagged one line of whitespace in the previous commit's InfPropagation rewrite (a continuation operator was indented one space too far). Applied clang-format 20.1.8 (matching the version pinned in requirements-lintrunner.txt) directly to this file; no functional change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…tration Independent review found that registering GetMlasPlatform().LayerNormF32Kernel on macOS arm64 causes the pre-existing, unmodified test_layernorm.cpp suite (previously a guaranteed no-op on this architecture, since MlasLayerNormF32 always returned false and every case hit its own `if (!used) return;`) to actually execute for the first time -- and fail for norm_size=7 (both bias=0 and bias=1, non-simplified/full LayerNorm). Root cause: test_layernorm.cpp's ScalarLayerNorm oracle uses a single-pass "E[x^2] - mean^2" formula, while this PR's kernel deliberately uses a centered two-pass formula (mean first, then sum of centered squares) to avoid catastrophic cancellation. Both are valid IEEE-754 evaluation orders of the same mathematical quantity, but they diverge measurably at small norm_size for this test's fixed input pattern: confirmed via independent Python fp32 simulation, norm_size=7 diverges by ~2.19e-3 absolute in 1/std_dev (~22x over the previous 1e-4f absolute ASSERT_NEAR tolerance), while the *relative* divergence is only ~4.4e-5 (0.0044%) -- three orders of magnitude below what a genuine kernel bug (wrong scale, dropped bias, misplaced epsilon) would produce. Changing the oracle's formula to "centered" was considered and rejected: onnxruntime/core/mlas/lib/riscv64/layernorm_kernel_rvv.cpp's kernel uses the same uncentered formula as this shared oracle, so "fixing" the reference would risk silently breaking RVV's own currently-passing test instead of resolving anything -- this environment has no RVV hardware to validate that risk. Fix: replace the oracle's absolute ASSERT_NEAR(..., 1e-4f) checks with a relative-tolerance comparison (ASSERT_NEAR_REL, rel_tol=1e-3, abs floor 1e-4), which accepts either valid evaluation order while still failing loud on genuine regressions. Verified against the exact norm_size=7 case (passes with large margin) and against a synthetic 5% magnitude bug (still fails). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Remove/rephrase comments claiming a nonexistent "AVX2 LayerNorm kernel" exists in this tree (no such kernel exists anywhere in this repository; only RVV registers a non-Apple LayerNormF32Kernel). Also rephrase a comment implying a sibling "Apple Accelerate Tanh kernel" file exists in this same file family -- it does not; PR microsoft#32036, which explored that idea, was closed without merging. - Remove unused #include <algorithm>. - Replace the heap-fallback scratch buffer (NormSize > kApplePerRowStackScratch) from std::vector<float>::resize() to a raw std::unique_ptr<float[]> allocation. vector::resize() value-initializes (zero-fills) every element; every byte of this scratch buffer is fully overwritten by the vDSP calls before it is ever read in both the Simplified and full-LayerNorm branches, so the zero-fill was pure overhead on every heap-fallback call (NormSize > 8192, e.g. GPT-3-class 12288 hidden dim). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… sizes Independent review found two issues with bench_layernorm_apple_accelerate.cpp: 1. BM_LayerNormDispatch/BM_RMSNormDispatch called MlasLayerNormF32 directly without checking whether a kernel is actually registered. On any platform/config where GetMlasPlatform().LayerNormF32Kernel is nullptr (everything except this option on ARM64, or RVV), MlasLayerNormF32 returns false and does nothing at all -- it does not fall back to a scalar implementation itself, contrary to what a comment here claimed. The benchmark would therefore silently time a no-op and report a fake, near-zero-cost throughput number labeled (correctly, per DispatchPathLabel) "no_kernel_registered", which is easy to miss. Fixed: RunLayerNormBenchmark now checks HasRegisteredKernel() up front and calls state.SkipWithMessage(...) for the two dispatch benchmarks when no kernel is registered, rather than silently benchmarking nothing. 2. No benchmark size exceeded kApplePerRowStackScratch (8192), so the heap-fallback scratch path's real performance (relevant for e.g. GPT-3-class 12288 hidden dim) was completely unverified -- the claimed speedup numbers in the PR description only covered the on-stack path. Added 12288 and 16384 to the benchmark size list. Also fixed a comment claiming an "AVX2 (x86)" kernel registers LayerNormF32Kernel; no such kernel exists in this repository (only RVV does), and rephrased the PR microsoft#32036 reference to make clear it is a closed, unmerged PR, not code present in this tree. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Remove comments claiming test_tanh_apple_accelerate.cpp and a "sibling AVX2 LayerNorm kernel's test suite" exist in this repository. Neither exists here: PR microsoft#32036 (which added an Apple Accelerate Tanh kernel and its own test file) was closed without merging, and no AVX2 LayerNorm kernel/test exists anywhere in this tree. - Add NormSize == 8192 to ForcedReachability. This is the exact kApplePerRowStackScratch boundary (see layernorm.cpp): the on-stack vs. heap-fallback scratch condition is strictly greater-than, so 8192 takes the on-stack path and was previously untested (LargeNormSizeHeapFallback only covers 8193/16384, the heap side of the boundary). - Harden NearEnough against a latent (not currently reachable by any caller in this file) Inf-vs-Inf comparison trap: fabs(inf - inf) is NaN, which would make the function incorrectly reject two equal same-sign infinities. Added an explicit isinf check before the subtraction. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- test_layernorm.cpp: fix clang-format violation in NearRelative() signature (continuation lines had one extra leading space), introduced by the previous relative-tolerance fix commit. Verified clean against pinned clang-format 20.1.8 with --style=file. - test_layernorm_apple_accelerate.cpp: update stale comment that still referenced the removed std::vector<float> heap scratch buffer; the heap fallback now uses std::unique_ptr<float[]> as of the layernorm.cpp cleanup commit. Both flagged by a fresh independent Opus rubber-duck review as the only remaining issues after the prior blocking bug and all medium/nit findings were confirmed resolved. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Justin Chu (justinchuby)
marked this pull request as ready for review
August 13, 2026 19:46
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an Apple Accelerate (vDSP)
MLAS_LAYERNORM_F32_KERNELfor native macOSarm64, dispatched via
GetMlasPlatform().LayerNormF32Kernel, covering bothfull LayerNorm (mean/variance + bias) and RMSNorm (simplified, mean-of-squares
only).
Depends on #32001 (
onnxruntime_USE_APPLE_ACCELERATECMake option,open, green, Opus-reviewed and Ready per prior convergence work) for the
opt-in build infrastructure. This PR adds no new CMake plumbing — it reuses
MLAS_USE_APPLE_ACCELERATEexactly as #32001 defines it and stacks directlyon its branch (merge-base
44a752ef84). This PR cannot merge before #32001does, and will need a rebase once #32001 lands.
Status: Ready for review. Real Apple Silicon CI evidence collected (all
tests passing, material speedup confirmed) across 3 CI rounds. Two rounds of
independent Opus rubber-duck review complete, all findings resolved — see
"Review history" below.
Why LayerNorm/RMSNorm (and not another Apple Accelerate candidate)
This is a follow-up survey after #32036 (Tanh via vForce) was closed with a
negative result: on native arm64,
vvtanhfwas consistently slower than theexisting 4-wide NEON polynomial kernel already compiled unconditionally for
Erf/Sigmoid/Tanh. I surveyed the other MLAS/CPU-EP hot paths before picking a
new candidate:
on ARM64 — ruled out (this is exactly why [MLAS] Add Apple Accelerate (vForce) Tanh kernel (depends on #32001) #32036 lost).
GetMlasPlatform().LayerNormF32Kernelisnullptron ARM64 today — there is no SIMD kernel at all here. Thereal fallback (
onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc,ComputeJob) is a scalar Welford/sum-of-squares loop. This is agenuine, currently-unfilled acceleration gap, and one with high model
relevance: LayerNorm/RMSNorm is invoked roughly twice per transformer
layer in every GenAI decoder model this runtime targets.
Since RVV already has a
LayerNormF32Kernelregistration using theuncentered
E[x²] − mean²variance formula (numerically risky under fp32due to catastrophic cancellation for inputs with large mean-to-variance
ratio), I used a centered two-pass formula (mean first, then
sum-of-centered-squares) here for numerical safety. This kernel is otherwise
fully independent — no shared dispatch-threshold/helper infrastructure with
any other branch — and registers orthogonally alongside the existing RVV
kernel.
Implementation notes
MlasLayerNormKernelAppleAccelerate(onnxruntime/core/mlas/lib/layernorm.cpp):uses
vDSP_meanv/vDSP_measqv/vDSP_svesq/vDSP_vsadd/vDSP_vsmul/vDSP_vmul/vDSP_vmafor the two reduction passes and the finalnormalize+scale+bias pass. A stack-allocated scratch buffer is used up to
8192 floats (covers every representative transformer hidden size), falling
back to a heap allocation (
std::unique_ptr<float[]>, intentionallyuninitialized since every element is unconditionally written by a vDSP
call before any read — avoids the zero-fill cost
std::vectorwould add)above that, tested explicitly (
LargeNormSizeHeapFallback, sizes 8193 and16384) and benchmarked explicitly (sizes 12288/16384, see below).
#include <Accelerate/Accelerate.h>. On Xcode 26.3 / MacOSX26.2 SDK,vecLib's cblas headers break C++ compilation (this broke [MLAS] Add Apple Accelerate (vForce) Tanh kernel (depends on #32001) #32036 mid-review
and had to be fixed by forward-declaring only the needed symbols). This PR
starts from that fix directly: only the specific vDSP symbols used are
forward-declared (
extern "C", rawlong/unsigned longparameter typesto avoid any typedef collision risk), with no umbrella header include.
platform.cppinside the existingMLAS_TARGET_ARM64block,alongside the other Apple-Accelerate-gated kernels from [MLAS] Add Apple Accelerate (vForce) Tanh kernel (depends on #32001) #32036.
Input == Output, the pattern real ORT LayerNorm/SkipLayerNorm call sites use) is supported and tested.
Tests (
onnxruntime/test/mlas/unittest/test_layernorm_apple_accelerate.cpp)All gated identically to the kernel (
MLAS_USE_APPLE_ACCELERATE && __APPLE__ && MLAS_TARGET_ARM64) — this file is an empty no-op translation uniteverywhere else. 11 cases against an independent fp64-accumulated reference:
ForcedReachability/LargeNormSizeHeapFallback: direct kernel callsacross representative sizes (1 .. 8192 including the exact 8192 stack/heap
boundary, plus real transformer hidden dims 768/1024/1536/2048/3072/4096,
and the >8192 heap-fallback path at 8193/16384), both LayerNorm and RMSNorm
modes, with/without bias.
PublicDispatchMatchesDirectKernelCall: proves the publicMlasLayerNormF32entry point (what real ORT callers use) actuallyresolves to this kernel.
InPlaceAliasing,NoWriteDetectionPoisonBuffer(poison-fills the outputwith
1e30fbefore the call so a silent no-write bug cannot coincidentallypass),
ZeroVarianceConstantRow,NanPropagation,InfPropagation,DenormalInputs,ZeroNormSizeNoWrite,NullMeanAndInvStdDevOutputs.Pre-existing shared test fix: registering this kernel causes
onnxruntime/test/mlas/unittest/test_layernorm.cpp(a pre-existing,platform-independent test file, previously never exercised on ARM64 because
no kernel was registered there) to run for the first time on this platform.
One case,
norm_size=7(bothbias=0/bias=1, non-simplified), failedunder the file's original absolute tolerance (
ASSERT_NEAR(..., 1e-4f)).Root cause: the shared test oracle uses an uncentered
E[x²] − mean²reference formula, which suffers catastrophic-cancellationerror at small
norm_size— this is oracle evaluation-order noise, not akernel bug (confirmed via independent fp32/fp64 Python simulation: absolute
divergence was ~22x the old tolerance, but relative divergence was only
~4.4e-5, three orders of magnitude below what an actual bug would produce).
The oracle's formula was deliberately not changed to "centered", because
onnxruntime/core/providers/riscv64/mlas/layernorm_kernel_rvv.cpp's existing,currently-passing test relies on the same uncentered oracle, and altering it
is unvalidatable in this environment (no RVV hardware) and out of scope.
Instead,
test_layernorm.cppgained a relative-tolerance helper(
NearRelative/ASSERT_NEAR_REL, rel_tol=1e-3 with a 1e-4 absolute floorfor near-zero expected values) replacing the 3 affected absolute-tolerance
assertions. This was verified (both locally and independently by Opus
review, via a standalone harness with deliberately injected bugs — dropped
bias, dropped scale, sign flip, misplaced epsilon, 1% systematic error) to
still correctly fail on any real regression while tolerating only the
sub-tolerance floating-point noise.
The CI test-discovery filter (see below) was also broadened for the same
reason: the original
--gtest_filter="LayerNormAppleAccelerate.*"wouldnever have caught this, since it only matches this PR's own 11 tests, not
the pre-existing 36-test shared suite this PR's registration newly activates
on this platform.
Benchmark (
onnxruntime/test/mlas/bench/bench_layernorm_apple_accelerate.cpp)A/B: public dispatch (
BM_LayerNormDispatch/BM_RMSNormDispatch) vs.BM_LayerNormPortableScalarBaseline/BM_RMSNormPortableScalarBaseline— astandalone reproduction of the actual
layer_norm_impl.ccscalar fallbackthis kernel replaces on ARM64 (not a hypothetical baseline), across sizes
{1, 8, 64, 256, 768, 1024, 1600, 2048, 3072, 4096, 8192, 12288, 16384} (the
last two exercise the heap-fallback path with real timing data, added after
review). The dispatch benchmarks are guarded with a
HasRegisteredKernel()check +
state.SkipWithMessage(...)so they cleanly skip (rather thansilently reporting fake near-zero-cost numbers for a no-op) on any platform
where
MlasLayerNormF32has no kernel registered —MlasLayerNormF32returns
falseand does nothing in that case, it does not fall back to ascalar loop itself.
Real Apple Silicon CI evidence
Latest run, after all review-response fixes including the clang-format fix:
MacOS CI Pipeline run 31732593546,
apple_acceleratejobs: Release(94556560112),
Debug
(94556560187).
Both succeeded. Lint run
31732593185
also green (confirms the clang-format fix resolved the prior red run).
(An earlier run, 31726519495,
already confirmed the blocking-bug fix — both
apple_acceleratejobs green,47/47 — before the clang-format-only follow-up commit; the numbers below are
from the latest, fully-clean run.)
Tests: 47/47 passed in both Debug and Release
(
--gtest_filter="*LayerNorm*"→ 2 suites:LayerNormAppleAccelerate.[11tests, this PR] +
LayerNorm.[36 tests, pre-existing shared oracle suite,now actually exercised on ARM64 for the first time] — including the
previously-failing
norm_size7/simplified0/bias{0,1}cases now passingunder the relative-tolerance fix above).
Benchmark (Release job, real M-series hardware, ns/call, lower is
better):
(A prior run's numbers were consistent within normal micro-benchmark
run-to-run noise, e.g. LayerNorm/768 measured 16.26x vs 26.07x here — both
runs agree on the qualitative conclusion below; noise at intermediate sizes
is expected on shared CI runners and does not affect the interpretation.)
Interpretation: below
norm_size≈64, vDSP call overhead dominates andthe vDSP path is slower than the trivial scalar loop (expected, and honestly
reported — not hidden). At every representative transformer hidden size
(256 and above, which covers essentially all real GenAI decoder model
dimensions: 768/1024/1536/2048/3072/4096, plus larger MoE/vocab-adjacent
sizes up to 16384), the vDSP kernel is 10x–27x faster for LayerNorm and
3.7x–6.8x faster for RMSNorm, including the heap-fallback path (>8192),
which still retains an ~11x/~3.7x advantage.
CI wiring
Neither
mainnor #32001 has anapple_accelerateCI job — it only existedon #32036's now-closed branch. This PR adds it back (
mac.yml: native arm64Debug + Release;
macos-ci-build-and-test-workflow.yml: newuse_apple_accelerateinput/flag + fail-loud steps), since it is the firstsurviving candidate that needs it:
onnxruntime_mlas_test --gtest_filter="*LayerNorm*"invocation, list-then-verify with two independent fail-loud checks
(
::error::on zero-match): one forLayerNormAppleAccelerate(substring)and one for the exact
^LayerNorm\.$suite-header line — the lattercannot be satisfied by the Apple-only suite alone, so this genuinely
proves both suites ran. This target is not registered with ctest, so the
existing
build.py --teststep never runs it (same underlying gap [MLAS] Run onnxruntime_mlas_test on native macOS arm64 CI lane #32050filled generically for
onnxruntime_mlas_teston native arm64, but thatlane does not set
--use_apple_accelerateso would not run these testsanyway).
Releaseonly): buildsonnxruntime_mlas_benchmarkvia
--build_micro_benchmarks, runs the LayerNorm/RMSNorm A/B filter,list-then-verify fail-loud, and asserts the dispatch benchmarks report the
apple_accelerate_vdsplabel (not some other silently-substituted path).CI cost impact: reuses the existing native arm64
apple_accelerateDebug+Release jobs (the same lane structure #32036 introduced and this PR
restores); adds a few seconds of extra
onnxruntime_mlas_test/onnxruntime_mlas_benchmarkinvocation time on top of an already-scheduledjob. No new job/matrix entries, no additional machine-hours beyond that.
Review history
unmodified
test_layernorm.cppshared oracle suite on ARM64 for thefirst time (see "Pre-existing shared test fix" above) — fixed via
NearRelative/ASSERT_NEAR_REL+ broadened CI filter, verified in realCI (47/47 passing, including the previously-failing cases).
faked measurement when no kernel is registered — fixed with
HasRegisteredKernel()+SkipWithMessage.added sizes 12288/16384 with real collected data (see table above).
and a nonexistent
test_tanh_apple_accelerate.cppsibling file; missingexact-8192-boundary test case; latent
NearEnoughInf-vs-Inffabs(inf-inf)=NaNtrap; unused#include <algorithm>— all fixed.above fixes were pushed and re-validated in CI: REQUEST CHANGES for
one new issue only — a clang-format violation introduced by the first
fix's commit (2 continuation lines in
test_layernorm.cpp'sNearRelativesignature, one space over-indented;onnxruntime/test/**is in clang-format lint scope even though
onnxruntime/core/mlas/**isnot). All previously-blocking/medium/nit items were independently
re-verified as correctly resolved with no new concerns (including a
from-scratch correctness check of
ASSERT_NEAR_REL's macro expansion andthe
std::unique_ptr<float[]>no-uninitialized-read safety argument, bothconfirmed sound). One additional non-blocking nit (stale comment still
saying "
std::vector<float>heap path" after the unique_ptr change) wasalso noted. Both were fixed in a follow-up commit, re-verified clean
against the exact pinned
clang-format==20.1.8, and re-validated in CI(Lint green,
apple_accelerateDebug+Release still 47/47 passing).Validation performed
after the review-1 fixes, and after this review-2 clang-format/comment
fix) —
apple_accelerateDebug/Release jobs green every round, 47/47tests every round; full benchmark table collected each time (final table
above is from the latest, fully-clean run 31732593546).
onnxruntime_mlas/onnxruntime_mlas_test/onnxruntime_mlas_benchmarkwith the option OFF (default) — succeeds,confirming zero regression to the non-Apple path.
aarch64-linux-gnu-g++cross-compiler syntax checks (-fsyntax-only,forced
-D__APPLE__=1 -DMLAS_USE_APPLE_ACCELERATE=1, real__aarch64__)of every modified file, all clean.
clang-format --style=file(exact pinned20.1.8) verified clean on bothmodified test files after the lint fix.
yaml.safe_load) andbash -nsyntax-checked every modifiedrun:block in both workflow files.CI convergence (final)
Two infra-only failures were observed and resolved via rerun (both diagnosed
as unrelated to this PR's code before retrying):
Build Linux arm64 Debug / build_test_pipeline(Linux CI workflow) failedon an earlier run. Both
test_layernorm.cppandtest_layernorm_apple_accelerate.cpp(the only files this PR modifies thatare compiled for this platform-independent config, with
use_apple_accelerate=Falseconfirmed in the build log) had alreadycompiled successfully — the build reached 1456/1462 targets (linking
onnxruntime_mlas_test) before stopping with no ninja/compiler error inthe log. The identical job was independently failing on
mainat the sametime across 3+ consecutive runs (e.g. run 31733035193, job 94557865363) —
a contemporaneous, shared upstream/infra issue, not a regression caused by
this branch. Rerun (
gh run rerun --failedon run 31732593568)completed the full build (all 1462 targets) and test suite successfully:
job 94575553312.
web_Release / build_onnxruntime_web(Web CI Pipeline) subsequently failedwith an explicit self-hosted-runner infra annotation ("The self-hosted
runner lost communication with the server") — not a code/test failure (all
prior steps, including all
ort-webtest suites, had already passed).Rerun (
gh run rerun --failedon run 31732593306) completed cleanlythrough the E2E package-consuming test and artifact upload: job
94596754825.
Final PR check rollup: 85/89 SUCCESS, 0 FAILURE, 3 pending, 1 neutral.
The 3 pending checks (
Test Linux CUDA Plugin EP x64 Release,Test Linux CUDA x64 Release,Test Linux TensorRT x64 Release) are GPU/CUDA/TensorRTlanes this PR does not touch (no CUDA/TensorRT code in this change), queued
since ~19:30–19:39 UTC on shared GPU runners — a capacity/queue matter, not a
failure.
CodeQLreportsNEUTRAL(informational security scan, not afailure state). No FAILURE checks remain, and no branch-caused failure was
found at any point in this convergence.
Commits
11 commits total on this branch (since merge-base with #32001 at
44a752ef84): initial kernel + tests + benchmark + CI wiring (4 commits),a CI-round
InfPropagation/clang-format fix (2 commits), then the 4review-response fixes (test regression, kernel cleanup, benchmark guard,
test nits) plus this review's clang-format/comment follow-up (1 commit).
Every commit includes the Copilot co-author trailer.