Repository navigation
feat(rocke): Initial drop for tiling primitives - #12606
Merged
Merged
Conversation
…tiling) Introduce rocke.helpers.tiling -- an N-D tiling + MMA primitives layer that lets an author express tiled GEMM/MMA kernels and custom thread/register distributions without hand-writing raw TileDistributionEncoding integers, resolving the concrete hardware intrinsic for the bound target (target-agnostic authorship). - TileMma resolves the intrinsic and owns the whole wave-tile subtile grid (m/n/k_iter + order); make_tile_desc authors any distribution as a quantity-major geometric table; make_tensor_desc/make_window give ptr-free memory + auto-clipping; load/store/fill_fragment are the b-first IR verbs; describe/render_* reflect layouts. - Dense MMA GEMM is bit-exact on gfx90a (single-atom, wave-tile subtiling, clipping); c_transpose / interleaved / sparse / MX are reserved seams. - Committed traits table (mma_traits.json); torch-free numpy golden-ref tests (95/95 pass) mirror the package under tests/helpers/tiling/. - Adds an AGENTS.md pointer so agents can discover and drive the layer. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds two reusable analysis layers under the tiling package plus the core-emit changes they build on. - visualization/: composable layout-viz components (register file, logical tile, LDS bank view, MMA tee) on a shared canvas + a Pipeline spine, with phase recipes and epilogue auto-detection. - lds_conflict.py: the arch-parameterized LDS write-port bank-conflict model + simulator + rocprof harness + 3-panel dataflow renderer (gfx90a-validated). The gate reconciles on scale-invariant conflicts/access, so a per-served-group sim can be checked against whole-run counters without a scale error. Supporting core changes: - emit: thread-driven coordinate emission (one path for single-wave and block-wide wave_dist), wide + LDS-targeted load/store, and the b128 XOR swizzle. - transforms: now hosts transform_fragment (moved out of emit so the memory verbs never import the layout solver; the verb duck-types the IRBuilder -- no IR import). - mma: K-alignment validation + a per-atom accumulator that keeps C in AGPRs across the K-loop; a_desc/b_desc become methods with a broken-interleaved guard. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
describe_edge names the five layout-edge kinds (identity, reorder, cross_lane, reposition, relabel) by routing through the classifier; reorder_between/ReorderPlan price a same-data register reorder into dword vs sub-dword lane moves. Labels stay invariant except on an explicit relabel. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Pure, tested backend (no matplotlib) that the visualization package and coalescing skill ground on: vector_transactions models the b128-capped hardware transaction pattern; analyze_coalescing reports cross-lane cache-line fusion (fused vs scattered). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Derives the load descriptor and gates the coalescing model against the real compiled ASM, so a render only ships when the assembly backs it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Records the transaction/op stream at the verb boundary (contextmanager + verb decorator + registry), with a byte-identical-IR gate and a b-witness completeness check, so recorded kernels round-trip. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
layout_render gains the register-file/LDS-bank coalescing address-space renderer (per-SIMD data-range grids) and in-register reorder panels; kernel_stages labels flow edges through the transform classifier and threads lds_base_bytes through the LDS recipes; _canvas gains a fit-to-cell font sizer. gate() takes an rtol for measured-vs-sim tolerance. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Interleaved GEMM (AOS load + in-register reorder) and the CRC free-dim-contiguous demo, exercising the tiling primitives end to end. Tests cover the interleaved bit-exact path plus the recorder and LDS-conflict use of these kernels. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
block_diagram builds the Level-0 loop-body block diagram (parallel prefetch/compute flows, edges from the recorded Value graph); auto_pipeline segments those flows and drives Pipeline renders end to end. layout_optimizer prices layout choices via the transforms reorder-cost model. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Add the label_flow_and_transforms and visualization_api_surface SOTs and a kernel design-report template; expand mma_is_machinery on per-operand soundness and point mma_operation's comment at it. tiling_gemm_demo gains an interleave_a flag proving the in-register reorder path bit-exact. Remove the superseded design proposal. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Mirror the self-contained skill/expert set the tiling layer depends on (layout-viz, bank-conflict, rocke-tiling-api skills; MMA/LDS/Tiling/ profiling/GPU experts; dispatch table and shared policies) so the tooling travels with the code. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replace the measured TFLOPS/multiplier with a qualitative statement of the store-transaction win; the proposal's rationale stands without the numbers. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
A parameter that changes a computed number (arch, dtype, wave size, strides, origin, swizzle, K-alias depth) or that is printed on a figure as fact (kernel/macro/operand/dims labels) no longer has a default. The failure mode being closed is specific: a default does not error, it silently analyzes a different kernel than yours and stamps your kernel's name on the answer. Concretely: pad0_depth's None-fallback assumed depth 8 and returned a pad the GPU does not agree is conflict-free for depth-16 aliases; element size was `2 if f16 else 4`, which gave bf16 4 bytes and a wrong bank map; wave size was a literal 64 in four places; and the figure title hardcoded one kernel's name and macro-tile. Wave size now lives on ArchLDS.WAVE and element size in dtype_bytes_of -- one home each, derived rather than re-defaulted per call site, and both raise on an unknown target instead of guessing. verify_roundtrip likewise derives pipeline.wave_size rather than assuming 64, which would replay half a wave on a 32-lane target and "verify" a round-trip that never happened. Two silently-wrong constants fell out and are fixed here. recommend_pad was passed per_dword where it wants dtype_bytes; these are equal for f16 and differ for everything else, so the pad was right only by accident. _locate_collision hardcoded bank 0, but which bank a store piles on is a property of its address map, so a fixed 0 reports an empty collision for any layout that piles elsewhere. Neither failed loudly -- both returned a plausible wrong number, which is the class of result this module exists to prevent. Presentation-only knobs (max_banks, full, verbose, n_iter, grid_ctas) keep their defaults; they cannot change a verdict. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Until now the only way to get a conflict number was a container + GPU run gated against rocprof counters. That is right for "what does this kernel actually do", but it makes design-time layout comparison prohibitively expensive, so the alternative in practice was hand-reasoning -- exactly what this module exists to replace. Simulate mode uses a model already validated on that arch: the gate is selftest(arch), i.e. the arch reproduces its OWN measured corpus. Seconds, no GPU. An arch with no registered model is a full stop, never an extrapolation from another target's constants. The honesty of the cheap path rests entirely on labelling, so provenance is structural rather than advisory: verdict gains a third state (SIMULATED, between VALIDATED and UNVALIDATED), the figure carries a diagonal SIMULATED watermark built from the same branch that gated it, facts_table reports the selftest instead of printing "?" where a counter would be, and the mode confusions (investigate without measure, simulate with it) are refused rather than silently down- or upgraded. The skill and the LDS expert were written for a single empirical-only path; both now fork on mode and lead with provenance. The mode is asked for, never inferred -- guessing simulate risks a modelled number read as measured, guessing investigate commits the user to a long container run. This also reverses the old "keep fixing the model silently until it matches" instruction. A simulator that disagrees with the counters is a finding: stop, report which config and which constant looks wrong, and offer to repair it. Patching quietly until the numbers line up buries a real disagreement between the model and the hardware. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The read-mode probe sized its output tensor as tile_k*warp_free while the wave_read fragment is register_count*wave_size elements. For tile_k not equal to the read descriptor's K extent these disagree: one measured geometry produced a fragment twice the window, so the kernel wrote past the allocation on the device. Besides the overrun, the compiler responds to the out-of-range store by bounds-guarding and scalarizing it, so the probe stops measuring the descriptor's access and prices a branchy mixed-width program instead -- a different experiment, not a bad data point. A partial-window golden still passes on the corner it covers, so nothing downstream catches it. Check the two sizes at build time and refuse the mismatch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
run_probe verified store mode only; read mode returned max_abs_diff=None, so every read counter came from a probe never checked to compute the right thing. This module's cardinal rule forbids trusting counters from an unverified probe. The probe stores one coop band (coop_free wide) while the wave reads warp_free wide, so only the first coop_free columns correspond to LDS this probe wrote; LDS holds (K, free) and the final store preserves that order, so the golden is the band transposed. Established empirically on gfx90a before wiring it in. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The read probe's loop was dead code -- n_iter 4, 64 and 256 produced byte-identical counters -- so its numbers carried the cooperative store as an unseparated term. Subtracting a store-only control is not available either: an LDS store with no consumer is deleted by the compiler (measured: zero LDS instructions), which is why the existing probe is a store-MIRROR. Emit N live wave reads instead, each to its own output slice and separated by barriers, so the distinct destinations defeat DCE and the barriers defeat CSE. The (n=2)-(n=1) slope is exactly one wave read and the single store cancels by construction. Check SQ_INSTS_LDS scales with n_reads; if it does not, the reads merged and the slope is invalid. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Only the write port was modelled. PORT_BANKS=8 and COMBINE=4 are write-side constants and every row of the validation corpus was a store-mirror measurement, so a read had no priced answer. Measured on gfx90a via rocprofv3 using the n_reads slope, the read port obeys a structurally different rule: served = max_bank_depth per (half-wave, phase), SUMMED over phases where writes MAX them, with productive = ceil(distinct_dwords/NB). So conflicts/access = max_depth - 1, with zero free constants: a general min(banks,P)*depth/min(banks,C) must satisfy min(b,P) == min(b,C) at b = 4, 8, 16, 32, so the terms cancel identically. PORT_BANKS and COMBINE are not features of the read path. At depth 2 the two rules agree numerically, which is why the depth-8 row is what decides it. max_depth is measured to be the sole cost determinant, not assumed: across pads 0/2/4/6 banks_used sweeps 16 -> 28 and mean depth falls 2.00 -> 1.14 while the cost holds constant, refuting mean-depth and accesses-per-bank rules alike. selftest enforces the structural minima that make MAX a measurement -- at least four distinct depths, a depth-1 floor row, and two rows sharing a depth while differing in banks_used -- and refuses a uniform-only corpus. analyze_read scans the whole wave rather than half-wave 0, returns the model's own input histograms so a reviewer can recompute, and raises on conflicts_per_access rather than returning a plausible number when the arch has no model or the access is out of envelope. dwords_per_lane is required and comes from the disassembly: the emit declares vw=1 and the backend merges to ds_read2_b32, so the emit cannot tell you. gfx90a ships registered at module scope, which re-runs selftest on import. Corpus, selftest block and registration are therefore one indivisible unit: split any finer and importing the module fails, not just the read tests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The skill stated that render_sweep runs verify_lds_roundtrip and verify_mma_soundness automatically. Repo-wide the only callers are tests, so following the skill produced renders over ungated recordings. Also document that both gates pass by omission -- verify_lds_roundtrip returns [] when a space has no store or no read, verify_mma_soundness returns 0 when no MMA was emitted -- which is exactly what an incomplete recording produces. Empty or zero is NOT RUN, not PASS. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A hand-filled CFG is a second source of truth for what the kernel does, sitting beside the emit and free to drift from it, and the natural place to copy one from is a sibling kernel. Its failure mode is silent: it analyzes a different kernel and stamps yours on the answer. Pull every field off the recorded transaction the way verify_lds_roundtrip does; where a field must be hand-supplied, assert it against the recording as soon as one exists and say so in the report. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A four-reviewer audit found the skill's five phases did not cover the common case of iterating on an existing kernel, and that its output contract was implicit enough that expert prose reached the user unsummarized. Redraw as four phases with both fork arms complete (NEW via Frame -> Design&Emit, EXISTING via Recover, both joining at Iterate -> Finish); add an explicit output contract governing what the user sees; add a Validation ladder distinguishing the structural CPU-side rungs from bit-exactness, since a structurally gated design that computes the wrong answer passes every check the skill can run without a GPU. Roughly twenty correctness fixes and a glossary alongside. This also folds in the earlier DERIVE-never-COPY edit to the same file; the restructure subsumed it and the two are no longer separable. tiling_expert.md's design-record path is corrected to match the kernel-folder layout this restructure adopts. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A seven-expert review of the read-port series found the rule itself sound -- every reviewer that checked the arithmetic reproduced all 8 corpus rows independently -- but the guardrails around it did not hold. Three of the eight critical findings were gates that could not fire. GATES THAT DID NOT BIND - The envelope check tested the caller's claimed instruction width, not the access: the same read priced at 2 dwords/lane and refused at 4. The width is genuinely NOT derivable from the address map (a lane with 8 contiguous dwords still issues 4 x ds_read2_b32; the merge follows register allocation), so rather than fake a derivation, price every instruction, require uniform per-instruction depth, and record the width on the report as caller-asserted provenance. - The structural-minima check demanded two rows sharing a max_depth with different banks_used, which all-uniform rows satisfy -- so a corpus where mean depth equals max depth everywhere passed, and MAX stayed an assumption. Now requires a genuinely non-uniform row and a pair whose MEAN depths differ. - The probe guard was a product equality bounding only the global window. Compensating per-axis errors passed, and nothing bounded the read against the LDS allocation, whose accesses are unclipped by design. Store mode had no guard at all. Now per-axis, both modes, plus the 64 KiB LDS budget. SILENT WRONG ANSWERS K and free were decoded with tile_free instead of the row stride, mislabelling every padded render. Cross-run dword dedup could shift one lane's phases against the others. _locate_collision raised KeyError on an inactive lane. The verdict spliced max_nway, the piled-bank count and the located cells from three different served groups into one sentence. A failed re-registration destroyed a working model, and register_arch rewrote the corpus wholesale, deleting a read corpus measured separately. A zero productive floor reported 0.0 -- a confident "conflict-free" on no data. CLAIMS SCOPED TO THE EVIDENCE productive is unvalidated for reads: every corpus row has footprint 128, so a constant 4 fits identically and the denominator of the headline result is measured at one value. "Zero free constants" is narrower than stated -- the served-group size multiplies IDX and is retained by convention, not measured. conflicts/access = max_depth - 1 is an in-envelope corollary, not the rule. The n=0 intercept reproduces the write model at the 8B-aligned pads only. And this work's own data refutes the write model for odd-dword row strides: at pads 2 and 6 the intercepts measure 64/0 where §1.4 predicts 16/12 and 8/4, because the hardware splits the misaligned ds_write_b64. The write pad sweep only ever tested pads divisible by 8, so it was invisible. Recorded as a known limitation next to the mechanism. COVERAGE read_provenance was dead data -- nothing read it, so it could drift from the corpus unnoticed. selftest now re-derives every slope from it and checks SQ_INSTS_LDS scales by exactly instrs_per_read, which is the condition that makes the slope design valid at all. The broadcast arm of the envelope had no test (two mutants survived); the structural minima masked each other so three could be deleted silently; the read golden verified only the first slice, where the extra reads are the whole basis of the slope. Registration at import is now tested in a fresh interpreter -- asserting in-process was itself vacuous, because sibling tests re-register and repair the state under test. The module header advertised a write-only module and omitted all five read symbols; lds_expert.md, which owns the model, had no read-port entry at all. 259 tests pass, up from 241. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two blockers from a value review of the read-port series, plus three prose
corrections a blind A/B of the skill files turned up.
FOOTPRINT WAS DOCUMENTED, NOT GATED
Every read corpus row has footprint 128 dwords, so `productive` is 4
throughout and a constant 4 fits the data identically. That limitation was
written into lds_banks.md 1.5 and treated as handled -- but the code would
still price an access at any other footprint, extrapolating the denominator
of conflicts/access under a SIMULATED verdict. Documenting a gap is not
gating it: a gate that cannot fire is not a gate, and this axis had no gate
at all. analyze_read now refuses unless the per-instruction footprint is the
validated NB x n_groups. Note the broadcast check cannot stand in for this --
it is per-(half-wave, phase), so two lanes of DIFFERENT half-waves sharing a
dword lower the footprint without tripping it, which is how the new test
drives the gate.
THE IMPORT TEST ONLY PASSED ON A SHELL THAT EXPORTS PYTHONPATH
test_gfx90a_read_model_is_registered_at_import spawns a subprocess to check
import-time registration -- correctly, since asserting in-process is vacuous
when sibling tests re-register and repair the state under test. But it did
not propagate sys.path, so from any cwd other than platform/python the child
raised ModuleNotFoundError and the test reported it as "fresh import did not
register the read model": red on CI, and for the wrong reason. It now passes
the parent's sys.path through the environment.
PROSE CORRECTED AGAINST THE CODE
The read-probe description still said `loop {read; sync}` and "robust even if
some iterations are optimized out" -- the loop body is dead code and n_iter
does not move the counters, so an agent following it would have measured
nothing. It now describes the n_reads slope. conflicts/access is stated as
BC/productive, which reduces to max_depth - 1 only for a full-wave
distinct-dword access, matching what the LDS expert file already said. And
the container recipe restores the known-good fmha-build-a:rocm714 /
rocprofv3 1.3.2 configuration verified working on this host, replacing an
unsourced "ROCm 7.10+" floor that nobody measured.
260 tests pass, up from 259.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two findings from a blind A/B of this skill against its predecessor. The restructure named the three layout styles (canonical / interleaved / custom) without explaining any of them, where the previous version spelled out the trade-off. That made the user's decision harder than it had been: a user with no preference was handed a bare three-way choice. Restored, with the MMA Expert still owning the detailed comparison and the encoding. Round 1 listed six required items under a single AskUserQuestion, which takes at most four questions per call -- the instruction could not execute as written. Split into 1a (goal, algorithm, ordered tensors, arch+wave) and 1b (kernel name, caller-side constraints), and clarified that the two-round cap applies to the optional levers, not to the required set. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
- lds_conflict: multi-instruction stores were folded into one, reporting 0.00 (or negative) conflicts/access. Now priced per instruction and summed, matching the read path. Single-instruction results unchanged. - lds_conflict: store_datum restarted phase numbering per run, so the second instruction overwrote the first in the address map and the 3-panel figure drew one instruction under the other's caption. - layout-viz: block-diagram captions were hardcoded (every store "col-major", every reorder "C epilogue", every global load "prefetch"). Now derived from the recorded strides and dataflow. - layout-viz: localization assumed waves split the free axis and printed a K extent on the M axis. Now derives which axis the waves split. - layout-viz: epilogue panel picked the first reorder, not the C-shuffle, then asserted native == store order. Now selects by reachability and raises instead of rendering an unknown classification. - layout-viz: operand reorders were routed to the epilogue, hiding the LDS->MMA bridge and dropping the localization readback stage. - layout-viz skill: a sweep needs a recording; no recording is a stop, not a cue to hand-build one from single panels. Workflow B renders one named view; composing the list is not the skill's call. - tiling-api skill: expert disagreement becomes a knob, not an arbitration round. Report adds Design parameters and Knobs tables. - Expert notes: record that LDS footprint can outrank arithmetic intensity, that the load width is a bank-map lever, and that a multi-patch accumulator is usable interleaved. 14 regression tests added. 280 pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
- Add the interleaving design doc and residency model - Refresh the label-flow, LDS-bank, and MMA docs - Update the tiling/bank-conflict skills and expert personas - Extend the recorder and auto-pipeline viz, with tests - Add the gfx-support audit doc Co-Authored-By: Claude <noreply@anthropic.com>
- Rework the interleaving definition and TL;DR - Generalize off MFMA/wave-64/16x16 to any MMA / atom / arch - Complete the §9 accumulator construction (N-axis, general formulas) - Fix the swapped interleave_idx stride direction (K->free = DPT), verified against code - Scrub flip&zip/gather from the doc and code comments; 32x32 is a valid interleaving target - Dedup and clarity passes from the expert reviews Co-Authored-By: Claude <noreply@anthropic.com>
Reorganize rocke.helpers.tiling into a coherent, glass-box public surface so authors reach for primitives with one home each, grouped by audience. - Re-tier __init__.__all__ into front door / toolbox / (excluded) machinery, and add a contract doc stating the promotion gate every public symbol must pass. - Split transforms.py into a transforms/ package: the shared solver core, the read-only observers (operand soundness, edge classification, C derivation), and the transform_fragment verb; surface the observers at the root. - Extract TileMma into a TileMmaPlan (target-resolved design) + TileMmaDriver (stateless atom-grid iteration) behind the unchanged front-door TileMma, and add a LayoutStyle seam (canonical/interleaved) that owns the operand layout and the memory bridge; the driver validates every operand (soundness + K-match + atom contiguity) unconditionally, so custom layouts are checked identically to derived ones. Remove the broken interleaved= path. - Add N-D descriptor primitives: author-declared axis roles, the typed at_index/squeeze rank-reducing slice, the rank-2 (free, contraction) operand gate, and TileDesc.swap_dims/reorder_registers. - Gate the wide unmasked global load on fully-in-bounds tiles so a ragged/odd tensor scalarizes to the masked path instead of reading out of bounds. - Remove the superseded TilingGemmSpec demo. Co-Authored-By: Claude <noreply@anthropic.com>
Lift the two verification gates -- verify_lds_roundtrip and verify_mma_soundness -- plus the shared geometry helpers (resolve_origin, _lane_span, _arch_wave, _elem_bytes) out of the 1265-line render file visualization/auto_pipeline.py into the analysis/ package (new geometry.py / roundtrip.py / soundness.py), so the gates sit alongside the other analysis (coalescing, vectorization) instead of inside the visualization layer. auto_pipeline.py re-imports these for its own renderers and re-exports them so existing callers keep resolving -- a behavior-neutral move (the moved bodies are byte-identical; the full tiling suite is unchanged). Co-Authored-By: Claude <noreply@anthropic.com>
Small coherence and safety hardening across the tiling library, all behaviour-neutral for existing kernels (the byte-equivalence golden is unchanged): - Add LayoutStyle.lds_bridge() to the base protocol as the declared extension point for a style that stages its operand through LDS -- it returns (lds_read_landing, mma_ready) or None for a direct-load style (canonical). The interleaved style implements it and operand_desc delegates to it, so a kernel composes the bridge off the protocol, never a concrete subclass. - Fire the N-D operand gate (assert_mma_operand in load_fragment) only for a tensor that declares a contraction axis, so a role-tagged (free, free) output tensor -- a C reload -- is no longer wrongly rejected. Operand loads are unaffected (they carry a contraction axis). - Guard a compile-time mid-tile origin in the wide-load gate and the clip mask-skip: a tile placed off the tile grid can overhang the tensor even when the extent is a tile multiple, so it falls to the scalar masked path. A runtime origin is trusted grid-aligned (every caller uses block_id * tile); arbitrary mid-tile origins are handled when that support lands. - Lock "masking an LDS access is unrepresentable" with a signature test (the LDS store/read verbs carry no bounds parameter), and record clipped-through-LDS as correct by construction in the contract. Co-Authored-By: Claude <noreply@anthropic.com>
cgmillette
changed the title
cgmillette
marked this pull request as ready for review
September 30, 2026 17:55
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Public API safety, RDNA target resolution, replicated transforms, cooperative analysis, and exports contain unresolved correctness issues.
Review effort: Balanced
Findings: 5
Open (9)
What changed in this PR
Introduces rocKE’s Python tiling toolkit for describing MMA layouts, fragment movement, memory behavior, analysis, visualization, and kernel experiments.
Changes:
- Adds tile descriptors, register mapping, MMA layout styles, scheduling, transforms, and public APIs.
- Adds coalescing, LDS round-trip, soundness analysis, visualization, and demo infrastructure.
- Adds extensive CPU/GPU-gated tests and architecture/design documentation.
| File | Description |
|---|---|
tests/helpers/tiling/visualization/test_layout_visualizer.py |
Tests text layout rendering. |
tests/helpers/tiling/visualization/__init__.py |
Initializes visualization tests. |
tests/helpers/tiling/traits/test_mma_traits.py |
Tests traits loading and validation. |
tests/helpers/tiling/traits/__init__.py |
Initializes traits tests. |
tests/helpers/tiling/test_wide_load_predication.py |
Tests wide-load safety gates. |
tests/helpers/tiling/test_tiling_recorder.py |
Tests recording and verification. |
tests/helpers/tiling/test_scheduling.py |
Tests scheduling counts and masks. |
tests/helpers/tiling/test_scaffold_smoke.py |
Tests package imports. |
tests/helpers/tiling/test_register_mapper.py |
Tests register-coordinate mapping. |
tests/helpers/tiling/test_layout_optimizer.py |
Tests layout optimization. |
tests/helpers/tiling/test_descriptors.py |
Tests tensor descriptors and slicing. |
tests/helpers/tiling/test_coalescing_probe.py |
Tests ASM coalescing gates. |
tests/helpers/tiling/test_analysis_coalescing.py |
Tests coalescing analysis. |
tests/helpers/tiling/mma/test_warp_encoding.py |
Tests MMA encodings. |
tests/helpers/tiling/mma/test_mma_operation.py |
Tests TileMma resolution. |
tests/helpers/tiling/mma/test_layout_styles.py |
Tests layout-style contracts. |
tests/helpers/tiling/mma/test_layout_oracle.py |
Compares layouts with rocKE oracles. |
tests/helpers/tiling/mma/test_emit_coordinates.py |
Tests emitted coordinates. |
tests/helpers/tiling/mma/__init__.py |
Initializes MMA tests. |
tests/helpers/tiling/layouts/__init__.py |
Initializes layout tests. |
tests/helpers/tiling/kernels/test_tiling_gemm_interleaved_demo.py |
Tests interleaved GEMM demos. |
tests/helpers/tiling/kernels/__init__.py |
Initializes kernel tests. |
tests/helpers/tiling/__init__.py |
Initializes tiling tests. |
tests/helpers/__init__.py |
Initializes helper tests. |
tests/__init__.py |
Initializes the test suite. |
helpers/tiling/visualization/layout_visualizer.py |
Implements text layout reflection. |
helpers/tiling/visualization/__init__.py |
Exposes visualization APIs. |
helpers/tiling/transforms/verb.py |
Implements fragment transformation. |
helpers/tiling/transforms/__init__.py |
Exposes transform APIs. |
helpers/tiling/traits/__init__.py |
Exposes MMA traits APIs. |
helpers/tiling/traits/mma_traits.py |
Loads and resolves MMA traits. |
helpers/tiling/scheduling.py |
Derives scheduling barrier counts. |
helpers/tiling/register_mapper.py |
Maps lane/register slots to coordinates. |
helpers/tiling/coalescing_probe.py |
Connects recorded accesses to ASM analysis. |
helpers/tiling/emit.py |
Implements fragment memory verbs. |
helpers/tiling/README.md |
Documents package usage. |
helpers/tiling/mma/styles/interleaved.py |
Defines interleaved layouts. |
helpers/tiling/mma/styles/canonical.py |
Defines canonical layouts. |
helpers/tiling/mma/styles/base.py |
Defines the layout-style protocol. |
helpers/tiling/mma/styles/__init__.py |
Exposes layout styles. |
helpers/tiling/mma/mma_operation.py |
Implements the front-door MMA object. |
helpers/tiling/mma/__init__.py |
Exposes MMA APIs. |
helpers/tiling/layouts/__init__.py |
Exposes tile authoring. |
helpers/tiling/kernels/tiling_gemm_crc_demo/__init__.py |
Exposes CRC GEMM demos. |
helpers/tiling/kernels/__init__.py |
Initializes demo kernels. |
helpers/tiling/fragments.py |
Defines tile and fragment values. |
helpers/tiling/encoding.py |
Defines warp distribution encoding. |
helpers/tiling/docs/kernels/_TEMPLATE.md |
Adds kernel report template. |
helpers/tiling/docs/api_proposals/README.md |
Documents proposal workflow. |
helpers/tiling/docs/api_proposals/PROPOSAL_c_transpose_crossed_ab.md |
Proposes transposed C support. |
helpers/tiling/docs/api_proposals/_TEMPLATE.md |
Adds API proposal template. |
helpers/tiling/analysis/vectorization.py |
Models vector transactions. |
helpers/tiling/analysis/soundness.py |
Verifies MMA operand soundness. |
helpers/tiling/analysis/roundtrip.py |
Verifies LDS round trips. |
helpers/tiling/analysis/geometry.py |
Resolves recorded geometry. |
helpers/tiling/analysis/coalescing.py |
Models cache-line coalescing. |
helpers/tiling/analysis/__init__.py |
Exposes analysis APIs. |
helpers/tiling/.gitignore |
Excludes temporary artifacts. |
.claude/team_members/rocm_compute_profiler_expert.md |
Adds profiling guidance. |
.claude/team_members/amdgpu_expert.md |
Adds AMD GPU review guidance. |
.claude/skills/shared/temporary_file_policy.md |
Defines temporary-file policy. |
.claude/skills/shared/relationship_discovery.md |
Defines dependency discovery workflow. |
.claude/skills/shared/prerequisites.md |
Defines skill prerequisites. |
helpers/tiling/__init__.py |
Exposes the public tiling API. |
dsl_docs/architecture/multi_arch_data_layout.md |
Documents residency and target resolution. |
dsl_docs/architecture/gfx_support_audit.md |
Audits architecture support gaps. |
platform/AGENTS.md |
Adds tiling development guidance. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Contributor
Author
Thanks for the thorough pass. Filed as follow-ups under epic AICK-1902:
|
…er the target-mismatch guard The _arch_wave failure message pointed at arch=/wave_size=, but those kwargs are forwarded to the build fn -- the declaring call is declared_arch=/ declared_wave_size=. Name the working call, and add recorder tests that a declared target disagreeing with the recorded TileMma fails loud while an agreeing one is accepted. Co-Authored-By: Claude <noreply@anthropic.com>
tenpercent
reviewed
Oct 1, 2026
The driver checked A/B but never the accumulator, so a mislabeled C (its register->(m,n) map != the machine's) was silently stored to wrong coordinates. - collapse the overlapping transform observers into a layered soundness API: per-operand soundness, per-atom K-match, new accumulator-consistency check - driver runs one issue-time validation (+ SOA contiguity) -- no issue site can check a subset - add a package-internal canonical_layouts helper (single canonical reference); route driver, plan oracles, recorder, viz onto it - thread per-atom free-atom counts through the soundness gate so rectangular wave tiles validate per atom Co-Authored-By: Claude <noreply@anthropic.com>
- new style/API_DESIGN_STANDARD.md: language-neutral rules for how functions are layered, named, exported, and tested, worked example = the MMA validation layer - PYTHON_STYLE.md / CPP_STYLE.md: add an "API design" section that points to the standard with each language's mechanics; grow PYTHON_STYLE §11 with test mechanics - KERNEL_AUTHORING.md: list the standard in the Style links Co-Authored-By: Claude <noreply@anthropic.com>
…cate descriptors - a_desc/b_desc/c_desc and a_layout/b_layout/c_layout now return the CONFIGURED style's descriptor/layout (previously returned canonical, ignoring the set style) - delete the duplicate a_operand_desc/b_operand_desc/c_native_desc; the canonical reference comes only from the internal canonical_layouts() helper - validators + MmaTee.from_mma source canonical from the helper, so the viz default stays canonical and is decoupled from the configured style - add a guard test: from_mma on a styled mma still renders canonical by default Co-Authored-By: Claude <noreply@anthropic.com>
…ccessors - a_layout/a_desc now follow the configured style (canonical only by default), not canonical-by-nature - the canonical reference is the canonical_layouts() helper, never mma.a_layout - MmaTee.from_mma default-canonical behaviour is preserved and restated Co-Authored-By: Claude <noreply@anthropic.com>
- reformat the committed tiling Python files with black 25.12.0 (the pinned pre-commit version) to clear the pre-commit CI formatting gate - formatting only, no logic change — tiling suite green (414), goldens unchanged Co-Authored-By: Claude <noreply@anthropic.com>
- reformat the committed demo state to clear the pre-commit gate; content unchanged (revamp-in-progress edits kept out of this commit) Co-Authored-By: Claude <noreply@anthropic.com>
- clears the pre-commit trailing-whitespace gate on mma_is_machinery.md and tiling_interleaving_design.md (whitespace only, no content change) Co-Authored-By: Claude <noreply@anthropic.com>
- fill_fragment: reject nonzero/runtime scalar instead of silently zero-filling - TensorWindow: reject negative compile-time origins (OOB via the clip fast path) - visualization __all__: drop the unbound block_diagram so `import *` works (callable stays at its submodule home; matplotlib stays lazy) - test_register_mapper: scope the atoms importorskip to the oracle test so the structural tests always run - add fill_fragment, negative-origin, and star-export regression tests Co-Authored-By: Claude <noreply@anthropic.com>
vidyasagar-amd
approved these changes
Oct 7, 2026
Contributor
TheRock Submodule Bump ActivityNewest first
|
cgmillette
added a commit
that referenced
this pull request
Oct 10, 2026
## Motivation #12606 source some feedback: Follow-up to AICK-2472 (epic AICK-1902): closes the remaining gap where a Fragment accepted a register vector whose length did not match its TileDesc, which lowered to out-of-range extractelement (LLVM poison) or silently dropped registers. ## Technical Details `Fragment` (the tiling layer's per-lane register object) now guarantees that its registers match its layout and element type. - `tile_desc` and `dtype` are read-only; `dtype` must be a scalar IR type and `tile_desc` must be a `TileDesc`. - Every write to `value`, construction included, must be a `vec<dtype x tile_desc.register_count>`. Anything else raises `ValueError` naming the expected and actual type and the fix; a failed write keeps the old value. - Reading an unfilled fragment raises instead of returning `None`. Without the check, a too-short register vector still lowered and every slice past its end extracted LLVM poison (a silent wrong answer); a too-long one silently dropped registers. The MMA driver and `transform_fragment` docstrings and the tiling API docs are updated to match. ## Test Plan - New `tests/helpers/tiling/test_fragment.py`; `test_fill_fragment.py` now uses a real `IRBuilder`. - LLVM IR for 50 tiling kernel configurations is byte-identical to before the change (valid kernels emit the same IR). - black and pre-commit pass. ## Test Result - Full tiling suite on gfx90a: 489 passed. ## Submission Checklist - [ x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
ISSUE ID : AICK-1902
Motivation
rocKE has no first-class model of how a tile's logical elements map onto lanes, registers, and LDS for MMA/WMMA kernels. That mapping is today implicit in hand-written kernels, so layout choices can't be inspected, reused, or checked for correctness/performance before lowering. This PR lands the initial "M1" tiling layer: a standalone Python toolkit that makes tile distribution, register mapping, LDS bank behavior, and memory coalescing explicit and analyzable — the foundation for the interleaved-layout work tracked under epic AICK-1902.
Tiling primitives can help accelerate development by packaging in design elements and giving access to common optimization knobs found in today's performant kernels: e.g., double buffering, MMA A/B swap, MMA set prio, epilogue shuffles / chunking, scheduling barriers, interleaved layouts, buffer_load/store clipping etc. This library is intended to enable kernel authors to decompose algorithms into tiling fragments and have access to the common performance knobs. Meanwhile, the glass-door design means they have all the power to modify or manually specify parts of the kernel themselves if they so choose.
Tiling primitives (public surface)
The layer is a small set of composable primitives with a single dataflow —
load_fragment → mma → store_fragment— that holds from "say almost nothing" to "pin the exact intrinsic and loop-nest order." The organizing idea: you state intent; the resolved MMA atom supplies every default, and every default is a value you can inspect and replace. The public names are grouped by audience into three tiers (front door / toolbox / machinery), and a name is public only if it is composable, has one home, makes misuse unrepresentable (or can say "unknown"), and names its audience.Value objects (thin, ptr-free / dtype-free):
TensorDesc— a pointer-free strided memory descriptor frommake_tensor_desc(lengths, strides, dtype). Strides are the data layout (one generic verb serves A/B/C — the operand difference is strides + layout, never the verb);lengthsare the valid extent and double as the default clip bound. Rank-generic with author-declared axis roles (free/contraction/batch), a typedat_index/squeezethat reduces only batch axes to reach the rank-2(free, K)MMA view, andpermutefor logical transposes.TensorWindow— a positioned, bounded view frommake_window(tensor, origin, bounds=None). A bound auto-enables clipping (OOB load contributes 0, OOB store is dropped), per-axis, with a compile-time fast path that emits no compare when an axis is a known tile multiple.TileDesc— a dtype-free, memory-free logical→register layout (shape +WarpDistributionEncoding). It is interchangeable at the verb regardless of source:mma.a/b/c_desc(resolved from the atom),make_tile_desc(...)(author a non-MMA tile as a geometric struct-of-arrays:thread_dist,thread_tile,block_repeat,wave_dist, orders, broadcasts,wave_size— no raw encoding integers), orTileDesc(shape, layout)(self-composed).Fragment— realized per-lane registers frommake_fragment(tile_desc, dtype); dtype binds here, filled by the verbs.Verbs (
emit.py):fill_fragment,load_fragment,store_fragment, andtransform_fragment— builder-first, with honest fail-fast casts (identity orf32→{f16,bf16}only) and an optionallds_swizzlebank-swizzle policy for LDS accesses.Driver + knobs:
TileMma(shape, *, a, b, c, target, tiling=, style=)is the intrinsic resolver and subtile driver. The waveshapeplus dtypes and target resolveop_id,atom_shape,subtiles,wave_size, and the A/B/C layouts from the SSOT catalogtraits/data/mma_traits.json(TileMmaPlan/TileMmaDriverunderneath, withoperand_soundness+validate_operandsrun unconditionally so custom operand overrides are validated identically to derived ones).Tiling(atom_shape=, order=)exposes exactly two knobs: atom selection (None= single MMA, a shape tuple = target-agnostic, an intrinsic name = pinned escape hatch) and the M/N/K subtile loop order (any of 6 permutations, all bit-exact since accumulation is commutative).Layout styles:
CanonicalStyle/InterleavedStyle(theLayoutStyleseam inmma/styles/) pick the operand-layout profile; a style composes public primitives (cooperative_load_desc,TileDesc.swap_dims/.reorder_registers, thelds_bridgeprotocol) and influences the accumulator only through the K-distribution its operands present — C is always derived, guarded by a runtime C-oracle.Reflection: every resolved default is a value you can dump —
describe(layout),render_forward_map(...)(lane × register → row/col),render_inverse_map(...), and themma.op_id/atom_shape/subtiles/wave_sizefields.Status: the M1 surface is BUILT and bit-exact on gfx90a (f16→f32), including wave-tile subtiling, arbitrary-shape clipping/predication, and the interleaved style over dense square atoms. Designed-in but reserved seams: LDS-space descriptors, per-load
coherency(CACHE_STREAM/NON_TEMPORAL),sparse/scaledMMA kinds,c_transpose, and theTilePipelinefront-door stage.Why interleaved layouts
The MMA is a fixed vector machine. The ISA documents a canonical data layout as a teaching baseline, but canonical is mechanical ordering, not a correctness requirement. The only physical constraint the hardware imposes is a sound multiply-accumulate: each operand's M/N assignment is fixed and the two K-distributions must match positionally (
A.K-dist == B.K-dist); where M and N land is free and merely routes the output through to accumulator. Canonical layouts satisfy that constraint but are frequently poor memory access patterns. They force one whole canonical atom's worth of elements through the load/store path at a time and don't consider a broader wave-tile.An interleaved layout is a different rule for which thread (and which register) owns which elements of the A/B/C matrices across the entire wave-tile. It computes the identical result as canonical — it only rearranges who-owns-what so the memory side is more efficient. Two properties are the point:
Each issued interleaved atom covers a small slice of all the canonical atoms of a wave-tile at once, instead of just one whole canonical atom at a time (i.e., interleaved). Iterating those slices covers the same wave-tile, but the per-step memory footprint is regular and wide, which makes global loads and LDS stores far more coalescing-friendly (and, on the LDS side, bank-conflict-friendly).
Re-ordering between memory order and MMA order can be simplified to an in-register move. Because the interleaved per-thread patches are regular, switching between a wide coalesced-memory order (e.g.,
free-dim-contiguous) and the MMA-native input order (e.g.,K-contiguous) is a permutation of register slots within each lane; no LDS round-trip and no cross-lane exchange. One layout therefore can serve the wide coalesced access and sits at most one cheap re-order away from what the next stage needs, on both sides of the MMA: input (memory → MMA) and output (MMA → store). The C accumulator is store-ready for free when issue order already matches store order, otherwise it pays a single intra-lane re-order.Canonical is the degenerate case of this model (a tile with one atom along the free dimension has nothing to interleave), not a separate path. The shipped construction (
mma/styles/interleaved.py) builds this reordering tier cross-lane-free by adjusting the static distribution read from each target's MMA traits, and it is designed to cover every dense square atom any AMDGPU architecture (CDNA directly, RDNA/gfx11 via operand replication), any dtype including f64 with no per-atom special-casing.Note: The single non-square catalog atom (
wmma_f32_32x16x128_f4) over-subscribes the wave and is deliberately deferred. The census of exactly which primitives interleave is enforced by thetest_interleave_all_dense.pyproof sweep rather than asserted in prose.Technical Details
New self-contained package
rocke/helpers/tiling/(no changes to existing runtime; nothing imports it yet). Major pieces:mma/):TileMmadense driver, warp/register distribution encoding, per-atom lowering plan, and canonical vs interleaved layout styles.traits/): machine-readable table of ~73 CDNA/RDNA MMA primitives plus the loader/query surface.transforms/): the interleaving engine, built on a lane-uniform ("reorder") tier that is cross-lane-free, with observers for tracing.descriptors.py(tile descriptors),register_mapper.py,layouts/tile_distribution.py,encoding.py,fragments.py,emit.py,memory.py.analysis/+lds_conflict.py,coalescing_probe.py): LDS bank-conflict model, memory coalescing / address-space views, vectorization, geometry, round-trip soundness.visualization/): renders the three layout contexts (logical→register, register→LDS-bank, MMA-atom A/B/C) via an auto pipeline, block diagrams, and kernel-stage views.kernels/): interleaved GEMM, CRC-checked interleaved GEMM, and a manual tiling demo.docs/,dsl_docs/): API contract/surface, interleaving design, LDS-bank notes, gfx support audit, plus a scoped.claude/skills workspace for the subtree.Related tickets
Epic: AICK-1902
Test Plan
Bundled pytest suite under
rocke/tests/helpers/tiling/covering descriptors, register mapping, transforms/interleaving, LDS conflicts, coalescing, the recorder, MMA operation/warp encoding/layout oracle, traits, and the visualization renderers. Notablytest_interleave_all_dense.pyasserts every dense square MMA primitive in the catalog is usable under interleaving.Run:
pytest rocke/tests/helpers/tiling/Test Result
Standalone Python suite; no C++ build or GPU numerical run required (layer is not yet on the emit path).
Submission Checklist