[ARM] MLAS: SVE i8mm (svmmla) int8 QGEMM kernels, portable machine code - #31146
Draft
mirounga wants to merge 1 commit into
Draft
[ARM] MLAS: SVE i8mm (svmmla) int8 QGEMM kernels, portable machine code#31146mirounga wants to merge 1 commit into
mirounga wants to merge 1 commit into
Conversation
Add S8S8 and U8S8 QGEMM compute kernels using the SVE i8mm svmmla
instructions, selected at runtime whenever the processor supports SVE with
the I8MM extension (HasArmSVE_I8MM). The kernels consume the exact packed
A/B panels of the existing NEON smmla/ummla kernels (byte-identical packing
code, PackedK=8, same RowSum/ColumnSum zero-point-correction layout), so
results are bit-identical to the NEON path; only the inner compute differs
(12x8 M-tiles with an 8x12 second shape, svmmla_s32 / svmmla_u32).
The compute kernels ship as portable machine code in the style of Arm's
KleidiAI library: every instruction is a raw word (GAS ".inst" / armasm64
"DCD" via aarch64/kai_asm_macros.h, macro set adopted verbatim from
KleidiAI), so one generated file assembles under both the GNU assembler and
Microsoft armasm64 with no SVE toolchain support required. The SVE
intrinsics reference implementation (sve/qgemm_mmla_sve_impl.cpp) remains
the regeneration source: sve/gen_sve_asm.py freezes it, verifying the code
is fully self-contained (no relocations, no adrp/bl, no literal pools;
compiled -fno-stack-protector) and the emitted words are byte-identical to
the compiler's object code. The driver/pack/dispatch translation units use
no SVE intrinsics and compile with plain AArch64 flags on any compiler, so
the dispatch is not OS-gated; Windows runtime detection uses the
PF_ARM_SVE_INSTRUCTIONS_AVAILABLE / PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE
feature constants (SDK-#ifdef-guarded), and the sources are wired through
the existing cl /P + armasm64 pipeline. cmake option
onnxruntime_SVE_QGEMM_ASM (default ON) selects the frozen machine code;
OFF builds the intrinsics reference instead.
M == 1 (GEMV-shaped) operations delegate to the NEON mmla dispatch on
Linux, where those kernels exist: a single packed row cannot amortize the
2-row mmla pair structure, and the packed-B interchangeability makes the
delegation valid for both the unpacked and prepacked paths (measured: M=1
prepacked cells go from up to 1.07x slower to parity).
Measured on Cortex-X925/A725 (SVE VL=128, 5 big + 5 mid cores), full MLAS
QGEMM benchmark (76 median cells, all four SignedA/UnsignedA x PackB/
NoPackB variants), pinned to the big cluster, versus the NEON baseline
built from the same base commit (real-time medians, lower is better):
Threads:4 (fits the big cluster) 0.84-0.86x (~1.17x speedup)
Threads:16 (oversubscribed on 5c) 0.93-0.97x
Threads:1, M=1 rows 0.98-1.02x (parity, after the
M==1 delegation)
All 76 cells 0.932x
The frozen machine code measures at parity with the intrinsics build
overall (0.932x vs 0.931x); the Threads:4 band pays ~1% for the kernel now
being an out-of-line call rather than force-inlined into the operation
loop.
Correctness: full onnxruntime_mlas_test with the SVE dispatches active by
default: 33213 passed, 0 failed. Results are bit-identical to the NEON
kernels (exact integer accumulation over identical packed operands).
mirounga
marked this pull request as draft
July 29, 2026 21:54
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| @@ -0,0 +1,181 @@ | |||
| #!/usr/bin/env python3 | |||
| @@ -0,0 +1,181 @@ | |||
| #!/usr/bin/env python3 | |||
|
|
||
|
|
||
| def run(cmd): | ||
| result = subprocess.run(cmd, capture_output=True, text=True) |
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.
Add S8S8 and U8S8 QGEMM compute kernels using the SVE i8mm svmmla instructions, selected at runtime whenever the processor supports SVE with the I8MM extension (HasArmSVE_I8MM). The kernels consume the exact packed A/B panels of the existing NEON smmla/ummla kernels (byte-identical packing code, PackedK=8, same RowSum/ColumnSum zero-point-correction layout), so results are bit-identical to the NEON path; only the inner compute differs (12x8 M-tiles with an 8x12 second shape, svmmla_s32 / svmmla_u32).
The compute kernels ship as portable machine code in the style of Arm's KleidiAI library: every instruction is a raw word (GAS ".inst" / armasm64 "DCD" via aarch64/kai_asm_macros.h, macro set adopted verbatim from KleidiAI), so one generated file assembles under both the GNU assembler and Microsoft armasm64 with no SVE toolchain support required. The SVE intrinsics reference implementation (sve/qgemm_mmla_sve_impl.cpp) remains the regeneration source: sve/gen_sve_asm.py freezes it, verifying the code is fully self-contained (no relocations, no adrp/bl, no literal pools; compiled -fno-stack-protector) and the emitted words are byte-identical to the compiler's object code. The driver/pack/dispatch translation units use no SVE intrinsics and compile with plain AArch64 flags on any compiler, so the dispatch is not OS-gated; Windows runtime detection uses the PF_ARM_SVE_INSTRUCTIONS_AVAILABLE / PF_ARM_SVE_I8MM_INSTRUCTIONS_AVAILABLE feature constants (SDK-#ifdef-guarded), and the sources are wired through the existing cl /P + armasm64 pipeline. cmake option onnxruntime_SVE_QGEMM_ASM (default ON) selects the frozen machine code; OFF builds the intrinsics reference instead.
M == 1 (GEMV-shaped) operations delegate to the NEON mmla dispatch on Linux, where those kernels exist: a single packed row cannot amortize the 2-row mmla pair structure, and the packed-B interchangeability makes the delegation valid for both the unpacked and prepacked paths (measured: M=1 prepacked cells go from up to 1.07x slower to parity).
Measured on Cortex-X925/A725 (SVE VL=128, 5 big + 5 mid cores), full MLAS QGEMM benchmark (76 median cells, all four SignedA/UnsignedA x PackB/ NoPackB variants), pinned to the big cluster, versus the NEON baseline built from the same base commit (real-time medians, lower is better):
Threads:4 (fits the big cluster) 0.84-0.86x (~1.17x speedup)
Threads:16 (oversubscribed on 5c) 0.93-0.97x
Threads:1, M=1 rows 0.98-1.02x (parity, after the
M==1 delegation)
All 76 cells 0.932x
The frozen machine code measures at parity with the intrinsics build overall (0.932x vs 0.931x); the Threads:4 band pays ~1% for the kernel now being an out-of-line call rather than force-inlined into the operation loop.
Correctness: full onnxruntime_mlas_test with the SVE dispatches active by default: 33213 passed, 0 failed. Results are bit-identical to the NEON kernels (exact integer accumulation over identical packed operands).