Repository navigation
[MIOpen] Use a single GEMM for point-output conv in all three directions (2D/3D) - #10912
Conversation
jfactory07
changed the title
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
Patch-embedding shapes (stride == kernel, zero pad, unit dilation, output spatially 1x1 or 1x1x1) reduce to a plain GEMM, but only the 3D forward path took advantage of that. Backward-data fell back to the strided IGEMM kernel, which decomposes stride=14 into 196 micro-GEMMs, and backward-weights had no GEMM path at all. Unify the 2D and 3D point-output detection into a single IsPointOutputStrideEqFilter() and extend the fast paths to cover both: - GemmBwdRest merges its 2D and 3D branches into one GEMM over N. 2D writes dx in place, since a per-batch dx slice enumerates elements in the same order as a row of w for NCHW ((c,y,x)) and NHWC ((y,x,c)) alike, so it needs no workspace. 3D permits a larger dx and keeps the batched Col2Im scatter. - GemmWrwUniversal gains the point-output GEMM for 2D and 3D, with fp32 accumulation via workspace for bf16 batches. The fp32-accumulation overload of CallGemm is rocBLAS-only, so request that backend explicitly. - The TunaNet heuristic promotes these GEMM solvers for all three directions, so the find modes that skip benchmarking also land on the fast path. Direct dx/dw writes are only a valid reinterpretation when the image spatial extent equals the filter extent, so the 2D bwd and the wrw checks require that; 3D bwd scatters through Col2Im and does not. Also fix double-counted profiling in the 2D bwd path, which reported twice the kernel time by passing GetKernelTime() to AccumKernelTime(). Measured on MI355X (gfx950), n=9500 c=3 14x14 k=1280 stride=kernel=14, bf16: fprop 0.147ms -> 0.027ms and dgrad 210.9ms -> 0.028ms. wgrad had no GEMM path before; GemmWrwUniversal now offers one at 0.124ms, which the in-tree IGEMM wrw still edges out at 0.098ms once Find benchmarks the two against each other.
9812e05 to
237d50d
Compare
The point-output branch of GemmBwdRest::IsApplicable() does its own checks instead of deferring to GemmBwdBase, and accepted NHWC for both 2D and 3D. That is only sound in 2D, where the GEMM writes dx in place and dx and w enumerate their elements in the same order. 3D scatters the GEMM result with Col2Im3dGPUBatched, which addresses dx as channel-first, so an NDHWC problem produced incorrectly ordered dx. Reject it the same way GemmWrwUniversal already does. The failure is easy to miss because the scatter degenerates. With col_d/h/w all 1 the kernel loop collapses to a single iteration and reads im_ch*wei_d*wei_h*wei_w + (im_d*wei_h + im_h)*wei_w + im_w while writing im_ch*D*H*W + im_d*H*W + im_h*W + im_w, which coincide exactly when the dx spatial extent equals the filter. Point-output only forces out == 1, and integer division admits any dx spatial extent in [wei, 2*wei-1], so a larger dx is what exercises the real scatter. Verified on gfx950 with n=8 c=4 5x5x5, k=32, 4x4x4 filter, stride 4: NDHWC failed verification before this change and falls back to a correct solver after it, while NCDHW keeps the fast path. Add that larger-dx shape to the GemmBwdRest tests, since the existing 3D case sits in the degenerate regime and never covered the scatter.
… can The previous commit kept 3D point-output on the default layout because the Col2Im3d scatter addresses dx as channel-first. That is stricter than needed: the scatter is only required when dx is spatially larger than the filter. When the two match, a per-batch dx slice enumerates its elements in the same order as a row of w for channel-first ((c,z,y,x)) and channel-last ((z,y,x,c)) alike, so the GEMM can write dx in place and the layout stops mattering. That is already how 2D works. Key the decision off that property via IsBwdDataPointOutputDirectWritable() instead of off the spatial rank. This admits 3D NDHWC when dx spatial equals the filter, still rejecting the larger-dx case that genuinely needs the channel-first scatter, and additionally drops the workspace round trip and the Col2Im launch for 3D default-layout problems in the same regime, where the scatter was only performing an element-wise copy. Cover both channel-last shapes in the GemmBwdRest tests, with C > 1 so the layout is distinguishable from the default one.
The K=1280 FP32 cases exceed the default 1.0*eps RMS threshold on non-TF32 GPUs.
|
This PR also adds fwd and wrw support for NHWC but there are no test cases added for this new ability. |
Point-output wrw accumulates dw in an fp32 workspace and casts the result back to bf16 once the batch size exceeds 1. Running the point-output cases as fp16 only left that branch unexecuted, so add bf16 cases to all three directions and cover both sides of the wrw batch size condition. FP32 stays out: at K=1280 its RMS error lands just under the 1.0*eps threshold that non-TF32 GPUs apply, which is too little margin to rely on. The new cases skip gfx90a, where rocBLAS does not support the BF16->BF16 GEMM that a bf16 point-output shape issues.
GetWti was pinned to 10.0 for point-output shapes, which asserts that the GEMM solver should win every shape matching the case. It never changed a decision. WTI is only consulted by the Immediate mode fallback, and there the solvers that could compete return wti_approximate_worst and are skipped, leaving the naive solvers at 0.01 as the only other candidates. For a 14x14 stride-14 K=1280 shape on gfx950 the unforced estimates are 0.63 for wrw and 0.855 for bwd, and Immediate mode still picks GemmWrwUniversal and GemmBwdRest without the override. Backward data also charged itself a Col2Im launch unconditionally. Count it only when dx cannot be written directly, matching what the invoker does.
Point-output wrw accepts NHWC in 2D, but no case exercised it. 3D still requires the default layout, so only a 2D case is added. Backward data already had its NHWC and NDHWC cases.
A point-output shape is a single GEMM with no Im2Col, and it stays correct channel-last: a per-batch x slice and a row of w both enumerate C*Z*Y*X, in (z,y,x,c) order rather than (c,z,y,x), so the dot product still pairs the same elements, and y is contiguous over K because its spatial extent is 1. GemmFwdBase confines every forward GEMM solver to the default layout, which is right for the Im2Col path it also guards but needlessly excludes this case. Accept it in GemmFwdRest rather than relaxing the shared check, since that would also admit the 1x1 solvers, whose GEMM does assume channel-first. Verified on gfx950: NHWC and NDHWC point-output are correct in fp16 and bf16 and run one GEMM with no Im2Col, while a non point-output NHWC shape stays inapplicable.
The point-output GEMM contracts the whole batch in a single call with beta = 0, so it overwrites all K * C*Z*Y*X elements of its destination without reading them, which makes the preceding SetTensor dead work. Its destination is exactly dw, or the fp32 accumulation buffer of the same element count, because filter_col_size reduces to C*Z*Y*X once the output spatial extent is 1, and ldc equals it. The accumulating per-batch loop in the general path relies on beta = 1 and keeps its own zeroing. Saves one kernel launch per invocation. Verified with the driver NaN pre-fill (-N 1), which the gtest cases cannot show since they zero dw themselves: fp16, bf16 with and without the fp32 accumulation buffer, and channel-last all return results identical to before.
I will add NHWC test |
|
Really nice result — collapsing the stride=14 dgrad from 196 micro-GEMMs into one, and giving wgrad a GEMM path where it had none, are both clear wins, and I appreciate the honesty about wgrad still being edged out by the in-tree IGEMM once Find compares them. The layout-aware "point-output is just a GEMM" argument you make here seems to generalize a bit further than the PR currently applies it, all in the same spirit and without new compute kernels. A few directions that look like cheap follow-ons (sharing as ideas, not asking for changes to this PR): Closest to what's already here: 3D NHWC point-output. The 2D path handles NHWC via the (y,x,c) ordering argument, but the predicate keeps 3D to the default layout (and 3D still uses Col2Im). If the same ordering reasoning holds for the 3D in-place case, that might extend NHWC to 3D too — does it, or is there a reason 3D has to keep the scatter? Slightly broader, same technique: NHWC 1×1 stride-1. The simplest case of all — a flat [N·H·W, C] × [C, K] GEMM, no batching/transpose/workspace — but the base still gates it to NCHW today. None of this is blocking and the PR stands on its own — just flagging adjacent possibly low hanging fruit in case any of it is a cheap next step. If useful, I've written up the per-case rationale with file:line anchors and a validation/timing plan in an attached doc (pr10912-gemm-conv-fruit.md) that someone (or an agent) could work from directly. |
|
Here is the file... |
|
Edited to remove points already covered by Brad's comment ... Advanced PyTorch users can reshape the input tensor and weights to convert the "point-output conv" to a linear layer. The same trick can be used for point-wise conv2d. This is useful information for anybody waiting for this PR to land. |
The guard rejecting a non-default layout in 3D was carried over from backward data, where Col2Im addresses dx channel-first and so cannot serve a channel-last layout. Point-output wrw launches no such kernel: it is a single GEMM whose extents come from the logical dimensions, and x and dw both enumerate C*Z*Y*X in (z,y,x,c) order, so the contraction pairs the same elements. The bf16 fp32 accumulation buffer lines up as well, since the GEMM writes element k*C*Z*Y*X + z*Y*X*C + y*X*C + x*C + c, which is how packed channel-last dw strides address it. Verified on gfx950 over asymmetric filters, C of 1, 3 and 5, and a ViT-like K=768, in fp16 and bf16 both with and without the fp32 accumulation buffer, also with dw pre-filled with NaN, and confirmed the trace is a single GEMM.
regarding 3D NHWC point-output, I tested it, and the guard can be removed directly. In this case, we can use a single GEMM. I also added the test case. For other cases, I agree that they are areas we need to address. I currently have two hipConv tickets on my plate, which involve supporting output > 1×1 patchify and NHWC 1×1 stride-2 in hipConv. |
Good idea! |
|
@BradPepersAMD I think the CI is clean now |
|
@jfactory07 wrote:
The slower WGRAD (wrw) performance is because the GEMM is not using a workspace. Enabling a 32 MB workspace would make the deep reduction significantly faster. Benchmarking the hiptBLASLt GEMM on gfx950 (m=1280 × n=588 × k=9500, bf16), I get 135 µs with no workspace and 40 µs with 32 MiB workspace. |
## Summary Backfills the **MIOpen 3.6.1 / ROCm 10.1.0** changelog section with externally-visible changes that landed since the `release/therock-10.0` branch but weren't yet documented. All 43 MIOpen commits in that range were reviewed; CI/build-infra, test-only, docs, version-bump, and already-documented changes were excluded. 21 entries added: ### Added * [RNN] Fused single-layer LSTM forward-inference kernel, gated behind `MIOPEN_DEBUG_RNN_FUSED_INFERENCE` (#8915) * [Conv] Naive reference kernel enabling 3D INT8 forward convolution (#10051) * [Conv] `ConvHipConv` solver backed by vendored hipconv kernels (#9746) * [Conv] float8 vector type support + deduplicated math kernels (#10130) * [Conv] Validated gfx1151 and gfx1200 (Navi) SystemDBs (#10050) * [Conv] `MIOPEN_DEBUG_DISABLE_SYSTEM_DB` / `MIOPEN_DEBUG_DISABLE_USER_DB` env vars (#10728) * [Conv] CK depthwise convolution on RDNA wave32 GPUs (#9373) ### Changed * [Conv] Winograd Rage RxS f2x3 on gfx950 + gfx942 v4_6_1/v4_9_1 kernels (#9014) * [Conv] Updated 3D AI heuristics for gfx942 & gfx950 (#8954) ### Optimized * [Conv] Single GEMM for point (1x1) output convolution, all directions, 2D/3D (#10912) * [Conv] Naive backward-data kernel work group widened to 1024 (#10684) ### Resolved Issues * [Conv] Batch-split miscalculation on large WrW tensors — ROCM-26011 (#11188) * [Conv] Conv failure on XNACK-unsupported gfx9 APUs (gfx90c) (#11211) * [BatchNorm] Device heap-overflow in fwd spatial `FinalMeanVariance` — ALMIOPEN-2424 (#11093) * [Conv] Deterministic `SetNextValue` polarity for grouped Bwd/Wrw CK solvers (#10240) * [LRN][LayerNorm] Grid alignment + OOB threads on gfx1250 (#10251) * [Conv] Stale gfx950 grouped-wrw SystemDB entries (#11130) * [Conv][Windows] System-DB path resolution from the MIOpen module (#10838) * [Conv] System-DB fallback so MI308X (gfx942) resolves a usable DB (#8802) * JSON perf logs (`MIOPEN_PERFORMANCE_LOGS`) dropping the last solver during Find (#11166) * Build failure against GCC 15 libstdc++ under C++20 (#9642) ## Notes Deliberately excluded: internal/default-off changes (#10219, #10088, #10649, #10231), a near-duplicate of an existing 10.0.0 entry (#10188), and CI/build-infra/test-only/docs commits. Docs-only change; no code impact. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Motivation
JIRA ID: ROCM-25757
Patch-embedding shapes (stride == kernel, zero pad, unit dilation, output spatially 1x1 or 1x1x1) reduce to a plain GEMM, but only the 3D forward path took advantage of that. Backward-data fell back to the strided IGEMM kernel, which decomposes stride=14 into 196 micro-GEMMs, and backward-weights had no GEMM path at all.
Technical Details
Unify the 2D and 3D point-output detection into a single IsPointOutputStrideEqFilter() and extend the fast paths to cover both:
Test Plan
Measured on MI355X (gfx950), n=9500 c=3 14x14 k=1280 stride=kernel=14, bf16: fprop 0.147ms -> 0.027ms and dgrad 210.9ms -> 0.028ms. wgrad had no GEMM path before; GemmWrwUniversal now offers one at 0.124ms, which the in-tree IGEMM wrw still edges out at 0.098ms once Find benchmarks the two against each other.
Test Result
Submission Checklist