[Kernel] conv3d: parameterize tile size + add autotuner#820
Draft
jiacao-amd wants to merge 7 commits into
Draft
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
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>
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_8wave(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(behavior verified bit-identical over the full input space).All of the above is applied to both the BF16 (
conv3d_implicit_8wave.py) and FP8 (_fp8.py) kernels.Changes
conv3d_implicit_8wave.pyacc[MI_M*MI_N]MFMA grid; generic double-buffered 2-stage loop (replaces the hand-scheduled 2x2-half pipeline); 3x1x1 temporal fast path; vec4 epilogue; BIG_IN/BIG_OUT 64-bit addressing; unified rank-dispatch entry + private_conv{1,2,3}d_implhelpers; merged_resolve_splitkconv3d_implicit_8wave_fp8.pyconv3d_autotune.pydo_benchover a candidate list, cached in-memory + JSON). The@autotunedecorator doesn't fit — the tile is baked into the@lru_cache'd compile factorytests/kernels/test_conv3d_implicit_8wave.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. 1D/2D/3D all go through the sameconv3d_implicit_8wave(x, weight, ...)entry (dispatched by filter rank).Results (MI355X gfx950, BF16)
Benchmarked across 10 VAE-style conv shapes (0.9 GB - 14.2 GB activations; 1x1x1, 3x1x1, 1x3x3, 3x3x3 filters), comparing this PR (autotune) against
main(fixed 128x128 tile) and MIOpen:main: faster on all 10/10 shapes, 1.26x - 2.60x (geomean 1.46x). The gain comes from autotune selecting the 256x256 tile vs main's fixed 128x128; smallest shapes and the narrow 1x1x1 projection benefit most. Fixed-tilemainplateaus at ~500-600 TFLOPS; this PR reaches 677-847 TFLOPS.main). Validated end-to-end against real WAN 2.1 and HunyuanVideo VAE decode (rel-err ~1e-3).Test plan
pytest tests/kernels/test_conv3d_implicit_8wave.py -v— CDNA4 GPUpytest tests/kernels/test_conv3d_implicit_8wave_fp8.py -v— gfx95x only