[ARM CPU] SVE support for Sgemm kernel - #26027
Aruna K (akote123) wants to merge 9 commits into
Conversation
@microsoft-github-policy-service agree company=“Fujitsu Research of India Private Ltd” |
7ebfed8 to
ff2a6a4
Compare
6ee2ffc to
5dd2161
Compare
|
Please address the build failures and please run onnxruntime_mlas_test in your env with your change and report any failures if any here. |
|
|
/azp run Linux QNN CI Pipeline, Win_TRT_Minimal_CUDA_Test_CI, Windows ARM64 QNN CI Pipeline, Windows GPU Doc Gen CI Pipeline, Windows x64 QNN CI Pipeline |
|
Azure Pipelines successfully started running 4 pipeline(s). |
|
Have you tried this on multiple Gemm problem shapes - Please submit comprehensive micro-benchmarks for all SVE platforms. I tried taking this change on a Graviton4 and for a Conv heavy model (which uses Im2Col + SGemm for the Conv implementation), it slows down the model very much - when I build with |
Yes. I have tried on different shapes ,and for large shapes tuning is required. This PR is added as patch to enable sve for gemm kernel and for larger shapes this gemm implementation has to be tuned and we are currently working on that. |
|
Hi Aruna K (@akote123): Any plan to resume this work ? Or at the very least, can you please guide what remains to be done for perf tuning ? |
|
Hariharan Seshadri (@hariharans29) ,Thank you. We will resume this work .This PR we planned as SVE sgemm enablement work and do performance tuning as future task and contribute |
7f3180f to
2e2c6db
Compare
2e2c6db to
e21b7b5
Compare
|
Hariharan Seshadri (@hariharans29) ,Thank you Below table shows shape wise timing: GEMM Performance Comparison (time in ms)
|
|
Thanks for this. I have a few PRs ahead of this in the review queue and I will get to this soon ! Meanwhile, I am happy to kick off CI and Copilot reviewing in the background. Also there is another SVE based Gemm implementation coming in (from KleidiAI) - Need to study the overlap and synergize |
|
/azp run Linux QNN CI Pipeline, Win_TRT_Minimal_CUDA_Test_CI, Windows ARM64 QNN CI Pipeline, Windows GPU Doc Gen CI Pipeline, Windows x64 QNN CI Pipeline |
|
No pipelines are associated with this pull request. |
Sorry - re-visiting this just now. Do we have perf data across a large bunch of representative M,N,K shapes for SVE 256 and 512 - or have you measured perf only across the 4 models listed in the PR description ? Personally, I am leaning towards including SVE SGemm support in MLAS and defaulting it to off (via the Mlas backend kernel selector) and allowing the user to opt-in to using this feature - atleast for 1 or 2 releases. We can announce the feature via release notes and get perf feedback. We can turn it on my default in a future release once we are satisifed it hasn't broken anything. |
|
Could you please re-base with main (and integrate with the Mlas backend kernel selector) ? Once done, I will prioritize merging this - thanks |
Sure .I will update with patch to integrate with the Mlas backend kernel selector. |
759bcce to
6c52be6
Compare
|
Hariharan Seshadri (@hariharans29),
|
Review: PR #26027 — [ARM CPU] SVE support for Sgemm kernelRequest changes. The feature is legitimately useful and the branch has done the right work in reintegrating with the Mlas backend kernel selector so it defaults off. However, one substantive design gap remains: the Blocking issue — default-off contract is broken for non-packed callersThe dispatch guard used in every SVE branch is (paraphrasing the diff on lines around 274, 325, 515, 608, 635, 690, 1086, 1096 of if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSVE() &&
(!BackendKernelSelectorConfig || BackendKernelSelectorConfig->enable_sve_sgemm)) {
// SVE path
}When
Which callers pass
Net effect: enabling SVE support at build time ( Two acceptable fixes:
Please pick one and apply it uniformly. This is the one issue that would prevent merging as-is. Correctness observations
Style and naming
Testing gaps
Perf data observations (non-blocking)The 14-hour-ago perf table on g3E is a nice sanity check but is narrow:
Meaningful but modest wins in the 1.0×–1.4× range on square/tall shapes. Consider adding:
Since the feature is opt-in, this is not a blocker, but it will inform whether it's worth turning on by default in a future release. Config wiring
SummarySubstantive design fix needed on the |
Performance Comparison: MLAS SVE vs MLAS NEONMetric: Time in milliseconds
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (4)
onnxruntime/core/mlas/lib/sve/sgemm_sve.cpp:1156
MlasSveTransposePackBNx4only handles VL==16/8/4. For any other valid SVE vector length, this loop performs no stores, leaving the packed buffer uninitialized and producing incorrect GEMM results.
Add a safe fallback path for unsupported VL values (e.g., scalar transpose matching MlasSgemmTransposePackBNx4’s D[0]/D[16]/D[32]/D[48] layout).
for (unsigned n = 0; n < N / 4; n++) {
if (VL() == 16) {
Transpose_SVE512_4x4(&D[0], &B[0], ldb);
} else if (VL() == 8) {
Transpose_SVE256_4x4(&D[0], &B[0], ldb);
onnxruntime/core/mlas/lib/sve/sgemm_sve.cpp:1223
MlasSveScatterStorewrites to d[15]/d[30]/d[45], but the packed-B layout used by the existing transpose pack routine stores rows at offsets 0/16/32/48 (seeMlasSgemmTransposePackBNx4in sgemm.cpp). These off-by-one (and off-by-two/three) offsets will corrupt the packed buffer layout.
Update the offsets to match the standard 16-float stride.
MlasSveStoreFloat32(pb_first_half, &d[0], vec0);
MlasSveStoreFloat32(pb_second_half, &d[15], vec0);
MlasSveStoreFloat32(pb_third_half, &d[30], vec0);
MlasSveStoreFloat32(pb_fourth_half, &d[45], vec0);
onnxruntime/core/mlas/lib/sve/sgemm_sve.cpp:1232
MlasSveLoadStoreusessvptrue_b32()and increments byVL(). Ifsvcntw()is larger than (or does not evenly divide)MLAS_SGEMM_STRIDEN_THREAD_ALIGN(16), this will read/write past the intended 16-float block.
Use a predicate based on the remaining element count so this stays in-bounds for any SVE vector length.
for (int i = 0; i < MLAS_SGEMM_STRIDEN_THREAD_ALIGN; i += VL()) {
svfloat32_t vec0 = MlasSveLoadFloat32(svptrue_b32(), b + i);
MlasSveStoreFloat32(svptrue_b32(), D + i, vec0);
}
onnxruntime/core/mlas/lib/sve/sgemm_sve.cpp:1241
MlasSveZeroInitializestores withsvptrue_b32()in chunks ofsvcntw(). Forsvcntw() > kMlasSvePackedBBlockWidth(16) this will write past the 16-float block; and forsvcntw()values that don’t evenly divide 16 it can also overrun on the final iteration.
Use a predicate capped to kMlasSvePackedBBlockWidth so this always zeros exactly 16 floats.
svfloat32_t zero = svdup_f32(0.0f);
for (int i = 0; i < kMlasSvePackedBBlockWidth; i += svcntw()) {
MlasSveStoreFloat32(svptrue_b32(), d + i, zero);
}
|
Hariharan Seshadri (hariharans29), mirounga |
|
There is a substantial amount of work required to bring this up to the current upstream state. You can take a look at how the Elementwise kernels have evolved and refactor SGEMM in the same key |
Hi mirounga, thanks for checking the PR and sharing your feedback |
Update session option for packing kernel
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
6836115 to
48d101a
Compare
|
mirounga Sir please support this PR further. Thankyou. |
This PR ports the SGEMM kernel and associated packing kernels to the ARM SVE (Scalable Vector Extension) backend. Specifically, it introduces new wrapper implementations for SVE functions in lib/sve/sgemm_sve.cpp and integrates these wrappers within the existing kernel implementations.
Motivation and Context
This work is part of an ongoing effort to enhance ONNX Runtime's performance and architecture-awareness on ARM platforms. By leveraging ARM SVE, we aim to unlock better computational efficiency and scalability on modern ARM hardware.
This PR builds upon and extends the SVE work introduced in PR #25238
Performance Analysis:

Results are captured from sve 256,128 and SVE 512 supported machines.
This PR is a joint contribution by: