Shape-blind pack-lh promotion regresses batch-1 GEMV on Arm SME2 (pqs8/pf32/qp8 slower than plain neondot/NEON paths at M=1)
- Dominant language
- C
- Stars
- 2.5k
- Forks
- 560
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 163
Description
### Summary
The subgraph optimizer's pack-lh promotion (`src/subgraph.c`, fully-connected / batch-matrix-multiply rewrite) switches an FC onto the packed GEMM config whenever `xnn_init_{pqs8_qc8w,pf32,qp8_f32_qc8w}_gemm_config()` is available, **with no shape (M) heuristic**. On Arm SME2 hardware this regresses batch-1 GEMV substantially: the packed SME2 kernels are 1.4-3x slower than the plain `neondot`/NEON paths that would otherwise be selected, while the same promotion is a large win for M>>1 GEMM.
### Environment
- XNNPACK @ `3131afead790c5c69a9aa12273dfc40399789ad7` (the commit pinned by ExecuTorch v1.1.0), via the ExecuTorch XNNPACK delegate
- Device: vivo X300 Pro (MediaTek Dimensity 9500, Arm C1 cores with SME2), Android; `/proc/cpuinfo` exposes `sme`/`sme2`
- A/B: identical `.pte` files run on the same device with two runner builds - NDK r26/clang 17 (SME2 kernels compiled out by the CMake `Clang < 18` guard, so the plain paths run) vs NDK r27+/clang 18+ (packed SME2 paths run). Kernel identity verified per run via ExecuTorch ETDump per-op events; times are averages over 30 executions (first run dropped).
### Measurements (`nn.Linear(2048x2048)` and `4096x4096`, statically quantized qs8_qc8w unless noted)
| Case | plain path (clang17 build) | packed SME2 path (clang18 build) | packed vs plain |
|---|---|---|---|
| M=1, 2048^2, qs8_qc8w | `qs8_qc8w` (neondot) 30-40 us | `pqs8_qc8w` (sme2) **92 us** | **2.3-3x slower** |
| M=1, 2048^2, dynamic qp8_f32_qc8w | 60 us (neondot variant) | **106 us** (sme2 variant) | **1.8x slower** |
| M=1, 2048^2, fp32 | `f32` 209-275 us | `pf32` **309 us** | ~1.3x slower |
| M=1, 4096^2, fp32 | 1177 us | **1701 us** | **1.4x slower** |
| M=1, 4096^2, qs8_qc8w | 347 us | 309 us | ~parity (DRAM-bound) |
| M=64, 2048^2, fp32 | 3499 us | **1118 us** | 3.1x faster |
| M=64, 2048^2, qs8_qc8w | 669 us | **348 us** | 1.9x faster |
So the promotion is clearly correct for GEMM shapes but counterproductive at M=1 (except when purely DRAM-bound): the 32-wide SME2 tile kernels underutilize on a single row and the plain neondot GEMV is significantly faster. The `MR_TO_INDEX(1)` sme2 gemm ukernel exists and is what runs, but it does not come close to the neondot GEMV.
### Impact
Decode-shaped workloads (LLM token generation: every linear is an M=1 GEMV) get slower when SME2 support is enabled in the build, while prefill/CNN/ViT-shaped workloads get faster - an unfortunate trade-off that a shape-aware fallback would avoid.
### Expected behavior
Kernel/config selection for FC (and BMM) should consider M - e.g. keep the non-packed config for M==1 (or below some threshold) instead of unconditionally inserting pack-lh when a packed config exists, or the packed 1xN SME2 kernels should be made competitive with the neondot GEMV.
### Repro
Any statically quantized (qs8_qc8w) or fp32 `nn.Linear` with batch 1, lowered through the ExecuTorch XNNPACK delegate on SME2 hardware, reproduces this; happy to share the exact `.pte` pairs and the ETDump event dumps if useful.
Contributor guide
Research direction
Start in src/subgraph.c at the fully-connected and batch-matrix-multiply pack-lh promotion, then inspect the pqs8_qc8w, pf32, and qp8_f32_qc8w GEMM configuration checks. Reproduce the batch-1 and larger-M cases through the ExecuTorch XNNPACK delegate on Arm SME2 hardware. Done means shape-aware selection avoids the batch-1 regression while retaining the packed path's gains for larger GEMM shapes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- c
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100