[MLAS] Add AArch64 SVE FP16 GEMM kernel and driver by nikhilfujitsu · Pull Request #32756 · microsoft/onnxruntime

Nikhil Sharma (nikhilfujitsu) · 2026-09-23T04:57:06Z

Adds an SVE implementation of the half-precision GEMM path:

hgemm.cpp - new sigmoid-style driver for HGEMM (operation, packing, threading, batch). It now owns MiasHgemmSupported and the MLAS_HGEMM_DATA_PARAMS overload of MiasGemmBatch, moved here from halfgemm.cpp. Adds a TransA path (A-panel transpose) and fast paths that skip B packing:
M=1 and M=2 gemv (NoTrans/NoTrans)
N=1 matrix-vector when B is contiguous
M=1 with TransB, which reuses the N=1 kernel with A and B swapped.
sve/halfgemm_kernel_sve.cpp - new SVE compute kernels: Zero/Add GEMM kernels, CopyPackB, TransposePackB, TransposeA and fp16 gemv (M=1, M=2). Vector-length agnostic (svctnh() 8/16/32). Streams B blocks and reduces register pressure for K-1 column operations.
sve/halfgemm_sve.h, sve/halfgemv_sve.h - extern "C" kernel declarations shared by the driver and both kernel implementations.
aarch64/halfgemm_sve_asm.S, aarch64/halfgemv_sve_asm.S - generated by sve/gen_sve_asm.py from the intrinsics sources; raw instruction words that assemble without SVE toolchain support and are byte-identical on every platform.
hlgemm.cpp - new MiasHgemmAccelSize / MiasHgemmPackfD and a MLAS_HGEMM_DATA_PARAMS: blsPacked field for pre-packed B. The TransA/TransB docs now list NoTrans and Trans.
platform.cpp - gate the Neon HGEMM dispatch on runtime fp16 vector acceleration.
cmake - build hgemm.cpp. New option omnxruntime_SVE_HGEMM_ASM (default ON) builds the frozen .S kernels. OFF builds the intrinsics sources with -march=armv8.2-a+sve+fp16.

Fixes in the existing NEON FP16 path:

fp16_common.h - MiasLoadPartialFloat16x4 / MiasStorePartialFloat16x4 did nothing for len=4. Any FP16 GEMM with M>3 and M&8=4 therefore got zeroes for that 8-column block.
halfgemm_kernel_neon_fp16.cpp
HGemm_Kernel_Complicated_read_C when beta=0, so a NaN in an uninitialized output buffer leaked into the result (alpha=1 only). The tile is now zeroed first.
HPack_B_kernel zero-fills the unused lanes of its N tail.

Tests / benchmarks:

test_halfgemm_sve.cpp - new: bit-exact unit test of the SVE TransA transpose primitive, plus an end-to-end sweep over all four (TransA, TransB) combinations, alpha/beta variants and padded strides. Gated on runtime SVE availability.
test_hgemm_neon.cpp - extend the shape sweep to TransA=Trans.
bench_hgemm.cpp - add TransA benchmark captures.

Contributer: nikhil.r.sharma@fujitsu.com

Adds an SVE implementation of the half precision GEMM path, structured like sgemm.cpp rather than the MLAS_HGEMM_DISPATCH function-pointer table. * hgemm.cpp: new sgemm-style driver for HGEMM (operation, packing, threading, batch). Owns MlasHGemmSupported and the MLAS_HGEMM_DATA_PARAMS overload of MlasGemmBatch, which move here out of halfgemm.cpp. Adds a TransA path (A-panel transpose) and an M == 1 gemv fast path, and keeps the MlasFp16AccelerationSupported() guard the old entry point had. * sve/halfgemm_kernel_sve.cpp: new SVE compute kernels -- Zero/Add GEMM kernels, CopyPackB, TransposePackB, TransposeA and an FP16 gemv. Vector length agnostic (svcnth() 8/16/32); the B-pack and A-transpose tiles use svzip cascades verified against a scalar reference. Follows the layout of the other SVE FP16 translation units: mlasi_sve.h, then the __ARM_FEATURE_FP16_VECTOR_ARITHMETIC && MLAS_F16VEC_INTRINSICS_SUPPORTED guard with <arm_sve.h> inside it, and SVE intrinsics called directly rather than through a wrapper layer. * mlasi.h: MLAS_HGEMM_DISPATCH moves here from halfgemm.h and gains HTransposeA_Fn and TransAStrideM. * platform.cpp: gate the Neon HGEMM dispatch on runtime FP16 vector acceleration. * cmake: build hgemm.cpp always; build the SVE kernel with -march=armv8.2-a+sve+fp16 under MLAS_USE_SVE. MlasHgemmKernelLoop() marks its arguments unreferenced on the non-SVE branch. Without MLAS_USE_SVE the SVE block is preprocessed away and the parameters would otherwise trip -Werror=unused-parameter, breaking non-SVE ARM64 builds. The branch is unreachable in practice because MlasHGemmSupported() returns false without SVE. Two FP16 NEON defects found while testing this path, both invisible to the existing tests: * MlasLoadPartialFloat16x4()/MlasStorePartialFloat16x4() decompose their length into its 1- and 2-bits, so len == 4 matched nothing: the load returned zeros and the store wrote nothing. HPackB_B_Kernel() hits this in its CountN <= 4 tail, so B packed as all zeros whenever the N tail was exactly 4, and any FP16 GEMM with M >= 3 (the packed-B path) and N % 8 == 4 silently produced zeros for that 8-column block. A sweep over M 1-4 x N 1-20 x K 1-6 found 36 such shapes; the existing NEON HGEMM tests use N values 1, 15, 17, 31, 32, 33, 79 and 263, none congruent to 4 mod 8. Both helpers now treat len >= 4 as a full 4-lane access, and the pack tail zero-fills lanes 4-7 of each 8-wide row rather than leaving them uninitialized. * beta == 0 makes C write-only, so a caller may pass an uninitialized output buffer. HGemm_B_Kernel_Complicated folds beta in as fma(mul(C, beta), accu, alpha) at every store and so reads C whatever beta is, letting a stale NaN survive 0 * NaN. HGemm_B_Kernel dispatches alpha == 1 with beta == 0 or 1 to HGemm_B_Kernel_Simple, which overwrites, so only alpha != 1 with beta == 0 reaches the unguarded path. A probe that fills C with NaN leaked in 160 of 480 (M<=4, N<=10, K<=4, alpha in {0,1,2}) combinations, all with alpha != 1. The tile is now zeroed up front when beta == 0, rather than threading a beta_behavior template through the five store sites the way HGemm_PackedB_Kernel_Impl already does. Tests and benchmarks * test_halfgemm_sve.cpp: new. A bit-exact unit test of the SVE TransA transpose primitive against a scalar reference, plus an end-to-end sweep over all four (TransA, TransB) combinations, alpha/beta variants and padded strides. Both suites are gated on runtime SVE availability. * test_hgemm_neon.cpp: extend the shape sweep to TransA = Trans. * bench_hgemm.cpp: add TransA benchmark captures.

…rnel Follows the approach of the SVE elementwise and QGEMM kernels: the intrinsics translation unit stays as the reference and regeneration source, and sve/gen_sve_asm.py freezes it into a portable machine-code variant where every instruction is a raw word. The frozen file assembles with toolchains that have no SVE support and the bytes are identical on every platform, so compilers can no longer disagree about how to optimise this code. * sve/halfgemm_sve.h: the extern "C" declarations the driver and both implementations agree on, plus PACKED_B_BLOCK_WIDTH_FP16 (hgemm.cpp static_asserts it against MLAS_HGEMM_STRIDEN_THREAD_ALIGN). * sve/halfgemm_kernel_sve.cpp: made self-sufficient (no mlasi.h) as the generator requires, entry points switched from extern "C++" to extern "C", helpers forced inline -- the generator rejects any call in a frozen function, and gcc had declined to inline the 32x32 transpose. * aarch64/halfgemm_sve_asm.S: generated, 6 functions. * cmake: onnxruntime_SVE_HGEMM_ASM (default ON) selects the frozen variant; OFF builds the intrinsics reference. -fno-tree-loop-distribute-patterns is applied to the reference so the zero-pad loop in HPackB_B_Kernel is not lowered to a memset call, which the generator rejects as a relocation. New matrix-vector kernel for N == 1 The general path packs B into 32-wide blocks, so a K x 1 operand pays for a whole block. Measured on Graviton3 at M = K = 1023, N = 1 through N = 16 all cost the same ~404 us, i.e. 21x the per-column rate reached at N >= 32. That is structural, not a tuning problem, so N == 1 now uses a kernel that skips the pack and streams B once per four-row group (four independent accumulator chains, vector length agnostic, predicated K tail). M == 1 with TransB reuses the same kernel with the operands swapped: with TransB, C[n] = dot(B[n, :], A), which is a matrix-vector product with B as the matrix. Previously that case fell into the packed path. B must be contiguous -- true for TransB, and for NoTrans when ldb == 1. SVE has no 16-bit gather, so a strided B stays on the general path. Measured against SGEMM on Graviton3 (8 threads), HGEMM >= SGEMM on: general GEMM 26/27 shapes, median 1.41x, up to 1.73x GEMV shapes 25/27 shapes, median 3.53x, up to 8.2x (was 10/27, 0.91x) The two remaining losses are K == 1, an outer product with no reduction, which neither kernel targets. Frozen and intrinsics builds are verified to agree: 261/261 SVE tests and 7051 HGEMM/HalfGemm tests pass in both configurations, and nm confirms the toggle actually swaps which object supplies the symbols.

…river

Copilot

Copilot review overview

🟡 Changes recommended

Critical build, compatibility, correctness, and memory-safety issues remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 5 · 1

Open (6)

Add HGEMM assembly sources to Windows ARM64 builds · New
Avoid reading C when zero-beta GEMM has K equal to zero · New
Restore NEON fallback for non-SVE HGEMM · New
Restore unconditional MLAS_HGEMM_DISPATCH declaration · New
Handle zero counts in the SVE HGEMM packer · New
Gate transposed-A tests on supported dispatch · New

What changed in this PR

Adds AArch64 SVE FP16 HGEMM kernels and driver support, including transpose and packed-B paths, alongside NEON correctness fixes.

Changes:

Adds SVE GEMM/GEMV kernels, packing, transpose paths, and generated assembly.
Adds packed-B APIs and updates runtime/build integration.
Expands tests and benchmarks and fixes NEON FP16 edge cases.

File	Summary	Review result
`onnxruntime/test/mlas/unittest/test_hgemm_neon.cpp`	Extends transpose coverage.	Moderate (2 votes): TransA tests must be gated because NEON-only systems do not support that path.
`onnxruntime/test/mlas/unittest/test_halfgemm_sve.cpp`	Adds SVE transpose and GEMM tests.	No direct issue identified.
`onnxruntime/test/mlas/bench/bench_hgemm.cpp`	Adds TransA benchmarks.	No issue identified.
`onnxruntime/core/mlas/lib/sve/halfgemv_sve.h`	Declares SVE GEMV kernels.	No issue identified.
`onnxruntime/core/mlas/lib/sve/halfgemv_kernel_sve.cpp`	Implements SVE GEMV kernels.	No issue identified.
`onnxruntime/core/mlas/lib/sve/halfgemm_sve.h`	Declares SVE GEMM kernels.	No issue identified.
`onnxruntime/core/mlas/lib/sve/halfgemm_kernel_sve.cpp`	Implements SVE GEMM, packing, and transpose kernels.	Critical (1 vote): Packing loops can underflow and corrupt memory when either count is zero; return early for zero dimensions.
`onnxruntime/core/mlas/lib/platform.cpp`	Adds runtime HGEMM dispatch gating.	No direct issue identified.
`onnxruntime/core/mlas/lib/mlasi.h`	Relocates and extends HGEMM dispatch definitions.	Critical (3 votes): Conditional removal of the dispatch declaration breaks x64/WebAssembly parsing; retain an unconditional forward declaration.
`onnxruntime/core/mlas/lib/hgemm.cpp`	Implements the HGEMM driver and packed-B operations.	Critical (3 votes): Direct HGEMM loses the existing non-SVE NEON fallback. Critical (3 votes): `K == 0 && beta == 0` reads `C`, preserving NaNs instead of zeroing output (also lines 323 and 588). Moderate (1 vote): The packed-B branch lacks end-to-end coverage for offsets, TransA, and N partitioning.
`onnxruntime/core/mlas/lib/halfgemm.h`	Removes the relocated dispatch definition.	No direct issue identified.
`onnxruntime/core/mlas/lib/halfgemm.cpp`	Routes eligible half-GEMM calls through HGEMM.	No direct issue identified.
`onnxruntime/core/mlas/lib/halfgemm_kernel_neon_fp16.cpp`	Fixes NEON packing and zero-beta behavior.	No issue identified.
`onnxruntime/core/mlas/lib/fp16_common.h`	Fixes four-element partial FP16 loads/stores.	No issue identified.
`onnxruntime/core/mlas/lib/aarch64/halfgemv_sve_asm.S`	Provides generated SVE GEMV assembly.	No issue identified.
`onnxruntime/core/mlas/inc/mlas.h`	Exposes packed-B APIs and transpose metadata.	Moderate (1 vote): Add end-to-end packed-B tests covering both transpose modes, tails, and multiple K panels.
`cmake/onnxruntime_mlas.cmake`	Integrates the HGEMM driver and SVE kernels into builds.	Critical (1 vote): Windows ARM64 enables SVE but omits both generated assembly sources, causing link failures.

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+          if (onnxruntime_SVE_HGEMM_ASM)
+            list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/aarch64/halfgemm_sve_asm.S)
+          else()
+            list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/sve/halfgemm_kernel_sve.cpp)
+            set_source_files_properties(${MLAS_SRC_DIR}/sve/halfgemm_kernel_sve.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+sve+fp16 -fno-tree-loop-distribute-patterns ${ORT_SVE_ABI_FLAGS} ")
+          endif()
+          if (onnxruntime_SVE_HGEMM_ASM)
+            list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/aarch64/halfgemv_sve_asm.S)


+    if (beta == 1.0f) {
+        return;
+    }
+    auto* c = reinterpret_cast<_mlas_fp16_*>(C);
+    for (size_t m = 0; m < CountM; ++m) {
+        _mlas_fp16_* row = c + m * ldc;
+        for (size_t n = 0; n < CountN; ++n) {
+            row[n] = MLAS_Float2Half(MLAS_Half2Float(row[n]) * beta);


+            MLAS_UNREFERENCED_PARAMETER(CountN);
+            MLAS_UNREFERENCED_PARAMETER(alpha);
+            MLAS_UNREFERENCED_PARAMETER(ZeroMode);
+            MLAS_THROW_EX(std::runtime_error, "HGEMM requires SVE in this build");


 // half gemm dispatch structure
 //
-struct MLAS_HGEMM_DISPATCH;
+#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) || defined(MLAS_TARGET_ARM64)


+{
+    const size_t vl = svcnth();
+
+    while (CountX >= PACKED_B_BLOCK_WIDTH_FP16) {


+    // TransA=Trans, TransB=Trans
+    TestHGemm<2, 1, 1, true, true>(MLAS_FP16(1.0f), MLAS_FP16(0.0f));
+    TestHGemm<1, 128, 512, true, true>(MLAS_FP16(0.5f), MLAS_FP16(1.0f));
+    TestHGemm<2, 128, 513, true, true>(MLAS_FP16(1.5f), MLAS_FP16(0.5f));
+    TestHGemm<1, 128, 511, true, true>(MLAS_FP16(1.0f), MLAS_FP16(0.0f));
+    TestHGemm<2, 129, 512, true, true>(MLAS_FP16(0.5f), MLAS_FP16(1.0f));
+    TestHGemm<1, 127, 512, true, true>(MLAS_FP16(1.5f), MLAS_FP16(0.5f));
+    TestHGemm<127, 513, 1023, true, true>(MLAS_FP16(1.0f), MLAS_FP16(0.0f));
+    TestHGemm<129, 511, 1025, true, true>(MLAS_FP16(0.5f), MLAS_FP16(1.0f));
+    // TransA=Trans, TransB=NoTrans
+    TestHGemm<2, 1, 1, true, false>(MLAS_FP16(1.0f), MLAS_FP16(0.0f));
+    TestHGemm<1, 128, 512, true, false>(MLAS_FP16(0.5f), MLAS_FP16(1.0f));
+    TestHGemm<2, 128, 513, true, false>(MLAS_FP16(1.5f), MLAS_FP16(0.5f));
+    TestHGemm<1, 128, 511, true, false>(MLAS_FP16(1.0f), MLAS_FP16(0.0f));
+    TestHGemm<2, 129, 512, true, false>(MLAS_FP16(0.5f), MLAS_FP16(1.0f));
+    TestHGemm<1, 127, 512, true, false>(MLAS_FP16(1.5f), MLAS_FP16(0.5f));
+    TestHGemm<127, 513, 1023, true, false>(MLAS_FP16(1.0f), MLAS_FP16(0.0f));
+    TestHGemm<129, 511, 1025, true, false>(MLAS_FP16(0.5f), MLAS_FP16(1.0f));


- Restore the NEON HGEMM dispatch as the fallback on CPUs without SVE (A x B and A x B'), and keep MlasHalfGemmBatch forwarding SVE-only. - Do not read C when beta == 0 in the K == 0 and beta scaling paths. - Define MLAS_HGEMM_DISPATCH on all targets so non-ARM builds compile. - Return early from MlasHGemmPackB when N or K is zero. - Add the HGEMM/HGEMV SVE assembly sources to Windows ARM64 builds. - Only run the transposed-A NeonHGemm cases when they are supported. - Add SveHGemmEdge tests for K == 0 with NaN-filled C and empty PackB.

Aruna K (akote123) and others added 5 commits July 7, 2026 08:53

Merged PR 997: Rebase fj-develop branch with OSS main

f23a93d

Merge remote-tracking branch 'upstream/main' into fj-develop

6fc4ce5

[MLAS] Forward FP16 x FP16 MlasHalfGemmBatch calls to the SVE HGEMM d…

5db455c

…river

Copilot AI balanced review requested due to automatic review settings September 23, 2026 04:57

Copilot started reviewing on behalf of Nikhil Sharma (nikhilfujitsu) September 23, 2026 04:58 View session

Copilot AI reviewed Sep 23, 2026

View reviewed changes

abhijain1204fujitsu force-pushed the nikhil/hgemm_sve_support branch from 26bdfa3 to 200beff Compare September 25, 2026 05:50

Merge branch 'main' into nikhil/hgemm_sve_support

8fefcf7

Navigation Menu

[MLAS] Add AArch64 SVE FP16 GEMM kernel and driver - #32756

[MLAS] Add AArch64 SVE FP16 GEMM kernel and driver#32756
Nikhil Sharma (nikhilfujitsu) wants to merge 7 commits into
microsoft:mainfrom
MonakaResearch:nikhil/hgemm_sve_support

Nikhil Sharma (nikhilfujitsu) commented Sep 23, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Nikhil Sharma (nikhilfujitsu) commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Nikhil Sharma (nikhilfujitsu) commented Sep 23, 2026 •

edited

Loading