You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
feat(rocke): stream-K for implicit-GEMM conv wgrad - #13151
Wgrad GEMMs usually have a filter-sized output (few output tiles) and a very long reduction (NHoWo). Split-K cuts every tile the same number of ways, so on these shapes it leaves CUs idle. This PR ports CK Tile's stream-K to rocKE wgrad. Stream-K balances MAC iterations instead of tiles:
Tiles that fill the device run data-parallel.
Only the remainder is spread across a pool of CTAs.
Stream-K is opt-in. The grouped-convolution dispatcher and its benchmark are unchanged (they are due for a refactor), so no production path picks stream-K on its own.
Technical Details
Core IR: three new leaf ops in both engines: memref.fence, memref.global_flag_store and memref.global_flag_wait_eq.
Each lowers to an LLVM fence, or to an atomic store/load with release/acquire ordering at a chosen sync scope.
The Python and C++ verifiers check arity, a global i32 flag pointer, i32 index/value operands, and the allowed scope/ordering values.
All four lowerers reject a missing or unknown scope/ordering with the same message.
Partitioner:StreamKIterPartition (Python) and rocke_streamk_iter_partition_t (C++) mirror CK's DP + SK split.
iters_per_sk_cta is clamped to at least 1, so an all-DP partition is safe.
flags_bytes(n_waves) sizes the per-(CTA, wave) flag table.
The device-side emitters read the partition from runtime values.
Wgrad kernel:
Launch modes: streamk = dp_sk | persistent.
Reductions: linear and tree are deterministic and hand off through per-(CTA, wave) flags. atomic (fp32 dW only) and workspace (f32 scratch + Stage-2 reduce) are the other two.
Conv groups are folded into GEMM-M per group, so the fold is exact for any wg_M. The C++ engine now builds grouped wgrad on the stream-K path.
AOT: the partition is passed as an sk_* kernarg block, appended after ks_count so existing kernarg offsets don't change.
Python side: conv_abi and ConvArgs.streamk_partition.
C++ side: rocke_conv_arg_names_ex.
The CTA pool is a launch parameter, so one binary serves any problem shape and any pool. streamk_ctas = -1 resolves to the arch default pool, in both engines.
Sweep / benchmark (opt-in):benchmark_implicit_gemm_conv.py (also reachable through benchmark_conv_compare.py) gains --streamk {off,dp_sk,persistent,all} (default off), --streamk-reductions and --streamk-ctas.
Stream-K wgrad candidates are added to the JIT sweep, the --compile-all AOT cache (kernel_sweep.py), and --run-from-cache.
With stream-K off, the cache identity and kernel labels are unchanged.
The stream-K launcher allocates the reduction's workspace and zeroes what each reduction accumulates into. It shares the Stage-2 reduce with the two-stage path.
Constraints: gfx942 and gfx950 only, MFMA wave64. Both engines reject the following with identical messages: split-K, two-stage, async DMA, unroll, cshuffle, chiplet swizzle and group merge.
Test Plan
Byte identity across both engines at llvm20, llvm22 and llvm23.
Differential parity: new wgrad configs 21–30 and negative configs 104–108.
Representative-IR golden: all 8 {dp_sk, persistent} × {linear, tree, atomic, workspace} combos (dense gfx950 plus grouped gfx942).
CPU tests:
partitioner invariants
validator agreement across both engines, including multi-cause rejects and the arch gate
admits-implies-builds matrix
host plan and default-pool agreement between Python and C++
stream-K ABI and launch values
sweep launcher, including all-DP partitions
dual-engine lowering
C++ ctests: rocke_streamk_iter_partition, rocke_wgrad_streamk_host, rocke_future_intrinsic_lowering and rocke_ir_serialize_roundtrip.
GPU tests:
numerics for both launch modes × all four reductions on dense, strided and grouped shapes, including a partial last M tile
one binary launched on several shapes (AOT)
20-launch bit-exact repeat test for linear and tree, with partials and dW poisoned before each launch
flag ordering pinned in IR (acquire spin / release publish) and ISA (cache invalidate / writeback)
sweep launcher on all-DP partitions (linear and workspace)
Test Result
Byte identity GREEN for all families at llvm20, llvm22 and llvm23. 31/31 wgrad configs pass differential parity.
Representative-IR golden OK at all three flavors after merging develop.
CPU tests, C++-binding tests and ctests pass.
GPU tests pass on gfx950. gfx942 is covered by the IR, golden and CPU gates only.
Add three leaf ops that let one workgroup hand data to another through
global memory, needed by the deterministic stream-K fixups:
- memref.fence: LLVM fence at workgroup / agent / system scope.
- memref.global_flag_store: scoped atomic i32 store of a flag.
- memref.global_flag_wait_eq: self-contained spin loop around an
acquire atomic load until the flag matches.
The protocol is expressed through the AMDGPU memory model rather than
raw cache-policy asm, so an agent-scope release writes back and an
acquire invalidates the per-XCD L2 on multi-XCD CDNA parts. The spin is
a leaf op, so structured control flow stays scf.for / scf.if and no
region op or recipe ABI change is needed.
Both engines implement the builders, verifier attrs, LLVM lowering and
HIP lowering. The new opcodes are appended after the control-flow block
so no existing opcode index moves. TestNewTargetIntrinsics and its C++
twin pin the emitted text and cross-check the engines.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add StreamKIterPartition, the CK Tile StreamKTilePartitioner resolved at
build time: output tiles that divide the CTA pool run data-parallel and
the remainder is spread over max_active_wgs stream-K CTAs, each owning
a contiguous, balanced range of MAC iterations that may straddle output
tiles. The existing StreamKPartition used by streamk_gemm is unchanged.
The helper exposes the host-side plan (DP/SK split, grid size, Linear
partner and Tree round bounds, flag-table size), start/end iteration of
a stream-K CTA, and SSA emitters that map a CTA onto its iteration
range. The C port returns the resolved plan by value so no derived
quantity can be recomputed with different rounding.
Reference vectors and an invariant sweep (contiguous coverage, at most
one non-owner tile per stream-K CTA, contributor bounds) run in both
engines; the C test is registered as rocke_streamk_iter_partition.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Port CK Tile's stream-K to the backward-weight convolution in both
engines. Split-K gives every output tile the same number of K slices,
which leaves CUs idle on a wgrad: the output is filter-sized while the
reduction N*Ho*Wo is long, so the tile count can be smaller than the
device no matter how the slices are cut. Stream-K balances MAC
iterations instead: tiles that divide the CTA pool run data-parallel
with no reduction at all, and only the remainder is spread over the
pool, each CTA owning a contiguous, balanced iteration range that may
straddle tile boundaries.
The whole partition is resolved on the host, since rocke compiles one
kernel per shape, so the device only maps its linear CTA index onto an
iteration range and walks it with one bounded loop.
Spec: streamk ("off" | "dp_sk" | "persistent"), streamk_reduction
("linear" | "tree" | "atomic" | "workspace") and streamk_ctas. All
default off or to the existing behaviour, so every current config keeps
its kernel name and emits byte-identical IR.
linear and tree are deterministic: a contributor publishes its partial
tile, releases, and sets a per-wave flag; the consumer acquires that
flag and folds in a fixed order, so dW needs no pre-zeroing. atomic and
workspace reuse the existing split-K and two-stage epilogues.
Conv groups fold into GEMM-M rather than riding block_id_z. Defining
the tile count per group keeps the fold exact when the per-group M is
not a multiple of the tile, which the upstream fold does not cover.
This is the one path on which the C port accepts groups > 1; the
grid-per-group path stays rejected there as before.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Add parity configs 18-28 covering both launch modes, all four
reductions, both MFMA arches, grouped shapes with a partial last M tile
per group, the all-data-parallel partition and the K-outer tile. The
byte-identity gate is GREEN for all 29 wgrad configs at llvm20, llvm22
and llvm23.
Add three representative-IR cases and re-bless that golden; the diff is
append-only, no existing hash moves.
The L5 differential anchor is deliberately left alone. It carries no
conv_implicit_gemm_wgrad entry today, and on a box whose clang is a
different LLVM vintage than the one it was blessed on it already
reports most of its families as changed -- so entries recorded here
would be anchored against a different toolchain state than the rest of
the file. Bless that family from the canonical toolchain instead.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
CPU: the spec surface and kernel-name tags, every validator gate with
validate() and is_valid_wgrad_spec agreeing on the reason string, the
C port rejecting with that same string, admits-implies-builds over
mode x reduction x arch x groups, the host launch plan cross-checked
against the C port through a new binding, and byte-identical lowering
between the engines.
GPU: numerics against a float32 torch reference on shapes chosen to
exercise the partitioner rather than the GEMM -- up to four CTAs per
tile (so the linear fold and both tree rounds run), a data-parallel
plus stream-K mix, and a grouped problem whose last M tile is partial.
A row-coverage guard reports whole unwritten output channels as such
instead of as an opaque tolerance miss. The linear and tree fixups are
launched repeatedly and must be bit-identical, which is what would
catch a missing release or acquire on the partial hand-off.
Both files are registered in the CTest lists so CI runs them.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Make stream-K reachable from the shipped selector instead of only from
a hand-built spec: ConvGroupedRequest gains streamk and
streamk_reduction, which pin split_k=1 and the direct epilogue and flow
into the instance spec. The launch grid routes through the instance's
own stream-K grid, so the spec the kernel is built from and the grid it
is launched with cannot disagree about the CTA pool. The gfx1250
candidate turns a stream-K request away rather than silently serving
the grid-per-tile kernel.
The grouped-conv benchmark gains --streamk / --streamk-reduction and
resets per launch exactly what each reduction's contract requires.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Describe why wgrad is the direction that wants stream-K, the two launch
modes, the four reductions and their determinism/zeroing contracts, the
workspace layout, the validator constraints and the group fold. Add the
three knobs to the runbook catalog, record the support matrix rows, and
point the forward-conv future-work note at what landed.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Resolve streamk_ctas=-1 from the reduction instead of always using one
CTA per CU. A CU-sized pool leaves a single workgroup per CU, so a
stream-K CTA stalled on memory has nothing to hide behind, while the
split-K grid it competes with keeps several resident.
The workspace and atomic reductions never wait on another CTA, so their
pool can be sized for throughput: the same waves-per-CU target the
split-K heuristic already sizes its grid for, divided by the waves per
CTA. linear and tree keep one CTA per CU: a CTA there waits on later
CTAs, so the whole pool must be co-resident, and those fixups get
slower with more contributors per tile because the owner folds each one.
Only the -1 default changes. Explicit pools, every parity config and
every golden are unaffected.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ConvGroupedRequest.streamk now defaults to "auto". The dispatcher picks
stream-K when the candidate's tiling yields few output tiles over a long
reduction, and keeps the grid-per-tile split-K kernel everywhere else.
In that regime split-K can only cut the reduction into equal slices per
tile, while stream-K balances iterations across the whole CTA pool.
With many tiles or a short reduction, split-K is already balanced and
stream-K's per-tile fixup is overhead. Both thresholds are coarse on
purpose, so the rule does not encode the corpus it was fitted on.
streamk_reduction defaults to "auto", which resolves to the workspace
reduction: it never waits on another CTA, so it gets the
occupancy-sized pool. "off" restores the previous selection exactly;
"dp_sk"/"persistent" and explicit reductions still force a choice.
One resolver feeds every wgrad candidate's instance spec, dispatch spec
and grid, so they cannot disagree on whether the kernel is stream-K.
gfx1250 treats "auto" as off and still rejects a forced stream-K
request. The benchmark follows the dispatcher default and accepts the
new values.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Size the default pool of the workspace and atomic reductions at four
CTAs per CU instead of deriving it from a waves-per-CU target. Each
resident CTA streams its own K range, so what hides memory latency is
the number of independent reduction streams per CU, not the number of
waves. A wider CTA adds no streams, and the waves-based rule halved the
pool for 8-wave CTAs, where the larger pool is measurably better. For
4-wave CTAs the two rules agree.
linear and tree keep their residency-safe pool of one CTA per CU.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Adapts stream-K conv wgrad to the AOT conv kernels develop switched to.
The problem shape and the CTA pool are no longer compile-time constants, so
the host-computed iteration partition travels as kernargs:
* conv_abi gains a `streamk` reduction selector on the wgrad list. It appends
the reduction's buffers (`sk_flags`/`sk_partials`, or the `ws_ptr` pair the
two-stage path already had) and then the partition block, after `ks_count`
and so after an unchanged prefix. C++ twin: `rocke_conv_arg_names_ex`.
* ConvArgs gains `streamk_partition` / `streamk_grid` and a `StreamKLaunch`
record, and fills the block from the launch shape and the pool.
* The device emitters take the partition as SSA values (`StreamKIterArgs` /
`rocke_streamk_iter_args_t`). The tree fixup's round count is a kernarg, so
its rounds became a loop carrying the partials, the pair stride and a
dropped-out flag instead of a Python-unrolled chain.
* The CTA pool leaves the kernel name, like the split-K degree, and the C++
builder no longer rejects `streamk_ctas=-1`.
Conflicts resolved in the wgrad builder (both engines), its parity emitters
and the dispatch test. Parity configs renumbered to 21-30 around develop's
new 18-20. The representative-IR golden is re-blessed for the three stream-K
cases only.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Address review feedback on the stream-K wgrad port:
- Reject a stream-K launch shape that changes what the kernel baked in
(grouping, 2-D/3-D, pointwise addressing, channel vector widths),
including on the workspace-size path.
- Carry group_merge into the C++ spec and gate it with the same reason
as Python, so both engines reject group_merge > 1 alike.
- Make the C++ HIP lowerer reject unknown fence/flag scope and ordering
values with ROCKE_ERR_VALUE instead of substituting a default.
- Correct the docs: a linear/tree pool need not be resident, since CTAs
wait only on later CTAs and publish before waiting; and the C builder
accepts streamk_ctas=-1.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This threshold can select a workspace Stream-K kernel that the partitioner immediately reduces to all-DP. On gfx942 the automatic workspace pool is 1216 CTAs, so a one-tile shape with 1024–1215 iterations passes this predicate but has tiles * iters < pool, yielding sk_ctas == 0; it still pays workspace atomics plus the Stage 2 reduction. Include the resolved pool/fallback condition in auto selection so such shapes remain on the direct path.
Address the second round of review feedback on the stream-K wgrad port:
- Give memref.fence, memref.global_flag_store and
memref.global_flag_wait_eq verifier contracts in both engines:
operand/result counts, a global flag pointer, i32 index and value,
and scope/ordering values from the builder's own sets (now shared
constants). Messages match between Python and C++.
- Make the C++ LLVM lowerer reject flag ops without exactly three
operands, since portable-IR import can reach it without the verifier.
- Have streamk="auto" build the partition the launch will use and keep
split-K when it has no stream-K CTAs. On a larger default pool a
single long tile would otherwise run its whole reduction on one CTA.
The default pool is now one shared helper for selection and launch.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1. Division-by-zero landmine in _sk_cta_of for all-data-parallel launches
library/kernels/common/conv_implicit_gemm_wgrad.py:2746-2754 (_sk_cta_of), C++ mirror conv_implicit_gemm_wgrad.cpp:2885/2925
StreamKIterPartition.iters_per_sk_cta (streamk.py:315-316) returns raw 0 when sk_ctas == 0 (an all-DP partition). _sk_cta_of unconditionally computes tail = extra_iters + (rel - boundary) / iters_per_sk_cta via b.div before the b.select — the divide is emitted unconditionally, not behind a branch, and for an all-DP launch the owner/started branch is taken for every tile, not a corner case.
The only guard is ConvArgs._streamk_values() (conv_args.py:586): "sk_iters_per_cta": max(part.iters_per_sk_cta, 1). Today's one shipped call path (dispatch_conv_grouped → launch_values_for → wgrad_streamk_launch_values → ConvArgs.to_launch_values) always routes through this clamp, so no active bug today.
Landmine: StreamKIterPartition.iters_per_sk_cta is a public property with no guard of its own, directly reachable by any future caller that builds kernargs from the partition without going through ConvArgs.to_launch_values (e.g. a future AOT path or a second kernel family reusing the helper). Worth a defensive max(..., 1) (or an assertion) at the source (streamk.py) rather than relying on the one call site remembering to clamp.
Note: the C++ conv_wgrad_streamk_host pybind binding does not expose this kernarg at all — one verification pass initially over-stated that as a second exposure route; it isn't.
2. StreamKIterPartition.flags_bytes is wrong for its only real use case, and silently bypassed
streamk.py:358-361: flags_bytes = ceil_div(4 * sk_ctas, 256) * 256 — "one i32 flag per stream-K CTA."
wgrad's actual flag scheme is one flag per (CTA, wave): _flag_idx(sk_idx) = sk_idx * n_waves + warp_id where n_waves = warp_m * warp_n (conv_implicit_gemm_wgrad.py:2709-2716).
wgrad_streamk_workspace_layout (conv_implicit_gemm_wgrad.py:1333-1349) therefore recomputes its own 4 * sk_ctas * warp_m * warp_n formula and never reads part.flags_bytes — and the C++ mirror does the exact same bypass (rocke_wgrad_conv_streamk_workspace_layout, conv_implicit_gemm_wgrad.cpp:203-222, line 218 recomputes flags locally instead of using plan.flags_bytes).
flags_bytes (both the Python property and the C++ struct field) is exercised only by its own unit test in both engines — dead in the shipped path. If a future caller (e.g. a forward-conv port trusting the "family-agnostic" partitioner API) calls part.flags_bytes directly, it under-allocates the flag table by a factor of n_waves.
Fix options: delete the property/field (if truly nobody should call it), or parameterize it by n_waves so wgrad's workspace layout can actually use it instead of duplicating the formula.
3. _wgrad_streamk_auto hand-rolls geometry instead of calling the authoritative accessor
library/dispatch/grouped_convolution.py:496-531 (the auto-stream-K-selection heuristic) manually derives wg_M/wg_N/wg_K → tiles/iters_per_tile from raw ConvProblem fields (spatial size, p.K // p.groups, etc.), duplicating the group-fold/3D-depth-fold logic that ConvArgs.from_problem(...).streamk_partition(pool) already encapsulates (the same accessor the real kernel build/launch path uses via _wgrad_streamk_args → ConvArgs.from_problem).
This duplication was avoidable: _wgrad_streamk_auto already receives req/tile_m/tile_n/tile_k as plain parameters before any spec is built — exactly what ConvArgs.from_problem needs — so it could call the real accessor directly instead of reinventing the math.
Risk: a future change to group-fold or 3D-depth-fold semantics only needs to be made in ConvGeometry/ConvArgs; this heuristic has its own copy and will silently fall out of sync, picking or skipping stream-K based on stale geometry math. A forward-conv port would likely copy this pattern, propagating it to a second family.
4. to_wgrad_spec reimplements _wgrad_streamk_fields instead of calling it
grouped_convolution.py:678-689 hand-rolls the identical 3-key dict (streamk/streamk_reduction/streamk_ctas: -1) and the identical off → {} branch that _wgrad_streamk_fields(mode, reduction) (grouped_convolution.py:558-567) already returns — and _wgrad_streamk_fields is already called elsewhere in the same file (line ~1269).
Trivial fix: streamk = _wgrad_streamk_fields(self.streamk, self.streamk_reduction). A future 4th field would otherwise need editing in two places with no signal connecting them.
5. Missing C++ port of CTA-pool auto-sizing (streamk_ctas == -1 default)
Python's wgrad_streamk_default_ctas/wgrad_streamk_default_pool (conv_implicit_gemm_wgrad.py:506-559, 1198-1234) auto-resolve streamk_ctas == -1 to a CU-based pool before grid/workspace sizing — and the dispatcher's _wgrad_streamk_fields always ships specs with streamk_ctas: -1 by default.
C++'s rocke_wgrad_conv_streamk_partition (conv_implicit_gemm_wgrad.cpp:166-183) has no such auto-sizing branch — it hard-rejects streamk_ctas <= 0 with ROCKE_ERR_VALUE, even though rocke_implicit_gemm_conv_wgrad_spec_default() itself defaults streamk_ctas = -1.
The pybind cross-check binding conv_wgrad_streamk_host (whose docstring says it exists "for cross-checking the Python wgrad_streamk_* helpers") throws immediately on a default-constructed spec, while the Python helper it's meant to agree with succeeds silently. Confirmed untested: the one test exercising this binding (test_cpp_host_plan_agrees) always passes streamk_ctas=7, never the default -1.
Scope note: this binding has no other caller in the repo (no AOT path, no production dispatch) — so it's a test/cross-check-tool-only gap today, but a real asymmetry that will bite the moment anyone adds a test or caller exercising the default.
6. Stream-K validity gate (wgrad_streamk_available / C++ twin) has zero C++ test coverage
wgrad_streamk_gates / rocke_implicit_gemm_conv_wgrad_is_valid_spec (conv_implicit_gemm_wgrad.cpp:~428-490) mirrors Python's wgrad_streamk_available string-for-string, but has no pybind binding and no C++ unit test anywhere in the diff.
The differential parity harness's own comment states the negative-config cases (100-108: split_k+streamk, gfx1250+streamk, cshuffle+streamk, unroll_k+streamk, atomic+fp16) "are not exercised by the differential gate" because "the C emitter defines only cases 0-30."
A logic/boundary slip in the C++ port of this gate (e.g. the streamk_ctas == 0 || streamk_ctas < -1 boundary, or the atomic+dtype_d=fp32 check) would silently accept/reject a different set of specs than Python with no CI signal. Note: the C++ rejection path is tested indirectly via the full conv_wgrad_lower_llvm entry point in test_conv_wgrad_streamk.py::test_cpp_rejects_with_the_same_reason for most negative configs — the real gap is narrower than "fully untested": it's the standalone validity-gate function itself (no direct binding) and the gfx1250 arch-gate case specifically (only exercised Python-side).
7. Golden byte-identity hash file doesn't cover all 8 new mode×reduction combos
platform/tests/golden/rocke_representative_ir_sha256.json gained exactly 3 new entries per LLVM flavor (persistent+tree, dp_sk+atomic, dp_sk+linear) out of the full 2×4 matrix ({dp_sk, persistent} × {linear, tree, atomic, workspace}) this PR introduces.
streamk_reduction="workspace" — the default that "auto" resolves to — has zero golden-hash entries in any LLVM flavor. 5 of 8 combos are uncovered by the golden lock.
Mitigating factor: the separate differential parity harness (conv_implicit_gemm_wgrad_emit.{c,py}, configs 21-30) does cover all 8 combos with real IR build/diff — so cross-engine correctness is checked, just not locked against future emission drift via the lightweight golden-hash regression test. Given CLAUDE.md's "re-bless the golden in the same change" rule, this is a gap worth closing — at minimum add a workspace-reduction golden entry.
_streamk_launcher (benchmark_grouped_conv.py:1184-1190) branches only on workspace vs. else, so "atomic" falls into the sk_flags/sk_partials branch — but atomic's ABI (conv_abi.py:238-247) has neither kernarg, so ConvArgs.to_launch_values's "supplied-but-unused" guard would raise.
Confirmed unreachable today: both the benchmark CLI's --streamk-reduction choices and grouped_convolution.py's request validation (line 389, with an explicit comment "'atomic' needs an fp32 dW, which a fp16/bf16 request never has") reject "atomic" before it reaches the launcher. Landmine only if atomic reduction is ever exposed through this harness for an fp32-dW case.
9. _streamk_launcher duplicates the two-stage branch's reset/launch/cast lifecycle
benchmark_grouped_conv.py:1138-1249 (new stream-K launcher) and the pre-existing two-stage branch at :1412-1481 independently implement the same four-step lifecycle (realloc workspace → build kernargs → _launch(fence) closure that memsets + launches main kernel + optional "+cast" reduce kernel). The PR had the natural opportunity to generalize _streamk_launcher (already extracted as a function) to also serve the two-stage path, but added a second parallel implementation instead.
One real complication noted on verification: the two paths differ in launcher shape, not just kernarg content (stream-K builds two independent KernelLaunchers called sequentially; two-stage gets one fused launcher called once with a config pair) — so a shared helper isn't pure copy-paste, but the lifecycle duplication (and the pre-existing wg_M_v/wg_N_v/s2_grid hand-computation in the two-stage branch, left un-consolidated) is real.
Also: the new path correctly calls the shared wgrad_reduce_grid/wgrad_reduce_signature for its Stage-2 grid, while the untouched two-stage branch still hand-computes the equivalent — so the duplication preserves a pre-existing drift risk rather than fixing it.
10. New flag/slot machinery reimplements (with an incompatible protocol) existing dead scaffolding
The new _flag_idx/_slot_base/_store_partial/_load_partial closures in _emit_streamk_tiles (and C++ mirror) use per-(CTA,wave) global_flag_store/global_flag_wait_eq with multi-component [sk_cta][slot][tid] addressing.
platform/python/rocke/helpers/streamk.py already ships emit_streamk_partial_store/emit_streamk_partial_load_accumulate (lines 184-229), explicitly scaffolded for this job per the module docstring ("per-tile reduction pass... lands with the StreamK GEMM kernel itself" later) — but this PR did not extend or call them. Verified: both functions have zero callers anywhere in the repo, before and after this PR (confirmed by repo-wide grep) — they were already fully dead, not "used by an older kernel family" as initially suspected. instances/common/streamk_gemm.py has its own, third, independent inline atomic-reduction implementation.
Net effect: three independently-invented "CTA publishes a partial + flag, another CTA waits and folds it in" protocols now exist in the codebase, none of them sharing code. Not a correctness bug (the new one is self-consistent and mirrored in both engines), but worth flagging before someone builds a 4th on top of the wrong one, or "fixes" the dead pair thinking it's live.
Lower-priority / judgment calls (not blocking)
_wgrad_streamk_choice + _wgrad_grouped_overrides computed twice per dispatch (once in support()/_build_instance_spec, once in select()) — deterministic/pure, so no correctness risk, just redundant CPU work consistent with this file's existing "recompute everything in both functions" convention. Could be threaded through as a return value instead.
StreamKIterPartition's cascading @propertys recompute sk_tiles ~5× per full read (7 properties on one part) — negligible (a handful of integer ops), flagged only because the C++ struct twin explicitly resolves everything once "so the C side cannot recompute with different rounding," a design-consistency gap rather than a bug (the Python dataclass is frozen, so no actual inconsistency hazard).
Duplicated started/ended predicate construction in the linear and tree branches of _emit_streamk_tiles (conv_implicit_gemm_wgrad.py:2843-2844, 2874-2875) — could be hoisted above the reduction-mode branch at zero cost for the other two modes.
Ambiguous shared error label "cpg" used for both B-operand and D-operand vector-width checks in _wgrad_streamk_launch_problem's chans tuple (conv_implicit_gemm_wgrad.py:1282-1294) — a rejected launch's error message can't tell the user which operand actually failed.
Stream-K partition recomputed 3-4× per dispatch+launch across _wgrad_grid/launch_values_for/wgrad_streamk_workspace_nbytes (and C++ mirrors) instead of memoized/threaded through once — negligible per-call cost (~10 int div/mod ops); only matters if any of these call sites turn out to be hot (per-launch) rather than cold (per-dispatch-setup), which couldn't be confirmed from this repo alone.
Four independent spec.streamk != "off" branch points across _wgrad_grid/_signature_for/launch_values_for/dispatch_conv_grouped — consistent with this file's pre-existing hand-branching style (not a new anti-pattern), but adds a 4th special-case site that a future stream-K variant (or forward-conv port) will need to touch again.
groups > 1 wg_M/wg_N formula change in C++ (rocke_wgrad_conv_spec_wg_M/wg_N switched from dense problem.K/problem.C to per-group kpg/cpg): identical for groups==1 (the only case previously exercised in C++), so no existing config changes output — but it's the first time grouped wg_M/wg_N is exercised in the C++ engine at all; worth a sanity nod from a C++-side reviewer even though it's in-scope for this PR.
The reason will be displayed to describe this comment to others. Learn more.
PR #13151 — stream-K for implicit-GEMM conv wgrad (jakpiase, AICK-1943) — 🟠 REJECT: sound core; gate the new auto default and make the race test able to fail
Prior item
Status
Copilot: unbounded explicit streamk_ctas can deadlock linear/tree
🟢 N/A 082f58426: no cap; docs argue in-order WG dispatch leaves ≤1 (linear) / ≤rounds (tree) CTAs blocked. Holds: waits only target higher CTAs (our sim, 5280 partitions)
Copilot: launch shape vs baked grouping / 3-D modes
🟢 FIXED cb6a47b8: arity, global ptr, i32 idx/value, scope/ordering sets in both. Threads still unresolved
Copilot (2nd review body): gfx942 auto picks all-DP workspace (sk_ctas=0)
🟢 FIXED cb6a47b8: auto builds the launch partition, requires sk_ctas > 0
Bots: policy (ISSUE ID), Math CI, base-freshness
🟢 pass; Math CI "skipped"
#
Sev
file:line
Net-new finding
N1
🟠
dispatch/grouped_convolution.py:286,549
streamk="auto" default enables stream-K where it never ran. Dispatch probe: gfx942 3x3, 1x1 and 3-D wgrad requests → dp_sk/workspace. Body: gfx942 is IR/CPU-only; no parity cfg, golden or GPU case covers 1x1 or 3-D stream-K; thresholds fitted on gfx950 only (:490). Gate auto to what the GPU tests cover, or add those cases.
N2
🟠
tests/test_conv_wgrad_streamk_gpu.py:15,181,439
20-launch bit-identity test can't catch a missing release/acquire, its stated purpose: only flags are re-zeroed, so an early read returns the previous launch's bit-identical partial (dW not reset either). Poison partials + dW (e.g. 0xFF) before each launch.
Public flags_bytes counts 1 flag/CTA; the kernel uses 1/wave (conv_implicit_gemm_wgrad.py:1346,2712): 1024 vs 4096 B at the gfx950 default pool. A workspace laid out with it puts partials over flags. Drop it or take n_waves.
N4
🟡
golden JSON, SUPPORT_MATRIX.md
mergeable_state: dirty: both conflict with develop. Rebase, re-run the all-flavor golden (3/3 green at head).
Utility: 🟢 Positive — CK-parity stream-K fills the few-tiles/long-K wgrad gap at the right layer (dual-engine IR ops, AOT kernargs, one resolver for spec/grid/kernargs); fixup protocol checked sound (fence + release/acquire, sim); host tests 62 pass (11 need C++/GPU). Merge after N1 + N2 + rebase.
- Clamp StreamKIterPartition.iters_per_sk_cta to >= 1 at the source
(Python and C++) so an all-data-parallel partition never feeds a zero
divisor to the kernel's owner lookup.
- Make flags_bytes take the wave count (one flag per CTA and wave) and
use it from both wgrad workspace layouts instead of a duplicated formula.
- Remove the unused emit_streamk_partial_store/_load_accumulate helpers;
the wgrad fixup is the reference flag/partial protocol.
- Port the streamk_ctas == -1 CTA-pool auto-sizing to C++ and bind the
C++ validity gate; test both against Python.
- Lock all eight stream-K mode x reduction combos in the golden IR.
- Poison partials and dW before each launch in the GPU repeat test so a
missing release/acquire can actually fail it.
- Build the linear/tree started/ended predicates once and name the
operand in the vector-width launch errors.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The grouped-conv dispatcher and its benchmark are due for a refactor, so
this change no longer extends them: the dispatcher stream-K selection and
the grouped-conv benchmark flags are reverted. Stream-K wgrad candidates
are instead opt-in in the implicit-GEMM sweep (--streamk,
--streamk-reductions, --streamk-ctas) for the JIT, AOT-cache and
run-from-cache paths. The cache identity and kernel label keep their
existing values when stream-K is off.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… engines
- global_flag_store / global_flag_wait_eq now require a global i32
pointer in the builders and verifiers of both engines, so the LLVM and
HIP lowerings cannot address different memory.
- All four lowerers reject a missing or unknown scope/ordering with the
same message instead of falling back to defaults or emitting invalid IR.
- Document the per-(CTA, wave) flag protocol and its forward-progress
assumption (waits only target higher-index CTAs).
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- The C host helpers (partition, grid, workspace layout/bytes) take the
target arch so the auto CTA pool matches the kernargs on gfx942.
- The C validator follows Python's gate order and ports the vector_size_c
and group_merge gates, so multi-cause rejects give the same reason.
- Python rejects streamk_ctas of 0 or below -1, as C does.
- Add a ctest with C-side host-plan and validator vectors so parity is
checked without the Python binding, and test a shape that actually
splits under the default pool.
- The GPU repeat test now claims release detection only; codegen tests
pin the acquire spin and release publish in IR and ISA.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…launch code
- The sweep's stream-K launcher always allocates a workspace, so a shape
whose partition has no stream-K CTAs no longer aborts the sweep.
- Stream-K builds reuse the regular wgrad build path, Stage 2 launches
share one helper, and the run-from-cache stream-K branch is additive.
- Drop the unused stream-K extras from kernel_sweep's launch values and
warn when stream-K options are given with --streamk off.
ISSUE ID : AICK-1943
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The reason will be displayed to describe this comment to others. Learn more.
P3 — Update the PR description to match the final scope
The Dispatch section still says ConvGroupedRequest.streamk defaults to auto and that the grouped-convolution dispatcher selects stream-K. The final series instead removes that implementation and moves stream-K to the opt-in implicit-GEMM sweep (--streamk, cache, and run-from-cache paths). Please update the description so it does not imply a production dispatcher-default behavior that is absent from this head.
Stream-K wgrad AOT kernels were built for every default-epilogue geometry
x launch mode x reduction x all 16 (vector_size_a, vector_size_b) pairs.
Cache them only where they are used:
- at the widest vector width W of the dtype (16-byte loads) for
(W, W) and (W, 1), with every requested reduction except "tree";
- fully scalar, dp_sk x {tree, workspace}, on tiles with tile_n >= 128,
for large-filter depthwise layers.
Shapes at other vector widths are served by the split-K kernels, which
are still cached at every width. The selection lives in
_wgrad_streamk_builds and count_jobs mirrors it.
Stream-K instances per dtype on gfx950: bf16/fp16 178560 -> 17076,
fp32 171720 -> 31374; gfx942: bf16/fp16 223200 -> 21234,
fp32 133704 -> 24192.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ke-aot-wgrad-perf
Conflicts with develop's stream-K wgrad (#13151) and dgrad AOT fixes
(#13494), resolved by keeping both sides. Follow-on fixes for how the two
interact:
- wgrad: prefetch_addr is rejected with stream-K (Python and C), since
stream-K never reaches the K loop it rewrites; the JIT sweep only sets
it on non-stream-K combos.
- wgrad parity emitters: the prefetch_addr configs move from 21-38 to
31-48, after develop's stream-K configs 21-30.
- dgrad: the uniform_k_decode A decode uses the record's magic pairs, as
develop's plain loop now does, and the C kernel's kpg divisor is c_kdiv
(kpg when grouped). The AOT sweep keeps develop's lds_k_outer predicate
(both sides had added an equivalent one).
- run-from-cache: _cached_bind registers its scratch in live_extras, so
the pruning/free logic keeps working; depthwise problems keep stream-K
kernels; an early stop also frees the stream-K workspace.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
## Motivation
Wgrad GEMMs usually have a filter-sized output (few output tiles) and a
very long reduction (N*Ho*Wo). Split-K cuts every tile the same number
of ways, so on these shapes it leaves CUs idle. This PR ports CK Tile's
stream-K to rocKE wgrad. Stream-K balances MAC iterations instead of
tiles:
- Tiles that fill the device run data-parallel.
- Only the remainder is spread across a pool of CTAs.
Stream-K is opt-in. The grouped-convolution dispatcher and its benchmark
are unchanged (they are due for a refactor), so no production path picks
stream-K on its own.
## Technical Details
- **Core IR:** three new leaf ops in both engines: `memref.fence`,
`memref.global_flag_store` and `memref.global_flag_wait_eq`.
- Each lowers to an LLVM fence, or to an atomic store/load with
release/acquire ordering at a chosen sync scope.
- The Python and C++ verifiers check arity, a global i32 flag pointer,
i32 index/value operands, and the allowed scope/ordering values.
- All four lowerers reject a missing or unknown scope/ordering with the
same message.
- **Partitioner:** `StreamKIterPartition` (Python) and
`rocke_streamk_iter_partition_t` (C++) mirror CK's DP + SK split.
- `iters_per_sk_cta` is clamped to at least 1, so an all-DP partition is
safe.
- `flags_bytes(n_waves)` sizes the per-(CTA, wave) flag table.
- The device-side emitters read the partition from runtime values.
- **Wgrad kernel:**
- Launch modes: `streamk = dp_sk | persistent`.
- Reductions: `linear` and `tree` are deterministic and hand off through
per-(CTA, wave) flags. `atomic` (fp32 dW only) and `workspace` (f32
scratch + Stage-2 reduce) are the other two.
- Conv groups are folded into GEMM-M per group, so the fold is exact for
any wg_M. The C++ engine now builds grouped wgrad on the stream-K path.
- **AOT:** the partition is passed as an `sk_*` kernarg block, appended
after `ks_count` so existing kernarg offsets don't change.
- Python side: `conv_abi` and `ConvArgs.streamk_partition`.
- C++ side: `rocke_conv_arg_names_ex`.
- The CTA pool is a launch parameter, so one binary serves any problem
shape and any pool. `streamk_ctas = -1` resolves to the arch default
pool, in both engines.
- **Sweep / benchmark (opt-in):** `benchmark_implicit_gemm_conv.py`
(also reachable through `benchmark_conv_compare.py`) gains `--streamk
{off,dp_sk,persistent,all}` (default `off`), `--streamk-reductions` and
`--streamk-ctas`.
- Stream-K wgrad candidates are added to the JIT sweep, the
`--compile-all` AOT cache (`kernel_sweep.py`), and `--run-from-cache`.
- With stream-K off, the cache identity and kernel labels are unchanged.
- The stream-K launcher allocates the reduction's workspace and zeroes
what each reduction accumulates into. It shares the Stage-2 reduce with
the two-stage path.
- **Constraints:** gfx942 and gfx950 only, MFMA wave64. Both engines
reject the following with identical messages: split-K, two-stage, async
DMA, unroll, cshuffle, chiplet swizzle and group merge.
## Test Plan
- Byte identity across both engines at llvm20, llvm22 and llvm23.
- Differential parity: new wgrad configs 21–30 and negative configs
104–108.
- Representative-IR golden: all 8 `{dp_sk, persistent} × {linear, tree,
atomic, workspace}` combos (dense gfx950 plus grouped gfx942).
- CPU tests:
- partitioner invariants
- validator agreement across both engines, including multi-cause rejects
and the arch gate
- admits-implies-builds matrix
- host plan and default-pool agreement between Python and C++
- stream-K ABI and launch values
- sweep launcher, including all-DP partitions
- dual-engine lowering
- C++ ctests: `rocke_streamk_iter_partition`,
`rocke_wgrad_streamk_host`, `rocke_future_intrinsic_lowering` and
`rocke_ir_serialize_roundtrip`.
- GPU tests:
- numerics for both launch modes × all four reductions on dense, strided
and grouped shapes, including a partial last M tile
- one binary launched on several shapes (AOT)
- 20-launch bit-exact repeat test for linear and tree, with partials and
dW poisoned before each launch
- flag ordering pinned in IR (acquire spin / release publish) and ISA
(cache invalidate / writeback)
- sweep launcher on all-DP partitions (linear and workspace)
## Test Result
- Byte identity GREEN for all families at llvm20, llvm22 and llvm23.
31/31 wgrad configs pass differential parity.
- Representative-IR golden OK at all three flavors after merging
develop.
- CPU tests, C++-binding tests and ctests pass.
- GPU tests pass on gfx950. gfx942 is covered by the IR, golden and CPU
gates only.
JIRA ID : AICK-1943
---------
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…12441)
ISSUE ID : AICK-2273
# summary
## What this branch is about
rocke's emitter had no way to know whether an intrinsic it declares
actually
exists on the architecture it is lowering for. An unknown
`llvm.amdgcn.*`
declare is not a parse error, not a verifier error, and not a `clang -S`
error —
it survives every check rocke ran and fails much later, in the backend,
as
`Cannot select`. Worse, the same name can be present at one LLVM flavor
and
absent at another, and present on one arch and absent on its neighbour.
#12269 put the measurement on develop: a generator that probes the
toolchain, an
emitted-IR validity gate that compiles *and links* every corpus case,
and three
committed columns of answers (one per LLVM flavor). Nothing read them.
This
branch is their first consumer, plus the fixes that building a consumer
forced:
1. **A reader and a warning lane** — the lowerer asks the table whenever
a kernel
demands an intrinsic, on both backends.
2. **Removal of the implicit `gfx950` target default** everywhere it had
spread:
both engines, the per-family bindings and the C wrappers, so every
lowering
names the target it is actually for. This is what lets the lane mean
anything.
3. **A corrected table and sturdier gates**: reading the columns as a
consumer
showed that one of them was wrong, and that the generator and the corpus
gates
each measured less than they reported.
Plus a long series of methodology fixes, all of which draw one line:
> **A non-answer must never be recorded as a negative.**
> "We could not ask" is not "the answer is no".
The lane is warn-only. Making it an error is a separate change
(AICK-2274).
The commit list includes commits that #12269 already merged by squash;
the diff
against `develop` is what this PR changes.
---
## 1. The warning lane
**Reader** `python/rocke/core/arch/domain.py` — validates the envelope
and raises
on a malformed or wrong-schema column rather than defaulting every field
(which
would have presented a corrupt artifact as a complete absence of data
and
switched the lane off in silence — the one failure the lane exists to
prevent). A
column that exists but cannot be read raises too; only a *missing* file
means
"no data".
**What it reads** —
`python/rocke/core/arch/data/intrinsic_arch_domain.<flavor>.json`,
schema `rocke.intrinsic_arch_domain/v1`, **171 keys × 7 wired arches =
1197 cells
per column**:
| Flavor | `ok` | `name_absent` | `arch_absent` | no-data |
|---|---|---|---|---|
| llvm20 | 535 | 294 | 194 | 174 |
| llvm22 | 669 | 147 | 340 | 41 |
| llvm23 | 740 | 21 | 391 | 45 |
One column per flavor because **a host can only measure its own flavor**
— flavor
is a runtime choice (`ROCKE_LLVM_FLAVOR`), so the generator can never
produce all
three in one run, and a gate must skip the flavors its host cannot
measure rather
than mark them failing.
The lowerer's declaration table has 27 more keys than the columns: the
26
declarations develop added for data prefetch, workgroup clusters and TDM
(#12857),
and `raw.ptr.buffer.load.f16` (#12710), have not been measured yet. An
unmeasured
key is no data, so the lane stays silent for those intrinsics until the
columns
are regenerated. That is a measurement on each of the three toolchains,
and it is
left to a follow-up.
**Status vocabulary** splits answers from non-answers:
- Answers: `ok`, `name_absent`, `arch_absent`
- No data: `target_unsupported`, `toolchain_crash`, `toolchain_timeout`,
`probe_error`
Only `arch_absent` is negative (`is_negative`). Everything else is
silence.
`name_absent` is silent on purpose: it is a flavor-axis fact that holds
for every
target, and the validity gate already reports a reachable one with the
compiler's own diagnostic.
**Warning lane** — `_Lowerer._need()` consults the table and emits
`ArchDomainWarning`. The default backend is the C++ engine, which
returns
finished IR without the Python lowerer running, so the same check also
runs over
its emitted `declare` lines; the declare → key mapping is exact because
both
engines emit the table's declares verbatim.
Controlled by `ROCKE_ARCH_DOMAIN`: `warn` (default) / `off`; any other
value
**raises**, because every way of getting this variable wrong is somebody
trying
to switch the lane off, and silently defaulting to `warn` hands them the
opposite of what they asked for. The value is read in
`lower_kernel_via_backend`,
which every lowering passes through (`lower_kernel_to_llvm`,
`compile_kernel` and
the per-family wrappers), so it raises for every kernel, entry point and
backend
rather than only for a kernel that happens to demand an intrinsic. The
LLVM
flavor is likewise resolved once per lowering and handed to each stage.
`error` mode is planned (AICK-2274).
**Gate** `tools/check_arch_domain.py` — wired into `run_all.py`. It
lowers the
194-case corpus at all three committed flavors (pure Python; it needs no
toolchain, so it rules on every flavor on any host) and compares the
warnings
with `EXPECTED_WARNINGS`, which only ever shrinks. The list holds nine
entries:
three intrinsics (`ds.read.tr16.b64`, `mfma.f32.16x16x32.bf16`,
`mfma.f32.16x16x32.f16`) on gfx942 at each flavor. All nine are one
defect in the
parity harness rather than in a kernel: it builds the gfx950 3D
attention kernel
for two gfx942 cases. The validity gate lists the same two cases as
`KNOWN_BAD`.
Fixing the harness is left to the change that makes the lane binding,
which needs
the list empty.
Each entry names the corpus cases allowed to demand it, so the list
cannot widen
in silence: a case an entry does not name is unexpected even when its
intrinsic is
already allowlisted, and an entry whose case stops demanding it, or no
longer
exists, is stale. A case that did not run (`--only`, `--arch`,
`--flavor`, or an
instance that declines the flavor) is judged neither way. The comparison
is a pure
function with a test per rule.
## 2. No implicit target default
Two commits, in order: first a pure refactor making all 158 implicit
lowering
call sites name their arch explicitly; then removal of the default
itself from
both engines. Omitting the target now raises.
This is what lets the warning lane mean anything — a kernel lowered at a
default
arch is a kernel whose arch nobody chose. The default had also spread
beyond the
lowerer, and the same series removes it from those places: the
per-family pybind
entry points (one shared `require_arch` replaces their per-family "usual
target"
fallbacks), the Python wrappers in `core/backend.py` that mirrored them,
and the
scalar-attention C wrappers below.
## 3. A corrected table and sturdier gates
- **Generator 4 → 6, and one column corrected.** The generator built its
probes
from rocke's hand-written declaration table, which records `immarg` only
where
an author happened to add it, and passed the unmarked operands as
variables
where the compiler demands constants. llvm22 and llvm23 reject that in
the
verifier, so the generator's rescue recorded the true answer; llvm20 has
no such
check, died at ISel with `Cannot select` — which is also what a genuine
negative
looks like — and recorded `arch_absent`. Probes are now built from the
declare
`opt` resolves, so the signature comes from the LLVM being measured. 33
llvm20
cells move: 27 `arch_absent` → `ok` and 6 `arch_absent` →
`toolchain_crash` (a
crash is no data, so the lane stays silent). llvm22 and llvm23 move no
status.
Two more ways to bank a non-answer as an answer are closed: a failed
name probe
(`opt` crashed or timed out) was recorded as `name_absent` for every
arch, and a
partial `--only`/`--arch` sweep could be filed as a whole column.
- **One meaning of "this flavor refuses the instance".** The arch-domain
gate
counted four gfx1250 instances that refuse an older LLVM as lowering
errors —
eight failures on every run of a clean tree — while the validity gate
already
treated them as unmeasurable. Both now import one predicate from
`tools/_flavor_gate.py`, and a test asserts they hold the same object.
- **Provenance.** The validity gate no longer trusts an `opt` from
another
install: an older `opt` rejecting a construct its parser predates reads
as "rocke
emitted invalid IR" when the IR is fine, so that lane is skipped on a
mismatch.
The generator's identity check covers `opt` as well as `clang`.
- **The arch-domain gate lowers through the native lowerer.** Through
the
dispatching entry point, any host with the C++ engine built would lower
every
case without consulting the table, harvest nothing, and call every
allowlist
entry stale.
## 4. What testing this turned up
**The scalar-attention C wrappers.** The three public `_lower_to_llvm`
wrappers
in `attention_unified_entry.cpp` are the only lowering entries whose
signature
carries no target, and nothing in the tree called them. They passed
`NULL`, which
the lowerer now rejects, so all three had turned into `ROCKE_ERR_VALUE`.
They now
name `gfx950` explicitly — a fixed target, not a fallback, and the
header says so;
a caller who wants another target builds with the `_new` entry and
lowers through
`rocke_lower_kernel_to_llvm_ex`.
**A crash in the C++ scalar reduce builder.** Testing the wrappers found
that
`rocke_build_unified_attention_reduce_scalar` died with SIGSEGV on every
input:
its loops read three tensor descriptors that nothing had built. It now
builds
them, in the order the Python builder does, and emits the kernel the
Python
builder emits (byte-identical). A new ctest lowers each wrapper and
requires text
identical to lowering the same kernel through the target-taking entry.
**Bindings that ignored their target.** A sweep over every per-family
binding
that takes a `(spec, arch)` found six that discarded `arch`
(layernorm2d,
elementwise and the scalar attention family's serialize/verify entry
points), two
fused-MoE e2e entry points that hard-coded gfx950 into the spec they
built, and a
ninth that kept a pybind default. All now call `require_arch` first. The
sweep is
a test now: an extension-gated check reads every `(spec, arch)` binding
off the
module and requires each to reject an empty target, so a binding added
later is
held to the same rule.
**Bindings that arrived with develop.** Three stream-K wgrad host-plan
bindings
(`conv_wgrad_is_valid_spec`, `conv_wgrad_streamk_default_pool` and
`conv_wgrad_streamk_host`, from #13151) resolved their target with
`arch_or_default`, the helper this branch removed, and declared `arch`
with a
gfx950 default. `rocke_engine.cpp` stopped compiling, so the pybind
extension
could not be built; CI builds no extension, so nothing turned red. They
now take
the target like every other binding, and the sweep covers the two that
take a
spec (pybind prints `arch: str = 'gfx950'` for a defaulted argument,
which is why
the sweep could not see a default before). One decision to look at:
develop's stream-K tests pinned
"without an arch both engines default to gfx950" for the host planner.
The Python
planning helpers keep their default; the C++ binding now refuses an
empty
target, and the test asserts that.
The same thing happened a second time with the gfx950 block-scaled GEMM
bindings
that arrived with #13450 (`gfx950_scaled_serialize_ir`,
`gfx950_scaled_verify` and
`gfx950_scaled_lower_llvm`: four uses of `arch_or_default`), and
`rocke_engine.cpp`
again stopped compiling without anything turning red. They call
`require_arch`
first now. Because the extension-gated sweep cannot run where no
extension is
built, a second test reads `cpp/bindings` and fails on either way a
binding guesses
its target: the removed `arch_or_default` helper, or a gfx target as the
default of
an `arch` argument. It lists every offending line, and it fails on the
previous
merge without the fix. It reads the source, so it runs from a source
checkout and
skips in an installed layout.
## 5. Review
Every review thread has a fix, one commit per concern:
- **Tracker IDs out of code** (Rule 1): four places cited IDs; each now
names the
defect. The allowlist's reason string was also wrong about the defect —
it
blamed two declaration-table items, but the nine warnings come from the
parity
harness — and now says so.
- **Trust boundaries** (Rule 4): a malformed or wrong-schema column, an
unreadable
column and an unknown `ROCKE_ARCH_DOMAIN` value each raise, instead of
reading
as "no data" or defaulting to `warn`.
- **One flavor per lowering**: with no flavor given, the C++ path and
the
arch-domain post-pass each resolved their own, and compiler discovery is
retryable, so the IR could be checked against the wrong flavor's column.
It is
now resolved once in the dispatcher, including for both engines under
`ROCKE_BACKEND=both`.
- **Tests lower at the arch they are about**: three substitutions of an
explicit
`arch=` had put the wrong arch in, and `test_legacy_equals_gfx950`,
which
compared a call with itself, is replaced by
`test_omitting_the_target_raises`.
- **`develop` merged four times** since the first review (111 commits,
then 99,
29 and 14). Only the first conflicted textually (four binding
translation units
and `test_compiler_detection.py`). The trouble in all four was semantic:
develop's new code was written against the default target this PR
removes, so
it called the lowerer, or a binding, without naming one. Each occurrence
is
fixed here: the compiler-detection tests (three of the calls sit inside
a
subprocess script string, and they had turned the gfx94X CI shards red)
and a
binding-selection probe; the flag-attribute test that turned the
hipkernelprovider jobs red (twelve subtests), a library gate test that
CI does
not run, and the FP64 elementwise and reduce tests; and the stream-K and
block-scaled GEMM bindings in section 4. A scan of the Python and C++
that the
merges brought in finds no other call that names no target.
- **Formatting after the merge**: this branch's parenthesized `with`
statements
make black treat `test_rocke.py` as Python 3.9+, so develop's three
multi-line
`with a, b, c:` blocks in that file now need reformatting. Each version
of the
file is black-clean alone; the merge is not, and the pre-commit job
checks every
file a PR changes.
- **The allowlist names its cases** (see the gate in section 1): it was
keyed by
intrinsic, target and flavor alone, so a second corpus case demanding an
allowlisted intrinsic left the gate green, and fixing one of the two
cases that
demand it could not make the entry stale.
- **Wording**: the status vocabulary in `TESTING.md` lists
`probe_error`,
durations are out of the generator's comments, and the documentation
examples
name the target.
---
## Verification
Run in the three toolchain images (one per LLVM flavor) against an
export of the
branch tip; no GPU. The pybind extension is built in each image and put
on
`PYTHONPATH`, so the tests that need it run (CI builds no extension and
skips
them).
| Check | llvm20 | llvm22 | llvm23 |
|---|---|---|---|
| byte identity, every family | GREEN | GREEN | GREEN |
| ctest | 20/20 | 20/20 | 20/20 |
| emitted-IR validity (194 cases) | GREEN (2 documented `KNOWN_BAD`, 8
unvalidated) | GREEN (2 documented `KNOWN_BAD`, 4 unvalidated) | GREEN
(2 documented `KNOWN_BAD`) |
| arch-domain corpus gate (194 cases, 3 flavors) | GREEN (exactly the 9
documented warnings) | same | same |
| platform pytest, default and under `ROCKE_BACKEND=both` | 2685 passed,
59 failed (both runs) | 2689 passed, 59 failed (both runs) | 2690
passed, 59 failed (both runs) |
| library pytest | 2598 passed, 1 failed | 2599 passed, 1 failed | 2599
passed, 1 failed |
Every failure is also a failure of `develop` itself, checked on an
export of its
tip (59 in the platform suite at every flavor, one in the library
suite). The
failing set is the same at the three flavors:
- 56 tests import `ml_dtypes`, which none of the images has: six in
`core/test_gfx1250_fp4_scaled_wmma.py` and fifty in
`instances/test_gfx950_mxfp_examples.py` (new with #13450);
- two tests in `core/test_arch_domain_artifact.py` (every declaration is
covered;
regenerating the host's column is a no-op) fail because of the 27
unmeasured
declarations in section 1. That file is excluded from the installed test
tree,
so CI does not run it;
-
`tools/test_run_all.py::test_existing_engine_archive_is_reused_by_pytest`
fails
whenever pytest runs under `tests/run_all.py`, which exports the
variable the
test expects to be unset;
- `library/tests/test_kda_hostpack.py::test_runner_geometry_without_gpu`
builds a
stub kernel without the `cluster_dims` attribute that #12857 made the
manifest
code read.
The key #13450 added to the three columns
(`mfma.scale.f32.32x32x64.f8f6f4.v8i32.v8i32`) was re-measured in each
image with
this branch's generator: every cell, status and evidence, matches the
committed one.
A note on the llvm20 run. At llvm20, `raw.ptr.buffer.load.async.lds`
runs past the
60 s probe budget on six targets (the committed column records
`toolchain_timeout`
for them), and the generator's timeout kills the `clang` driver but not
the
`clang -cc1` it started, so each regeneration (including the
regeneration test
above) leaves twelve runaway compiler processes behind. In a
two-CPU container they starve the rest of the suite: an earlier full run
timed out
`test_binding_compiler_selection.py::test_resolver_failure_releases_binding_builders`
(valgrind, a fixed 120 s limit), which passes in about 15 s alone, on
this branch
and on `develop`. The numbers above are from a run that killed the
leftover probes
as they appeared. The leak is in the generator already on `develop`, not
in this
PR, and is left to a follow-up.
Not run: a GPU, Windows.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 5 <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
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.
github-actions
Bot
added
documentation
project: hip-kernel-provider
rocke: platform
cgmillette
left a comment
CongMa13
left a comment
Motivation
Wgrad GEMMs usually have a filter-sized output (few output tiles) and a very long reduction (NHoWo). Split-K cuts every tile the same number of ways, so on these shapes it leaves CUs idle. This PR ports CK Tile's stream-K to rocKE wgrad. Stream-K balances MAC iterations instead of tiles:
Stream-K is opt-in. The grouped-convolution dispatcher and its benchmark are unchanged (they are due for a refactor), so no production path picks stream-K on its own.
Technical Details
memref.fence,memref.global_flag_storeandmemref.global_flag_wait_eq.StreamKIterPartition(Python) androcke_streamk_iter_partition_t(C++) mirror CK's DP + SK split.iters_per_sk_ctais clamped to at least 1, so an all-DP partition is safe.flags_bytes(n_waves)sizes the per-(CTA, wave) flag table.streamk = dp_sk | persistent.linearandtreeare deterministic and hand off through per-(CTA, wave) flags.atomic(fp32 dW only) andworkspace(f32 scratch + Stage-2 reduce) are the other two.sk_*kernarg block, appended afterks_countso existing kernarg offsets don't change.conv_abiandConvArgs.streamk_partition.rocke_conv_arg_names_ex.streamk_ctas = -1resolves to the arch default pool, in both engines.benchmark_implicit_gemm_conv.py(also reachable throughbenchmark_conv_compare.py) gains--streamk {off,dp_sk,persistent,all}(defaultoff),--streamk-reductionsand--streamk-ctas.--compile-allAOT cache (kernel_sweep.py), and--run-from-cache.Test Plan
{dp_sk, persistent} × {linear, tree, atomic, workspace}combos (dense gfx950 plus grouped gfx942).rocke_streamk_iter_partition,rocke_wgrad_streamk_host,rocke_future_intrinsic_loweringandrocke_ir_serialize_roundtrip.Test Result
JIRA ID : AICK-1943