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 all 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
90 changes: 0 additions & 90 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 @@ -1243,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]:
Expand Down Expand Up @@ -1626,53 +1583,6 @@ def _fwd_epilogue_with_acc_transform(accs_in):
current_accs = emit_mfma_phase(cur[0], cur[1], current_accs)
b.sync()

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.
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.
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
elif not spec.async_dma:
for_op = b.scf_for_iter(c0, c_K_gemm, c_block_k, accs, iv_name="k0")
Expand Down
111 changes: 7 additions & 104 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 @@ -1960,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]:
Expand Down Expand Up @@ -2207,51 +2155,6 @@ def emit_mfma_phase(
current_accs = emit_mfma_phase(cur[0], cur[1], current_accs)
b.sync()

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))
)
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)
b.sync()
if pending_staged is not None:
emit_lds_write(pending_staged, A_smem, B_smem)
pending_staged = None

final_accs = current_accs
elif not spec.async_dma:
for_op = b.scf_for_iter(k_lo, _k_upper, c_block_k, accs, iv_name="k0")
Expand Down
34 changes: 3 additions & 31 deletions ...p-kernel-provider/rocke/platform/cpp/include/rocke/instance_conv_implicit_gemm_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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. */
Expand All @@ -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
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
Loading
Loading