[MLAS] Route ARM64 S8U8 QGEMM to the UDOT kernel

[kjg0724]

Description

ARM64 quantized GEMM never routed the signed-A/unsigned-B (S8U8) combination to a SIMD dispatch — MLAS_PLATFORM didn't even declare a GemmS8U8Dispatch member on ARM64, unlike x86_64 (AVX2/VNNI), LoongArch64 (LSX), and RISC-V64, which all have one. It fell through to the generic scalar reference kernel.

This reuses the existing UDOT kernel (already shared by U8U8 and U8S8) instead of adding new assembly. UDOT only does unsigned x unsigned, so B is already shifted into the unsigned domain with an XOR 0x80 sign flip when signed; this does the same for A (MlasGemmQuantFixupZeroPointA, MlasGemmQuantCopyPackA). K-padding bytes are left untouched across all four row-count paths (8/4/2/1-row) so they keep contributing zero to the dot product and RowSumBuffer. platform.cpp initializes the new dispatch member to the scalar default and only overrides it to UDOT inside the existing dot-product-support check, so cores without dot-product stay on the scalar path exactly as before.

Also parameterized bench_qgemm.cpp to cover S8U8, and added a guard around the packed-B benchmark: MlasGemmPackBSize returns 0 for combinations the active dispatch doesn't support packing for, and calling MlasGemmPackB anyway was dereferencing a null CopyPackBRoutine.

Motivation and Context

QLinearMatMul, MatMulInteger, and QLinearConv register their signed-A kernels with T2 accepting both int8 and uint8, so S8U8 is reachable from real models, not just a synthetic case. Existing test_qgemm.cpp correctness tests already cover int8_t/uint8_t, they were just exercising the scalar fallback on ARM64.

Tested on Apple M1 (macOS) and Ampere Neoverse-N1 (Ubuntu), both with NEON dot-product:

• Existing MLAS qgemm suites (S8U8/U8U8/U8S8/S8S8) pass on both, no regressions.
• Benchmarked bench_qgemm (NoPackB, real_time) before/after by temporarily forcing the scalar dispatch for comparison:

Shape (M/N/K, threads)	M1	Neoverse-N1
1/512/512, 1	149µs → 29µs	152µs → 51µs
1/1024/1024, 1	610µs → 202µs	1544µs → 240µs
384/1024/1024, 4	4.96ms → 3.98ms	12.3ms → 1.34ms
384/1024/3072, 4	20.7ms → 4.68ms	37.2ms → 4.17ms
1536/1024/4096, 16	124.6ms → 37.6ms	201.6ms → 37.4ms
3072/4096/1024, 16	135.4ms → 53.4ms	388.8ms → 69.2ms

Packed-B for S8U8 works correctly on both machines now that they're on the UDOT dispatch. Without the guard mentioned above, the new packed-B benchmark case would crash on any target that still resolves to the scalar dispatch, since MlasGemmPackBSize returns 0 there and MlasGemmPackB would call a null CopyPackBRoutine.

ARM64 cores without dot-product support are unaffected by the dispatch change — they stay on MlasGemmQuantDispatchDefault (scalar) exactly as before, MlasGemmPackBSize still returns 0 for S8U8 there, and the benchmark guard keeps that case skipping cleanly instead of crashing. This PR doesn't add an optimized S8U8 path or packed-B support for those cores; that's left for a follow-up if there's interest.

The routing change in qgemm.h lives inside the #if defined(MLAS_TARGET_ARM64) guard, so native Windows ARM64 (MSVC, MLAS_TARGET_ARM64) picks up the same S8U8 → UDOT routing as macOS/Linux ARM64 — it wasn't separately benchmarked here. ARM64EC (MLAS_TARGET_ARM64EC) and 32-bit ARM stay on their existing logic, unaffected.

[azure-pipelines]

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

[kjg0724]

Followed up on whether the unconditional veorq/XOR-with-zero added to CopyPackA<UDOT> regresses the existing U8U8/U8S8 paths (both go through this same packer with AIsSigned == false).

Methodology: swapped qgemm_kernel_udot.cpp between the pre-PR version and this PR's version (nothing else changed, so U8S8 still dispatches to the same UDOT kernel either way), rebuilt onnxruntime_mlas_benchmark for each, and ran the existing UnsignedAPackB/UnsignedANoPackB (U8S8) cases interleaved (before, after, before, after...) for 3 rounds. Machine: Oracle Ampere Altra (Neoverse-N1), idle (load average 0.07) for the whole run, 4 vCPUs.

Median ratio (after / before) by thread count:

Threads	ratio range	notes
1	0.988 – 1.033	matches core count 1:1
4	0.988 – 1.005	matches vCPU count 1:1
16	0.827 – 1.324	4x oversubscribed on 4 vCPUs

At 1 and 4 threads (properly subscribed), the difference is within measurement noise in both directions — no regression. At 16 threads the numbers swing widely, but that's oversubscription scheduling noise rather than a code effect: the same binary against itself across the 3 rounds shows the same order of swing (e.g. UnsignedAPackB/M:1536/N:1024/K:1024/Threads:16, before-only across rounds 1/2/3: 15.1M / 24.1M / 15.3M ns).

Makes sense given what actually changed: a single pipelined NEON XOR against an all-zero constant, folded into an already memory-bound copy loop.

[kjg0724]

@microsoft-github-policy-service agree

[hariharans29]

Review: PR #29787 — [MLAS] Route ARM64 S8U8 QGEMM to the UDOT kernel

Author: @kjg0724. Target: microsoft:main. Single commit cdf7b51. Head fetched via patch-diff.

Verdict: LGTM — approve. Small, well-scoped, mathematically sound, and thoroughly benchmarked. No blockers.

---

What it does

Adds GemmS8U8Dispatch to MLAS_PLATFORM on ARM64 (parity with x86_64/LoongArch64/RISC-V64, which already had one), threads it through MlasGemmQuantGetDispatch in qgemm.h, and wires the ARM64 UDOT kernel to handle a signed A matrix by:

• XOR-flipping every real A byte with 0x80 inside MlasGemmQuantCopyPackA<MLAS_GEMM_U8X8_KERNEL_UDOT> — a new template specialization MlasGemmQuantFixupZeroPointA<...> applies the matching ^ 0x80 to ZeroPointA so both stay in the same shifted-unsigned domain.
• Leaving K-padding bytes as literal zero across all four row-count paths (8/4/2/1), so they contribute nothing to the UDOT accumulator or to RowSumBuffer.
• Initializing the new dispatch member to MlasGemmQuantDispatchDefault (scalar) in platform.cpp, then overriding to MlasGemmU8X8DispatchUdot only inside the existing HasDotProductInstructions block — so no-dot-product cores stay on the scalar path exactly as before.
• Parameterizing bench_qgemm.cpp for S8U8 and adding a SkipWithError guard when MlasGemmPackBSize returns 0 (which was previously making MlasGemmPackB deref a null CopyPackBRoutine on any target where the active dispatch didn't support packing).

---

Correctness — the signed→unsigned shift math

UDOT is unsigned × unsigned. The standard "flip bit 7" trick converts a signed byte v (int8) to v + 128 (uint8) via v ^ 0x80:

int8 v	v ^ 0x80 (as uint8)	v + 128
-128	0	0
-1	127	127
0	128	128
127	255	255

Full quantized-GEMM identity (dropping i,j indices, with _s for shifted):

$$ \begin{aligned} C &= \sum_k (A_k - Z_A)(B_k - Z_B) \\ &= \sum_k A_k B_k - Z_A \cdot \mathrm{ColSum}_B - Z_B \cdot \mathrm{RowSum}_A + K \cdot Z_A Z_B \end{aligned} $$

Substituting $A_{s} = A + 128$ and $Z_{A,s} = Z_A + 128$, the UDOT accumulator computes:

$$ \mathrm{MainAccum}_s = \sum_k A_{s,k} B_k = \sum_k A_k B_k + 128 \cdot \mathrm{ColSum}_B $$

and the correction step (using shifted values consistently) evaluates to Correction_original − 128·ColSum_B. The +128·ColSum_B from the accumulator and the −128·ColSum_B from the correction cancel, so the total equals the unshifted C. This is the same accounting already used for the pre-existing BIsSigned path in the same UDOT kernel via MlasGemmQuantFixupZeroPointB — the PR is a symmetric extension.

The RowSum_A term is computed after the flip (from the packed A bytes, which the diff leaves unchanged in that region), so it's naturally RowSum_{A,s} — matches the shifted Z_{A,s} used in the correction. Padding bytes stay literal zero, so their contribution to both the UDOT accumulator and RowSum_A is zero regardless of AIsSigned. Consistent.

The ZeroPointA ^ 0x80 in the fixup is applied via OffsetAType(...) (a uint8_t cast), which is correct even when ZeroPointA is a negatively-signed int32_t — only bit 7 flips, and the cast to uint8_t drops the sign-extension bits.

---

Blast radius on the existing U8U8 / U8S8 paths

Those two combinations also flow through MlasGemmQuantCopyPackA<UDOT>, and after this PR the pack loop unconditionally executes:

const uint8_t  AByteFlip     = AIsSigned ? 0x80 : 0;
const uint8x16_t AFlipVector = vdupq_n_u8(AByteFlip);
...
uint32x4_t v0 = veorq_u32(vld1q_u32(...), AFlipVector32);

For AIsSigned == false, AFlipVector32 is all-zero and every veorq is a no-op semantically — but the instruction still executes (the compiler cannot fold this away because AIsSigned is a runtime bool). The concern is a perf regression on the very common U8U8/U8S8 paths.

The author addressed this explicitly in the PR conversation: swapped the pre-PR and PR-PR versions of qgemm_kernel_udot.cpp in place, interleaved runs across 3 rounds on Neoverse-N1 (idle, load-avg 0.07), and reported median after/before ratios of 0.988–1.033 at 1 thread and 0.988–1.005 at 4 threads — well within measurement noise. The 16-thread numbers swing more (0.827–1.324), but the author showed that noise is present when running the pre-PR binary against itself across rounds (e.g., 15.1M / 24.1M / 15.3M ns), so that's oversubscription-scheduling noise, not a code effect. The reasoning matches the operation: one pipelined NEON XOR against a compile-time-splatted constant, folded into an already-memory-bound copy loop. Convincing.

---

Dispatch selection change

The qgemm.h diff collapses:

if (BIsSigned) {
    GemmQuantDispatch = AIsSigned ? GemmS8S8Dispatch : GemmU8S8Dispatch;
} else if (!AIsSigned) {
    GemmQuantDispatch = GemmU8U8Dispatch;
}

into:

if (BIsSigned) {
    GemmQuantDispatch = AIsSigned ? GemmS8S8Dispatch : GemmU8S8Dispatch;
} else {
    GemmQuantDispatch = AIsSigned ? GemmS8U8Dispatch : GemmU8U8Dispatch;
}

Semantically equivalent for the pre-existing three combinations, and previously S8U8 was falling through to GemmQuantDispatch = nullptr (implicit default) which the rest of the pipeline resolved to MlasGemmQuantDispatchDefault (scalar). Now S8U8 is bound explicitly, and on no-dot-product cores it's bound to MlasGemmQuantDispatchDefault — same behavior as before, but now discoverable via GetMlasPlatform().GemmS8U8Dispatch rather than by omission.

---

Benchmark guard

if (packed_b_size == 0) {
    state.SkipWithError("Packing is not supported for this A/B signedness combination on this target");
    return;
}

Small but valuable — before, MlasGemmPackBSize returning 0 on the scalar dispatch would cause the next call MlasGemmPackB to dereference a null CopyPackBRoutine. The guard makes bench_qgemm portable across ARM64 variants with and without dot-product, and would have also fixed a latent crash risk if anyone ran the pre-existing SignedAPackB case on a target without S8S8 packing support (the pre-existing #if !defined(MLAS_TARGET_AMD64) guard on SignedAPackB was doing this job by hand). Good defensive cleanup.

---

Scope observations (not blockers)

1. Windows ARM64EC path unchanged. qgemm.h #elif defined(MLAS_TARGET_ARM64EC) || (defined(MLAS_TARGET_ARM) && !defined(_MSC_VER)) is deliberately not touched. Native Windows ARM64 (MSVC, MLAS_TARGET_ARM64) does pick up the change — only the ARM64EC (x64-emulation-compat) and 32-bit ARM legs stay on their existing logic. Worth calling out in the PR description that Windows ARM64 native benefits too; ARM64EC follow-up is trivially straightforward if anyone asks.
2. No-dot-product ARM64 cores. Explicitly out of scope per the description. MlasGemmPackBSize returns 0 there, the benchmark skips cleanly, and correctness is preserved via the scalar fallback. Reasonable — those cores are the tail of the ARMv8 population and adding an assembly path there is a much larger PR.
3. Test coverage. Author notes existing test_qgemm.cpp covers int8_t/uint8_t already, previously exercising the scalar fallback. Now the same tests exercise the UDOT path on any ARM64 with dot-product. That is sufficient for correctness verification, but a very small explicit test that constructs an S8U8 GEMM and asserts equality against a reference scalar computation (parameterized over AIsSigned=true, BIsSigned=false) would make the coverage self-documenting rather than implicit. Not a blocker.
4. MLAS_UNREFERENCED_PARAMETER(AIsSigned); removal. Correct — AIsSigned is now used. Good hygiene.
5. Constants recomputed per call. AByteFlip, AFlipVector, AFlipVector32, AFlipWord32 are computed at the top of MlasGemmQuantCopyPackA<UDOT> and then used across the four row-count blocks. Compiler will hoist and constant-splat these into registers cheaply; no restructuring needed.

---

Nits

• The four flip constants (AByteFlip, AFlipVector, AFlipVector32, AFlipWord32) all encode the same information — the same "flip bit 7 iff AIsSigned" — at four register widths. A brief comment near the declaration block already explains why (the 8/4/2/1-row paths use different register widths). Fine as-is.
• The new template specialization signature int32_t MlasGemmQuantFixupZeroPointA<MLAS_GEMM_U8X8_KERNEL_UDOT>(int32_t ZeroPointA, bool AIsSigned) matches the primary template signature — this is a pure specialization addition, not a signature change. No effect on other kernels.

---

CI

PR page shows 1 / 1 checks OK on the single commit at fetch time. No stalled or failing legs.

---

Bottom line

Approve. The math is sound and symmetric to the existing signed-B UDOT handling, the perf regression on the unsigned-A paths has been empirically ruled out on the two most relevant ARM64 microarchitectures, the no-dot-product path is correctly preserved, and the benchmark null-deref fix is a nice bonus. Merge when Microsoft-side signoff arrives.

[Copilot]

Pull request overview

This PR improves ARM64 MLAS quantized GEMM performance by enabling SIMD dispatch for the A=int8 / B=uint8 (S8U8) case, routing it through the existing UDOT kernel path instead of falling back to the scalar reference implementation. It also expands the MLAS QGEMM benchmark coverage for signedness combinations and prevents a packed-B benchmark crash when packing is unsupported by the active dispatch.

Changes:

• Add a new ARM64 platform dispatch slot for S8U8 QGEMM and route ARM64 S8U8 to the UDOT dispatch when dot-product is available.
• Extend the UDOT QGEMM kernel’s A-side fixup/packing to correctly handle signed A by XOR sign-bit flipping and matching zero-point adjustment.
• Update bench_qgemm to parameterize B signedness, add S8U8 coverage, and guard packed-B benchmarks when packing is unsupported.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
onnxruntime/test/mlas/bench/bench_qgemm.cpp	Parameterizes QGEMM benchmark by A/B signedness; adds S8U8 cases and a packed-B unsupported guard.
onnxruntime/core/mlas/lib/qgemm.h	Updates ARM64 dispatch selection to use the new S8U8 dispatch for AIsSigned && !BIsSigned.
onnxruntime/core/mlas/lib/qgemm_kernel_udot.cpp	Adds signed-A handling for UDOT via A zero-point fixup and XOR sign-bit flipping during A packing while preserving padding semantics.
onnxruntime/core/mlas/lib/platform.cpp	Initializes and conditionally overrides ARM64 GemmS8U8Dispatch to UDOT under the existing dot-product capability check.
onnxruntime/core/mlas/lib/mlasi.h	Adds GemmS8U8Dispatch to the ARM64 MLAS_PLATFORM layout.

[kjg0724]

Pushed 12e6e89 for the S8U8 test coverage point — turned out MatrixGuardBuffer's generic fill only ever produces bytes in [21,63], so the existing S8U8 tests never actually put A at the signed extremes (-128/-1/0/127) the XOR 0x80 flip has to handle. This pins those values across the K-padding and M row-block boundaries, 840 cases, all passing alongside the existing suite.

Also added a note to the description about native Windows ARM64 picking up the same routing via the MLAS_TARGET_ARM64 guard, and fixed the SkipWithError/SkipWithMessage nit Copilot caught on the benchmark guard.

[hariharans29]

Re-review: PR #29787 — [MLAS] Route ARM64 S8U8 QGEMM to the UDOT kernel (head 12e6e89)

Three commits since pass-1 (was cdf7b51, now 12e6e89). Delta summary:

• e14074a — bench-guard changed from SkipWithError to SkipWithMessage (Copilot review nit). Correct — the "packing not supported for this A/B combo" case is a runtime capability skip, not a failure, so SkipWithMessage is the right benchmark disposition.
• 12e6e89 — new QgemmS8U8SignedInputTest fixture in test_qgemm_fixture.h with explicit signed-A boundary coverage. Registered in both single-thread and threaded variants from test_qgemm.cpp.
• PR description now explicitly notes native Windows ARM64 (MLAS_TARGET_ARM64, MSVC) picks up the same S8U8→UDOT routing, matching pass-1 scope observation Set up CI with Azure Pipelines #1. ARM64EC and 32-bit ARM stay on their existing logic.

CI: 86/86 checks OK on 12e6e89. All green.

Merge status: "At least 1 approving review is required to merge." Copilot AI review + resolved. Waiting on human approver.

New test analysis

The boundary-coverage commit is genuinely important, and turned out to catch a latent hole in the pre-existing suite that my pass-1 review incorrectly waved away. Author diagnosis (verbatim in the PR thread):

"MatrixGuardBuffer's generic fill only ever produces bytes in [21,63], so the existing S8U8 tests never actually put A at the signed extremes (-128/-1/0/127) the XOR 0x80 flip has to handle."

Reading the fixture — this is correct. Every byte in [21,63] reinterpreted as int8 is a non-negative small positive; the sign-bit-flip path was completely unexercised for negative-A. Post-fix, the fixture pins:

• A values: -128, -1, 0, 1, 127, -64, 64, -33 — covers INT8_MIN, negative-one (all-bits-set), zero, positive-one, INT8_MAX, and mid-range signeds on both sides. Every representative bit pattern the ^ 0x80 shift must handle.
• B values: 0, 1, 127, 128, 255, 63, 200, 17 — covers UINT8_MIN/MAX, both sides of the signed/unsigned boundary at 127/128, and a few mid-range unsigneds.
• K ∈ {1, 3, 4, 5, 7, 8, 9, 15, 16, 17} — exercises K < 4 (byte-tail only), K = 4 (word boundary), K = 8, K = 16 (vector boundary), and K = 17 (past-vector remainder). This hits every K-padding path across the four row blocks.
• M ∈ {1, 2, 3, 4, 7, 8, 9} — dispatches through every row-count block (8/4/2/1) and their boundaries.
• offa ∈ {0, 128, 255} and offb ∈ {0, 255} — spans the zero-point range at endpoints and the signed/unsigned crossover midpoint.

Total: 10 × 7 × 3 × 2 = 420 cases per thread mode × 2 = 840 cases. Matches author's count. All pass on the reported test hardware.

Coverage note (not a blocker)

The new fixture uses MlasQgemmTest<int8_t, uint8_t, int32_t, false /*Packed*/, Threaded> — it covers NoPackB only. The packed-B S8U8 path relies on MlasGemmQuantCopyPackA<UDOT> (A-side packing) for the sign flip, which is byte-for-byte the same code path regardless of whether B is packed or not, so the signed-A boundary logic is fully covered by the NoPackB fixture. The pre-existing QgemmShortExecuteTest<int8_t, uint8_t, int32_t, true, false> still handles packed-B correctness broadly (just without signed-A extremes). Reasonable scope call; the packed-B and signed-A dimensions are orthogonal here.

If we wanted defense-in-depth, adding a QgemmS8U8SignedInputTest<Threaded> variant with Packed=true (gated on MlasGemmPackBSize(...) > 0 like the existing packed tests) would cost ~20 lines and would exercise the same 840 cases through the packed-B pipeline. Genuinely a nice-to-have, not a merge blocker.

Verdict

LGTM, approve. Pass-1 verdict stands, and both Copilot and pass-1 review comments have been addressed. Notably, the boundary-test commit uncovered and fixed a real gap in the existing suite that my pass-1 review misjudged — good catch by the author.

Merge when a Microsoft-side approver signs off.