[ARM] MLAS: portable machine-code SVE elementwise kernels, FEXPA exp - #31145
Draft
mirounga wants to merge 1 commit into
Draft
[ARM] MLAS: portable machine-code SVE elementwise kernels, FEXPA exp#31145mirounga wants to merge 1 commit into
mirounga wants to merge 1 commit into
Conversation
Rework the SVE elementwise kernels (f32: Erf, Logistic, Exp, SumExp, Softmax, LogSoftmax, ReduceMaximum, ReduceMinimumMaximum; fp16: Erf, Gelu, Tanh) so they no longer require an SVE-capable compiler or a Linux-only build, and speed them up while doing so. Structure: each kernel becomes a relocation-free *Impl function that receives its constants through table pointers (no literal pools, no global data references), fronted by a plain C++ dispatch layer (sve/elementwise_sve_dispatch.cpp) that owns the public entry points and fastpaths. The production kernels ship as portable machine code in the style of Arm's KleidiAI library (aarch64/elementwise_sve_asm.S: raw instruction words via KAI_ASM_INST, which is GAS ".inst" on Linux/macOS and armasm64 "DCD" on Windows via aarch64/kai_asm_macros.h, macro set adopted verbatim from KleidiAI). The SVE intrinsics translation units remain in-tree as the reference implementation and regeneration source, selectable with cmake option onnxruntime_SVE_ELEMENTWISE_ASM=OFF. The exp path uses the SVE FEXPA instruction (base SVE, no extension required) for the SumExp/exp evaluation, with the input clamped at -88.0f: below -127*ln2 the FEXPA index underflows into the exponent field and would produce NaN. Two kernels gain platform-dispatch coverage on AArch64/SVE that previously existed only on AMD64/RISCV64: ComputeExpF32Kernel and ReduceMinimumMaximumF32Kernel (MLAS_PLATFORM fields moved to a shared guard; call sites in compute.cpp/quantize.cpp extended with MLAS_USE_SVE). Windows enablement: runtime detection via PF_ARM_SVE_INSTRUCTIONS_AVAILABLE (SDK-#ifdef-guarded, Windows 11 24H2+), sources wired through the existing cl /P + armasm64 pipeline, the fp16 kernel dispatch un-gated from !_WIN32 (the NEON fp16 fallbacks remain POSIX-only; on Windows without SVE the routines stay null and callers use their scalar fallbacks). Structurally complete; awaiting validation on Windows-on-ARM SVE hardware. Measured on Cortex-X925/A725 (SVE VL=128), big-core pinned, 104 median cells, versus the SVE intrinsics kernels currently shipping (real-time medians, lower is better): ComputeSoftmaxInplace (32 cells) 0.772x (1.30x speedup) GeluErf fused/unfused (36 cells) 0.795-0.797x Silu (36 cells) 0.927-0.930x All 104 cells 0.832x (1.20x speedup) Correctness: full onnxruntime_mlas_test with the machine-code kernels linked: 33213 passed, 0 failed. The code is fully position-independent and vector-length agnostic (runtime cntw/cnth + whilelo loops).
mirounga
marked this pull request as draft
July 29, 2026 21:53
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
1 participant
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.
Rework the SVE elementwise kernels (f32: Erf, Logistic, Exp, SumExp, Softmax, LogSoftmax, ReduceMaximum, ReduceMinimumMaximum; fp16: Erf, Gelu, Tanh) so they no longer require an SVE-capable compiler or a Linux-only build, and speed them up while doing so.
Structure: each kernel becomes a relocation-free *Impl function that receives its constants through table pointers (no literal pools, no global data references), fronted by a plain C++ dispatch layer (sve/elementwise_sve_dispatch.cpp) that owns the public entry points and fastpaths. The production kernels ship as portable machine code in the style of Arm's KleidiAI library (aarch64/elementwise_sve_asm.S: raw instruction words via KAI_ASM_INST, which is GAS ".inst" on Linux/macOS and armasm64 "DCD" on Windows via aarch64/kai_asm_macros.h, macro set adopted verbatim from KleidiAI). The SVE intrinsics translation units remain in-tree as the reference implementation and regeneration source, selectable with cmake option onnxruntime_SVE_ELEMENTWISE_ASM=OFF.
The exp path uses the SVE FEXPA instruction (base SVE, no extension required) for the SumExp/exp evaluation, with the input clamped at -88.0f: below -127*ln2 the FEXPA index underflows into the exponent field and would produce NaN.
Two kernels gain platform-dispatch coverage on AArch64/SVE that previously existed only on AMD64/RISCV64: ComputeExpF32Kernel and ReduceMinimumMaximumF32Kernel (MLAS_PLATFORM fields moved to a shared guard; call sites in compute.cpp/quantize.cpp extended with MLAS_USE_SVE).
Windows enablement: runtime detection via
PF_ARM_SVE_INSTRUCTIONS_AVAILABLE (SDK-#ifdef-guarded, Windows 11 24H2+), sources wired through the existing cl /P + armasm64 pipeline, the fp16 kernel dispatch un-gated from !_WIN32 (the NEON fp16 fallbacks remain POSIX-only; on Windows without SVE the routines stay null and callers use their scalar fallbacks). Structurally complete; awaiting validation on Windows-on-ARM SVE hardware.
Measured on Cortex-X925/A725 (SVE VL=128), big-core pinned, 104 median cells, versus the SVE intrinsics kernels currently shipping (real-time medians, lower is better):
ComputeSoftmaxInplace (32 cells) 0.772x (1.30x speedup)
GeluErf fused/unfused (36 cells) 0.795-0.797x
Silu (36 cells) 0.927-0.930x
All 104 cells 0.832x (1.20x speedup)
Correctness: full onnxruntime_mlas_test with the machine-code kernels linked: 33213 passed, 0 failed. The code is fully position-independent and vector-length agnostic (runtime cntw/cnth + whilelo loops).
Description
Motivation and Context