[MLAS][Kleidiai] Sve Gemm and IMatmul Integration by JonathanC-ARM · Pull Request #27643 · microsoft/onnxruntime

Jonathan Clohessy (JonathanC-ARM) · 2026-03-13T15:21:00Z

Description

Adds initial Arm SVE enablement for the KleidiAI MLAS backend, including SVE
ukernel wiring, SGEMM dispatch/packing support, and an SVE convolution path
with runtime selection (prefer SME/SME2 when available, otherwise use SVE).
Also updates the KleidiAI dependency to a newer release to pick up the
required SVE kernels KleidiAI 1.22

Motivation and Context

Enables KleidiAI acceleration on Arm systems that expose SVE but not SME/SME2,
reducing fallback to the generic MLAS implementations and broadening hardware
coverage. This is an initial bring-up focused on correctness and integration,
with some configuration limitations (e.g., SVE SGEMM currently targets non-
transposed inputs; SVE conv has capability constraints).

Hariharan Seshadri (hariharans29) · 2026-03-13T17:46:10Z

/azp run Linux QNN CI Pipeline,Win_TRT_Minimal_CUDA_Test_CI,Windows ARM64 QNN CI Pipeline,Windows GPU Doc Gen CI Pipeline

azure-pipelines · 2026-03-13T17:46:29Z

Azure Pipelines successfully started running 4 pipeline(s).

Copilot

Pull request overview

Adds initial Arm SVE enablement for the KleidiAI MLAS backend, wiring in SVE SGEMM packing/dispatch and an SVE IMATMUL-based convolution path with runtime selection (prefer SME/SME2 when available, otherwise SVE). Also bumps the KleidiAI dependency to a version that provides the required SVE kernels.

Changes:

Extend MLAS platform dispatch to select KleidiAI overrides on SVE-only CPUs (and include SME2 in SME selection).
Add SVE variants for KleidiAI SGEMM (pack + batched GEMM dispatch) and convolution (SVE IMATMUL indirection path).
Expand FGEMM unit test coverage for batched short-path shapes and non-trivial alpha/beta cases.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
onnxruntime/test/mlas/unittest/test_fgemm_fixture.h	Adds small batched FGEMM short-path coverage with varied alpha/beta.
onnxruntime/core/mlas/lib/sgemm.cpp	Removes TransA gating so KleidiAI override can decide support; minor cleanup.
onnxruntime/core/mlas/lib/platform.cpp	Enables KleidiAI override selection for SME2 and SVE-only runtime.
onnxruntime/core/mlas/lib/kleidiai/sgemm_kleidiai.cpp	Adds SVE SGEMM packer + batched GEMM path; refactors alpha/beta application helpers.
onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h	Introduces `UseSVE` runtime feature flag.
onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp	Adds SVE IMATMUL convolution path and runtime selection between SME/SME2 vs SVE.
onnxruntime/core/mlas/lib/kai_ukernel_interface.h	Adds SVE matmul/imatmul wrapper typedefs + getter declarations; small comment edits.
onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp	Registers SVE kernels and implements SVE getter functions.
cmake/deps.txt	Updates KleidiAI dependency to v1.22.0.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+    // Match SME alpha/beta behavior: apply beta per-batch when alpha==0 or K==0.
+    if (Data->alpha == 0.0f || K == 0) {
+        if (BatchSize == 1) {
+            ApplyBetaToC(Data->C, Data->ldc, M, N, Data->beta);
+        } else {
+            for (size_t batch = 0; batch < BatchSize; ++batch) {
+                ApplyBetaToC(Data[batch].C, Data[batch].ldc, M, N, Data[batch].beta);
+            }
+        }
+        return true;
+    }


+            const size_t tile_elems = TileSizeM * TileSizeN;
+            g_kai_tls.output_tile.resize(tile_elems);
+            out_tile = g_kai_tls.output_tile.data();
+            out_row_stride_bytes = TileSizeN * sizeof(float);


 #include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
+#include "kai/ukernels/matmul/pack/kai_rhs_pack_kxn_x32p4vlx1b_x32_x32_sve.h"


+const KaiF32SveIMatmulKernel GetKleidiAISveImatmulUKernel();
+
+const KaiF32SveKernel GetKleidiAISveSGemmUKernel();



 #include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi4c32p/kai_matmul_clamp_f32_qai8dxp_qsi4c32p_interface.h"

+// matmul inferfaces


Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>

Hariharan Seshadri (hariharans29) · 2026-03-23T17:04:15Z

/azp run Linux QNN CI Pipeline,Win_TRT_Minimal_CUDA_Test_CI,Windows ARM64 QNN CI Pipeline,Windows GPU Doc Gen CI Pipeline

azure-pipelines · 2026-03-23T17:04:35Z

Azure Pipelines successfully started running 4 pipeline(s).

Copilot

Pull request overview

This PR adds initial Arm SVE enablement for the KleidiAI-backed MLAS paths, wiring up SVE ukernels for SGEMM/IMATMUL and enabling runtime selection between SME/SME2 and SVE where applicable. It also bumps the KleidiAI dependency to a newer release containing the required SVE kernels.

Changes:

Enable KleidiAI SGEMM overrides to be considered more broadly (remove TransA gating in MLAS dispatch/pack helpers) while letting KleidiAI return “unsupported” for unsupported transpose cases.
Add SVE runtime selection for KleidiAI SGEMM packing/dispatch and an SVE convolution path (with capability checks).
Update KleidiAI dependency to v1.22.0.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
onnxruntime/core/mlas/lib/sgemm.cpp	Removes TransA gating around KleidiAI SGEMM override/packing hooks so KleidiAI can decide support internally.
onnxruntime/core/mlas/lib/platform.cpp	Registers KleidiAI overrides for SME/SME2 and (when enabled) for SVE-only systems.
onnxruntime/core/mlas/lib/kleidiai/sgemm_kleidiai.cpp	Adds SVE SGEMM pack-size/packer helpers and an SVE GEMM batch implementation with SME/SME2 preference.
onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h	Adds a cached `UseSVE` capability flag for KleidiAI runtime selection.
onnxruntime/core/mlas/lib/kleidiai/convolve_kleidiai.cpp	Adds an SVE IMATMUL-based convolution path with runtime selection and new capability checks.
onnxruntime/core/mlas/lib/kai_ukernel_interface.h	Introduces SVE kernel wrapper typedefs and declares SVE kernel getters.
onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp	Adds SVE ukernel includes, wrapper instances, and getter implementations.
cmake/deps.txt	Updates KleidiAI dependency tarball from v1.20.0 to v1.22.0.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-23T17:12:07Z

+    if (Parameters->Beta != 0.f) {
+        KLEIDIAI_DEBUG_LOG("CheckCapabilitiesSve rejected: unsupported Beta=" << Parameters->Beta);
+        return false;
+    }
+


CheckCapabilitiesSve rejects non-zero Parameters->Beta, but the new ConvolveSve implementation explicitly supports beta by preserving the original output and applying beta after the kernel run. As written, the beta handling in ConvolveSve is unreachable for the SVE path. Either allow Beta != 0 in CheckCapabilitiesSve (and keep the post-processing), or remove the beta code and keep the restriction (but then drop the extra copy/memcpy logic).

Suggested change

if (Parameters->Beta != 0.f) {

KLEIDIAI_DEBUG_LOG("CheckCapabilitiesSve rejected: unsupported Beta=" << Parameters->Beta);

return false;

}

Copilot · 2026-03-23T17:12:07Z

+// SVE kernels
+//   IMATMUL
+#include "kai/ukernels/matmul/imatmul_clamp_f32_f32_f32p/kai_imatmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla.h"
+//   GEMM
+#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla.h"


The SVE ukernel headers and wrappers are included/defined unconditionally here. Since MLAS_USE_SVE is only enabled when onnxruntime_USE_SVE is set and the compiler supports -march=armv8.2-a+sve, this can introduce build failures (or unwanted SVE codegen/link requirements) in builds where KleidiAI is enabled but SVE is not. Consider wrapping the SVE includes + imatmul_sve/sgemm_gemm_sve definitions + the GetKleidiAISve*UKernel() getters in #if defined(MLAS_USE_SVE) to match the rest of the integration.

Suggested change

// SVE kernels

// IMATMUL

#include "kai/ukernels/matmul/imatmul_clamp_f32_f32_f32p/kai_imatmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla.h"

// GEMM

#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla.h"

// SVE kernels

#if defined(MLAS_USE_SVE)

// IMATMUL

#include "kai/ukernels/matmul/imatmul_clamp_f32_f32_f32p/kai_imatmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla.h"

// GEMM

#include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla.h"

#endif // defined(MLAS_USE_SVE)

Copilot · 2026-03-23T17:12:08Z


+const KaiF32SveIMatmulKernel imatmul_sve = 
+    KAI_WRAP_UKERNEL_RUN_IMATMUL_6_NO_LHS_PACKED_OFFSET(imatmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla);

+const KaiF32SveKernel sgemm_gemm_sve = 
+    KAI_WRAP_UKERNEL_RUN_MATMUL_10_LHS_OFFSET(matmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla);


These SVE kernel wrapper instances are defined even when MLAS_USE_SVE is not enabled. If SVE support is meant to be optional (gated by onnxruntime_USE_SVE), these should be conditionally compiled to avoid pulling in SVE symbols/codepaths in non-SVE builds.

Suggested change

const KaiF32SveIMatmulKernel imatmul_sve =

KAI_WRAP_UKERNEL_RUN_IMATMUL_6_NO_LHS_PACKED_OFFSET(imatmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla);

const KaiF32SveKernel sgemm_gemm_sve =

KAI_WRAP_UKERNEL_RUN_MATMUL_10_LHS_OFFSET(matmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla);

#if defined(MLAS_USE_SVE)

const KaiF32SveIMatmulKernel imatmul_sve =

KAI_WRAP_UKERNEL_RUN_IMATMUL_6_NO_LHS_PACKED_OFFSET(imatmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla);

const KaiF32SveKernel sgemm_gemm_sve =

KAI_WRAP_UKERNEL_RUN_MATMUL_10_LHS_OFFSET(matmul_clamp_f32_f32_f32p4vlx1b_6x4vl_sve_mla);

#endif // MLAS_USE_SVE

Copilot · 2026-03-23T17:12:08Z


+#if defined(MLAS_USE_SVE)
+// SVE convolution kernel (indirection-LHS IMATMUL).
+const KaiF32SveIMatmulKernel imatmul_conv_sve = GetKleidiAISveImatmulUKernel();


imatmul_conv_sve is stored by value, whereas the SME/SME2 kernel (imatmul_conv) is stored as a const&. For consistency (and to avoid an extra static copy/initialization), consider making this a const KaiF32SveIMatmulKernel& as well.

Suggested change

const KaiF32SveIMatmulKernel imatmul_conv_sve = GetKleidiAISveImatmulUKernel();

const KaiF32SveIMatmulKernel& imatmul_conv_sve = GetKleidiAISveImatmulUKernel();

Copilot · 2026-03-23T17:12:08Z

 // Wrapper for FP32 IMATMUL kernels used by the KleidiAI convolution implementation.
 using KaiF32IMatmulKernel = KaiMatmulKernel<kai_imatmul_clamp_f32_f32p_f32p_ukernel>;

+// Wrapper for BF16 based IMatmul


The comment says "Wrapper for BF16 based IMatmul", but KaiBF16SBgemmKernel is a wrapper for a BF16 matmul (kai_matmul_clamp_f32_bf16p_bf16p_ukernel), not an IMATMUL ukernel. Please adjust the comment to match the actual type to avoid confusion when adding more IMATMUL variants.

Suggested change

// Wrapper for BF16 based IMatmul

// Wrapper for BF16 SBGEMM kernels

Hariharan Seshadri (hariharans29) · 2026-03-31T19:26:55Z

/azp run Linux QNN CI Pipeline,Win_TRT_Minimal_CUDA_Test_CI,Windows ARM64 QNN CI Pipeline,Windows GPU Doc Gen CI Pipeline

azure-pipelines · 2026-03-31T19:27:16Z

Azure Pipelines successfully started running 4 pipeline(s).

Hariharan Seshadri (hariharans29) requested a review from Copilot March 13, 2026 17:45

Copilot started reviewing on behalf of Hariharan Seshadri (hariharans29) March 13, 2026 17:46 View session

Copilot AI reviewed Mar 13, 2026

View reviewed changes

Added initial SVE Sgemm changes

6bd3238

Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>

Jonathan Clohessy (JonathanC-ARM) force-pushed the jonclo01/sve_kleidi_integration branch from f352bfd to 6bd3238 Compare March 23, 2026 15:42

Jonathan Clohessy (JonathanC-ARM) added 2 commits March 23, 2026 15:50

Fix SVE ukernel declarations and include guards

128ab78

Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>

Consolidate gemm function checks and fix ukernel function signature

7fd7778

Signed-off-by: Jonathan Clohessy <Jonathan.Clohessy@arm.com>

Hariharan Seshadri (hariharans29) requested a review from Copilot March 23, 2026 17:04

Copilot started reviewing on behalf of Hariharan Seshadri (hariharans29) March 23, 2026 17:04 View session

Copilot AI reviewed Mar 23, 2026

View reviewed changes

Merge branch 'microsoft:main' into jonclo01/sve_kleidi_integration

07dc282

Hariharan Seshadri (hariharans29) mentioned this pull request Apr 24, 2026

[ARM CPU] SVE support for Sgemm kernel #26027

Open

Hariharan Seshadri (hariharans29) marked this pull request as draft May 19, 2026 16:36

Navigation Menu

[MLAS][Kleidiai] Sve Gemm and IMatmul Integration - #27643

[MLAS][Kleidiai] Sve Gemm and IMatmul Integration#27643
Jonathan Clohessy (JonathanC-ARM) wants to merge 4 commits into
microsoft:mainfrom
JonathanC-ARM:jonclo01/sve_kleidi_integration

Jonathan Clohessy (JonathanC-ARM) commented Mar 13, 2026

Uh oh!

Hariharan Seshadri (hariharans29) commented Mar 13, 2026

Uh oh!

azure-pipelines Bot commented Mar 13, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Hariharan Seshadri (hariharans29) commented Mar 23, 2026

Uh oh!

azure-pipelines Bot commented Mar 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Copilot AI Mar 23, 2026

Uh oh!

Hariharan Seshadri (hariharans29) commented Mar 31, 2026

Uh oh!

azure-pipelines Bot commented Mar 31, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
		#include "kai/ukernels/matmul/pack/kai_rhs_pack_kxn_x32p4vlx1b_x32_x32_sve.h"

		const KaiF32SveIMatmulKernel GetKleidiAISveImatmulUKernel();

		const KaiF32SveKernel GetKleidiAISveSGemmUKernel();


		#include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi4c32p/kai_matmul_clamp_f32_qai8dxp_qsi4c32p_interface.h"

		// matmul inferfaces

	const KaiF32SveIMatmulKernel imatmul_conv_sve = GetKleidiAISveImatmulUKernel();
	const KaiF32SveIMatmulKernel& imatmul_conv_sve = GetKleidiAISveImatmulUKernel();

	// Wrapper for BF16 based IMatmul
	// Wrapper for BF16 SBGEMM kernels

Conversation

Jonathan Clohessy (JonathanC-ARM) commented Mar 13, 2026

Description

Motivation and Context

Uh oh!

Hariharan Seshadri (hariharans29) commented Mar 13, 2026

Uh oh!

azure-pipelines Bot commented Mar 13, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Hariharan Seshadri (hariharans29) commented Mar 23, 2026

Uh oh!

azure-pipelines Bot commented Mar 23, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Hariharan Seshadri (hariharans29) commented Mar 31, 2026

Uh oh!

azure-pipelines Bot commented Mar 31, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants