From 4a943d0f24f1f16bc4bf6ae1ec82528c71ed0513 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Tue, 22 Sep 2026 04:41:22 -0500 Subject: [PATCH 1/4] feat(rocke): disable unroll k for basic pipeline --- .../kernels/common/conv_implicit_gemm.py | 66 +++------------- .../common/conv_implicit_gemm_wgrad.py | 78 ++++--------------- .../conv_implicit_gemm_conv_build_glue.cpp | 20 ++--- .../common/conv_implicit_gemm_wgrad.cpp | 51 +++++------- 4 files changed, 50 insertions(+), 165 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py index 0614dfeddb76..3b1286e24a04 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py @@ -401,19 +401,6 @@ def is_valid_spec_for_problem( f"N_gemm={problem.N_gemm} tile_n={spec.tile_n}" ) - # Reject pipeline="basic" configs that would Python-unroll the K loop - # beyond this threshold — above it IR size explodes and comgr compilation - # time grows unacceptably. - _MAX_BASIC_K_ITERS = 10 - if spec.pipeline == "basic": - _k_iters = (problem.K_gemm + spec.tile_k - 1) // spec.tile_k - if _k_iters > _MAX_BASIC_K_ITERS: - return False, ( - f"pipeline='basic' K-loop would unroll to {_k_iters} iterations " - f"(K_gemm={problem.K_gemm} tile_k={spec.tile_k}), " - f"exceeding the {_MAX_BASIC_K_ITERS}-iteration limit" - ) - return True, "ok" @@ -1628,52 +1615,17 @@ def _fwd_epilogue_with_acc_transform(accs_in): final_accs = current_accs elif spec.pipeline == "basic": - # CK pipeline_basic: single-buffer, global-read/compute overlap. - # - # The buffer_load_vN for tile k+1 is issued before the sync+mfma for - # tile k so VMEM latency is hidden behind compute. The LDS write - # (smem_store_vN) is deferred until AFTER the second sync (after all - # ds_reads for tile k have drained), using the split emit_global_read / - # emit_lds_write helpers. Only one LDS buffer is needed. - # - # Per-iteration instruction order: - # emit_global_read(k+1) buffer_load_vN (VMEM, in flight) - # sync() s_waitcnt(lgkmcnt=0) + s_barrier - # (drains prior ds_write; tile k RAW-safe) - # k_off_capture = k (descriptor uses tile k's offset) - # emit_mfma_phase ds_read(A_smem,B_smem) + mfma - # sync() s_waitcnt(lgkmcnt=0) + s_barrier - # (drains ds_reads; A_smem WAR-safe) - # emit_lds_write(staged_k+1) smem_store_vN (now safe to write) - K_iters = (p.K_gemm + block_k - 1) // block_k - current_accs = [v for _, v in accs] - - # Prologue: global read for tile 0 then immediately write to LDS. - # (No prior ds_reads to drain, so lds_write can follow immediately.) - staged0 = emit_global_read(b.const_i32(0)) - emit_lds_write(staged0, A_smem, B_smem) - - pending_staged = None # staged tuple for the tile whose ds_write is next - - for it in range(K_iters): - # Issue buffer_load for tile it+1 BEFORE the sync. The VMEM latency - # (~300-600 cycles) overlaps with the mfma stream that follows. - if it + 1 < K_iters: - pending_staged = emit_global_read(b.const_i32((it + 1) * block_k)) - # Drain the current tile's ds_write (prologue or previous iter's - # emit_lds_write), then barrier all waves. + # Runtime K-loop (single buffer, no Python unroll). Uses the same + # emit_load_phase + sync + emit_mfma_phase + sync structure as the + # plain "mem" path so IR size stays bounded regardless of K_gemm. + for_op = b.scf_for_iter(c0, c_K_gemm, c_block_k, accs, iv_name="k0") + with for_op as (k0, iter_vars): + emit_load_phase(k0, A_smem, B_smem) b.sync() - # Set k offset so descriptors address tile it during mfma. - k_off_capture[0] = b.const_i32(it * block_k) - current_accs = emit_mfma_phase(A_smem, B_smem, current_accs) - # Drain ds_reads before the next ds_write can overwrite A_smem/B_smem. + new_accs = emit_mfma_phase(A_smem, B_smem, iter_vars) b.sync() - # Now safe to commit the next tile's staged VGPRs to LDS. - if pending_staged is not None: - emit_lds_write(pending_staged, A_smem, B_smem) - pending_staged = None - - final_accs = current_accs + b.scf_yield(*new_accs) + final_accs = for_op.results elif not spec.async_dma: for_op = b.scf_for_iter(c0, c_K_gemm, c_block_k, accs, iv_name="k0") with for_op as (k0, iter_vars): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index bdbd6f0886f3..a4bb89f1a6a5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -697,9 +697,7 @@ def validate(self) -> None: "global->LDS load needs a stride-1 reduction axis, which wgrad " "only has once the tile is stored K-outer" ) - if self.split_k == 0 and ( - self.async_dma or self.unroll_k or self.pipeline == "basic" - ): + if self.split_k == 0 and (self.async_dma or self.unroll_k): # split_k == 0 means the split degree is a launch-time kernel # argument, so the K-slice length is not known at build time. The # async and unrolled k-loops both need a compile-time trip count to @@ -709,8 +707,7 @@ def validate(self) -> None: # sweep drivers swallow into a silent skip. raise ValueError( "wgrad split_k=0 (runtime degree) is incompatible with " - "async_dma/unroll_k/pipeline='basic': those pipelines need a " - "compile-time " + "async_dma/unroll_k: those pipelines need a compile-time " "iteration count. Use a fixed split_k >= 1." ) if self.lds_k_outer: @@ -1153,13 +1150,10 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool "epilogue='cshuffle' is invalid (use epilogue='default')" ) - if spec.split_k == 0 and ( - spec.async_dma or spec.unroll_k or spec.pipeline == "basic" - ): + if spec.split_k == 0 and (spec.async_dma or spec.unroll_k): return False, ( "wgrad split_k=0 (runtime degree) is incompatible with " - "async_dma/unroll_k/pipeline='basic': those pipelines need a " - "compile-time iteration " + "async_dma/unroll_k: those pipelines need a compile-time iteration " "count. Use a fixed split_k >= 1." ) if spec.lds_k_outer and spec.lds_k_pad is not None: @@ -1234,21 +1228,13 @@ def is_valid_wgrad_spec(spec: WgradConvSpec, arch: str = "gfx950") -> Tuple[bool if spec.pipeline == "basic" and spec.async_dma: return False, "pipeline='basic' is incompatible with async_dma=True" - if spec.pipeline == "basic" or spec.async_dma: - # Both loops are unrolled in Python, one full load+mfma body per K - # iteration, so a deep reduction explodes compile time and code size. - # The cap is a build-practicality bound, not a hardware one. - # - # This used to guard 'basic' only, which left async uncapped. A deep - # reduction at a low split-K degree then unrolled five figures of - # bodies into one kernel; a sweep that reached those specs exhausted - # host memory during the IR build rather than failing validation. - _label = "pipeline='basic'" if spec.pipeline == "basic" else "async_dma" + if spec.async_dma: + # async_dma is Python-unrolled; a deep reduction explodes compile time. _slice_k = spec.wg_K_padded() // max(spec.split_k, 1) _k_iters = (_slice_k + spec.tile_k - 1) // spec.tile_k if _k_iters > _MAX_UNROLLED_K_ITERS: return False, ( - f"{_label} would unroll to {_k_iters} K iterations " + f"async_dma would unroll to {_k_iters} K iterations " f"(slice_k={_slice_k}, tile_k={spec.tile_k}), over the " f"{_MAX_UNROLLED_K_ITERS} limit; raise split_k or tile_k" ) @@ -2209,50 +2195,16 @@ def emit_mfma_phase( final_accs = current_accs elif spec.pipeline == "basic": - # CK pipeline_basic: single LDS buffer, global-read/compute overlap. - # - # The buffer_load_vN for tile it+1 is issued BEFORE the sync+mfma for - # tile it, so the VMEM latency hides behind the MFMA stream. The LDS - # write is deferred past the second sync -- once every ds_read for tile - # it has drained -- which is what makes one buffer sufficient: the tile - # it+1 data waits in VGPRs, not in a second LDS tile. unroll_k needs - # A_smem2/B_smem2 precisely because it moves the *store* inside the - # read window; this does not. - # - # Per iteration: - # emit_global_read(k_lo + (it+1)*block_k) buffer_load, in flight - # sync() drain the previous ds_write -> tile it RAW-safe - # k_off_capture = k_lo + it*block_k descriptors address tile it - # emit_mfma_phase ds_read + mfma - # sync() drain the ds_reads -> buffer WAR-safe - # emit_lds_write(staged it+1) ds_write - # - # Each k offset is materialised once and carried in the staged tuple: - # wgrad's offsets are add(k_lo, const), so re-deriving one in - # emit_lds_write would emit a stray add after the barrier. - slice_k = wg_K if k_hi is None else (spec.wg_K_padded() // spec.split_k) - K_iters = (slice_k + block_k - 1) // block_k - current_accs = [v for _, v in accs] - - # Prologue: read tile 0 and commit it immediately -- no prior ds_read - # exists to drain, so the write needs no barrier in front of it. - emit_lds_write(emit_global_read(k_lo), A_smem, B_smem) - - pending_staged = None - for it in range(K_iters): - if it + 1 < K_iters: - pending_staged = emit_global_read( - b.add(k_lo, b.const_i32((it + 1) * block_k)) - ) + # Runtime K-loop (single buffer, no Python unroll). Mirrors the plain + # "mem" path so IR size stays bounded regardless of K_gemm / split_k. + for_op = b.scf_for_iter(k_lo, _k_upper, c_block_k, accs, iv_name="k0") + with for_op as (k0, iter_vars): + emit_load_phase(k0, A_smem, B_smem) b.sync() - k_off_capture[0] = b.add(k_lo, b.const_i32(it * block_k)) - current_accs = emit_mfma_phase(A_smem, B_smem, current_accs) + new_accs = emit_mfma_phase(A_smem, B_smem, iter_vars) b.sync() - if pending_staged is not None: - emit_lds_write(pending_staged, A_smem, B_smem) - pending_staged = None - - final_accs = current_accs + b.scf_yield(*new_accs) + final_accs = for_op.results elif not spec.async_dma: for_op = b.scf_for_iter(k_lo, _k_upper, c_block_k, accs, iv_name="k0") with for_op as (k0, iter_vars): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp index c4c8a1929394..f2462f369256 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp @@ -714,22 +714,18 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv(rocke_ir_builder_t* b, /* ---- K-loop driver selection (1276-1347) ---- * * The driver is chosen by K-loop *structure*, not by pipeline name. - * "mem", "compv3", and "compv4" all use the same scf.for_iter shape - * (load->sync->mfma->sync per tile) and therefore share kloop_simple. - * Their behavioural differences (scheduling hints, double-buffering) are - * encoded in ctx->schedule and ctx->double_buffer, which are consulted - * inside emit_mfma_phase and the LDS allocation respectively -- no - * K-loop structural change is needed for those pipelines. + * "mem", "compv3", "compv4", and "basic" all use the same scf.for_iter + * shape (load->sync->mfma->sync per tile) and therefore share kloop_simple. + * Their behavioural differences (scheduling hints) are encoded in + * ctx->schedule, which is consulted inside emit_mfma_phase -- no K-loop + * structural change is needed for those pipelines. * * A new driver is only justified when the K-loop itself has a different * shape that cannot be expressed inside kloop_simple: * * unroll_k -> kloop_unroll (Python-unrolled prologue+ping-pong; * 2 LDS buffers; no scf.for_iter) - * pipeline="basic"-> kloop_basic (split global_read/lds_write; - * single LDS buffer; no scf.for_iter; - * VMEM/compute overlap without double-buf) - * else no async -> kloop_simple (scf.for_iter; mem/compv3/compv4 all + * else no async -> kloop_simple (scf.for_iter; mem/compv3/compv4/basic * share this; pipeline string only * affects schedule hints inside mfma) * else (async) -> kloop_async (SoftwarePipeline.run_ping_pong over @@ -738,10 +734,6 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv(rocke_ir_builder_t* b, { rocke_conv_emit_kloop_unroll(&ctx); } - else if(spec->pipeline != NULL && strcmp(spec->pipeline, "basic") == 0) - { - rocke_conv_emit_kloop_basic(&ctx); - } else if(spec->pipeline != NULL && strcmp(spec->pipeline, "wavelet") == 0) { rocke_conv_emit_kloop_wavelet(&ctx); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index 3d0fa06b29f3..c66a99748338 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -489,14 +489,13 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv /* split_k == 0 puts the split degree in a kernel argument, so the K-slice * length is unknown at build time; the async and unrolled k-loops both need * a compile-time trip count. Mirrors the Python validator. */ - if(s->split_k == 0 - && (s->async_dma || s->unroll_k || (s->pipeline && strcmp(s->pipeline, "basic") == 0))) + if(s->split_k == 0 && (s->async_dma || s->unroll_k)) { if(reason && reason_cap) snprintf(reason, reason_cap, "wgrad split_k=0 (runtime degree) is incompatible with " - "async_dma/unroll_k/pipeline='basic': those pipelines need a " + "async_dma/unroll_k: those pipelines need a " "compile-time iteration count. Use a fixed split_k >= 1."); return false; } @@ -570,34 +569,26 @@ bool rocke_implicit_gemm_conv_wgrad_is_valid_spec(const rocke_implicit_gemm_conv return false; } - /* Both loops are unrolled at build time, one full load+mfma body per K - * iteration, so a deep reduction explodes compile time and code size. A - * build-practicality bound, not a hardware one. This used to guard 'basic' - * only, which left async uncapped and let a low split-K degree unroll five - * figures of bodies into one kernel. Mirrors Python. */ + /* async_dma is Python-unrolled; a deep reduction explodes compile time. + * Mirrors Python is_valid_wgrad_spec. */ + if(s->async_dma) { - const bool is_basic = s->pipeline && strcmp(s->pipeline, "basic") == 0; - if(is_basic || s->async_dma) + const int spk = (s->split_k > 1) ? s->split_k : 1; + const int slice_k = rocke_wgrad_conv_spec_wg_K_padded(s) / spk; + const int k_iters = (slice_k + s->tile_k - 1) / s->tile_k; + if(k_iters > ROCKE_MAX_UNROLLED_K_ITERS) { - const char* label = is_basic ? "pipeline='basic'" : "async_dma"; - const int spk = (s->split_k > 1) ? s->split_k : 1; - const int slice_k = rocke_wgrad_conv_spec_wg_K_padded(s) / spk; - const int k_iters = (slice_k + s->tile_k - 1) / s->tile_k; - if(k_iters > ROCKE_MAX_UNROLLED_K_ITERS) - { - if(reason && reason_cap) - snprintf(reason, - reason_cap, - "%s would unroll to %d K iterations " - "(slice_k=%d, tile_k=%d), over the %d limit; " - "raise split_k or tile_k", - label, - k_iters, - slice_k, - s->tile_k, - ROCKE_MAX_UNROLLED_K_ITERS); - return false; - } + if(reason && reason_cap) + snprintf(reason, + reason_cap, + "async_dma would unroll to %d K iterations " + "(slice_k=%d, tile_k=%d), over the %d limit; " + "raise split_k or tile_k", + k_iters, + slice_k, + s->tile_k, + ROCKE_MAX_UNROLLED_K_ITERS); + return false; } } @@ -2488,8 +2479,6 @@ rocke_kernel_def_t* rocke_build_implicit_gemm_conv_wgrad( /* --- K-loop --- */ if(spec->unroll_k) rocke_conv_emit_kloop_unroll(&ctx); - else if(spec->pipeline && strcmp(spec->pipeline, "basic") == 0) - rocke_conv_emit_kloop_basic(&ctx); else if(!spec->async_dma) rocke_conv_emit_kloop_simple(&ctx); else From 2ab2c9c45f67bf9abdfc37f01ba3f5dca7795965 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Wed, 23 Sep 2026 11:47:57 +0000 Subject: [PATCH 2/4] fixes --- .../instances/common/conv_implicit_gemm_wgrad.cpp | 12 ++++++++---- .../platform/tests/core/test_pipeline_basic.py | 15 +++++++-------- 2 files changed, 15 insertions(+), 12 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp index c66a99748338..87eb7ba04a52 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp @@ -1955,12 +1955,16 @@ static bool wgrad_build_ctx_init(rocke_conv_build_ctx_t* ctx, ctx->kloop_k_lo = k_lo; /* Python: slice_k = wg_K if k_hi is None else wg_K_padded() // split_k * K_iters = ceil(slice_k / block_k) - * k_hi is None exactly when split_k == 1. */ + * k_hi is None exactly when split_k == 1. + * split_k == 0 (runtime degree): the trip count is unknown at build time; + * kloop_simple uses c_K_gemm (the runtime upper bound) directly via + * scf_for_iter so kloop_num_iters is not consulted -- leave it 0. */ { const int wgk = rocke_wgrad_conv_spec_wg_K(spec); - const int slice_k - = (k_hi_v == NULL) ? wgk : (rocke_wgrad_conv_spec_wg_K_padded(spec) / spec->split_k); - ctx->kloop_num_iters = (slice_k + ctx->block_k - 1) / ctx->block_k; + const int slice_k = (k_hi_v == NULL) ? wgk + : (split_k == 0) ? 0 /* runtime: no static trip count */ + : (rocke_wgrad_conv_spec_wg_K_padded(spec) / split_k); + ctx->kloop_num_iters = (slice_k > 0) ? (slice_k + ctx->block_k - 1) / ctx->block_k : 0; } /* Chiplet swizzle */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_pipeline_basic.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_pipeline_basic.py index 0f1e8527dce5..9bb40a9203da 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_pipeline_basic.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_pipeline_basic.py @@ -8,7 +8,8 @@ - CoalescedTileLoader.load_global() / .store_lds() IR structure - is_valid_spec rejects pipeline="basic" with async_dma=True - build_implicit_gemm_conv with pipeline="basic" produces a valid kernel - whose IR contains both buffer_load_vN and smem_store_vN ops + whose IR uses a bounded runtime scf.for_iter K-loop and contains both + buffer_load_vN and smem_store_vN ops """ from __future__ import annotations @@ -330,15 +331,13 @@ def test_build_cshuffle_epilogue(self): ] assert c_allocs, "cshuffle epilogue must allocate a C_smem tile" - def test_build_does_not_use_scf_for_iter(self): - """Unlike the simple K-loop (which uses scf.for_iter), pipeline='basic' - Python-unrolls the K loop — there should be no scf.for_iter in the IR.""" + def test_build_uses_scf_for_iter(self): + """pipeline='basic' now emits a bounded runtime scf.for K-loop + (same structure as 'mem') so IR size is independent of K_gemm.""" kernel = self._build() ops = _flatten(kernel.body, []) - for_iters = [op for op in ops if "for_iter" in op.name] - assert ( - not for_iters - ), "pipeline='basic' K-loop is Python-unrolled; scf.for_iter must not appear" + for_ops = [op for op in ops if op.name == "scf.for"] + assert for_ops, "pipeline='basic' K-loop must use scf.for (runtime loop)" if __name__ == "__main__": # pragma: no cover From 5c8430b1d978c95a99f2d1fecd8dc18ec188aea3 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Wed, 23 Sep 2026 11:52:55 +0000 Subject: [PATCH 3/4] ir --- .../rocke_representative_ir_sha256.json | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index 83f2ff5ccb49..c1b51916a6dc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -412,18 +412,18 @@ }, "conv/gfx950/n1h8c16k32r3/basic_cshuffle": { "arch": "gfx950", - "bytes": 50395, + "bytes": 22838, "family": "conv", "kernel_name": "irhash_conv_950_basic_b_N1H8W8C16_K32Y3X3_t64x64x32_w2x2_a32x32x16_basic_cshuffle", - "sha256": "8bbdf4a1910e00d6eb4590098ad320bbd896ee11fbb9d7b0754d38078ce6000a", + "sha256": "38ee311fea9a4e5c00e4a220c6a230ee989db5fef4df975eeb78b6fe1dab1b86", "status": "ok" }, "conv/gfx950/n1h8c16k32r3/basic_default": { "arch": "gfx950", - "bytes": 64236, + "bytes": 35938, "family": "conv", "kernel_name": "irhash_conv_950_basic_a_N1H8W8C16_K32Y3X3_t64x64x32_w2x2_a32x32x16_basic_default", - "sha256": "75fca036909dd993cc811cb05c8d06ec2d9e205e875a8366687f54cd03e1c1bb", + "sha256": "0a0bd6674508dc25315639a775fea5c202fe33dee8611d2e084c61f47c3718e6", "status": "ok" }, "conv/gfx950/n2h16c32k32r1": { @@ -1617,18 +1617,18 @@ }, "conv/gfx950/n1h8c16k32r3/basic_cshuffle": { "arch": "gfx950", - "bytes": 50534, + "bytes": 22971, "family": "conv", "kernel_name": "irhash_conv_950_basic_b_N1H8W8C16_K32Y3X3_t64x64x32_w2x2_a32x32x16_basic_cshuffle", - "sha256": "80da37bd539fa495c55f6c7bec0ebaa3a16d8a207dcf1e18d519f61a1eb1b789", + "sha256": "07df00c797d0664453a56b00691ef59150c9b8512a46ba9b00c464441a12abda", "status": "ok" }, "conv/gfx950/n1h8c16k32r3/basic_default": { "arch": "gfx950", - "bytes": 64375, + "bytes": 36077, "family": "conv", "kernel_name": "irhash_conv_950_basic_a_N1H8W8C16_K32Y3X3_t64x64x32_w2x2_a32x32x16_basic_default", - "sha256": "e712f2a45210cfb28be9e28bf43ef28b05042f2cfae87aa1c4d83b99ec859d31", + "sha256": "ef97e3df68a627ba468cf4e9a66592cc529bd1521b3dd20e1422995381364c9d", "status": "ok" }, "conv/gfx950/n2h16c32k32r1": { @@ -2822,18 +2822,18 @@ }, "conv/gfx950/n1h8c16k32r3/basic_cshuffle": { "arch": "gfx950", - "bytes": 50538, + "bytes": 22975, "family": "conv", "kernel_name": "irhash_conv_950_basic_b_N1H8W8C16_K32Y3X3_t64x64x32_w2x2_a32x32x16_basic_cshuffle", - "sha256": "b30fed8bd47d12db2a05694a4a0d02338898461bf76d9864468c04d1f902e533", + "sha256": "c226cd0c1239454e8928b95a9e847825bf314c3f29a18237ad9aea3db3a645da", "status": "ok" }, "conv/gfx950/n1h8c16k32r3/basic_default": { "arch": "gfx950", - "bytes": 64379, + "bytes": 36081, "family": "conv", "kernel_name": "irhash_conv_950_basic_a_N1H8W8C16_K32Y3X3_t64x64x32_w2x2_a32x32x16_basic_default", - "sha256": "a0b6acb3748dae3ee5e3961503cb12b7a08bbf03d232f7baf24b726a0ba74bd3", + "sha256": "7a64ce1c297e125855241f895544397f90b56b956c43cf3d98873bee535a6e94", "status": "ok" }, "conv/gfx950/n2h16c32k32r1": { From d42c4ea52f13f75f7d5fc2f957fa248bbfe0d960 Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Wed, 23 Sep 2026 13:44:38 +0000 Subject: [PATCH 4/4] fixes --- .../kernels/common/conv_implicit_gemm.py | 42 --------- .../common/conv_implicit_gemm_wgrad.py | 49 ---------- .../instance_conv_implicit_gemm_internal.h | 34 +------ .../conv_implicit_gemm_conv_kloop_drivers.cpp | 89 ------------------- .../conv_implicit_gemm_conv_load_phase.cpp | 54 ----------- .../dsl_docs/instances/convolution.md | 4 +- .../optimization/optimization_runbook.md | 4 +- 7 files changed, 7 insertions(+), 269 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py index 3b1286e24a04..6519661267d0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py @@ -1230,36 +1230,6 @@ def emit_load_phase(k_off: Value, A_dst: Value, B_dst: Value) -> None: rsrc=b_rsrc, ) - def emit_global_read(k_off: Value) -> tuple: - """Issue only the global memory reads (buffer_load_vN) for one K tile. - - Returns ``(k_off, a_staged, b_staged)`` — the tile offset and the two - lists of ``(row, col, v)`` triples from :meth:`CoalescedTileLoader.load_global`. - The caller must later call :func:`emit_lds_write` to commit these values - to LDS. Only valid on the sync (non-async-DMA) path; CK pipeline_basic - uses this to overlap VMEM latency with MFMA compute. - """ - k_off_capture[0] = k_off - a_staged = a_sync_loader.load_global( - b, tid=tid, descriptor=a_descriptor, rsrc=a_rsrc - ) - b_staged = b_sync_loader.load_global( - b, tid=tid, descriptor=b_descriptor, rsrc=b_rsrc - ) - return k_off, a_staged, b_staged - - def emit_lds_write(staged_tuple: tuple, A_dst: Value, B_dst: Value) -> None: - """Commit previously-staged VGPR values to LDS (smem_store_vN). - - ``staged_tuple`` is the value returned by :func:`emit_global_read`. - Restores ``k_off_capture`` so the descriptor sees the correct k offset - even though the global read and LDS write happen in different loop positions. - """ - k_off, a_staged, b_staged = staged_tuple - k_off_capture[0] = k_off - a_sync_loader.store_lds(b, smem_dst=A_dst, staged=a_staged) - b_sync_loader.store_lds(b, smem_dst=B_dst, staged=b_staged) - def emit_wmma_phase( A_src: Value, B_src: Value, iter_vars: Sequence[Value] ) -> List[Value]: @@ -1614,18 +1584,6 @@ def _fwd_epilogue_with_acc_transform(accs_in): b.sync() final_accs = current_accs - elif spec.pipeline == "basic": - # Runtime K-loop (single buffer, no Python unroll). Uses the same - # emit_load_phase + sync + emit_mfma_phase + sync structure as the - # plain "mem" path so IR size stays bounded regardless of K_gemm. - for_op = b.scf_for_iter(c0, c_K_gemm, c_block_k, accs, iv_name="k0") - with for_op as (k0, iter_vars): - emit_load_phase(k0, A_smem, B_smem) - b.sync() - new_accs = emit_mfma_phase(A_smem, B_smem, iter_vars) - b.sync() - b.scf_yield(*new_accs) - final_accs = for_op.results elif not spec.async_dma: for_op = b.scf_for_iter(c0, c_K_gemm, c_block_k, accs, iv_name="k0") with for_op as (k0, iter_vars): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py index a4bb89f1a6a5..5c0b9e931a97 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py @@ -1946,44 +1946,6 @@ def _x_kouter(b_, row, col): return _dy_kouter, _x_kouter return dy_descriptor, x_descriptor - def emit_global_read(k_off: Value) -> tuple: - """Issue only the global reads (buffer_load_vN) for one K tile. - - Returns ``(k_off, a_staged, b_staged)`` -- the tile offset plus the two - lists of ``(row, col, v)`` triples from - :meth:`CoalescedTileLoader.load_global`. The caller commits them later - with :func:`emit_lds_write`. Sync path only; this is what lets the - CK pipeline_basic loop overlap VMEM latency with MFMA compute. - - ``k_off`` is carried in the tuple because wgrad's offsets are - ``add(k_lo, const)`` -- a real emitted op, unlike the forward conv's - cached ``const_i32`` -- so re-deriving it in :func:`emit_lds_write` - would strand an extra add after the barrier. - """ - k_off_capture[0] = k_off - a_desc_fn, b_desc_fn = _split_desc_fns() - a_staged = a_sync_loader.load_global( - b, tid=tid, descriptor=a_desc_fn, rsrc=dy_rsrc - ) - b_staged = b_sync_loader.load_global( - b, tid=tid, descriptor=b_desc_fn, rsrc=x_rsrc - ) - return k_off, a_staged, b_staged - - def emit_lds_write(staged_tuple: tuple, A_dst: Value, B_dst: Value) -> None: - """Commit previously staged VGPR values to LDS. - - Restores ``k_off_capture`` so any descriptor consulted here sees the - offset the values were read at, even though the read and the write sit - in different loop positions. ``store_lds`` funnels through the same - ``_store_tile`` as the fused loader, so ``vector_axis="row"`` replays - the transpose-on-store scatter unchanged. - """ - k_off, a_staged, b_staged = staged_tuple - k_off_capture[0] = k_off - a_sync_loader.store_lds(b, smem_dst=A_dst, staged=a_staged) - b_sync_loader.store_lds(b, smem_dst=B_dst, staged=b_staged) - def emit_wmma_phase( A_src: Value, B_src: Value, iter_vars: Sequence[Value] ) -> List[Value]: @@ -2194,17 +2156,6 @@ def emit_mfma_phase( b.sync() final_accs = current_accs - elif spec.pipeline == "basic": - # Runtime K-loop (single buffer, no Python unroll). Mirrors the plain - # "mem" path so IR size stays bounded regardless of K_gemm / split_k. - for_op = b.scf_for_iter(k_lo, _k_upper, c_block_k, accs, iv_name="k0") - with for_op as (k0, iter_vars): - emit_load_phase(k0, A_smem, B_smem) - b.sync() - new_accs = emit_mfma_phase(A_smem, B_smem, iter_vars) - b.sync() - b.scf_yield(*new_accs) - final_accs = for_op.results elif not spec.async_dma: for_op = b.scf_for_iter(k_lo, _k_upper, c_block_k, accs, iv_name="k0") with for_op as (k0, iter_vars): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h index 63f8d31bfdd2..bbd8be27bcb5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h @@ -466,29 +466,6 @@ void rocke_conv_emit_load_phase(rocke_conv_build_ctx_t* ctx, rocke_value_t* A_dst, rocke_value_t* B_dst); -/* Split-load helpers for CK pipeline_basic (sync path only). - * - * emit_global_read: issue only buffer_load_vN for A and B into VGPR staging. - * Sets ctx->k_off_capture = k_off so descriptors address the correct tile. - * Fills *a_staged and *b_staged (caller-allocated); these are consumed later - * by emit_lds_write. Mirrors Python emit_global_read() -> (k_off, a_staged, b_staged). - * - * emit_lds_write: commit the staged VGPRs to LDS via smem_store_vN. - * Restores ctx->k_off_capture = k_off (from the staged tuple) so any - * descriptor that re-reads k_off_capture sees the correct value. - * Mirrors Python emit_lds_write(staged_tuple, A_dst, B_dst). */ -void rocke_conv_emit_global_read(rocke_conv_build_ctx_t* ctx, - rocke_value_t* k_off, - rocke_ctl_staged_t* a_staged, - rocke_ctl_staged_t* b_staged); - -void rocke_conv_emit_lds_write(rocke_conv_build_ctx_t* ctx, - rocke_value_t* k_off, - const rocke_ctl_staged_t* a_staged, - const rocke_ctl_staged_t* b_staged, - rocke_value_t* A_dst, - rocke_value_t* B_dst); - /* emit_wmma_phase(ctx, A_src, B_src, iter_vars[n], out_accs[n]): one K-tile of * WMMA atoms, fully MMA-contract driven (gfx1151). Reads iter_vars (length * ctx->num_accs), writes the new accs into out_accs. */ @@ -513,20 +490,15 @@ void rocke_conv_emit_mfma_phase(rocke_conv_build_ctx_t* ctx, /* ----- K-loop drivers (ctx-driven; write ctx->final_accs) ----- * * Exactly one is called per build, chosen as Python does: * unroll_k -> rocke_conv_emit_kloop_unroll - * pipeline=="basic" -> rocke_conv_emit_kloop_basic - * else not async_dma -> rocke_conv_emit_kloop_simple + * else not async_dma -> rocke_conv_emit_kloop_simple (mem/compv3/compv4/basic) * else (async_dma) -> rocke_conv_emit_kloop_async */ /* spec.unroll_k branch (lines 1276-1310): double-buffered Python-unrolled * software pipeline (ping-pong A_smem/A_smem2). */ void rocke_conv_emit_kloop_unroll(rocke_conv_build_ctx_t* ctx); -/* pipeline=="basic" branch: CK pipeline_basic single-buffer, global-read/compute - * overlap. Global read for tile k+1 is issued before the sync+mfma for tile k - * so VMEM latency is hidden behind compute. Single LDS buffer, no double-buf. */ -void rocke_conv_emit_kloop_basic(rocke_conv_build_ctx_t* ctx); - -/* not-async branch (lines 1311-1319): single scf.for_iter load+sync+mfma+sync. */ +/* not-async branch (lines 1311-1319): single scf.for_iter load+sync+mfma+sync. + * Covers mem, compv3, compv4, and basic (basic now uses the same runtime loop). */ void rocke_conv_emit_kloop_simple(rocke_conv_build_ctx_t* ctx); /* async_dma branch (lines 1320-1347): SoftwarePipeline.run_ping_pong over the diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp index 9db0d09e97f8..f97508f863d7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_kloop_drivers.cpp @@ -114,95 +114,6 @@ void rocke_conv_emit_kloop_unroll(rocke_conv_build_ctx_t* ctx) rocke_conv_set_final_accs(ctx, current_accs, num_accs); } -/* ===================================================================== * - * rocke_conv_emit_kloop_basic - * - * CK pipeline_basic: single-buffer, global-read/compute overlap. - * Mirrors the Python ``elif spec.pipeline == "basic":`` branch. - * - * Per-iteration order (byte-identical to Python): - * emit_global_read(it+1) buffer_load_vN (VMEM, in flight) - * sync() s_waitcnt(lgkmcnt=0) + s_barrier - * (drains prior ds_write; tile it RAW-safe) - * k_off_capture = it descriptor uses tile it's offset for mfma - * emit_mfma_phase ds_read(A_smem,B_smem) + mfma - * sync() s_waitcnt(lgkmcnt=0) + s_barrier - * (drains ds_reads; A_smem WAR-safe) - * emit_lds_write(staged) smem_store_vN (safe to write now) - * ===================================================================== */ -void rocke_conv_emit_kloop_basic(rocke_conv_build_ctx_t* ctx) -{ - rocke_ir_builder_t* b = ctx->b; - const rocke_conv_problem_t* p = ctx->p; - int block_k = ctx->block_k; - int K_iters = (ctx->kloop_num_iters > 0) - ? ctx->kloop_num_iters - : (rocke_conv_problem_k_gemm(p) + block_k - 1) / block_k; - int num_accs = ctx->num_accs; - int it, i; - - rocke_value_t* current_accs[ROCKE_CONV_MAX_ACCS]; - rocke_value_t* new_accs[ROCKE_CONV_MAX_ACCS]; - - /* pending_staged_{a,b}: staging buffers for the next tile's VGPRs. - * has_pending mirrors Python `pending_staged is not None`. */ - rocke_ctl_staged_t pending_a; - rocke_ctl_staged_t pending_b; - rocke_value_t* pending_k_off; - int has_pending; - - /* staged_{a,b}_0: prologue staging buffers (tile 0). */ - rocke_ctl_staged_t staged_a0; - rocke_ctl_staged_t staged_b0; - rocke_value_t* k0_val; - - for(i = 0; i < num_accs; ++i) - current_accs[i] = ctx->acc_inits[i]; - - /* Prologue: global read for tile 0 then write to LDS immediately. - * staged0 = emit_global_read(b.const_i32(0)) - * emit_lds_write(staged0, A_smem, B_smem) */ - /* wgrad slices the reduction, so the prologue tile starts at kloop_k_lo, - * not 0. kloop_k_off returns the bare const when kloop_k_lo is NULL, so - * the forward conv is unchanged. */ - k0_val = (ctx->kloop_k_lo != NULL) ? ctx->kloop_k_lo : rocke_b_const_i32(b, 0); - rocke_conv_emit_global_read(ctx, k0_val, &staged_a0, &staged_b0); - rocke_conv_emit_lds_write(ctx, k0_val, &staged_a0, &staged_b0, ctx->A_smem, ctx->B_smem); - - has_pending = 0; - pending_k_off = NULL; - - for(it = 0; it < K_iters; ++it) - { - /* Issue buffer_load for tile it+1 BEFORE the sync. */ - if(it + 1 < K_iters) - { - pending_k_off = kloop_k_off(b, ctx, (int64_t)(it + 1) * block_k); - rocke_conv_emit_global_read(ctx, pending_k_off, &pending_a, &pending_b); - has_pending = 1; - } - /* Drain prior ds_write then barrier. */ - rocke_b_sync(b); - /* Set k offset for mfma descriptors. */ - ctx->k_off_capture = kloop_k_off(b, ctx, (int64_t)it * block_k); - /* ds_read + mfma for the current tile. */ - rocke_conv_emit_mfma_phase(ctx, ctx->A_smem, ctx->B_smem, current_accs, num_accs, new_accs); - for(i = 0; i < num_accs; ++i) - current_accs[i] = new_accs[i]; - /* Drain ds_reads before the next smem_store_vN. */ - rocke_b_sync(b); - /* Commit the next tile's staged VGPRs to LDS. */ - if(has_pending) - { - rocke_conv_emit_lds_write( - ctx, pending_k_off, &pending_a, &pending_b, ctx->A_smem, ctx->B_smem); - has_pending = 0; - } - } - - rocke_conv_set_final_accs(ctx, current_accs, num_accs); -} - /* ===================================================================== * * rocke_conv_emit_kloop_simple (Python lines 1311-1319) * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp index 6ab0f0d44f00..de7bd3c35631 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_load_phase.cpp @@ -313,57 +313,3 @@ void rocke_conv_emit_load_phase(rocke_conv_build_ctx_t* ctx, b, &ctx->b_sync_loader, ctx->tid, B_dst, b_fn, ctx, ctx->b_rsrc, NULL); } } - -/* ===================================================================== * - * rocke_conv_emit_global_read -- split load, phase 1 of 2. - * - * Mirrors Python emit_global_read(k_off) -> (k_off, a_staged, b_staged). - * Sets ctx->k_off_capture = k_off (so the descriptor closures see the - * correct tile offset during load_global), then issues only the - * buffer_load_vN ops for A and B into VGPR SSA values stored in - * *a_staged / *b_staged. The smem_store_vN ops are deferred to - * rocke_conv_emit_lds_write. Only valid on the sync (non-async-DMA) path. - * ===================================================================== */ -void rocke_conv_emit_global_read(rocke_conv_build_ctx_t* ctx, - rocke_value_t* k_off, - rocke_ctl_staged_t* a_staged, - rocke_ctl_staged_t* b_staged) -{ - rocke_ir_builder_t* b = ctx->b; - /* Honour the per-family descriptor overrides, as rocke_conv_emit_load_phase - * already does. wgrad installs its own (including the lds_k_outer - * coordinate-swap variants); without the fallback a wgrad basic kernel - * would emit forward-conv addressing. Both are NULL on the forward path, - * so its emission is unchanged. */ - rocke_loads_descriptor_fn a_fn - = (ctx->a_descriptor_fn != NULL) ? ctx->a_descriptor_fn : rocke_conv_a_descriptor; - rocke_loads_descriptor_fn b_fn - = (ctx->b_descriptor_fn != NULL) ? ctx->b_descriptor_fn : rocke_conv_b_descriptor; - /* k_off_capture[0] = k_off */ - ctx->k_off_capture = k_off; - rocke_coalesced_tile_loader_load_global( - b, &ctx->a_sync_loader, ctx->tid, a_fn, ctx, ctx->a_rsrc, NULL, a_staged); - rocke_coalesced_tile_loader_load_global( - b, &ctx->b_sync_loader, ctx->tid, b_fn, ctx, ctx->b_rsrc, NULL, b_staged); -} - -/* ===================================================================== * - * rocke_conv_emit_lds_write -- split load, phase 2 of 2. - * - * Mirrors Python emit_lds_write(staged_tuple, A_dst, B_dst). - * Restores ctx->k_off_capture = k_off (from the staged tuple's recorded - * offset), then emits smem_store_vN for the VGPRs staged by a prior - * rocke_conv_emit_global_read call. - * ===================================================================== */ -void rocke_conv_emit_lds_write(rocke_conv_build_ctx_t* ctx, - rocke_value_t* k_off, - const rocke_ctl_staged_t* a_staged, - const rocke_ctl_staged_t* b_staged, - rocke_value_t* A_dst, - rocke_value_t* B_dst) -{ - rocke_ir_builder_t* b = ctx->b; - ctx->k_off_capture = k_off; - rocke_coalesced_tile_loader_store_lds(b, &ctx->a_sync_loader, A_dst, a_staged); - rocke_coalesced_tile_loader_store_lds(b, &ctx->b_sync_loader, B_dst, b_staged); -} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md index 19ea2a3dbbce..f390fd90b758 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/instances/convolution.md @@ -426,9 +426,9 @@ M-outer default for dgrad, and matches the fp32 reference for wgrad. ### Backward Knob Changes -- `pipeline="basic"` (CK `pipeline_basic`): single LDS buffer, global-read / compute overlap. wgrad and forward only; dgrad has no `basic` branch. +- `pipeline="basic"`: now uses a bounded runtime `scf.for` K-loop (load → sync → MFMA → sync per tile), identical to the `"mem"` path. No K-trip-count limit, no `split_k=0` restriction on wgrad. dgrad has no `basic` branch. - `async_dma` is a swept axis on the wgrad sweep driver, not a run-level flag. Unlike `lds_k_outer` it is not deducible from `(arch, spec)`: it removes the register staging of the tile, but it also forces the K-outer row pad to 0 and coarsens the load-width ladder to the widths the intrinsic accepts, and both terms are functions of tile width and channel run, which are themselves sweep axes. -- `_MAX_UNROLLED_K_ITERS = 128` (wgrad) now caps **both** statically-unrolled loops, `pipeline="basic"` and `async_dma`. It previously guarded only `"basic"`, leaving async uncapped: a deep reduction at a low split-K degree then unrolled five figures of load+MFMA bodies into one kernel and exhausted host memory during the IR build rather than failing validation. Mirrored as `ROCKE_MAX_UNROLLED_K_ITERS` in the C engine. +- `_MAX_UNROLLED_K_ITERS = 128` (wgrad) caps the statically-unrolled `async_dma` loop only. `pipeline="basic"` is no longer unrolled and is not bounded by this constant. Mirrored as `ROCKE_MAX_UNROLLED_K_ITERS` in the C engine. - Removed: the `--lds-k-outer`, `--lds-k-pad` and `--dtype-d` CLI flags and the `ROCKE_WGRAD_LDS_K_OUTER` env override. All replaced by deduction or dropped. ## Direct Grouped Convolution diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index 33f3ecbe3e6c..2621e20254eb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -1849,10 +1849,10 @@ matches → zero re-pack. See §17.4 for the quantitative reduction. | Knob | Spec | Values | Effect | |---|---|---|---| | `pipeline` | GEMM `TraitSpec` | `"mem"` / `"compv3"` / `"compv4"` | `mem` = single-buffer; `compv4` = double-buffered async DMA + MFMA overlap. Compv3 / compv4 trade LDS for latency hiding | -| `pipeline` | conv_implicit_gemm, conv_implicit_gemm_wgrad | `"mem"` / `"compv3"` / `"compv4"` / `"wavelet"` / `"basic"` | The GEMM names above plus `"basic"` = CK `pipeline_basic`: a single LDS buffer, with the global read of tile `k+1` issued before the `sync` + MFMA of tile `k` and the matching LDS write deferred past the second `sync`, so VMEM latency overlaps MFMA compute and tile `k+1` waits in VGPRs rather than in a second LDS tile. Statically unrolls the K loop — bounded by `_MAX_UNROLLED_K_ITERS` on wgrad and by the parallel `_MAX_BASIC_K_ITERS` in `conv_implicit_gemm.py` on forward conv — is mutually exclusive with `async_dma`, and needs a compile-time trip count, so on wgrad it rejects `split_k=0` | +| `pipeline` | conv_implicit_gemm, conv_implicit_gemm_wgrad | `"mem"` / `"compv3"` / `"compv4"` / `"wavelet"` / `"basic"` | The GEMM names above plus `"basic"` = same runtime `scf.for` K-loop as `"mem"` (load → sync → MFMA → sync per tile); it is mutually exclusive with `async_dma` but otherwise imposes no K-trip-count limit and supports `split_k=0` on wgrad | | `scheduler` | GEMM `TraitSpec` | `"intrawave"` / `"interwave"` | Where the scheduler injects waits. Intrawave keeps producers and consumers in one wave; interwave splits them | | `async_dma` | conv_implicit_gemm, conv_implicit_gemm_wgrad, conv_implicit_gemm_dgrad | False / True | Enable direct global → LDS DMA (`raw_ptr_buffer_load_lds`). It **pins the schedule**: forward conv and wgrad both build `SchedulePolicy.for_pipeline("async_dma")` (interwave, `s_setprio 1`) and ignore `spec.pipeline` on that leg, which is why the sweep pins the pipeline to `"mem"` there instead of compiling one body under five kernel names. Incompatible with `pipeline="basic"` (all conv families) and with `pipeline="wavelet"` (forward conv, dgrad). A **swept axis**, not a run-level flag — unlike `lds_k_outer` it is not deducible from `(arch, spec)`: it removes the register staging of the tile, but it also forces the LDS pad to 0 (the intrinsic writes a packed lane-contiguous tile) and coarsens the load-width ladder to the chunk widths the intrinsic accepts, and both terms are functions of tile width and channel run, which are themselves sweep axes. On wgrad it *requires* `lds_k_outer=True`; on dgrad it is *rejected* under `lds_k_outer` (the tilde builder has no direct global→LDS path). Python-unrolls the K loop, so it is bounded by `_MAX_UNROLLED_K_ITERS` | -| `_MAX_UNROLLED_K_ITERS` | `instances/common/conv_implicit_gemm_wgrad.py` module constant (mirrored as `ROCKE_MAX_UNROLLED_K_ITERS` in the C engine) | 128 | Cap on `ceil((wg_K_padded / split_k) / tile_k)` for **both** statically-unrolled wgrad loops — `pipeline="basic"` **and** `async_dma`. A build-practicality bound, not a hardware one: over the cap `is_valid_wgrad_spec` rejects the spec, so raise `split_k` or `tile_k` rather than raising the constant. It previously guarded `basic` only, leaving async uncapped — a deep reduction at a low split-K degree then unrolled five figures of load+MFMA bodies into one kernel and exhausted host memory during the IR build instead of failing validation. Forward conv keeps its own parallel `_MAX_BASIC_K_ITERS = 128` in `conv_implicit_gemm.py`, which still guards `pipeline="basic"` only | +| `_MAX_UNROLLED_K_ITERS` | `instances/common/conv_implicit_gemm_wgrad.py` module constant (mirrored as `ROCKE_MAX_UNROLLED_K_ITERS` in the C engine) | 128 | Cap on `ceil((wg_K_padded / split_k) / tile_k)` for the statically-unrolled wgrad `async_dma` loop. A build-practicality bound, not a hardware one: over the cap `is_valid_wgrad_spec` rejects the spec, so raise `split_k` or `tile_k` rather than raising the constant. `pipeline="basic"` no longer applies — it uses a runtime `scf.for` loop and is not bounded by this constant | | `unroll_k` | conv_implicit_gemm | False | Python-time unroll of the K loop. Bigger code, fewer waits | | `use_early_v_schedule` | Attention 2D | False | Issue current-V async copy before QK so V overlaps QK + softmax (§8.1, §17.4). Use only on no-SW prefill | | `prefetch distance` (implicit) | `pipeline` choice | — | More stages ⇒ better latency hiding but more LDS |