Fix #31573, prevent ARM64 SymmQgemm int16 overflow

[Sylvester Kaczmarek (sylvesterkaczmarek)]

Widen each 8-lane int8 product group into the int32 accumulators before multiplying the second half of the packed K block. This avoids the signed int16 overflow that occurs when two -128 * -128 products share a halfword lane.

Description

• Update the plain-NEON ARM64/AArch64 SymmQgemm S8 kernels to reduce each 8-lane product group into int32 before the next multiply.
• Keep the ARM64 and AArch64 assembly implementations in sync.
• Remove the temporary non-dotprod test guard introduced in [MLAS] Add int8 extreme-value coverage for ARM64 SymmQgemm #31606 so the existing signed-input regression coverage runs on the plain-NEON path again.

Motivation and Context

Fixes #31573.

The previous smull + smlal sequence accumulated two int8 products in a signed int16 lane before widening. For the extreme case, (-128 * -128) + (-128 * -128) = 32768, which overflows int16. The revised sequence reduces each product group into the int32 accumulators before processing the second half of the packed K block.

Validation

The existing signed-input regression test is re-enabled for non-dotprod ARM64. The repository also contains onnxruntime_mlas_benchmark with SYMMQGEMM/SignedActivation; no representative non-dotprod Arm64 hardware was available for a trustworthy throughput comparison, so no performance numbers are claimed here.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Sylvester Kaczmarek (sylvesterkaczmarek)]

@microsoft-github-policy-service agree

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR fixes a known overflow issue in the plain-NEON (non-dotprod) SymmQgemm S8 kernel on Arm64/AArch64, and removes the corresponding unit-test skip so the signed-input test runs on non-dotprod hosts again.

Changes:

• Reworked the NEON kernel inner loops to avoid int16 accumulation overflow by reducing into 32-bit accumulators between half-vector multiplies.
• Removed the Arm NEON dotprod capability guard (and related include) that previously skipped the signed-input SymmQgemm test on non-dotprod hosts.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/test/mlas/unittest/test_symm_qgemm_fixture.h	Removes the runtime feature skip so the signed-input SymmQgemm test executes on plain-NEON hosts.
onnxruntime/core/mlas/lib/arm64/SymQgemmS8KernelNeon.asm	Updates the Arm64 NEON kernel to avoid int16 accumulator overflow by accumulating via 32-bit reductions.
onnxruntime/core/mlas/lib/aarch64/SymQgemmS8KernelNeon.S	Mirrors the Arm64 kernel fix in the AArch64 assembly source.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Sylvester Kaczmarek (sylvesterkaczmarek)]

Hariharan Seshadri (@hariharans29) Could you please authorize the CI runs and review this PR when convenient? The branch is now updated to current main and the diff is limited to the two NEON SymmQgemm kernels plus the regression-test guard.

The overflow explanation requested by Copilot is now in both assembly files and references #31573. For the performance comment, the repository already has onnxruntime_mlas_benchmark with SYMMQGEMM/SignedActivation. I do not have representative non-dotprod Arm64 hardware for a trustworthy throughput comparison; forcing the NEON path on dotprod-capable hardware would not provide representative performance data.

This PR also removes the guard added in #31606, as requested in #31573.

[Hariharan Seshadri (hariharans29)]

Review — PR #32057: Fix ARM64 SymmQgemm int16 accumulator overflow (#31573)

Summary of changes (3 files, 2 commits)

• SymQgemmS8KernelNeon.asm (Windows MASM) and SymQgemmS8KernelNeon.S (GCC/Clang) — same edit in both, kept in sync. Each inner-loop block that used the pattern

  smull v12.8h, v0.8b, v4.8b      ; product 1 → int16 lanes
  smlal v12.8h, v2.8b, v8.8b      ; += product 2 → still int16 lanes
  sadalp v16.4s, v12.8h           ; widen-and-add to int32
  

  is rewritten as

  smull v12.8h, v0.8b, v4.8b
  sadalp v16.4s, v12.8h           ; widen product 1 immediately
  smull v12.8h, v2.8b, v8.8b
  sadalp v16.4s, v12.8h           ; widen product 2 separately
  

  Applied to all four loops (M4_ComputeBlockLoop, M4_ComputeBlockLoopFinish, M2_ComputeBlockLoop/Finish, M1_ComputeBlockLoop/Finish) — the same pattern change every place smlal-into-.8h appeared. One explanatory comment added at the top of M4_ComputeBlockLoop in each file.

• test_symm_qgemm_fixture.h — removes the HasArmNeonDot() GTEST_SKIP that [MLAS] Add int8 extreme-value coverage for ARM64 SymmQgemm #31606 added specifically to hide this bug, and drops the now-unused #include "core/mlas/lib/mlasi.h".

The diagnosis is correct — I verified it against the pre-PR file

I looked at the current on-disk SymQgemmS8KernelNeon.S lines 143–150. The pattern is exactly:

smull v12.8h,v0.8b,v4.8b     // one int8×int8 product per int16 lane
...
smlal v12.8h,v2.8b,v8.8b     // add second product into same int16 lane

For the pathological case a = -128, b = -128, each product is -128 · -128 = 16384. 16384 + 16384 = 32768 = INT16_MAX + 1 → wraps to -32768, and sadalp then sign-extends the wrapped value into int32. The kernel silently produces a wildly wrong contribution. That's #31573 exactly.

The reordering is bitwise-identical for correct inputs

sadalp v16.4s, v12.8h is defined as, per int32 lane i:

$$v16[i] \mathrel{+{=}} v12[2i] + v12[2i+1]$$

• Old: v12[j] = a1[j]·b1[j] + a2[j]·b2[j] in int16, then v16[i] += v12[2i] + v12[2i+1].
• New: v16[i] += a1[2i]·b1[2i] + a1[2i+1]·b1[2i+1], then v16[i] += a2[2i]·b2[2i] + a2[2i+1]·b2[2i+1].

Both expand to the same int32 sum of four products. Because the accumulation is now in int32 throughout, the overflow window at int16 is closed. For any input where the old code did not overflow, the new code produces identical bits. ✓

The perf trade-off Copilot flagged is real but bounded

Old inner-block (four output cols, one A-row group): 4 smull + 4 smlal + 4 sadalp = 12 NEON ops.
New: 4 smull + 4 sadalp + 4 smull + 4 sadalp = 16 NEON ops.

That's +33% NEON-pipe instructions on the fused-multiply-then-widen path. On A55 (in-order, dual-issue, one NEON pipe), that's close to a 1:1 wall-clock hit for the multiply-widen stage — probably 10–20% overall kernel slowdown after amortizing load/loop overhead. Copilot's suggested sshll/sxtl + smlal into 32-bit-widened operands would recover most of that, but at 2× register pressure and a much more invasive rewrite.

I think this trade is right for merge as-is:

• The fix is for a known correctness bug in a shipping kernel. Correctness beats a 10–20% cost on the older-hardware path.
• The plain-NEON SymQgemmS8Kernel only runs when FEAT_DotProd is absent. That path serves pre-Cortex-A75 cores and pre-Apple-M1 (2021+) processors. Every modern deployment target either has FEAT_DotProd (and hits MlasGemmS8S8DispatchDot) or has FEAT_I8MM (and hits MlasGemmS8S8DispatchSmmla). This kernel is the correctness floor, not the perf hot path.
• The author explicitly acknowledges no hardware validation for perf and doesn't claim any. Filing a follow-up for the sshll+smlal-int32 rewrite as a perf-only PR would be the right shape if someone is motivated.

Test guard removal is correct

The removed guard from #31606 explicitly said "Skip on hosts that would fall back to it instead of failing; drop this guard once that kernel is fixed." That preconditioned drop is being honored here, and the CI ARM64 lanes will now actually exercise the fix on any non-dotprod code path they cover. Dropping the #include "core/mlas/lib/mlasi.h" is correct — it was added in #31606 solely for MLAS_CPUIDINFO::GetCPUIDInfo(), which the fixture no longer touches.

Minor comments

1. The explanatory comment lives only in M4_ComputeBlockLoop. Same pattern change is applied at M4_ComputeBlockLoopFinish, M2_ComputeBlockLoop / Finish, M1_ComputeBlockLoop / Finish. Not repeating the comment is fine (the pattern is now uniform across the file), but a one-liner cross-reference in each other loop — "same reason as M4_ComputeBlockLoop above; see [MLAS] ARM64 non-dotprod SymmQgemm NEON kernel produces wrong results for extreme int8 operands (int16 intermediate overflow) #31573" — would prevent someone from micro-optimizing one of the tail loops back to the fused pattern in isolation. Non-blocking.

2. Consider adding one M-loop-covering test to lock the fix in for M1/M2/M4 paths. The signed-input regression test exists but I don't know its M coverage off the top of my head. If the current test only exercises the M4 block, a future refactor could regress M1 / M2 silently. Cheap to add a shape or two if not already covered.

3. PR title. "Fix [MLAS] ARM64 non-dotprod SymmQgemm NEON kernel produces wrong results for extreme int8 operands (int16 intermediate overflow) #31573, prevent ARM64 SymmQgemm int16 overflow" — minor comma-vs-colon inconsistency with the ORT convention. Not a merge blocker; the maintainer can adjust on squash.

Recommendation

Approve. Correctness fix for a shipping bug, mirror-edited in both .asm and .S, bitwise-equivalent to the old code for non-overflowing inputs, and correctly re-enables the test that was gated on this fix landing. Perf hit on the plain-NEON path is real but affects only pre-dotprod hardware where correctness is the priority. The sshll-based faster rewrite is a legitimate follow-up if a maintainer wants to file it.

[Sylvester Kaczmarek (sylvesterkaczmarek)]

Thanks Hariharan, appreciate the detailed review and validation, and for getting the CI running. The follow-up suggestions make sense, particularly broader M-loop regression coverage and keeping any performance optimization as a separate follow-up.

[Hariharan Seshadri (hariharans29)]

Thanks Hariharan, appreciate the detailed review and validation, and for getting the CI running. The follow-up suggestions make sense, particularly broader M-loop regression coverage and keeping any performance optimization as a separate follow-up.

Thanks for this fix. Should we consider adding in the M-loop regression test for this PR ? What are your thoughts on this matter ?

[Sylvester Kaczmarek (sylvesterkaczmarek)]

Thanks Hariharan, appreciate the detailed review and validation, and for getting the CI running. The follow-up suggestions make sense, particularly broader M-loop regression coverage and keeping any performance optimization as a separate follow-up.

Thanks for this fix. Should we consider adding in the M-loop regression test for this PR ? What are your thoughts on this matter ?

I checked the existing signed-input regression coverage. It already registers M = {1, 2, 3, 4, 5, 7, 8, 9}, so it exercises the M1, M2 and M4 paths directly as well as mixed/tail combinations.

Given that, I don't think a separate M-loop regression test would add meaningful coverage here.

Happy to make that intent more explicit in the test if you prefer.

[Hariharan Seshadri (hariharans29)]

Thanks Hariharan, appreciate the detailed review and validation, and for getting the CI running. The follow-up suggestions make sense, particularly broader M-loop regression coverage and keeping any performance optimization as a separate follow-up.

Thanks for this fix. Should we consider adding in the M-loop regression test for this PR ? What are your thoughts on this matter ?

I checked the existing signed-input regression coverage. It already registers M = {1, 2, 3, 4, 5, 7, 8, 9}, so it exercises the M1, M2 and M4 paths directly as well as mixed/tail combinations.

Given that, I don't think a separate M-loop regression test would add meaningful coverage here.

Happy to make that intent more explicit in the test if you prefer.

Sounds good, we can leave it as such then, thanks.

There are a couple of unstarted CI checks blocking merge. Can you please rebase with main to see if that is mitigated with that ?

[Sylvester Kaczmarek (sylvesterkaczmarek)]

Thanks Hariharan, appreciate the detailed review and validation, and for getting the CI running. The follow-up suggestions make sense, particularly broader M-loop regression coverage and keeping any performance optimization as a separate follow-up.

Thanks for this fix. Should we consider adding in the M-loop regression test for this PR ? What are your thoughts on this matter ?

I checked the existing signed-input regression coverage. It already registers M = {1, 2, 3, 4, 5, 7, 8, 9}, so it exercises the M1, M2 and M4 paths directly as well as mixed/tail combinations.
Given that, I don't think a separate M-loop regression test would add meaningful coverage here.
Happy to make that intent more explicit in the test if you prefer.

Sounds good, we can leave it as such then, thanks.

There are a couple of unstarted CI checks blocking merge. Can you please rebase with main to see if that is mitigated with that ?

Done, I’ve rebased the branch onto the latest main and force-pushed the updated commits. The CI checks have been retriggered, although they’re currently showing as requiring approval.

Could you please authorize them again when convenient? Thanks!