Skip to content

[Rock][NonAccel] Vectorize the C writeback along gemmM for M-contiguous outputs - #2454

Open
pfultz2 wants to merge 1 commit into
developfrom
non-accel-output-swizzle
Open

[Rock][NonAccel] Vectorize the C writeback along gemmM for M-contiguous outputs#2454
pfultz2 wants to merge 1 commit into
developfrom
non-accel-output-swizzle

Conversation

@pfultz2

@pfultz2 pfultz2 commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Motivation

NHWC f32 convolutions on gfx1201 run ~7-9% slower than NCHW even with
exhaustive tuning (e.g. 512→1536 k=2 pad=1 @32x32: 408us vs 385us). These use
the non-accelerated GEMM path (no f32 WMMA), where the C writeback iterates
with n_thread innermost unconditionally. For NHWC outputs the
memory-contiguous dimension is gemmM (k innermost), so every thread issues
scalar b32 stores scattered 6KB apart across lanes. The same problem applies
to any --transC-style GEMM: forcing an M-contiguous C through this path
costs ~10% (450us → 404us on a 1536x1089x2048 f32 GEMM once fixed).

Additionally, the outputSwizzle tuning parameter (9th field of v3 perf
configs) was silently ignored on the non-accel path: the OutputSwizzle pass
reads it from a func attribute that only the accel and attention lowerings set.

Technical Details

  • In GridwiseGemmRewritePattern (GridwiseGemmToBlockwise.cpp), query
    getMaxVectorization on the C operand for gemmM and gemmN. When M
    vectorizes better than N (storeMFast):
    • Keep the accumulator n-major in registers: view the flat C register
      buffer as [n, m] and hand the blockwise gemm a transposed [m, n] view
      of it, so the register order matches the new store order.
    • Emit the writeback's iter merge as
      {n_repeat, n_thread, m_repeat, m_thread} instead of
      {m_repeat, m_thread, n_repeat, n_thread}, putting m_thread innermost
      so the threadwise write vectorizes along gemmM (b64/b128 stores, with
      m_cuwave lanes forming contiguous 32-64B clusters).
  • When N vectorizes at least as well as M, codegen is unchanged; NCHW kernels
    produce bit-identical ISA (their output vectorizes in neither dimension).
  • Set the outputSwizzle func attribute from GeneralGemmParams in the
    non-accel lowering, matching the accel/attention paths. Default
    (2 = heuristic) behavior is unchanged; this only makes the existing knob
    functional.

Test Plan

Test Result

Submission Checklist

@pfultz2
pfultz2 requested a review from causten as a code owner August 17, 2026 22:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant