diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index d55a4d49fb455..319c897b8349c 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -112,6 +112,15 @@ function(setup_mlas_source_for_windows) ${MLAS_SRC_DIR}/qnbitgemm_kernel_neon.cpp ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_fp32.cpp ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp + # Portable W2 pack helpers + scalar reference kernel. Misleadingly + # Avx512-named because the AVX-512 W2 path was the first consumer, but + # the TU contains no x86 intrinsics (see sqnbitgemm_kernel_avx512_2bit.cpp). + # ARM64 W2 dispatch reuses these for pack-size / pack / layout; the + # native compute kernel is the NEON DotProd TU listed just below. + ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.h + ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.cpp + # W2 CompInt8 DotProd kernel (NEON FEAT_DotProd backend). + ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8_2bit.cpp ${MLAS_SRC_DIR}/cast_kernel_neon.cpp ${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16.cpp ${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16_8bit.cpp @@ -517,6 +526,15 @@ else() ${MLAS_SRC_DIR}/qnbitgemm_kernel_neon.cpp ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_fp32.cpp ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp + # Portable W2 scalar pack / reference kernel. See ARM64 (Windows) branch + # above for the rationale; this is the matching entry for non-Windows + # ARM64 builds (Linux, macOS). + ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.h + ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.cpp + # W2 CompInt8 DotProd kernel (NEON FEAT_DotProd backend). Compiled + # with -march=armv8.2-a+dotprod on this branch -- see flag override + # further below alongside the W4/W8 dotprod TU. + ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8_2bit.cpp ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.h ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.cpp ${MLAS_SRC_DIR}/qkv_quant_kernel.h @@ -551,6 +569,8 @@ else() endif() set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+dotprod") + set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8_2bit.cpp + PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+dotprod") set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8_i8mm.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+i8mm ") diff --git a/docs/Arm64_w2_kernel_future_enhancements.md b/docs/Arm64_w2_kernel_future_enhancements.md new file mode 100644 index 0000000000000..50d0db064a859 --- /dev/null +++ b/docs/Arm64_w2_kernel_future_enhancements.md @@ -0,0 +1,435 @@ +# ARM64 W2 Kernel: Future Performance Enhancements + +This document captures two follow-up perf enhancements for the ARM64 W2 +(`SQNBIT_CompInt8`, BlkBitWidth=2) NEON kernel — the templated +R{1,2}×C{1,4,8} DotProd tile grid that ships in this PR. Both are +deferred to separate PRs because they each touch the packed-B layout or +add a new translation unit, and they are independent of the +correctness-and-tiling work that already shipped. + +This document is the authoritative reference for the current ARM64 W2 +DotProd layout and dispatch; earlier design and handoff drafts have +been folded into it. + +The W2 DotProd kernel as of the templated-grid commit is the recipe +shipped in the table below. Both enhancements below extend it; neither +replaces it. The DotProd TU stays the correctness baseline for every +ARM64 host without FEAT_I8MM and is reused by the I8MM dispatch for R1 +under the new SMMLA layout (see §1.3). + +| Path | A | B in-kernel | Dot | A-sign correction | B-zp correction | +|---|---|---|---|---|---| +| ARM DotProd W4 | int8 | int8 after in-kernel `vsubq_s8(B, zp_b)` | SDOT | none | folded into unpack | +| ARM DotProd W8 | uint8 (+128 offset) | uint8 raw | UDOT | extra 128-correction SGEMM | post-kernel BlkSum SGEMM | +| ARM I8MM W8 | int8 | uint8 raw | USDOT (mixed-sign) | none | post-kernel BlkSum SGEMM | +| **ARM DotProd W2 (current)** | int8 | uint8 weights in [0, 3] (no in-kernel `vsub`) | SDOT | none | fused into accumulator via ABlockSum × QuantBBlkSum | + + +## 1. SMMLA-based W2 I8MM TU + +### 1.1 Motivation + +FEAT_I8MM offers two new int8 instructions on top of DotProd: + +- `USDOT` — unsigned-A × signed-B 4-way dot, **same throughput as SDOT** + on every implementing core. ARM W8 uses it to skip the 128-offset + correction SGEMM that the W8 DotProd path needs. W2 cannot benefit + from USDOT alone because the in-kernel `vsubq_s8(B, 2)` step is + mandatory anyway (unpacking 2-bit weights into int8 lanes). +- `SMMLA` (`vsmmlaq_s32`) — signed 2×8 × 8×2 → 2×2 int32 + matrix-multiply-accumulate. **2× the int8 ops/cycle of SDOT.** This + is the actual reason to ship a W2 I8MM TU. + +This enhancement adds a SMMLA-based TU that targets the R2+ tiles of +the grid where two M-rows are processed together. + +### 1.2 B-pack layout + +SMMLA's right-hand operand is laid out as +`[B_col0[k=0..7], B_col1[k=0..7]]` — two columns interleaved at an +8-K-element boundary, total 16 bytes per SMMLA. The current W2 packed-B +layout (shared with the AVX-512 reference, single-column per int8x16 +vec) is incompatible with this. The I8MM path forks the layout the same +way W8 already does via `Q8BitGemmPackQuantBDataSize<8, IsI8MM>`: + +- DotProd-only host → current 1-col-per-vec pack + DotProd TU (no change). +- I8MM-capable host → new 2-col-interleaved pack at 8-K boundary + new + SMMLA TU. + +`Q2BitGemmPackQuantBDataSize` and `SQ2BitGemmPackQuantBDataAndBlkSum` +in `qnbitgemm_kernel_neon.cpp` become I8MM-aware (pick layout based on +`InitializeWithI8MMSupport` at dispatch construction). + +### 1.3 R1 viability under the SMMLA layout + +The concern: the SMMLA-friendly B-pack layout must also work for the R1 +tiles (M=1 batch) that share the same packed buffer on an I8MM host. + +Resolution: R1 uses **SMMLA with the A row replicated into both SMMLA +row slots** (e.g. via `vdupq_lane_s64` after loading the A row). SMMLA +then computes a 2×2 int32 block where the top two int32 lanes are the +two dots we want (one per output column) and the bottom two are +duplicates that get discarded. Throughput is identical to 2× SDOT on +the current DotProd layout — no win, no loss for R1. + +Alternatives considered and rejected: + +- Deinterleave B in-kernel via `vuzp1q_s8` / `vuzp2q_s8` then SDOT — + adds 2 insts per 16 K-bytes of B and ends up slower than R1 on the + DotProd layout. Defeats the point of the new pack. +- Keep both B layouts in memory side-by-side — 2× the packed-B + footprint with no perf benefit (only one layout is ever live on a + given host). + +So the dispatch model is: + +| Host capability | B pack | R1 path | R2 path | +|---|---|---|---| +| DotProd only | 1-col-per-vec | SDOT (current) | SDOT (current) | +| FEAT_I8MM | 2-col-interleaved, 8-K boundary | SMMLA + replicated A row | SMMLA (real 2× win) | + +M=1 workloads on I8MM hosts therefore land at the same speed as DotProd +hosts; the real SMMLA throughput win is M≥2 only. That is expected and +matches the W8 I8MM precedent (R1 on I8MM is no faster than R1 on +DotProd for W8 either). + +### 1.4 Files to add / touch + +New: + +- `onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8_i8mm_2bit.cpp` + — SMMLA W2 TU. Compiled with `-march=armv8.2-a+i8mm` on Linux/macOS + ARM64 (no flag on Windows ARM64). Mirrors the structure of + `sqnbitgemm_kernel_neon_int8_i8mm.cpp` (the W8 I8MM TU). +- Pack helper alongside the existing portable W2 pack in + `sqnbitgemm_kernel_avx512_2bit.{h,cpp}` (TU is portable C++ despite + the name) that produces the 2-col-interleaved-at-8-K layout. + +Touched: + +- `onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp` — make + `Q2BitGemmPackQuantBDataSize`, `SQ2BitGemmPackQuantBDataAndBlkSum`, + and the W2 kernel pointer I8MM-aware. +- `onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h` — forward decl + for `SQ2BitGemmKernel_BlkSum_CompInt8_NeonI8MM`. +- `cmake/onnxruntime_mlas.cmake` — add the new TU to both ARM64 + source lists with the `+i8mm` flag entry (mirrors the existing W8 + I8MM TU entries). +- `onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp` — + add `NeonI8MM_*` test variants gated on `HasArmNeon_I8MM()`. + +### 1.5 Expected gains + +Inner-loop dot throughput per cycle, holding scalar overhead constant: + +- R1: 1× SDOT-equivalent (no change). On par with current DotProd R1. +- R2 C1: ~2× the current DotProd R2 C1 inner loop. +- R2 C4 / R2 C8: same 2× factor at the inner loop; outer-loop scale / + BlkSum gather is shared with the DotProd kernel, so end-to-end speedup + is < 2× and shape-dependent. Microbench across (M=2/4/8, N=4/8/16, + K=512/1024/2048, BlkLen=64) is the right way to quantify before + shipping. + + +## 2. Lane-indexed SDOT for the C8 tile + +### 2.1 Motivation + +The current C8 tile (`Q2Int8GemmRxC_DotProd`) issues +one full `vdotq_s32(acc, A_vec, B_vec)` per (M-row, N-col) pair, which +reloads the A vec for each of the 8 N-cols. `vdotq_laneq_s32` lets one +A vec source four 4-lane dots against four different B vecs by indexing +a lane group of A. For C8 (= 2 lane groups of 4 N-cols), that halves +the A loads in the inner loop: + +``` +// Current C8 inner step (8x full SDOT, 8x A reloads): +acc[m][n] = vdotq_s32(acc[m][n], A[m], B[n]); // n = 0..7 + +// Lane-indexed equivalent (8x lane-SDOT, 2x A loads via lane groups): +A0 = A[m]; // lanes [0..3] feed n=0..3 +acc[m][0] = vdotq_laneq_s32(acc[m][0], B[0], A0, 0); +acc[m][1] = vdotq_laneq_s32(acc[m][1], B[1], A0, 1); +acc[m][2] = vdotq_laneq_s32(acc[m][2], B[2], A0, 2); +acc[m][3] = vdotq_laneq_s32(acc[m][3], B[3], A0, 3); +// lanes [4..7] of A0 reused for n=4..7 (one more vec load on A) +``` + +The exact win depends on the M-row count of the tile (R1 vs R2 share +the A loads, so R2 C8 amortises further), but on cores that don't +forward the A-vec load across cycles it's a measurable inner-loop win. + +### 2.2 B-pack layout change (C8 only) + +`vdotq_laneq_s32` requires the four B vecs that share a single A vec to +be laid out so that their `k` dimensions align lane-for-lane with the +A vec. In practice that means **4-N-interleaved B-pack for the C8 tile +only** — N-cols 0..3 of a K-block are stored as four adjacent int8x16 +vecs in B; the next four cols 4..7 are the next four vecs; etc. + +C1 and C4 keep the current 1-col-per-vec layout (no benefit to changing +those — C4 already uses one A vec per 4-N-col group via 4× full SDOT +and a lane-indexed C4 would not reduce A loads). + +So this is the first time we'd have a tile-specific B-pack within the +same W2 TU. Options: + +a. **Tile-conditional pack at prepack time** — emit both layouts; C8 + path consumes the interleaved buffer, C1/C4 paths consume the + current buffer. Doubles W2 packed-B memory. Rejected unless C8 wins + are large enough to justify. +b. **Single interleaved-layout pack, C1/C4 deinterleave in-kernel via + `vuzp1q_s8` / `vuzp2q_s8`** — same packed-B size, C1/C4 take a small + inner-loop hit. Acceptable if the hit is small relative to the C8 + win on a balanced (M, N) workload. +c. **Single interleaved-layout pack, C1/C4 read with strided gather** — + uses `vld4q_s8` to gather 4 strided lanes per load. Same packed-B + size, C1/C4 load cost similar to (b), no shuffle insts. + +Recommended starting point: **option (c)**, with microbench validation +on a (C1, C4, C8) mix workload before committing. If the C1/C4 hit +exceeds the C8 win on representative shapes, fall back to (a) with the +2× B-pack footprint accepted as the cost. + +### 2.3 Interaction with the SMMLA enhancement (§1) + +The two enhancements compose cleanly: + +- DotProd host: current B-pack OR new C8-interleaved B-pack (depending + on the §2.2 decision). Both R1/R2 use SDOT / lane-SDOT. +- I8MM host: SMMLA-layout B-pack (2-col-interleaved at 8-K boundary, + §1.2). Lane-indexed SDOT is **not relevant on this host** — SMMLA + already provides the M-row amortisation in a different way. + +So the C8 lane-SDOT work only affects the DotProd TU and the DotProd +pack helper. It does not touch the SMMLA TU. + +### 2.4 Files to touch + +- `onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_2bit.{h,cpp}` — + the portable W2 pack lives here. Add a C8-interleaved layout helper + (or switch the single layout per §2.2). +- `onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8_2bit.cpp` — + change `Q2Int8GemmRxC_DotProd` to use + `vdotq_laneq_s32`. C1 / C4 specialisations unchanged unless + option (b)/(c) is picked, in which case they pick up a small + deinterleave step in their inner loops. +- `onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp` — + the existing 12 R*xC4 tile-coverage shapes plus the C8 entries in + `kSimdShapes*` already exercise the C8 path; the + `PackUnpackRoundTrip` tests cover any pack-layout change. No new + tests needed unless a new public layout helper is exposed. + +### 2.5 Expected gains + +Inner-loop A-load count for a C8 K-step: + +- Current: 1 A load + 8 full SDOTs. +- Lane-indexed: 1 A load + 8 lane-SDOTs. + +The SDOT-vs-lane-SDOT instruction throughput is identical on +implementing cores; the win is purely the A-load reduction in +**multi-row tiles** (R2 C8 reuses the same A vec across both M-rows +under both schemes, so the gain there is the C1/C4 pack-layout overhead +being absorbed by the C8 path, not raw A-load savings). On R1 C8 the +savings are exactly the load reduction. + +Quantify with a microbench at C8-heavy shapes (N ≥ 8, mixed M, K=1024) +before deciding which pack option from §2.2 to ship. + + +## 3. Sequencing + +Both enhancements are independent and can land in either order. Suggested +order: + +1. **SMMLA W2 I8MM TU first.** Larger expected win on R2+ tiles; same + B-pack-fork pattern as W8 I8MM (low review-novelty risk). Validate + with the existing `NeonDotProd_*` test pattern reused as + `NeonI8MM_*`. +2. **Lane-indexed C8 SDOT second.** Smaller, more localised change but + touches the portable W2 pack (review touches AVX-512 file). Land + after §1 so the pack-layout-fork code added by §1 is the reference + for any further fork in §2 option (a). + +Either way, microbench the inner-loop change in isolation before +committing — both enhancements have shape regimes where the win is +small or negative, and a microbench harness around a fixed (M, N, K, +BlkLen) grid is the cheapest way to discover those before shipping. + +## 4. Observed cross-ISA W2 perf asymmetry (context, not a task) + +Reported end-to-end on a QA transformer with `MatMulNBits`, `seq_len` in +{32, 64, 128}, batch=1, `SQNBIT_CompInt8` (`accuracy_level=4`): + +- **ARM64** (Snapdragon X, DotProd, KleidiAI OFF in the build): W2 runs + ~20–40 % **faster** than W4, monotone across `seq_len`. +- **AVX-512** (VNNI host, previously measured): W2 runs ~10–15 % + **slower** than W4. + +Both platforms exercise the same `MatMulNBits` operator over the +same packed-B pipeline; the split is entirely inside the ISA-specific +inner dot-product loop. Recorded here so a future contributor comparing +W2 perf across CI hosts does not read the AVX-512 gap as a regression. + +### 4.1 Root-cause sketch + +Two facts to keep separate: + +**(a) Algorithmic bytes-per-MAC is set by the quant width, not by the +ISA.** For an int8×int8 MAC the algorithm streams 1 B byte per MAC at +W8, 0.5 B byte per MAC at W4, 0.25 B byte per MAC at W2. This is the +same on `VPDPBUSD` and on `SDOT`. Comparing "MACs/cycle" across ISAs +does *not* tell you how bandwidth-hungry the loop is; it tells you the +per-cycle absolute throughput of the dot pipe. + +**(b) Compute-bound vs bandwidth-bound is a *ratio* between the +algorithm's bytes-per-MAC and the ratio the hardware supports between +peak compute (MACs/sec) and peak byte-delivery from the level of the +cache hierarchy where B actually lives.** Wider dot units (Xeon SPR / +EPYC AVX-512 VNNI) ship with proportionally wider L1/L2 pipes and +larger private L2, so B tends to sit in L2 and the dot pipe is the +ceiling. Client ARM SoCs (Snapdragon X Oryon) have narrower dot units +*and* narrower cache pipes, but the two scale down at different rates; +in practice B spills out of L1D into SLC and the delivery pipe becomes +the ceiling for typical `MatMulNBits` shapes. + +W2 vs W4 trades two things against each other: + +- **W2 halves the B bytes streamed** (0.25 vs 0.5 byte/MAC). This shows + up as speedup only in proportion to how much of the runtime was + actually spent waiting on bytes. +- **W2 adds an extra bit-unpack pass** — two shift+mask+`SUB` rounds + per 4 output lanes instead of one — issued into the vector/SIMD pipes + that also issue the dot instruction. This shows up as slowdown only + in proportion to how tight those pipes already are. + +Per ISA the trade lands opposite ways: + +| ISA | int8 dot inst | typical M ≤ 128, K-loop regime | W2 vs W4 outcome | +|---|---|---|---| +| AVX-512 VNNI | `VPDPBUSD` | dot-pipe-limited: SPR/EPYC L1/L2 keep `VPDPBUSD` near its 2/cycle peak while B is L2-resident | small bytes-saving win, unpack uops steal directly from the dot pipe → W2 ~10–15 % slower | +| ARM64 NEON DotProd | `SDOT` | delivery-pipe-limited: on Oryon the L2/SLC pipes cannot keep 2 `SDOT`/cycle fed once B spills out of L1D | halved B translates almost linearly, dot pipe has slack to absorb the extra unpack uops → W2 ~20–40 % faster | + +Two clarifications this table hides: + +1. **NEON unpack does share the SIMD pipes with `SDOT`** — `USHR` / + `SSHR` / `AND` / `SUB` on vector regs are SIMD-issue ops, not scalar + ALU ops, and Oryon does not have a dedicated "unpack pipe" hiding + next to the dot unit. The reason the extra unpack cost does not + surface as a W2 slowdown on ARM64 is not that unpack runs on free + pipes — it is that the loop was not dot-pipe-limited to begin with, + so the SIMD pipes had cycles to spare. On AVX-512 the loop *is* + dot-pipe-limited, so the same class of extra uops pushes against the + ceiling. +2. Snapdragon X is a *client* SoC. The dot-pipe / cache-pipe balance + here is not representative of ARM server parts (Neoverse V2 / V3), + which have wider cache pipes and can look dot-pipe-limited on the + same kernel. The "W2 wins on ARM64" observation is defensible on + Snapdragon X and similarly balanced client parts; extending it to + Graviton3/4 is a separate measurement. + +Amplifier on the AVX-512 side: the AVX-512 W4 kernel already uses a +near-peak `VPDPBUSD` layout, so W4 is essentially dot-pipe-limited and +W2 has an even higher bar to clear. The observed 10–15 % gap is +consistent with the added unpack pressure and needs no other +explanation. + +### 4.2 Is this worth chasing? + +**Not right now.** The theory follows from published int8-dot throughput +numbers, the gap is small, and closing it would require redesigning the +AVX-512 W2 inner loop — most obviously by materialising 8× int8 lanes +into a scratch buffer per K-block ahead of the `VPDPBUSD` stream. That +trades unpack uops for extra memory traffic, i.e. it gives back the very +bandwidth advantage that made W2 attractive on ARM64. ARM64 is the +target platform for W2 shipping today; that is also the ISA where W2's +smaller B footprint actually helps. + +If a future PR does want to close it, two cheap experiments settle the +question without touching kernel code: + +1. On an AVX-512 VNNI host, run + `perf stat -e uops_dispatched_port.port_5,cycle_activity.stalls_l1d_miss,cycles` + over `onnxruntime_mlas_benchmark --benchmark_filter='QNBITGEMM.*2.*'` + and the equivalent 4-bit filter at a fixed (M, N, K). If W2 has + markedly more port-5 (shift/rotate) uops and similar or fewer L1 + miss stalls than W4, the compute-bound story is confirmed. +2. Sweep K upward at fixed M (e.g. `K ∈ {512, 2048, 8192}`). If the + W2/W4 ratio on AVX-512 narrows as K grows, bandwidth pressure is + starting to bite and points at large-K regimes where the redesign + would pay off. If the ratio is flat, the current inner loop has no + bandwidth regime in which W2 wins on AVX-512 and no redesign is + warranted. + +Either result is standalone; neither blocks the ARM64 enhancements in +§1 or §2. + +### 4.3 Alternative explanations to weigh before investing in a rewrite + +The §4.1 compute-vs-bandwidth-ratio story (plus the shared-SIMD-pipe +observation) is the most compact single-cause explanation, but it is +not the only plausible one, and the AVX-512 number by itself is not +tight enough to isolate a single mechanism. Ranked roughly by the +weight I would put on each: + +1. **Kernel-maturity asymmetry (probably the biggest factor).** The + AVX-512 W4 path in ORT is very mature — iterated across multiple + PRs, specialised VPDPBUSD B-layout, patterns copied from known-good + `q4_0` kernels. The AVX-512 W2 path is much younger and has had + far less hand-tuning attention. A 10–15 % gap can be entirely a + "how many perf-tuning passes has each kernel received" gap with no + per-instruction cause. Hard to distinguish from the port-pressure + story from numbers alone. +2. **Scale / zero-point amortisation at smaller `BlkLen`.** W2 usually + ships with a smaller default `BlkLen` (16 or 32) than W4 to keep + quantisation error in range. Each K-block boundary applies scale + + zp to a partial-sum vector; at `BlkLen=16` W2 does 2× as many + boundary applications per K-unit as W4 at `BlkLen=32`. If the + AVX-512 W2 kernel applies scales via scalar FMAs per block instead + of fusing them into a per-N-tile post-loop `BlkSum` gather (which + is what the ARM64 W4 / W8 paths do), that overhead can drown a real + bandwidth win. Source-checkable. +3. **Measurement-quality asymmetry.** The AVX-512 W2 / W4 gap was + collected earlier and separately from the ARM64 run. Same host? + Same thread count? Same `allow_spinning`? Same + `client_package_build`? Same seq_lens? Multiple runs with variance? + A single-run 10–15 % gap on AVX-512 is inside the typical noise + floor without NUMA pinning, turbo lock, and repeated runs. A clean + re-measurement under the same discipline as the ARM64 run might + shrink or eliminate the gap. +4. **AVX-512 W2 may not use `VPDPBUSD`.** If the kernel uses + `VPMADDUBSW + VPADDD` (non-VNNI fallback) or a lane-permute-heavy + sequence instead of pure `VPDPBUSD`, throughput headroom is much + lower than W4's and the port-pressure story from §4.1 does not + apply — the two paths are apples-to-oranges. Easy to eliminate by + reading `onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_2bit.{h,cpp}`. +5. **Register pressure / spills on AVX-512.** W2 unpack needs extra + ZMM registers for shift constants and intermediate masks. With only + 32 ZMMs, spilling one or two accumulators to stack per K-step adds + a small hidden load-forwarding cost. On ARM64 with 32 128-bit + V-regs the same working set fits without spills. A couple-percent + effect at most, not a headline one. + +Weaker candidates (real effects, but unlikely at this magnitude in this +regime): µop-cache / DSB pressure from a longer W2 inner loop, +prefetcher unfriendliness of a non-sequential W2 B-layout stride, +AVX-512 frequency-license downclocking on older SKUs, dependency-chain +length changes if W2 uses fewer accumulators. + +**Cheapest-first order for a future investigation:** + +1. Re-run the AVX-512 measurement under the same discipline as the ARM64 + run (item 3). If the gap narrows to < 5 %, stop. +2. Grep `sqnbitgemm_kernel_avx512_2bit.{h,cpp}` to confirm `VPDPBUSD` + is used and that the B-layout matches the AVX-512 W4 kernel's layout + in spirit (item 4). If it does not, that is the first thing to fix, + independent of any port-pressure work. +3. Check whether AVX-512 W2 applies scales per K-block or via a fused + post-loop `BlkSum` (item 2). If per-block, port the ARM64 pattern + and re-measure. +4. Only then run the `perf stat` port-pressure experiment from §4.2 + and the K-sweep. By that point items 1–3 have already ruled out the + easy explanations. + +None of this blocks the ARM64 enhancements in §1 or §2. diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc index 162d7257d0a4c..42dd26c20c98a 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc @@ -420,7 +420,14 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All #if defined(MLAS_TARGET_AMD64_IX86) return true; #else - return (nbits_ == 8); + // W2 on ARM64 (NEON DotProd) uses the same 3-call prepack pattern as + // AVX-512: B in the input_idx==B call, scales+ZP in the separate + // input_idx==scales / input_idx==zero_points calls. Without this, the + // ZP tensor is silently dropped and QuantBBlkSum bakes in the symmetric + // default ZP=2 for every block, producing wrong outputs whenever the + // model's real ZPs ≠ 2. W4 stays out because it goes through the + // KleidiAI scale-baked path above and does not use the 3-call pattern. + return (nbits_ == 2 || nbits_ == 8); #endif }(); @@ -549,12 +556,14 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All } } - // Pack zero_points separately only for 8-bit (matching standard SQNBIT_CompInt8 behavior). - // For 4-bit, zero_points are passed directly in data params or handled via KleidiAI BZpCorr. - if (input_idx == InputIndex::zero_points && packed_b_ != nullptr && nbits_ == 8 && !packed_b_finalized_) { + // Pack zero_points separately for W2; W4/W8 are folded into packed_b_ during the B PrePack. + // Pass scales_fp32_ so QuantBBlkSum = -scale*zp is recomputed (otherwise it keeps the symmetric + // default ZP=2 from the B-pack call and silently produces wrong outputs when real ZPs != 2). + if (input_idx == InputIndex::zero_points && packed_b_ != nullptr && !packed_b_finalized_ && nbits_ == 2) { auto zptr = tensor.Data(); - MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, SQNBIT_CompInt8, nullptr, packed_b_.get(), nullptr, - has_zp_input_, zptr, nullptr, &mlas_backend_kernel_selector_config_); + MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, SQNBIT_CompInt8, nullptr, packed_b_.get(), + scales_fp32_.get(), has_zp_input_, zptr, nullptr, + &mlas_backend_kernel_selector_config_); is_packed = false; } @@ -706,7 +715,7 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*ou #if defined(MLAS_TARGET_AMD64_IX86) return true; #else - return (nbits_ == 8); + return (nbits_ == 2 || nbits_ == 8); #endif }(); diff --git a/onnxruntime/core/mlas/lib/qnbitgemm.cpp b/onnxruntime/core/mlas/lib/qnbitgemm.cpp index 1337b798d3a8c..441cd7f0d3c70 100644 --- a/onnxruntime/core/mlas/lib/qnbitgemm.cpp +++ b/onnxruntime/core/mlas/lib/qnbitgemm.cpp @@ -1153,9 +1153,24 @@ InitializeWorkspace_CompInt8( } }); } else { - // TODO(hasesh): Clean-up the following logic so that it is clean AND it works as expected on all platforms + // TODO(hasesh): The (BlkBitWidth x A-layout x A-signedness) matrix below is + // resolved by a hand-rolled cascade because the dispatch exposes both an + // interleaved-scale (W4-style) and a separate-scale (W8-style) quantize fn, + // plus a W2-specific signed-A override. Ideally the dispatch itself would + // expose a single "correct quantize fn for this (bit-width, kernel)" pointer + // so this call-site would not have to reason about layout/sign compatibility. if (BlkBitWidth == 4 || BlkBitWidth == 2) { - if (QuantizeARow) { + // W2 requires the W8-style separate-scale layout produced by + // QuantizeARowComputeBlkSum_CompInt8 because the W2 SQ2BitGemm + // kernel reads QuantData (M*K int8 flat) and QuantScale + // (M*BlockCountK float) as independent buffers. The W4-style + // interleaved layout produced by QuantizeARow_CompInt8 packs the + // scale inside each Q8Blk and is therefore incompatible. On hosts + // that register both (e.g. NEON FEAT_DotProd, which uses + // QuantizeARow for the W4 path and QuantizeARowComputeBlkSum for + // W8), force W2 down the W8-compatible branch. + const bool prefer_compute_blksum = (BlkBitWidth == 2); + if (QuantizeARow && !prefer_compute_blksum) { MlasTrySimpleParallel(ThreadPool, BatchN, [&](ptrdiff_t gemm_idx) { const auto& data = DataParams[gemm_idx]; @@ -1169,6 +1184,14 @@ InitializeWorkspace_CompInt8( } }); } else if (QuantizeARow2) { + // W2 needs SIGNED int8 A (NEON DotProd-only hosts wire the + // shared QuantizeARowComputeBlkSum_CompInt8 to the UNSIGNED + // u8 = i8+128 W8-compatible variant). Prefer the W2-specific + // override when set so the W2 kernel always sees signed A. + const auto QuantizeARow2_W2 = + (BlkBitWidth == 2 && GetMlasPlatform().QNBitGemmDispatch->QuantizeARowComputeBlkSum_CompInt8_W2 != nullptr) + ? GetMlasPlatform().QNBitGemmDispatch->QuantizeARowComputeBlkSum_CompInt8_W2 + : QuantizeARow2; MlasTrySimpleParallel(ThreadPool, BatchN, [&](ptrdiff_t gemm_idx) { const auto& data = DataParams[gemm_idx]; const float* ARowPtr = data.A; @@ -1179,7 +1202,7 @@ InitializeWorkspace_CompInt8( float* QuantARowScalePtr = quant_a_data.QuantScale; float* QuantARowBlkSum = quant_a_data.BlockSum; for (size_t m = 0; m < M; ++m) { - QuantizeARow2(BlkLen, ARowPtr, K, QuantARowPtr, QuantARowScalePtr, QuantARowBlkSum); + QuantizeARow2_W2(BlkLen, ARowPtr, K, QuantARowPtr, QuantARowScalePtr, QuantARowBlkSum); ARowPtr += data.lda; QuantARowPtr += BlockCountK * BlkLen; QuantARowScalePtr += BlockCountK; diff --git a/onnxruntime/core/mlas/lib/qnbitgemm.h b/onnxruntime/core/mlas/lib/qnbitgemm.h index b9c88ab1c7ede..441da585c6611 100644 --- a/onnxruntime/core/mlas/lib/qnbitgemm.h +++ b/onnxruntime/core/mlas/lib/qnbitgemm.h @@ -679,6 +679,13 @@ struct MLAS_QNBIT_GEMM_DISPATCH { ); QuantizeARowComputeBlkSum_CompInt8_Fn* QuantizeARowComputeBlkSum_CompInt8 = nullptr; + // Optional W2-specific override of QuantizeARowComputeBlkSum_CompInt8. + // The W2 NEON DotProd kernel assumes SIGNED int8 A, but on DotProd-only hosts + // the shared field is wired to the UNSIGNED (u8 = i8+128) W8-compatible variant. + // When set, the W2 dispatch in InitializeWorkspace_CompInt8 uses this in + // preference to the shared field; otherwise it falls back to the shared one. + QuantizeARowComputeBlkSum_CompInt8_Fn* QuantizeARowComputeBlkSum_CompInt8_W2 = nullptr; + /** * @brief Multiply fp16 matrix A rows with fp16 matrix B columns. * Results are written to fp16 matrix C. diff --git a/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp b/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp index 5a3c8005d8318..4ff9b6f1a4d6f 100644 --- a/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp +++ b/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp @@ -27,6 +27,13 @@ Module Name: #include "qnbitgemm.h" #include "sqnbitgemm_q8_block.h" +// W2 block-group pack helpers and scalar reference kernel declared in the +// Avx512-named header are pure C++ (no x86 intrinsics). They serve as the +// cross-arch layout authority for W2 and are reused on ARM64 for pack-size / +// pack / layout; compute is provided by the native NEON DotProd kernel in +// sqnbitgemm_kernel_neon_int8_2bit.cpp. +#include "sqnbitgemm_kernel_avx512_2bit.h" + #ifdef USE_KLEIDIAI #include "kai/kai_common.h" #include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0.h" @@ -655,6 +662,47 @@ GetMlasQNBitGemmDispatchNeon( d.SQ8BitGemmKernel_BlkSum_CompInt8 = sqnbitgemm_neon::SQ8BitGemmKernel_BlkSum_CompInt8; } + // W2 native CompInt8 path. + // + // Pack-size, pack-and-blksum, and EffectiveBlockCountK use the + // portable AVX-512-namespaced helpers (the file is misleadingly + // named -- the TU contains no x86 intrinsics) which serve as the + // cross-arch layout authority for W2. + // + // SQ2BitGemmKernel_BlkSum_CompInt8 is wired to the native NEON + // DotProd kernel when FEAT_DotProd is available; the kernel + // handles BlkLen ∈ {32, 64, 128} natively via SDOT inner loops. + // The 2-bit weights are unpacked to unsigned values in [0, 3] and + // fed directly to SDOT; the B zero-point correction is fused into + // the accumulator via the ABlockSum × QuantBBlkSum term (where + // QuantBBlkSum bakes in -scale × zp per block), so no post-kernel + // zero-point correction SGEMM is needed. + // + // FEAT_I8MM hosts also take this path: FEAT_I8MM always implies + // FEAT_DotProd per the ARM spec, and USDOT and SDOT have identical + // throughput on every core that implements both -- a separate I8MM + // TU using USDOT would be pure duplication. The real I8MM-only + // throughput win lives in SMMLA (R2-tile 2×2 matrix-multiply, + // 2× the dots/cycle of SDOT), which is future work. + // + // The Scalar fall-back is defensive: an I8MM-without-DotProd host + // is unreachable on conformant ARMv8.2+ hardware. + if (InitializeWithDotSupport || InitializeWithI8MMSupport) { + d.Q2BitGemmPackQuantBDataSize = onnxruntime::mlas::sq2bit_avx512::Q2BitGemmPackQuantBDataSize_Avx512; + d.SQ2BitGemmPackQuantBDataAndBlkSum = onnxruntime::mlas::sq2bit_avx512::SQ2BitGemmPackQuantBDataAndBlkSum_Scalar; + d.SQ2BitGemmKernel_BlkSum_CompInt8 = InitializeWithDotSupport + ? sqnbitgemm_neon::SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd + : onnxruntime::mlas::sq2bit_avx512::SQ2BitGemmKernel_BlkSum_CompInt8_Scalar; + d.Q2BitGemmEffectiveBlockCountK = [](size_t BlockCountK) { + return MlasDivRoundup(BlockCountK, kSq2BitAvx512WeightKBlockGroup) * kSq2BitAvx512WeightKBlockGroup; + }; + // W2 NEON DotProd kernel uses vdotq_s32 over A and so requires + // SIGNED int8 A. The shared QuantizeARowComputeBlkSum is wired + // to the UNSIGNED W8 variant on DotProd-only hosts; route W2 + // through the signed variant explicitly. + d.QuantizeARowComputeBlkSum_CompInt8_W2 = sqnbitgemm_neon::QuantizeARowComputeBlkSum_CompInt8; + } + #if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64) d.HQ4BitGemmPackQuantBData = sqnbitgemm_neon::HQ4BitGemmPackQuantBData_CompFp16; d.HQ4BitBlkDequantBForHgemm_CompFp16 = sqnbitgemm_neon::HQ4BitBlkDequantBForHgemm_CompFp16; diff --git a/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h b/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h index 2c31017732a0a..b211bd06ee477 100644 --- a/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h +++ b/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h @@ -195,6 +195,28 @@ SQ4BitGemmKernel_CompInt8( const float* Bias ); +// W2 CompInt8 kernel entry point (DotProd backend). Implemented in +// sqnbitgemm_kernel_neon_int8_2bit.cpp. Signature matches the +// SQ4BitGemmKernel_BlkSum_CompInt8_Fn typedef in qnbitgemm.h. +size_t +SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd( + size_t BlkLen, + const std::byte* QuantA, + const float* QuantAScale, + const std::byte* QuantBData, + const float* QuantBScale, + const std::byte* QuantBZeroPoint, + float* C, + size_t CountM, + size_t CountN, + size_t CountK, + size_t BlockCountK, + const float* Bias, + size_t ldc, + const float* ABlockSum, + const float* QuantBBlkSum +); + #ifdef USE_KLEIDIAI void QuantizeA_Packed_CompInt8( diff --git a/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8_2bit.cpp b/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8_2bit.cpp new file mode 100644 index 0000000000000..552e2c12f404c --- /dev/null +++ b/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8_2bit.cpp @@ -0,0 +1,535 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. +Licensed under the MIT License. + +Module Name: + + sqnbitgemm_kernel_neon_int8_2bit.cpp + +Abstract: + + ARM NEON (FEAT_DotProd) implementation of the W2 SQNBIT_CompInt8 inner + kernel. + + This TU consumes the block-group packed B-data layout produced by the + portable W2 pack helpers in `sqnbitgemm_kernel_avx512_2bit.{h,cpp}`. The + file naming "avx512_2bit" is historical: those pack helpers are pure + C++ and serve as the cross-arch layout authority for W2. + + Scope: + * `BlkLen == 32`, `BlkLen == 64`, `BlkLen == 128`: native NEON dotprod + (SDOT) inner loops. The packed B layout is the same 4-K-block group + with bit positions {0..1, 2..3, 4..5, 6..7} per byte for all three + BlkLens (only the bytes-per-K-block scales with BlkLen). One block- + group is loaded as 2 / 4 / 8 int8x16 chunks respectively, the in- + block-group selector is applied with one shift + mask per 16-byte + chunk, and the per-block int8 dot is computed with 2 / 4 / 8 SDOTs. + + Tile grid (per BlkLen): + R2xC8 (M-pair, main 8-N), R1xC8 (M-tail, main 8-N), + R2xC4 (M-pair, main 4-N), R1xC4 (M-tail, main 4-N), + R2xC1 (M-pair, 1..3-N tail), R1xC1 (M-tail, 1..3-N tail). + All six variants are template instantiations of a single tile kernel + (`Q2Int8GemmRxC_DotProd`). The per-BlkLen + dispatcher splits CountN into 8-aligned, 4-aligned, and 1..3 tail + regions and CountM into M-pairs + odd-M-tail, then dispatches to up + to six tile-shape kernels. + + Deviation vs. W4/W8 DotProd ARM reference: W4/W8 use R1-only on + DotProd (R2 reserved for I8MM). W2 adds R2 to DotProd too because + W2's B-unpack (8 weights/byte, 2 shifts+masks per chunk) is the + heaviest of any width and amortising it + the per-K-block scale / + BlkSum gather across an M-pair is a meaningful win. + +--*/ + +#include + +#include +#include +#include + +#include "qnbitgemm.h" +#include "qnbitgemm_kernel_neon.h" +#include "sqnbitgemm_kernel_avx512_2bit.h" + +namespace sqnbitgemm_neon +{ + +namespace +{ + +namespace sq2 = onnxruntime::mlas::sq2bit_avx512; + +// ============================================================================ +// Shared templated tile machinery for the W2 DotProd kernel. +// +// One templated tile kernel `Q2Int8GemmRxC_DotProd` +// covers all 18 BlkLen x tile combinations. The per-BlkLen public entries +// each delegate to `SQ2BitGemm_Dispatch_NeonDotProd`, which splits +// the (M, N) plane into up to 6 sub-rectangles (R{1,2} x C{1,4,8}) and +// dispatches one instantiation per piece. +// +// Layout (see sqnbitgemm_kernel_avx512_2bit.h): +// * Packed B is grouped into `kNCols4` (=4) col groups in the main region +// (n < NMain = floor(CountN / kNCols4) * kNCols4) and column-major in +// the tail region (n >= NMain). C8/C4 tiles operate entirely in the +// main region; the C1 tile handles 1..3 trailing cols in the tail. +// * `BlockCountK` is padded to a multiple of `kBlockGroupBlks` (=4) for +// storage; padding contributes 0 to dot products and to the BlkSum +// correction. +// * QuantBBlkSum uses the existing width-16 chunked layout consumed by +// the SGEMM correction step. +// ============================================================================ + +template struct BlkLenTraits; +template <> struct BlkLenTraits<32> { static constexpr size_t kVecsPerBlk = 2; }; +template <> struct BlkLenTraits<64> { static constexpr size_t kVecsPerBlk = 4; }; +template <> struct BlkLenTraits<128> { static constexpr size_t kVecsPerBlk = 8; }; + +template inline constexpr size_t kBlockGroupBytesV = 0; +template <> inline constexpr size_t kBlockGroupBytesV<32> = sq2::kBlockGroupBytes32; +template <> inline constexpr size_t kBlockGroupBytesV<64> = sq2::kBlockGroupBytes; +template <> inline constexpr size_t kBlockGroupBytesV<128> = sq2::kBlockGroupBytes128; + +// Per-BlkLen packed-B byte offset (folds the BlkLen-specific overloads). +template +MLAS_FORCEINLINE size_t +PackedBOffsetBytes(size_t n, size_t blk_group, + size_t BlockGroupCountKPadded, size_t NMain) +{ + if constexpr (BlkLen == 32) { + return sq2::PackedQuantBOffsetBytes_W2_BlkLen32( + n, blk_group, BlockGroupCountKPadded, NMain); + } else if constexpr (BlkLen == 128) { + return sq2::PackedQuantBOffsetBytes_W2_BlkLen128( + n, blk_group, BlockGroupCountKPadded, NMain); + } else { + return sq2::PackedQuantBOffsetBytes_W2( + n, blk_group, BlockGroupCountKPadded, NMain); + } +} + +// Unpack one in-group block (BlkInGroup in [0,4)) into `kVecsPerBlk` +// uint8x16_t vectors of 2-bit weights in [0, 3]. +template +MLAS_FORCEINLINE void +UnpackBlockGroupSlice_DotProd(const std::byte* group, + uint8x16_t out[BlkLenTraits::kVecsPerBlk]) +{ + static_assert(BlkInGroup < sq2::kBlockGroupBlks, "BlkInGroup must be < 4"); + constexpr size_t kVecs = BlkLenTraits::kVecsPerBlk; + const uint8x16_t mask03 = vdupq_n_u8(0x03); + const uint8_t* p = reinterpret_cast(group); + for (size_t v = 0; v < kVecs; ++v) { + const uint8x16_t raw = vld1q_u8(p + v * 16); + if constexpr (BlkInGroup == 0) { + // vshrq_n_u8 immediate must be in [1, 8] -- skip shift here. + out[v] = vandq_u8(raw, mask03); + } else { + constexpr int kShift = 2 * static_cast(BlkInGroup); + out[v] = vandq_u8(vshrq_n_u8(raw, kShift), mask03); + } + } +} + +// One in-group block's contribution to NRows x NCols scalar accumulators. +// All A vecs for the K-block are pre-loaded once per row, then for each +// output col we unpack one col's B and run NRows int dot+reduce paths. +// Worst-case live NEON regs (R2xC*, BlkLen=128): 16 A + 8 B + 1 scratch = 25. +template +MLAS_FORCEINLINE void +AccumOneBlock_RxC_DotProd( + float (&acc)[NRows][NCols], + const std::byte* const group_cols[NCols], + const int8_t* const a_blk_bases[NRows], + const float* const a_scale_rows[NRows], + const float* const a_blksum_rows[NRows], + const float b_scales[NCols][sq2::kBlockGroupBlks], + const float b_blksums[NCols][sq2::kBlockGroupBlks], + size_t blk0) +{ + constexpr size_t kVecs = BlkLenTraits::kVecsPerBlk; + constexpr size_t kK = BlkLen; + + int8x16_t av[NRows][kVecs]; + for (size_t r = 0; r < NRows; ++r) { + const int8_t* a = a_blk_bases[r] + BlkInGroup * kK; + for (size_t v = 0; v < kVecs; ++v) { + av[r][v] = vld1q_s8(a + v * 16); + } + } + + for (size_t c = 0; c < NCols; ++c) { + uint8x16_t bw[kVecs]; + UnpackBlockGroupSlice_DotProd(group_cols[c], bw); + for (size_t r = 0; r < NRows; ++r) { + int32x4_t s = vdupq_n_s32(0); + for (size_t v = 0; v < kVecs; ++v) { + s = vdotq_s32(s, av[r][v], vreinterpretq_s8_u8(bw[v])); + } + const int32_t dot = vaddvq_s32(s); + const float a_s = a_scale_rows[r][blk0 + BlkInGroup]; + const float a_bs = a_blksum_rows[r][blk0 + BlkInGroup]; + acc[r][c] += a_s * b_scales[c][BlkInGroup] * static_cast(dot); + acc[r][c] += a_bs * b_blksums[c][BlkInGroup]; + } + } +} + +// Full block-group: 4 in-group blocks unrolled at compile time. +template +MLAS_FORCEINLINE void +AccumOneBlockGroup_RxC_DotProd( + float (&acc)[NRows][NCols], + const std::byte* const group_cols[NCols], + const int8_t* const a_blk_bases[NRows], + const float* const a_scale_rows[NRows], + const float* const a_blksum_rows[NRows], + const float b_scales[NCols][sq2::kBlockGroupBlks], + const float b_blksums[NCols][sq2::kBlockGroupBlks], + size_t blk0) +{ + AccumOneBlock_RxC_DotProd( + acc, group_cols, a_blk_bases, a_scale_rows, a_blksum_rows, b_scales, b_blksums, blk0); + AccumOneBlock_RxC_DotProd( + acc, group_cols, a_blk_bases, a_scale_rows, a_blksum_rows, b_scales, b_blksums, blk0); + AccumOneBlock_RxC_DotProd( + acc, group_cols, a_blk_bases, a_scale_rows, a_blksum_rows, b_scales, b_blksums, blk0); + AccumOneBlock_RxC_DotProd( + acc, group_cols, a_blk_bases, a_scale_rows, a_blksum_rows, b_scales, b_blksums, blk0); +} + +// K-tail block-group: only `blocks_in_tail` (1..3) in-group blocks are +// real K-blocks (pack helper zero-padded the rest). We still must not +// read A past BlockCountK, so the run-time gate stays. +template +MLAS_FORCEINLINE void +AccumLastBlockGroup_RxC_DotProd( + float (&acc)[NRows][NCols], + const std::byte* const group_cols[NCols], + const int8_t* const a_blk_bases[NRows], + const float* const a_scale_rows[NRows], + const float* const a_blksum_rows[NRows], + const float b_scales[NCols][sq2::kBlockGroupBlks], + const float b_blksums[NCols][sq2::kBlockGroupBlks], + size_t blk0, size_t blocks_in_tail) +{ + assert(blocks_in_tail >= 1 && blocks_in_tail <= sq2::kBlockGroupBlks); + if (blocks_in_tail >= 1) { + AccumOneBlock_RxC_DotProd( + acc, group_cols, a_blk_bases, a_scale_rows, a_blksum_rows, + b_scales, b_blksums, blk0); + } + if (blocks_in_tail >= 2) { + AccumOneBlock_RxC_DotProd( + acc, group_cols, a_blk_bases, a_scale_rows, a_blksum_rows, + b_scales, b_blksums, blk0); + } + if (blocks_in_tail >= 3) { + AccumOneBlock_RxC_DotProd( + acc, group_cols, a_blk_bases, a_scale_rows, a_blksum_rows, + b_scales, b_blksums, blk0); + } + // blocks_in_tail == 4 is the main-loop case, not invoked here. +} + +// Templated tile kernel for the sub-rectangle +// [m_start, m_start+m_count) x [n_start, n_start+n_count) +// Caller guarantees: m_count is a multiple of NRows; n_count is a multiple +// of NCols (for NCols==1 this is trivial -- the C1 tile iterates the 1..3 +// trailing cols one at a time). +template +MLAS_FORCEINLINE void +Q2Int8GemmRxC_DotProd( + const std::byte* QuantA, + const float* QuantAScale, + const std::byte* QuantBData, + const float* QuantBScale, + float* C, + size_t CountN, // full N (for NMain calc + BlkSum width-16) + size_t BlockCountK, + const float* Bias, + size_t ldc, + const float* ABlockSum, + const float* QuantBBlkSum, + size_t m_start, + size_t m_count, + size_t n_start, + size_t n_count) +{ + if (m_count == 0 || n_count == 0) return; + + constexpr size_t kK = BlkLen; + constexpr size_t kBlockGroupBytes = kBlockGroupBytesV; + // C8/C4 tiles operate fully in the main NCols4-grouped region; C1 tile + // operates fully in the column-major tail region. + constexpr bool kInMain = (NCols > 1); + constexpr size_t kGroupStrideBytes = + kInMain ? (sq2::kNCols4 * kBlockGroupBytes) : kBlockGroupBytes; + + const size_t BlockGroupCountKPadded = + MlasDivRoundup(BlockCountK, sq2::kBlockGroupBlks); + const size_t BlockCountKPadded = BlockGroupCountKPadded * sq2::kBlockGroupBlks; + const size_t NMainLocal = (CountN / sq2::kNCols4) * sq2::kNCols4; + const size_t MainBlockGroups = BlockCountK / sq2::kBlockGroupBlks; + const size_t TailBlocks = BlockCountK % sq2::kBlockGroupBlks; + const size_t lda = BlockCountK * kK; + const size_t lda_scale = BlockCountK; + + assert(m_count % NRows == 0); + + for (size_t m = m_start; m < m_start + m_count; m += NRows) { + const int8_t* a_rows[NRows]; + const float* a_scale_rows[NRows]; + const float* a_blksum_rows[NRows]; + float* c_rows[NRows]; + for (size_t r = 0; r < NRows; ++r) { + a_rows[r] = reinterpret_cast(QuantA + (m + r) * lda); + a_scale_rows[r] = QuantAScale + (m + r) * lda_scale; + a_blksum_rows[r] = ABlockSum + (m + r) * lda_scale; + c_rows[r] = C + (m + r) * ldc; + } + + for (size_t n = n_start; n < n_start + n_count; n += NCols) { + float acc[NRows][NCols]; + for (size_t r = 0; r < NRows; ++r) { + for (size_t c = 0; c < NCols; ++c) { + acc[r][c] = (Bias != nullptr) ? Bias[n + c] : 0.0f; + } + } + + // Per-col g=0 base pointers; per-group advance is `kGroupStrideBytes` + // (constant for this tile shape / region). + const std::byte* group_cols_g0[NCols]; + for (size_t c = 0; c < NCols; ++c) { + group_cols_g0[c] = QuantBData + PackedBOffsetBytes( + n + c, 0, BlockGroupCountKPadded, NMainLocal); + } + + for (size_t g = 0; g < MainBlockGroups; ++g) { + const size_t blk0 = g * sq2::kBlockGroupBlks; + + const std::byte* group_cols[NCols]; + for (size_t c = 0; c < NCols; ++c) { + group_cols[c] = group_cols_g0[c] + g * kGroupStrideBytes; + } + + float b_scales[NCols][sq2::kBlockGroupBlks]; + float b_blksums[NCols][sq2::kBlockGroupBlks]; + for (size_t c = 0; c < NCols; ++c) { + for (size_t i = 0; i < sq2::kBlockGroupBlks; ++i) { + b_scales[c][i] = QuantBScale[sq2::PackedQuantBScaleOffset_W2( + n + c, blk0 + i, BlockCountKPadded, NMainLocal)]; + const size_t off = + (((n + c) / 16) * BlockCountK + (blk0 + i)) * 16 + ((n + c) % 16); + b_blksums[c][i] = QuantBBlkSum[off]; + } + } + + const int8_t* a_blk_bases[NRows]; + for (size_t r = 0; r < NRows; ++r) { + a_blk_bases[r] = a_rows[r] + blk0 * kK; + } + + AccumOneBlockGroup_RxC_DotProd( + acc, group_cols, a_blk_bases, + a_scale_rows, a_blksum_rows, + b_scales, b_blksums, blk0); + } + + if (TailBlocks != 0) { + const size_t g = MainBlockGroups; + const size_t blk0 = g * sq2::kBlockGroupBlks; + + const std::byte* group_cols[NCols]; + for (size_t c = 0; c < NCols; ++c) { + group_cols[c] = group_cols_g0[c] + g * kGroupStrideBytes; + } + + float b_scales[NCols][sq2::kBlockGroupBlks] = {}; + float b_blksums[NCols][sq2::kBlockGroupBlks] = {}; + for (size_t c = 0; c < NCols; ++c) { + for (size_t i = 0; i < TailBlocks; ++i) { + b_scales[c][i] = QuantBScale[sq2::PackedQuantBScaleOffset_W2( + n + c, blk0 + i, BlockCountKPadded, NMainLocal)]; + const size_t off = + (((n + c) / 16) * BlockCountK + (blk0 + i)) * 16 + ((n + c) % 16); + b_blksums[c][i] = QuantBBlkSum[off]; + } + } + + const int8_t* a_blk_bases[NRows]; + for (size_t r = 0; r < NRows; ++r) { + a_blk_bases[r] = a_rows[r] + blk0 * kK; + } + + AccumLastBlockGroup_RxC_DotProd( + acc, group_cols, a_blk_bases, + a_scale_rows, a_blksum_rows, + b_scales, b_blksums, blk0, TailBlocks); + } + + for (size_t r = 0; r < NRows; ++r) { + for (size_t c = 0; c < NCols; ++c) { + c_rows[r][n + c] = acc[r][c]; + } + } + } + } +} + +// Per-BlkLen tile-grid dispatcher. Splits (CountM, CountN) into up to 6 +// sub-rectangles and dispatches one tile-shape instantiation per piece. +template +size_t +SQ2BitGemm_Dispatch_NeonDotProd( + const std::byte* QuantA, const float* QuantAScale, + const std::byte* QuantBData, const float* QuantBScale, + float* C, size_t CountM, size_t CountN, size_t BlockCountK, + const float* Bias, size_t ldc, + const float* ABlockSum, const float* QuantBBlkSum) +{ + if (BlockCountK == 0) return 0; + if (CountM == 0 || CountN == 0) return CountM; + + const size_t NMain8 = (CountN / 8) * 8; + const size_t NRem = CountN - NMain8; + const size_t NMain4 = (NRem / 4) * 4; // 0 or 4 + const size_t NTail = NRem - NMain4; // 0..3 + const size_t M_main = (CountM / 2) * 2; + const size_t M_tail = CountM - M_main; + const size_t n_tail_start = NMain8 + NMain4; + + if (NMain8 > 0) { + if (M_main > 0) + Q2Int8GemmRxC_DotProd( + QuantA, QuantAScale, QuantBData, QuantBScale, C, + CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum, + 0, M_main, 0, NMain8); + if (M_tail > 0) + Q2Int8GemmRxC_DotProd( + QuantA, QuantAScale, QuantBData, QuantBScale, C, + CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum, + M_main, M_tail, 0, NMain8); + } + if (NMain4 > 0) { + if (M_main > 0) + Q2Int8GemmRxC_DotProd( + QuantA, QuantAScale, QuantBData, QuantBScale, C, + CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum, + 0, M_main, NMain8, NMain4); + if (M_tail > 0) + Q2Int8GemmRxC_DotProd( + QuantA, QuantAScale, QuantBData, QuantBScale, C, + CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum, + M_main, M_tail, NMain8, NMain4); + } + if (NTail > 0) { + if (M_main > 0) + Q2Int8GemmRxC_DotProd( + QuantA, QuantAScale, QuantBData, QuantBScale, C, + CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum, + 0, M_main, n_tail_start, NTail); + if (M_tail > 0) + Q2Int8GemmRxC_DotProd( + QuantA, QuantAScale, QuantBData, QuantBScale, C, + CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum, + M_main, M_tail, n_tail_start, NTail); + } + + return CountM; +} + +// Thin per-BlkLen public entry. Pins BlkLen at compile time so the templated +// tile kernel constant-folds the BlkLen-specific layout. +size_t +SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd_BlkLen64( + const std::byte* QuantA, const float* QuantAScale, + const std::byte* QuantBData, const float* QuantBScale, + float* C, size_t CountM, size_t CountN, size_t BlockCountK, + const float* Bias, size_t ldc, + const float* ABlockSum, const float* QuantBBlkSum) +{ + return SQ2BitGemm_Dispatch_NeonDotProd<64>( + QuantA, QuantAScale, QuantBData, QuantBScale, C, CountM, CountN, + BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum); +} + +size_t +SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd_BlkLen128( + const std::byte* QuantA, const float* QuantAScale, + const std::byte* QuantBData, const float* QuantBScale, + float* C, size_t CountM, size_t CountN, size_t BlockCountK, + const float* Bias, size_t ldc, + const float* ABlockSum, const float* QuantBBlkSum) +{ + return SQ2BitGemm_Dispatch_NeonDotProd<128>( + QuantA, QuantAScale, QuantBData, QuantBScale, C, CountM, CountN, + BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum); +} + +size_t +SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd_BlkLen32( + const std::byte* QuantA, const float* QuantAScale, + const std::byte* QuantBData, const float* QuantBScale, + float* C, size_t CountM, size_t CountN, size_t BlockCountK, + const float* Bias, size_t ldc, + const float* ABlockSum, const float* QuantBBlkSum) +{ + return SQ2BitGemm_Dispatch_NeonDotProd<32>( + QuantA, QuantAScale, QuantBData, QuantBScale, C, CountM, CountN, + BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum); +} + +} // unnamed namespace + +// +// W2 CompInt8 kernel entry point (DotProd backend). +// +// All three BlkLens ({32, 64, 128}) route to native NEON dotprod kernels. +// +size_t +SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd( + size_t BlkLen, + const std::byte* QuantA, + const float* QuantAScale, + const std::byte* QuantBData, + const float* QuantBScale, + const std::byte* QuantBZeroPoint, + float* C, + size_t CountM, + size_t CountN, + size_t CountK, + size_t BlockCountK, + const float* Bias, + size_t ldc, + const float* ABlockSum, + const float* QuantBBlkSum) +{ + if (BlkLen == sq2::kBlkLen) { + return SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd_BlkLen64( + QuantA, QuantAScale, QuantBData, QuantBScale, + C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum); + } + if (BlkLen == sq2::kBlkLen128) { + return SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd_BlkLen128( + QuantA, QuantAScale, QuantBData, QuantBScale, + C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum); + } + if (BlkLen == sq2::kBlkLen32) { + return SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd_BlkLen32( + QuantA, QuantAScale, QuantBData, QuantBScale, + C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum); + } + + // Unsupported BlkLen for W2 -- defer to the portable scalar reference. + // Should be unreachable: the W2 dispatch advertises only BlkLen + // in {32, 64, 128} via MlasIsQNBitGemmAvailable. + return sq2::SQ2BitGemmKernel_BlkSum_CompInt8_Scalar( + BlkLen, QuantA, QuantAScale, QuantBData, QuantBScale, QuantBZeroPoint, + C, CountM, CountN, CountK, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum); +} + +} // namespace sqnbitgemm_neon diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index 8e133caa15d55..025c52b86f181 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -1167,6 +1167,93 @@ TEST(MatMul2Bits, MLFloat16_2b_Uint8ZP_Fallback) { .RunWithConfig(); } +// MLFloat16 activation + Uint8 ZP with accuracy_level=4 -- exercises the native MLAS W2 +// CompInt8 path with PrePack of constant scales/ZPs. On ARM64 with f16 intrinsics this +// goes through GetComputeType = HQNBIT_CompInt8 which redirects pack/compute +// through SQNBIT_CompInt8. The bug being guarded: the HQNBIT_CompInt8 PrePack branch +// originally only forwarded scales to the packer (zptr=nullptr), and the 3-call ZP +// follow-up was gated on compute_type_ == SQNBIT_CompInt8 (so it did not fire here). +// QuantBBlkSum then baked in the symmetric default ZP=2 for every block, giving wrong +// outputs whenever the model's real ZPs != 2. Shape: M=1, N=32, K=128, block_size=64 +// (K%BlkLen==0, BlkLen in {32,64,128}); large/varied uint8 ZPs. +TEST(MatMul2Bits, MLFloat16_2b_Uint8ZP_Accuracy4) { + RandomValueGenerator random{1234}; + const int64_t M = 1, N = 32, K = 128, block_size = 64; + std::vector input0_fp32_vals(random.Gaussian(AsSpan({M, K}), 0.0f, 0.25f)); + std::vector input1_fp32_vals(random.Gaussian(AsSpan({K, N}), 0.0f, 0.25f)); + + int q_rows, q_cols; + MlasBlockwiseQuantizedShape(static_cast(block_size), true, + static_cast(K), static_cast(N), + q_rows, q_cols); + size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; + MlasBlockwiseQuantizedBufferSizes(static_cast(block_size), true, + static_cast(K), static_cast(N), + q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); + + std::vector input1_vals(q_data_size_in_bytes); + std::vector scales(q_scale_size); + std::vector zero_points(q_zp_size_in_bytes); + + auto& ortenv = **ort_env.get(); + onnxruntime::concurrency::ThreadPool* tp = ortenv.GetEnvironment().GetIntraOpThreadPool(); + + MlasQuantizeBlockwise( + input1_vals.data(), scales.data(), zero_points.data(), + input1_fp32_vals.data(), static_cast(block_size), + true, static_cast(K), static_cast(N), + static_cast(N), tp); + + // Force non-symmetric ZPs so the bug actually surfaces. With all ZPs == 2 the + // missing-ZP code path coincidentally produces correct results; varying ZPs + // (1 / 2 / 3 packed two-per-byte) make the QuantBBlkSum correction non-trivial. + for (size_t i = 0; i < zero_points.size(); ++i) { + const uint8_t lo = static_cast(((i * 3) % 4)); // 0,3,2,1,0,3,... + const uint8_t hi = static_cast(((i * 5 + 1) % 4)); // 1,2,3,0,1,2,... + zero_points[i] = static_cast((hi << 4) | (lo & 0x0F)); + } + + // Reference dequant via MLAS using the same (varied) ZPs so the expected matmul + // uses the same B values that the kernel sees post-pack. + MlasDequantizeBlockwise( + input1_fp32_vals.data(), input1_vals.data(), scales.data(), zero_points.data(), + static_cast(block_size), true, + static_cast(K), static_cast(N), tp); + + std::vector expected_vals(M * N); + for (int64_t m = 0; m < M; m++) { + for (int64_t n = 0; n < N; n++) { + float sum = 0.0f; + for (int64_t k = 0; k < K; k++) { + sum += input0_fp32_vals[m * K + k] * input1_fp32_vals[n * K + k]; + } + expected_vals[m * N + n] = sum; + } + } + + int64_t k_blocks = (K + block_size - 1) / block_size; + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("block_size", block_size); + test.AddAttribute("bits", QBits); + test.AddAttribute("accuracy_level", static_cast(4)); + + test.AddInput("A", {M, K}, FloatsToMLFloat16s(input0_fp32_vals), false); + test.AddInput("B", {q_cols, k_blocks, q_rows / k_blocks}, input1_vals, true); + test.AddInput("scales", {N, k_blocks}, FloatsToMLFloat16s(scales), true); + test.AddInput("zero_points", + {N, static_cast(q_zp_size_in_bytes) / N}, zero_points, true); + + test.AddOutput("Y", {M, N}, FloatsToMLFloat16s(expected_vals)); + test.SetOutputAbsErr("Y", 0.15f); + test.SetOutputRelErr("Y", 0.03f); + + test.ConfigEp(DefaultCpuExecutionProvider()) + .RunWithConfig(); +} + TEST(MatMul2Bits, Float32_2b_Accuracy0) { TestMatMul2BitsTyped(); TestMatMul2BitsTyped(); diff --git a/onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp b/onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp index 9e5a83d4fad99..ad4a06e61c63a 100644 --- a/onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp +++ b/onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp @@ -36,8 +36,6 @@ Module Name: #include "core/mlas/lib/mlasi.h" #include "core/mlas/lib/sqnbitgemm_kernel_avx512_2bit.h" -#if defined(MLAS_TARGET_AMD64) - namespace { namespace sq2 = onnxruntime::mlas::sq2bit_avx512; @@ -159,9 +157,9 @@ using W2KernelFn = size_t(MLASCALL*)( const std::byte*, float*, size_t, size_t, size_t, size_t, const float*, size_t, const float*, const float*); -void RunW2Case(size_t M, size_t N, size_t K, bool WithBias, uint32_t seed, - bool WithZeroPoints, W2KernelFn kernel, - const char* kernel_name) { +[[maybe_unused]] void RunW2Case(size_t M, size_t N, size_t K, bool WithBias, uint32_t seed, + bool WithZeroPoints, W2KernelFn kernel, + const char* kernel_name) { const size_t BlockCountK = (K + kBlkLen - 1) / kBlkLen; ASSERT_EQ(K % kBlkLen, 0u) << "Test K must be a multiple of BlkLen=64"; // BlockCountK no longer required to be a multiple of kBlockGroupBlks -- @@ -302,6 +300,8 @@ void RunW2Case(size_t M, size_t N, size_t K, bool WithBias, uint32_t seed, } // namespace +#if defined(MLAS_TARGET_AMD64) + // // Scalar block-group test, no zero-points. Covers the same small synthetic // shapes + representative prefill sizes used by the production W2 tests. All @@ -393,6 +393,8 @@ TEST(MlasSq2BitTest, Scalar_BlkLen64_WithZeroPoints) { } } +#endif // defined(MLAS_TARGET_AMD64) -- kSimdShapes itself is cross-arch (also used by the ARM64 NEON DotProd W2 tests further down) + // // SIMD block-group shape coverage. Phase-3+K-tail kernel requires: // * BlkLen == 64 @@ -416,7 +418,7 @@ TEST(MlasSq2BitTest, Scalar_BlkLen64_WithZeroPoints) { // the right packed-B address regardless of K % 4. K=384 and the // synthetic K=320, K=448 shapes exercise this path. // -constexpr struct { +[[maybe_unused]] constexpr struct { size_t M, N, K; } kSimdShapes[] = { {1, 16, 256}, // R1 only @@ -469,8 +471,26 @@ constexpr struct { {1, 17, 384}, {4, 33, 384}, {128, 19, 448}, + // R*xC4 tile coverage (CountN % 8 in [4..7]). The per-BlkLen NEON tile + // dispatcher splits CountN into NMain8 + NMain4 + NTail; without these + // shapes the NMain4 region (R1xC4 / R2xC4 tiles) is never exercised. + // M={1,2,3} covers M-tail-only / M-pair-only / mixed M-pair+M-tail. + {1, 4, 256}, // R1xC4 only + {2, 4, 256}, // R2xC4 only + {3, 4, 256}, // R2xC4 + R1xC4 + {1, 5, 256}, // R1xC4 + R1xC1 + {2, 5, 256}, // R2xC4 + R2xC1 + {3, 5, 256}, // R2xC4 + R2xC1 + R1xC4 + R1xC1 + {1, 12, 256}, // R1xC8 + R1xC4 + {2, 12, 256}, // R2xC8 + R2xC4 + {3, 12, 256}, // R2xC8 + R2xC4 + R1xC8 + R1xC4 + {1, 13, 256}, // R1xC8 + R1xC4 + R1xC1 + {3, 13, 256}, // ALL 6 tiles (R2xC8 + R2xC4 + R2xC1 + R1xC8 + R1xC4 + R1xC1) + {3, 21, 384}, // ALL 6 tiles + K-tail }; +#if defined(MLAS_TARGET_AMD64) + // // AVX-512BW (non-VNNI) SIMD block-group kernel. // @@ -542,10 +562,15 @@ TEST(MlasSq2BitTest, BlkLen64_Avx512Vnni_WithZeroPoints) { } } +#endif // defined(MLAS_TARGET_AMD64) + // ============================================================================= // BlkLen=128 coverage. Mirrors the BlkLen=64 tests above. The helpers are // duplicated rather than templated to keep the BlkLen=64 path bit-identical; // the diff is purely additive. +// +// The helper namespace is intentionally cross-arch (not under MLAS_TARGET_AMD64) +// so the ARM64 NEON DotProd tests in this file can reuse RunW2Case_BlkLen128. // ============================================================================= namespace { @@ -624,9 +649,9 @@ void ReferenceGemm_W2_CompInt8_BlkLen128(size_t M, size_t N, size_t K, } } -void RunW2Case_BlkLen128(size_t M, size_t N, size_t K, bool WithBias, uint32_t seed, - bool WithZeroPoints, W2KernelFn kernel, - const char* kernel_name) { +[[maybe_unused]] void RunW2Case_BlkLen128(size_t M, size_t N, size_t K, bool WithBias, uint32_t seed, + bool WithZeroPoints, W2KernelFn kernel, + const char* kernel_name) { const size_t BlockCountK = (K + kBlkLen128 - 1) / kBlkLen128; ASSERT_EQ(K % kBlkLen128, 0u) << "BlkLen128 test K must be a multiple of 128"; @@ -761,13 +786,17 @@ void RunW2Case_BlkLen128(size_t M, size_t N, size_t K, bool WithBias, uint32_t s // * BlockCountK in {1, 2, 3, 4, 8, 16, 32} -- full + K-tail variants // * N-tail (NMain=0 and various NTail) combined with K-tail // -constexpr struct { +[[maybe_unused]] constexpr struct { size_t M, N, K; } kSimdShapes_BlkLen128[] = { {1, 16, 128}, // R1, BlockCountK=1 {1, 32, 256}, // R1, BlockCountK=2 (K-tail, no full group) {1, 1024, 1024}, // R1, BlockCountK=8 {1, 1024, 4096}, // R1, BlockCountK=32 (representative N) + // K < BlkLen coverage lives at the operator layer + // (MatMul2Bits.Float32_2b_BlkLen128_Accuracy4), which goes through the + // quantizer's K-pad-up path. This direct-kernel runner intentionally + // enforces K % BlkLen == 0. {2, 16, 128}, {2, 32, 256}, {2, 64, 512}, // BlockCountK=4 (one full group) @@ -807,10 +836,28 @@ constexpr struct { {1, 17, 384}, {4, 33, 384}, {128, 19, 640}, + // R*xC4 tile coverage (CountN % 8 in [4..7]). The per-BlkLen NEON tile + // dispatcher splits CountN into NMain8 + NMain4 + NTail; without these + // shapes the NMain4 region (R1xC4 / R2xC4 tiles) is never exercised. + // M={1,2,3} covers M-tail-only / M-pair-only / mixed M-pair+M-tail. + {1, 4, 256}, // R1xC4 only + {2, 4, 256}, // R2xC4 only + {3, 4, 256}, // R2xC4 + R1xC4 + {1, 5, 256}, // R1xC4 + R1xC1 + {2, 5, 256}, // R2xC4 + R2xC1 + {3, 5, 256}, // R2xC4 + R2xC1 + R1xC4 + R1xC1 + {1, 12, 256}, // R1xC8 + R1xC4 + {2, 12, 256}, // R2xC8 + R2xC4 + {3, 12, 256}, // R2xC8 + R2xC4 + R1xC8 + R1xC4 + {1, 13, 256}, // R1xC8 + R1xC4 + R1xC1 + {3, 13, 256}, // ALL 6 tiles (R2xC8 + R2xC4 + R2xC1 + R1xC8 + R1xC4 + R1xC1) + {3, 21, 384}, // ALL 6 tiles + K-tail }; } // namespace +#if defined(MLAS_TARGET_AMD64) + TEST(MlasSq2BitTest, Scalar_BlkLen128) { if (!GetMlasPlatform().Avx512Supported_) { GTEST_SKIP() << "AVX-512 not available on this host"; @@ -907,10 +954,15 @@ TEST(MlasSq2BitTest, BlkLen128_Avx512Vnni_WithZeroPoints) { } } +#endif // defined(MLAS_TARGET_AMD64) + // ============================================================================= // BlkLen=32 coverage. Mirrors BlkLen=128 above with per-block byte width 8 // and per-group bytes 32. Helpers duplicated (rather than templated) to keep // the BlkLen=64 hot path bit-identical. +// +// The helper namespace is intentionally cross-arch (not under MLAS_TARGET_AMD64) +// so the ARM64 NEON DotProd tests in this file can reuse RunW2Case_BlkLen32. // ============================================================================= namespace { @@ -989,9 +1041,9 @@ void ReferenceGemm_W2_CompInt8_BlkLen32(size_t M, size_t N, size_t K, } } -void RunW2Case_BlkLen32(size_t M, size_t N, size_t K, bool WithBias, uint32_t seed, - bool WithZeroPoints, W2KernelFn kernel, - const char* kernel_name) { +[[maybe_unused]] void RunW2Case_BlkLen32(size_t M, size_t N, size_t K, bool WithBias, uint32_t seed, + bool WithZeroPoints, W2KernelFn kernel, + const char* kernel_name) { const size_t BlockCountK = (K + kBlkLen32 - 1) / kBlkLen32; ASSERT_EQ(K % kBlkLen32, 0u) << "BlkLen32 test K must be a multiple of 32"; @@ -1121,7 +1173,7 @@ void RunW2Case_BlkLen32(size_t M, size_t N, size_t K, bool WithBias, uint32_t se // K-shape constraint: K multiple of 32 (BlkLen=32). Covers BlockCountK in // {1, 2, 3, 4, 8, 16, 32, 64} -- both K-tail variants (BlockCountK not a // multiple of 4) and exact block-group multiples. -constexpr struct { +[[maybe_unused]] constexpr struct { size_t M, N, K; } kSimdShapes_BlkLen32[] = { {1, 16, 32}, // R1, BlockCountK=1 @@ -1167,10 +1219,28 @@ constexpr struct { {1, 17, 96}, {4, 33, 96}, {128, 19, 160}, + // R*xC4 tile coverage (CountN % 8 in [4..7]). The per-BlkLen NEON tile + // dispatcher splits CountN into NMain8 + NMain4 + NTail; without these + // shapes the NMain4 region (R1xC4 / R2xC4 tiles) is never exercised. + // M={1,2,3} covers M-tail-only / M-pair-only / mixed M-pair+M-tail. + {1, 4, 128}, // R1xC4 only + {2, 4, 128}, // R2xC4 only + {3, 4, 128}, // R2xC4 + R1xC4 + {1, 5, 128}, // R1xC4 + R1xC1 + {2, 5, 128}, // R2xC4 + R2xC1 + {3, 5, 128}, // R2xC4 + R2xC1 + R1xC4 + R1xC1 + {1, 12, 128}, // R1xC8 + R1xC4 + {2, 12, 128}, // R2xC8 + R2xC4 + {3, 12, 128}, // R2xC8 + R2xC4 + R1xC8 + R1xC4 + {1, 13, 128}, // R1xC8 + R1xC4 + R1xC1 + {3, 13, 128}, // ALL 6 tiles (R2xC8 + R2xC4 + R2xC1 + R1xC8 + R1xC4 + R1xC1) + {3, 21, 160}, // ALL 6 tiles + K-tail }; } // namespace +#if defined(MLAS_TARGET_AMD64) + TEST(MlasSq2BitTest, Scalar_BlkLen32) { if (!GetMlasPlatform().Avx512Supported_) { GTEST_SKIP() << "AVX-512 not available on this host"; @@ -1267,6 +1337,133 @@ TEST(MlasSq2BitTest, BlkLen32_Avx512Vnni_WithZeroPoints) { } } +#endif // defined(MLAS_TARGET_AMD64) + +#if defined(MLAS_TARGET_ARM64) + +#include "core/mlas/lib/qnbitgemm_kernel_neon.h" + +// +// W2 NEON DotProd kernel direct-call tests. Mirror the AVX-512 layout +// above: BlkLen 64 / 128 / 32, each with a no-zero-points and a +// with-zero-points variant, all driven from the shared kSimdShapes* +// shape tables and the cross-arch RunW2Case* helpers (which build the +// packed B and quantized A in the layout the kernel expects). +// +// End-to-end coverage through MlasQNBitGemmBatch is provided by the +// MatMulNBits operator tests; this file only verifies the inner kernel. +// +// Skipped on ARM64 hosts that lack FEAT_DotProd -- the NEON DotProd +// kernel emits SDOT instructions and would SIGILL on a pre-armv8.2 core. +// + +// +// NEON DotProd SIMD block-group kernel, BlkLen=64. +// +TEST(MlasSq2BitTest, BlkLen64_NeonDotProd) { + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot()) { + GTEST_SKIP() << "ARM NEON FEAT_DotProd not available on this host"; + } + for (uint32_t seed : {0xC0FFEEu, 0xBADC0DEu}) { + for (const auto& s : kSimdShapes) { + for (bool bias : {false, true}) { + RunW2Case(s.M, s.N, s.K, bias, seed + (bias ? 1u : 0u), + /*WithZeroPoints=*/false, + sqnbitgemm_neon::SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd, + "NEON-DotProd"); + } + } + } +} + +TEST(MlasSq2BitTest, BlkLen64_NeonDotProd_WithZeroPoints) { + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot()) { + GTEST_SKIP() << "ARM NEON FEAT_DotProd not available on this host"; + } + for (uint32_t seed : {0xC0FFEEu, 0xBADC0DEu}) { + for (const auto& s : kSimdShapes) { + for (bool bias : {false, true}) { + RunW2Case(s.M, s.N, s.K, bias, seed + (bias ? 1u : 0u), + /*WithZeroPoints=*/true, + sqnbitgemm_neon::SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd, + "NEON-DotProd"); + } + } + } +} + +// +// NEON DotProd SIMD block-group kernel, BlkLen=128. +// +TEST(MlasSq2BitTest, BlkLen128_NeonDotProd) { + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot()) { + GTEST_SKIP() << "ARM NEON FEAT_DotProd not available on this host"; + } + for (uint32_t seed : {0xC0FFEEu, 0xBADC0DEu}) { + for (const auto& s : kSimdShapes_BlkLen128) { + for (bool bias : {false, true}) { + RunW2Case_BlkLen128(s.M, s.N, s.K, bias, seed + (bias ? 1u : 0u), + /*WithZeroPoints=*/false, + sqnbitgemm_neon::SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd, + "NEON-DotProd"); + } + } + } +} + +TEST(MlasSq2BitTest, BlkLen128_NeonDotProd_WithZeroPoints) { + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot()) { + GTEST_SKIP() << "ARM NEON FEAT_DotProd not available on this host"; + } + for (uint32_t seed : {0xC0FFEEu, 0xBADC0DEu}) { + for (const auto& s : kSimdShapes_BlkLen128) { + for (bool bias : {false, true}) { + RunW2Case_BlkLen128(s.M, s.N, s.K, bias, seed + (bias ? 1u : 0u), + /*WithZeroPoints=*/true, + sqnbitgemm_neon::SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd, + "NEON-DotProd"); + } + } + } +} + +// +// NEON DotProd SIMD block-group kernel, BlkLen=32. +// +TEST(MlasSq2BitTest, BlkLen32_NeonDotProd) { + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot()) { + GTEST_SKIP() << "ARM NEON FEAT_DotProd not available on this host"; + } + for (uint32_t seed : {0xC0FFEEu, 0xBADC0DEu}) { + for (const auto& s : kSimdShapes_BlkLen32) { + for (bool bias : {false, true}) { + RunW2Case_BlkLen32(s.M, s.N, s.K, bias, seed + (bias ? 1u : 0u), + /*WithZeroPoints=*/false, + sqnbitgemm_neon::SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd, + "NEON-DotProd"); + } + } + } +} + +TEST(MlasSq2BitTest, BlkLen32_NeonDotProd_WithZeroPoints) { + if (!MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot()) { + GTEST_SKIP() << "ARM NEON FEAT_DotProd not available on this host"; + } + for (uint32_t seed : {0xC0FFEEu, 0xBADC0DEu}) { + for (const auto& s : kSimdShapes_BlkLen32) { + for (bool bias : {false, true}) { + RunW2Case_BlkLen32(s.M, s.N, s.K, bias, seed + (bias ? 1u : 0u), + /*WithZeroPoints=*/true, + sqnbitgemm_neon::SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd, + "NEON-DotProd"); + } + } + } +} + +#endif // defined(MLAS_TARGET_ARM64) + // // Availability contract test for W2 + SQNBIT_CompInt8. // @@ -1277,9 +1474,18 @@ TEST(MlasSq2BitTest, BlkLen32_Avx512Vnni_WithZeroPoints) { // availability return true for shapes that Q2BitGemmPackQuantBDataSize_Avx512 // would refuse to size (returning 0). // +// The contract is platform-agnostic: it must hold wherever a W2 dispatch is +// installed (AVX-512 on x86_64, NEON+DotProd or i8mm on ARM64). We skip on +// hosts where no W2 dispatch is wired in (e.g. pure-NEON-only ARM, or x86 +// without AVX-512) -- there is nothing to assert about a feature that is +// uniformly unavailable. +// TEST(MlasSq2BitTest, AvailabilityContract_BlkLens) { - if (!GetMlasPlatform().Avx512Supported_) { - GTEST_SKIP() << "W2 native dispatch is AVX-512-only on x86_64"; + // Probe a representative supported BlkLen to detect whether ANY W2 dispatch + // is installed on this host. This mirrors the dispatch-pointer-identity + // guard pattern that the W4/W8 tests use. + if (!MlasIsQNBitGemmAvailable(2, 64, SQNBIT_CompInt8)) { + GTEST_SKIP() << "No W2 native dispatch on this host"; } // Supported BlkLens. @@ -1295,5 +1501,3 @@ TEST(MlasSq2BitTest, AvailabilityContract_BlkLens) { EXPECT_FALSE(MlasIsQNBitGemmAvailable(2, 64, SQNBIT_CompFp32)); EXPECT_FALSE(MlasIsQNBitGemmAvailable(2, 64, HQNBIT_CompFp16)); } - -#endif // defined(MLAS_TARGET_AMD64)