Skip to content

feat(rocke): disable unroll k for basic pipeline #12409

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bartekxk merged 4 commits into from
Sep 24, 2026
Merged

feat(rocke): disable unroll k for basic pipeline #12409

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4a943d0
feat(rocke): disable unroll k for basic pipeline
bartekxk Sep 22, 2026
2ab2c9c
fixes
bartekxk Sep 23, 2026
5c8430b
ir
bartekxk Sep 23, 2026
d42c4ea
fixes
bartekxk Sep 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 9 additions & 57 deletions dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,19 +401,6 @@ def is_valid_spec_for_problem(
f"N_gemm={problem.N_gemm} tile_n={spec.tile_n}"
Comment thread
bartekxk marked this conversation as resolved.
Resolved
)

# 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"


Expand Down Expand Up @@ -1628,52 +1615,17 @@ def _fwd_epilogue_with_acc_transform(accs_in):

final_accs = current_accs
elif spec.pipeline == "basic":
Comment thread
bartekxk marked this conversation as resolved.
Resolved
# 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)
Comment thread
bartekxk marked this conversation as resolved.
Resolved
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):
Expand Down
78 changes: 15 additions & 63 deletions dnn-providers/hip-kernel-provider/rocke/library/kernels/common/conv_implicit_gemm_wgrad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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"
)
Expand Down Expand Up @@ -2209,50 +2195,16 @@ def emit_mfma_phase(

final_accs = current_accs
elif spec.pipeline == "basic":
Comment thread
bartekxk marked this conversation as resolved.
Resolved
# 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):
Expand Down
20 changes: 6 additions & 14 deletions ...ernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_conv_build_glue.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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);
Expand Down
63 changes: 28 additions & 35 deletions ...ders/hip-kernel-provider/rocke/platform/cpp/instances/common/conv_implicit_gemm_wgrad.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Comment thread
bartekxk marked this conversation as resolved.
Resolved
{
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;
}
Expand Down Expand Up @@ -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;
}
}

Expand Down Expand Up @@ -1964,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 */
Expand Down Expand Up @@ -2488,8 +2483,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)
Comment thread
bartekxk marked this conversation as resolved.
Resolved
rocke_conv_emit_kloop_basic(&ctx);
else if(!spec->async_dma)
rocke_conv_emit_kloop_simple(&ctx);
else
Expand Down
15 changes: 7 additions & 8 deletions dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_pipeline_basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading