[Rock][NonAccel] Vectorize the C writeback along gemmM for M-contiguous outputs - #2454
Open
pfultz2 wants to merge 1 commit into
Open
[Rock][NonAccel] Vectorize the C writeback along gemmM for M-contiguous outputs#2454pfultz2 wants to merge 1 commit into
pfultz2 wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_threadinnermost unconditionally. For NHWC outputs thememory-contiguous dimension is gemmM (
kinnermost), so every thread issuesscalar
b32stores scattered 6KB apart across lanes. The same problem appliesto any
--transC-style GEMM: forcing an M-contiguous C through this pathcosts ~10% (450us → 404us on a 1536x1089x2048 f32 GEMM once fixed).
Additionally, the
outputSwizzletuning parameter (9th field ofv3perfconfigs) 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
GridwiseGemmRewritePattern(GridwiseGemmToBlockwise.cpp), querygetMaxVectorizationon the C operand for gemmM and gemmN. When Mvectorizes better than N (
storeMFast):buffer as
[n, m]and hand the blockwise gemm a transposed[m, n]viewof it, so the register order matches the new store order.
itermerge as{n_repeat, n_thread, m_repeat, m_thread}instead of{m_repeat, m_thread, n_repeat, n_thread}, puttingm_threadinnermostso the threadwise write vectorizes along gemmM (
b64/b128stores, withm_cuwavelanes forming contiguous 32-64B clusters).produce bit-identical ISA (their output vectorizes in neither dimension).
outputSwizzlefunc attribute fromGeneralGemmParamsin thenon-accel lowering, matching the accel/attention paths. Default
(
2= heuristic) behavior is unchanged; this only makes the existing knobfunctional.
Test Plan
Test Result
Submission Checklist