[Kernel] conv3d: parameterize tile size + add autotuner#820
Open
jiacao-amd wants to merge 13 commits into
Open
[Kernel] conv3d: parameterize tile size + add autotuner#820jiacao-amd wants to merge 13 commits into
jiacao-amd wants to merge 13 commits into
Conversation
Generalize the BF16 and FP8 implicit-GEMM conv3d kernels from a single hardcoded 128x128/2x4-wave shape to a compile-time-parameterized tile (TILE_M, TILE_N, WAVE_M, WAVE_N) and add a manual per-shape autotuner. - Replace the hand-scheduled 2x2-half BF16 pipeline (phase_*_prefetch + peeled prologue/main/epilogue, welded to QM_STEPS==2/QN_STEPS==1) with a generic double-buffered 2-stage loop over a flat acc[MI_M*MI_N] MFMA grid. - Generalize gather/commit/read/store to range_constexpr over MI_M/MI_N and per-thread LDG counts; add tile legality asserts (MFMA divisibility, LDG vectorization, LDS budget). - FP8: same flat-acc generalization, g2s halves -> row-block loop, Mfma16x16x128(MI_M, MI_N), flat_work_group_size from BLOCK_THREADS. - New kernels/conv/conv3d_autotune.py: benchmark a small candidate list with do_bench, cache the best tile per shape (in-memory + JSON disk), reusing flydsl.autotune. Enabled via autotune=True or FLYDSL_CONV3D_AUTOTUNE=1; default path stays (128,128,2,4) for zero behavior change. - TILE_K and filter size (kt/kh/kw) stay compile-time constants. - Add tile-sweep + autotune-cache tests and scripts/bench_conv3d_tiles.py.
Buffer load offsets are 32-bit, so a single buffer resource can address
at most ~2 GB. The existing BIG_IN / BIG path rebased per image (nbase =
m_offset // dhw), which is a no-op for n=1 and fails entirely when a
single image exceeds 2^31 elements.
Fix: per-block time-plane rebase for bf16 and fp8:
* conv3d_implicit_8wave (bf16):
- compile_transpose_ncdhw_ndhwc: BIG = n*c*s > 2^31 path adds per-block
(nb, c0, s0) rebase so read/write residuals are tile-local (<900M <2^31).
- compile_conv3d_implicit_8wave: BIG_IN path computes base_t (block's
first time-plane, shifted down by pt) and rebases x_rsrc so gather_a
residual (di*d + in_t - base_t)*h*w*c stays within int32 (<200M).
* conv3d_implicit_8wave_fp8:
- compile_pack_activation_ncdhw_bf16_to_ndhwc_fp8: same BIG per-block
(nb, c0, s0) rebase for read (BF16 input) and write (FP8 byte output).
- compile_conv3d_implicit_8wave_fp8: BIG_IN path adds
_make_fp8_buffer_tensor_from_addr helper and computes a rebased x_div
so g2s_a_block's g_elem residual stays within int32.
Verified on gfx950 (MI355X):
- 1x1024x240x160x90 conv (1,3,3) and (3,1,1): PASS (was crash/wrong)
- 1x1024x120x160x90, 1x2048x60x80x45: no regression
- All 14 bf16 unit tests pass; 6/7 fp8 tests pass (1 is pre-existing flaky)
For n==1 (no split-K), pack a lane's 4 contiguous, 8-byte-aligned output values into one vec4 bf16 buffer_store instead of 4 scalar stores. The epilogue VMEM-store was ~12% of kernel stalls on the short-reduction 3x1x1 path; this brings the 3x1x1 core (256x256x4x4 tile) from 3.18ms to 2.85ms (854->951 TFLOPS), beating MIOpen (2.92ms). Gated on n==1, not split-K, and dhw % 4 == 0; scalar fallback otherwise. Add conv2d_implicit / conv1d_implicit as thin wrappers that reshape to the degenerate 3D form (D=T=1, and H=R=1 for 1D) and reuse the same kernel and fast paths -- zero kernel duplication. Guard temporal_only_fast off when BIG_IN (>2^31 elems): the large-tensor rebase uses the generic n/t/h/w decode, which the temporal fast decode does not carry, so >2^31 3x1x1 inputs fall through to the correct generic path. Tests: add conv2d/conv1d vs torch coverage (26 total, all pass).
Previously the temporal-only fast decode was disabled under BIG_IN (>2^31 elements), falling back to the generic path. Instead, rebase the temporal fast-path address against the same x_base_elem origin the BIG_IN x_rsrc is built from, so the residual element offset stays within int32. This keeps the faster temporal decode for large 3x1x1 inputs while staying correct. Verified vs torch.nn.functional.conv3d on a >2^31-element input (1x640x240x160x90, 3x1x1): allclose, maxdiff 1.0 (matches the generic path). Adds a large_shape regression test.
…upport Port the BF16 conv3d optimizations to the FP8 kernel: vectorized (vec4) epilogue store for n==1, the temporal-only 3x1x1 fast decode path, 64-bit output addressing for >2^31-element outputs (BIG_OUT), and a split-K overflow guard. Add conv2d_implicit_fp8 / conv1d_implicit_fp8 thin wrappers mirroring the BF16 API, plus 2D/1D FP8 correctness tests. Tidy the BF16 kernel comments.
jiacao-amd
force-pushed
the
jiacao/conv3d-multi-tile-autotune-v2
branch
from
July 12, 2026 08:01
31e5866 to
1a740c3
Compare
Make conv3d_implicit_8wave(_fp8) the single public entry that dispatches
1D/2D/3D by filter rank; move the rank-specific bodies to private
_conv{1,2,3}d_impl(_fp8) helpers so the 2D/1D reshape wrappers no longer
recurse through the public name. Merge _choose_splitk into _resolve_splitk
(behavior verified bit-identical over the full input space). Also drop
stale explanatory comments and trailing whitespace.
Co-Authored-By: Claude <noreply@anthropic.com>
Port hgemm_splitk's async-copy deep-pipeline recipe into the bf16 conv3d kernel. Replace the sync global->VGPR->LDS hop (gather_a/commit_a) with raw_ptr_buffer_load_lds direct global->LDS DMA, freeing the A-tile VGPRs and letting the software pipeline go from 2 to 4 stages. Padding is masked by OOB-routing: the buffer resources are rebuilt with the real num_records and invalid im2col taps are routed past the bounds so the hardware bounds check writes 0 to LDS (the DMA path has no VGPR step for a select mask). Gated by USE_ASYNC = not BIG_IN and X/W byte sizes <= 2^31; BIG_IN and >2^31 tensors keep the proven sync path. Measured on MI355X (gfx950), tile 256x256x4x4, controlled back-to-back A/B: c2048 1x3x3 M216k 842 -> 924 TFLOPS (+9.7%) c2048 3x3x3 M216k 848 -> 967 TFLOPS (+13.9%) 26/26 conv tests pass.
The buffer_load_lds voffset is unsigned 32-bit (verified), so no 64-bit soffset split is needed to reach past the old 2^31-byte limit. Extend the async path to two more regimes: - not-BIG_IN tensors up to ~4.29GB: raw resource with real num_records, padding taps routed to an OOB sentinel just under 2^32. - BIG_IN with n==1: reuse the existing per-block rebase (relative offsets <~0.3GB) and give the rebased resource a fixed 2GB num_records, between the max legal tap and the OOB sentinel, so padding still zeroes. n>1 BIG_IN and >~4.29GB weights keep the sync fallback (di can jump a whole batch past 2^32). All real video-VAE shapes now take the async path. Measured on MI355X (gfx950), tile 256x256x4x4: 512 [240,320,180] 1x3x3 728 -> 784 TFLOPS (+7.7%, 14GB) 1024 [240,160,90] 1x3x3 799 -> 887 TFLOPS (+11.0%, 7GB) 1024 [120,160,90] 1x3x3 804 -> 902 TFLOPS (+12.2%, 3.5GB) 2048 [60,80,45] 3x3x3 847 -> 977 TFLOPS (+15.4%) 26/26 conv tests pass; large shapes verified correct (allclose).
jiacao-amd
force-pushed
the
jiacao/conv3d-multi-tile-autotune-v2
branch
3 times, most recently
from
July 15, 2026 00:29
e05dd6c to
5ad796b
Compare
Rename the conv modules so the 8-wave implementation detail is not in the name: conv3d_implicit_8wave -> conv3d_implicit conv3d_implicit_8wave_fp8 -> conv3d_implicit_fp8 conv3d_autotune -> conv3d_implicit_autotune (+ their tests, internal symbols, and imports) The bf16 conv3d kernel + test are synced to the jiacao/conv3d-gemm-port versions (async-only pipeline, sync dead code removed, async naming dropped, PIPE_STAGES=4). The fp8 kernel + test keep this PR's base content unchanged (file + symbol renamed only, no logic diff); the fp8 rewrite lands in a separate PR.
jiacao-amd
force-pushed
the
jiacao/conv3d-multi-tile-autotune-v2
branch
from
July 15, 2026 00:34
5ad796b to
5311d0f
Compare
jiacao-amd
marked this pull request as ready for review
July 15, 2026 06:46
…ernel - wgm L2-swizzle: grouped-M block remap keeps the weight tile (B matrix) hot in L2 across wgm consecutive m-tiles, reducing TCC MISS by ~40% on large spatial shapes (measured: L2 hit rate 71.8% -> 83.5% with wgm=4) - XOR bank-swizzle on LDS a_lds_off/b_lds_off: spreads MFMA ds_read across banks using row//2 as the block selector (~0% bank conflict) - autotune extended to sweep (tile, wgm) pairs: WGM_VALUES=[1,4,8] combined with BF16_CANDIDATES; cache schema bumped to v3 (nested [tile, wgm] format) - autotune dry-run before do_bench to exclude compile latency from timing Performance (gfx950, BF16, vs image(3) reference FLY column): spatial 3x3: 1053-1148 TF/s (+45-67% vs reference, 2.0-2.1x MIOpen) temporal 1x3: 856-1154 TF/s (+43-84% vs reference, 2.3-3.5x MIOpen) Co-Authored-By: Claude <noreply@anthropic.com>
Previously BIG_IN required n==1 because the x_rsrc rebase only handled a single-sample base address. This lifts that restriction: - BIG_IN + n==1: unchanged -- rebase x_rsrc to block's temporal origin, all per-lane g_off values fit int32 within BIG_IN_NR (2 GB). - BIG_IN + n>1: x_rsrc stays at the tensor base; each DMA load in _load_a computes its row's sample index (n_idx) and rebases via create_buffer_resource_from_addr(x_base_addr + n_idx * X_SAMPLE_BYTES) so g_off stays int32-safe within a single sample (< 2 GB). Also reverts the XOR bank-swizzle on a_lds_off/b_lds_off: async DMA (raw_ptr_buffer_load_lds) writes LDS linearly so applying a read-side swizzle produces wrong results. The wgm L2-swizzle is kept. Fixes: shapes like (240,1024,162,92,1024) and (240,512,182,322,512) that previously triggered "BIG_IN requires n==1" assertion. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
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
1 participant
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.
Summary
Generalizes and hardens the two implicit-GEMM conv3d kernels (from #794):
128x128 / 2x4-wavetile becomes a compile-time(TILE_M, TILE_N, WAVE_M, WAVE_N), with a per-shape tile autotuner.TILE_Kand the filter(kt, kh, kw)stay compile-time constants (their im2col div/mod constant-folding is the perf source).Kt x 1 x 1convs, and a vec4 store that packs a lane's 4 accumulators into one aligned store instead of 4 scalar stores.BIG_IN) and outputs (BIG_OUT) so large tensors no longer overflow the i32 buffer offset; split-K disabled when its fp32 output would overflow.conv3d_implicit(and the FP8 variant) is now a single public entry that dispatches by filter rank: true 3D runs directly, 2D/1D reshape to the degenerate 3D case and reuse the same kernel. The rank-specific bodies live in private_conv{1,2,3}d_implhelpers; the two split-K helpers are merged into one_resolve_splitk.The kernels are renamed to drop the
_8wavesuffix:conv3d_implicit_8wave.py->conv3d_implicit.py,conv3d_implicit_8wave_fp8.py->conv3d_implicit_fp8.py,conv3d_autotune.py->conv3d_implicit_autotune.py. All of the above is applied to both the BF16 and FP8 kernels.Changes
conv3d_implicit.py(was_8wave.py)_resolve_splitkconv3d_implicit_fp8.py(was_8wave_fp8.py)conv3d_implicit_autotune.py(wasconv3d_autotune.py)do_benchover a candidate list, cached in-memory + JSON)tests/kernels/test_conv3d_implicit.py/_fp8.pyUsage
Default path unchanged --
(128, 128, 2, 4), zero behavior change. Opt in per call withautotune=True(ortile=(...)), or globally viaFLYDSL_CONV3D_AUTOTUNE=1.Results (MI355X gfx950, BF16)
Test plan
pytest tests/kernels/test_conv3d_implicit.py -v-- CDNA4 GPUpytest tests/kernels/test_conv3d_implicit_fp8.py -v-- gfx95x only