Repository navigation
feat(rocke): one exact MMA catalog row per instruction variant, and IR dtypes through the tiling layer - #15007
cgmillette wants to merge 24 commits into
Conversation
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
acab706 to
bb253bb
Compare
|
Code-quality pass on this PR's own code, folded into the commits that introduced it (no rebase onto develop):
Gates: byte identity green at llvm20/22/23/24; touched tests pass with both engines in strict mode; platform + tiling suites show no new failures; pre-commit passes. |
|
GPU check at head bb253bb on gfx950 (MI350X): test_mma_k_coupling.py (incl. the driven fp8 x fp4 tile) and the fragment storage round-trip selection -- 24 passed, 11 skipped (C++ engine comparisons without the binding, gfx90a-only round trips, and a target not yet run on hardware). Ready for review. |
The interleaved GEMM demo's launching tests now skip through rocke.tests.gpu.current_device_arch unless the current GPU is a gfx90a. Its old module-level check only caught an exception, so on a host whose HIP runtime loads but sees no device it ran the launches and failed. Building the demo kernel needs no GPU and still runs everywhere. The tests that draw a figure skip where matplotlib is not installed, through a requires_matplotlib fixture in the tiling tests' conftest.py: matplotlib is an optional dependency. The tests of the same modules that draw nothing still run. Three tests parse rocke's own Python sources (the range-marker caller check and the import check in test_accesses.py, the fragment-slicer check in test_mma_validation_recipe.py). They now read the sources as UTF-8: with the platform's default encoding, cp1252 on Windows, they stopped on the first source holding a non-ASCII character. Co-Authored-By: Claude <noreply@anthropic.com>
The tiling package keeps its tests inside the package (rocke/tests/helpers/tiling), and rocke_pytest collects only the installed tests/ tree, so no CTest entry ran them. Register rocke_tiling_pytest, which runs that tree from the installed package with the package's PYTHONPATH and the compiler tool locations rocke_pytest gets (not the MLA entry's compiler requirement), and add it to the quick and fast-feedback categories so the tier selection picks it up. The installed CTest environment contract test now checks this entry's environment too. It runs under the installed tests/pytest.ini, the configuration rocke_pytest uses, named with -c: no directory above the installed package holds a rocke pytest configuration, so pytest would otherwise take one from a pytest.ini or pyproject.toml above the install, as in a CI workspace whose pyproject.toml collects only *_test.py files. Co-Authored-By: Claude <noreply@anthropic.com>
The tiling layer's MMA catalog (helpers/tiling/traits/data/mma_traits.json)
keyed rows by op_id, so the five multi-format f8f6f4 builtins, which had one
record per format family, loaded "last one wins". Operand types were
instruction-name tokens (fp8bf8, bf16f32, iu8, ...), and the newer WMMA
target's mixed fp8/bf8 rows carried the wrong one. Rows the loader could not parse became
"reserved" entries.
The catalog is now schema version 2, one row per instruction variant:
- Per-operand dtypes a_dtype, b_dtype, c_dtype (accumulator in) and d_dtype
(result) in canonical core.dtypes names, read from the instruction name.
Integer rows state operand_signedness: "selectable" (integer WMMA, so u8/u4
also match) or "signed" (integer MFMA).
- Each f8f6f4 builtin has one row per (A format, B format), named
<stem>_<a_dtype>_<b_dtype>, with its formats' existing parameters.
- The wave64 RDNA rows are removed.
- Dims and layout parameters are integers. A multi-block row records its
block count in dims B. The parameters form one layout_params object, which
adds AMB (M blocks of A per lane; 1 on every row).
- Scaled rows carry a scale contract {block_k, a_dtypes, b_dtypes}. The
former scale16 builtins are family "scaled" with block_k 16.
- The sparse MFMA rows carry their index operand layout as index_layout;
the sparse WMMA rows hold null there.
- auto_select: false marks the pin-only rows: the gfx11 _tied variants, the
WMMA f8f6f4 8-bit x 8-bit variants, and swmmac_bf16f32_16x16x64_bf16.
- layout_source records where the layout values come from.
Targets and layout values are unchanged.
The loader is strict. A wrong schema version, missing or unexpected fields, a
repeated op_id, a non-integer dim or parameter, a non-canonical dtype, a
mismatched scale contract or replication counts that disagree (dims R and
the A and B repeats AR and BR) fails the load with a message naming the row. A
dtype is canonical when core.dtypes.dtype_info knows it and returns it as its
own name, so an alias (fp8), a typo or a retired token (iu8) fails. So do two
selectable rows that answer the same selection. MmaLayoutParams and
MmaScale are exported, and INTEGER_OPERAND_DTYPES and UNSIGNED_TWIN name
the integer operand dtypes and their unsigned twins.
select() now takes a=, b=, c= and an optional d= (default d = c). It also
takes scale_block_k=, which family "scaled" requires and every other family
refuses, so the block-16 and block-32 scaled instructions are never
auto-picked for each other. The selection key is (target, family, A, B, C,
D, shape, scale block size).
Multi-block rows now load. The encoder refuses them by name through
encoder_supports_layout(): "cannot be laid out by the tiling encoder ... the
block axis is not modelled".
TileMmaPlan keeps the author's dtype strings and normalises them where they
meet the catalog. On atom_override it takes dtypes from the row. It refuses
a row whose result dtype differs from its accumulator dtype, by intent, by
name and by atom_override, because the driver keeps the result in the
accumulator dtype; today that is only wmma_bf16f32_16x16x32_bf16. emit_op
passes the row's names to the backend, and the driver types operands with
core.ir.dtype_to_ir_type. The private alias map and token-to-IR-type map are
gone.
The MMA tee takes each operand's width from its dtype and labels it with
the IR type name. A and B are spelled "<a>/<b>" when they differ. The 16-
and 32-bit labels are unchanged. fp8/bf8 and the f8/f6/f4 formats now show
their canonical names (fp8e4m3, ...), and a bf16f32 row shows its result
name. C/D is still drawn at the result width. The width therefore follows
LLVM for swmmac_bf16f32_16x16x64_bf16, whose result is f32: it is now drawn
at 32 bits.
traits/README.md documents the schema, dtype names, layout parameters,
selection (including that TileMma refusal) and the edit procedure. The tiling docs no longer call the table
generated or give it a row count.
Co-Authored-By: Claude <noreply@anthropic.com>
…ource
Every MMA catalog row describes its register layout twice: as hardware-
notation strings and as layout parameters, which the tiling encoder lays out.
On many rows the two disagreed, and the parameters described only some K order
of the instruction rather than the order the hardware uses.
Each row now stores one layout, the same in both forms. The values are
transcribed from the row's source (the public Matrix Instruction
Calculator, the public CDNA4 ISA tables, rocke core's lane maps and its
scaled WMMA operand layout), else Composable Kernel's parameters, else the
earlier table's strings:
- gfx12 wave32 16-bit dense and sparse rows: AKN = BKN = 2, the two K runs
per lane the calculator shows.
- f8f6f4 rows: per-format K runs on each operand. On the WMMA rows that is 4 runs
for an 8-bit operand and 2 for a 6- or 4-bit one, from core's scaled layout;
the A and B strings are rewritten per format to match.
- The wave32 sparse WMMA rows of the newer WMMA target: AKN = BKN = 2 for 16-bit operands and 4 for 8-bit
ones, from their strings.
- The three 32x16x128_f4 rows: AKN = BKN = 2, and AMB = 2: A's 32 rows are
two blocks of 16, the second in the most significant registers.
- The newer WMMA target's dense rows outside f8f6f4: their strings were only another K
order of CK's parameters and are rewritten to CK's order.
- gfx950 16-bit and 8-bit SMFMAC rows: AKN = 2, the two K runs per lane
their A and index strings state and B holds, so A and B pair the same K.
CK's AKN = 1 does not match its own BKN = 2.
- gfx11 16-bit-result rows: the C string names the 8 logical results,
M{2} L{M0N} V{M1}, and dims S is 1. Their parameters were already exact.
- The f64 strings lose a footnote marker whose note no longer exists.
layout_source records each row's evidence.
The A encoding gains the register-block level AMB describes; at AMB = 1 the
encoding is unchanged. InterleavedStyle does not model AMB > 1 and refuses
such a row by name; the canonical style lays it out.
traits/README.md documents the layout-string grammar, where each row's
layout comes from and which of it the tests check mechanically, where the
catalog differs from CK, every row whose layout has no independent source
(and the sparse WMMA rows' index operand, which no source the catalog
follows describes), and why the WMMA f8f6f4 8-bit x 8-bit rows hold
another K order than the dedicated fp8 rows that answer the same selection.
tiling_interleaving_design.md states why the 32x16x128_f4 atom does not
interleave.
Tests: test_mma_catalog_layouts.py (with the string interpreter
layout_descriptor.py) requires every row's strings to equal the encoder's
layout element for element, one block at a time for multi-block rows and in
compressed K for sparse A; requires A and B of every equal-width row to hold
the same K in each lane; requires every pin-only row to lay out A, B and C
like its selectable twin, except that the WMMA f8f6f4 8-bit rows may
differ from the dedicated fp8 rows by one K permutation shared by A and B;
pins the scaled WMMA rows to core's scaled layout, two calculator point
facts, the sparse index strings to their A, and the README list to the data;
every sparse MFMA row must record its index layout.
New encoder and layout-style tests cover AMB = 2 and the interleaved refusal;
the interleave sweep treats that refusal as the 32x16x128_f4 row's boundary.
Co-Authored-By: Claude <noreply@anthropic.com>
… instruction
The MMA catalog's rows listed family names (gfx11, gfx12) that no other
part of rocke knows, left the gfx950 scaled f8f6f4 rows with no target at
all, and listed the xf32 instructions on gfx950, where the compiler does not
select them.
Each row's supported_targets now lists the real target ids, among rocke's
backend targets, on which the newest clang compiles and links the builtin:
gfx908, gfx90a, gfx942, gfx950, gfx11-generic and every gfx11 member rocke
lists (gfx1154 included), gfx12-generic, gfx1200, gfx1201 and a newer WMMA target. The
lists were measured, not written by hand. The gfx12 sparse wave32
instructions are also listed on that newer WMMA target, and the xf32 instructions on
gfx942 only.
Each row records the LLVM intrinsic its builtin lowers to (llvm_intrinsic,
type-overload suffix removed), read from the IR clang emits.
The scaled WMMA rows record the scale formats the compiler accepts: an
fp4e2m1 operand takes e8m0, e4m3 and e5m3, any other operand e8m0. The
README states the one pairing rule the per-operand sets do not express (both
operands fp4e2m1 on the 16x16x128 rows: the same scale format for A and B).
_meta.excluded_targets names gfx940 and gfx941 with the reason they are not
listed. The loader refuses a row that lists no target ("lists no targets --
list every id the compiler selects it on"), a row that lists an excluded id,
and an excluded id without a reason.
Docs: traits/README.md gains a Targets section and the new fields;
tiling_api_surface.md's example uses target="gfx1201"; the arch registry
and the catalog now share one target namespace, which gfx_support_audit.md
item 6 (now marked resolved) and multi_arch_data_layout.md's
catalog-resolution section say instead of describing family matching.
Tests: every listed target is a backend target, every row lists one, the
excluded ids are never listed, every row names an llvm.amdgcn intrinsic, the
scaled WMMA rows' scale formats follow the operand formats, the three new loader
errors, and selection on real ids (gfx1100, gfx1154, gfx11-generic, gfx1200,
gfx942), while the family name gfx11 selects nothing.
Co-Authored-By: Claude <noreply@anthropic.com>
tools/check_ir_validity.py finds the host's LLVM tools, compiles and links a module, and tells a refused target from a refused module. Other tools and the test helpers reached these through underscore-private names. They are now public: llvm_tool, compile_code_object, target_unsupported and first_diagnostic. tools/gen_arch_domain.py, rocke.tests.ir_validity and the artifact tests use the public names. Behaviour is unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
The MMA catalog is maintained by hand, and four of its facts are the compiler's to decide: on which targets a row's builtin compiles, which LLVM intrinsic it lowers to, which scale formats a scaled builtin's scale-format operands accept, and whether every matrix builtin that compiles on a backend target has a row. tools/check_mma_catalog.py asks the installed clang each of these and reports every disagreement. For every row it compiles a C call to the row's builtin (with the row's A/B format codes on an f8f6f4 builtin) on each of rocke's backend targets and links it the way tools/check_ir_validity.py does, since only the link proves the intrinsic exists. A target an older compiler does not accept by name is compiled as its generic target. It reads the emitted intrinsic with intrinsic_base_name(), which strips the type-overload suffix: the trailing name parts that mangle a type in the declare's own signature. Builtin prototypes come from clang's AST dump, and the list of all matrix builtins from the BuiltinsAMDGPU table installed beside clang. The scale-format codes come from core.arch.wmma_scale.SCALE_FORMATS, now public, the table the scaled WMMA lowering selects them from. It exits 1 on a target, intrinsic or scale-format mismatch and on a builtin with no row, and names what the clang cannot judge (a target it cannot compile for, a builtin it does not know, a missing builtin table) without failing. The catalog agrees with the newest clang available. The unit tests classify hand-written answers against real catalog rows, read real declare shapes and pin the parsers' refusals; they run no compiler and are kept out of the installed test tree. The catalog README's edit procedure gains the step that runs the tool. Co-Authored-By: Claude <noreply@anthropic.com>
…catalog gfx950's f8f6f4 MFMA has an unscaled form, v_mfma_f32_16x16x128_f8f6f4 and v_mfma_f32_32x32x64_f8f6f4, beside the scaled one. It has no clang builtin of its own: the compiler issues it for the scaled builtin when both scales are zero. The catalog listed only the scaled rows, so a kernel could not select the unscaled instruction by its formats. The catalog gains its 50 rows, one per (A format, B format) pair at each of the two shapes, named mfma_f32_<shape>_f8f6f4_<a>_<b> after the instruction. They are family dense with no scale contract, name the scaled builtin and its intrinsic, list gfx950, and carry the same per-operand layout parameters and strings as the scaled row of the same formats (layout_source isa_doc): the scales change no register. The op_id rule now drops 'scale_' from the stem of a dense row on a scaled builtin; the README states it and the builtin rule. Tests select the fp8 x fp8 and fp8 x fp4 rows by their formats, pin each new row's layout to its scaled twin, and the interleave sweep's population count covers the new dense rows. Co-Authored-By: Claude <noreply@anthropic.com>
…sues The MMA catalog lists every instruction variant on every target; nothing said which of them rocke's backend can issue. A committed record now does: traits/data/backend_coverage.json holds, for every (catalog row, target) pair, one measured class and the backend op the row was matched to, read through load_backend_coverage() with status(op_id, target) and backend_op(op_id, target). Only "issuable" means the backend issues the row's instruction there. A pair's candidates are the backend ops with the row's operand key (A, B, C and result dtypes, M, N, K and scale contract). Each is lowered alone with the operand form the tiling layer feeds it, the row's logical elements per lane, and the one whose module calls the row's LLVM intrinsic is the pair's backend op. The classes, first match wins: no_backend_op, other_instruction, multi_block, lowering_not_wired, operand_form_words_only, lane_map_differs, invalid_ir and issuable. A scaled row is recorded per scale format and a selectable integer row per operand dtype pair, since the backend holds one op for each. A lookup names the scale format when a scaled pair has more than one; without operand_dtypes, a selectable integer pair answers for the row's own signed operand dtypes, so (op_id, target) alone looks up every dense and sparse row. The record also lists the backend ops no row was matched to. A CoverageEntry checks its own fields when it is built (TypeError for a field of the wrong kind, ValueError for a wrong value); the reader adds the pair to any error. helpers/mma_io.mma_operand_ir_type is the one home of the operand-form rule: the 32-bit words that hold a sub-byte dtype's elements (core.storage.register_packing), the logical vector otherwise. test_mma_backend_coverage.py measures the record and requires it to equal the committed file. The classes are measured at the record's LLVM flavor, passed explicitly, and a check in fresh processes with ROCKE_LLVM_FLAVOR set to llvm20 and llvm23 shows the host's flavor does not move them. The compile test tells invalid_ir from issuable with rocke.tests.ir_validity.arbiter_verdict, lowering each module at the flavor of the arbiter's clang (the new arbiter_llvm_flavor), and skips by name a target that clang cannot judge. Other tests pin the operand counts in logical elements, the gfx950 scaled MFMA and the scaled WMMA pairs, the two WMMA bf16 pairs that differ only in result dtype, the lookup of a selectable integer row, the reader's errors and mma_operand_ir_type. The record is regenerated with the module's --write option. The catalog README gains a "Backend coverage" section, which says the measurement runs from a source checkout only and when and how to regenerate the record, and an edit step that re-runs the test. extending.md (section 5.5) and KERNEL_AUTHORING.md tell a core MMA change to regenerate it. Co-Authored-By: Claude <noreply@anthropic.com>
IRBuilder.zero_vec now takes u8, in the Python builder and the C++ builder alike. Both lowerings already emit any zero vector as zeroinitializer, and LLVM has no unsigned types, so a u8 zero vector lowers to <n x i8> zeroinitializer. Kernels that do not build a u8 zero vector lower to the same bytes as before. The zero_vec test builds a u8 zero vector and checks its type and the <8 x i8> zeroinitializer it lowers to. ir_model.md lists u8 among zero_vec's element types, and dtype_support.md says zero_vec builds a u8 zero vector. Co-Authored-By: Claude <noreply@anthropic.com>
core.storage.register_packing decided inline which dtypes rocke packs into registers by hand and how wide a register is. Both are now named: REGISTER_BITS is the 32-bit register width, and is_register_packed(dtype) is True for a sub-byte dtype (fp4, fp6, i4), which has no LLVM element type; a predicate raises as before. register_packing is built on the two, with unchanged results and messages. storage_dtype(dtype) is the one rule for the dtype that holds another in registers and in LDS: i32 words for a register-packed dtype, i32 for tf32, i8 for the scale encodings e8m0 and e5m3 (SCALE_ENCODINGS, which core.ir's buffer-access check now reads too), the dtype itself otherwise. storage_group(dtype) gives the fewest elements that fill whole storage elements and how many they fill: (1, 1) for a dtype held one to one, (8, 1) at 4 bits and (16, 3) at 6. elements_to_storage and storage_to_elements convert between element and storage-element counts (a partly filled last word counts, as an allocation needs), and element_bytes gives a dtype's byte width, refusing a sub-byte dtype. Tests pin the predicate on both kinds of dtype and on a predicate, the storage dtype of every kind, the groups against register_packing (whole registers, no padding), the conversions and their refusals, and the byte widths and the sub-byte refusal with its text. Co-Authored-By: Claude <noreply@anthropic.com>
…_int The constant folder read an integer arith.constant through a private helper. It is now public as core.passes.constant_int(value), so code built on the IR builder asks one function whether a value is known at compile time. It returns a Python int itself, reads an i1, i32 or i64 arith.constant at its width as the folder does (an i32 0x80000000 is -2**31), and returns None for a float constant, a bool or a runtime value. The folder uses it unchanged. Tests cover an int, i32 and i64 constants, the width wrap, and the three values that are not compile-time ints. Co-Authored-By: Claude <noreply@anthropic.com>
A Fragment's dtype is the logical number format; its value now holds the registers in Fragment.storage_type, derived from the dtype by core.storage.storage_dtype: the logical vector for every arithmetic dtype, i32 for tf32, and whole i32 words for a packed sub-byte dtype (fp4, fp6, i4), packed low bit first by core.storage.register_packing. helpers/mma_io.storage_element_type and storage_vector_type give that rule in IR types, and storage_ir_type, the unit a global pointer addresses, now reads it too (unchanged answers); mma_operand_ir_type takes its packed form from storage_vector_type. A value of another type is refused; a logical tf32 vector names the bitcast to i32 that fixes it, and a packed fragment names the verbs that build its words. load_fragment, store_fragment and fill_fragment move storage units, in global memory and LDS, a u8 fragment included. Element and unit widths come from core.dtypes.dtype_info, replacing the byte-width table. For every byte-or-wider dtype a unit is an element, so their IR is unchanged. A packed access moves whole 32-bit words: each lane's registers must form whole-word runs along the tensor's stride-1 axis, starting on a word boundary, and a clipped packed window must be bounded on whole words; a clipped global tile masks one whole-word group at a time. A compile-time origin or bound on the stride-1 axis (core.passes.constant_int) is checked; a runtime one is refused unless the caller passes packed_runs_aligned=True to state it is on a whole-word group. The flag is refused for a dtype that is not packed. A packed LDS access is no wider than the word alignment its start has. An LDS buffer holds the tensor dtype's storage element, as registers do: i32 words for packed data and i32 for tf32 (which an LDS buffer of tf32 elements held before; the LLVM IR is unchanged), i8 for e8m0 and e5m3. emit.lds_allocation(tensor, lds_swizzle) gives the buffer to allocate, (storage element, count); the verbs refuse any other buffer, naming lds_allocation, and an LDS swizzle of packed data. A packed access's word index is its element index divided by the elements of one whole-word group, then multiplied by that group's words, so it is exact and stays in the i32 range wherever the element index does. fragments.storage_slots is the one rule for which storage registers hold which logical registers, and core.storage.storage_group gives the registers that move together: one for a byte-or-wider dtype, a whole-word group for a packed one. The verbs, transform_fragment and the word arithmetic of the memory verbs use them, with core.storage's element and storage-element count conversions. Whether a tile axis can overhang its bound is one rule (the bound, the axis's tile extent and its origin) that the clip mask, the unmasked wide access and the packed-bound check share. register_moves.move_registers_into copies registers from one vector into another, and move_registers into a new zero vector; load_fragment places each loaded run with the first, and transform_fragment reorders with the second. The LDS verbs take their access width from one public function, emit.lds_access_units, which the bank-conflict address map uses too. transform_fragment reorders the storage registers: the elements of a byte-or-wider dtype (i32 for tf32) and the words of a packed fragment. A packed reorder must move whole word groups, in order, onto another group; one that would move elements within a word is refused, since repacking bits is not implemented. The byte-addressed views take element widths from core.storage.element_bytes and refuse a sub-byte dtype instead of mis-sizing it; the LDS bank model's address map refuses packed data by name. Tests cover the storage table, storage_slots and storage_group, lds_allocation for every kind of dtype, every refusal and its message (a float-constant origin counts as a runtime one), the alignment each packed access claims, the word index of the largest whole-group element index i32 holds, in i32 arithmetic, u8 fragments filled with u8 zeros and loaded from global memory, lowering identically in both engines, transform_fragment on tf32, u8, fp4 and fp6 fragments, and on gfx90a bit-exact round trips of a tf32 tile through LDS, of fp4, fp6 and i4 operand tiles in gfx950's canonical f8f6f4 layout through global memory and an i32 LDS buffer, checking the loaded words against FragmentPacking.pack, plus a clipped packed window, and of tf32, u8 and fp4 tiles moved to another register order by transform_fragment. Co-Authored-By: Claude <noreply@anthropic.com>
TileMma, TileMmaPlan and MmaTraitsCatalog.select take their dtypes as
scalar IR types from rocke.core.ir (F16, F32, FP8E4M3, ...) instead of
strings. A string, or a composite type such as a vector, is refused with
TypeError naming the constant to pass, e.g. "a= must be a scalar IR type such as
rocke.core.ir.F16, not the string 'f16' -- pass rocke.core.ir.F16". A
Fragment refuses a dtype that is not a scalar IR type, and a layout that
is not a TileDesc, with TypeError too, so the same mistake raises the same
error at either entry point.
A new d= names the result dtype; omitted, it equals the accumulator c,
as for every accumulate-in-place instruction. Selection keys on it, so
it tells apart two instructions that differ only in their result dtype.
The plan and TileMma expose a_dtype, b_dtype, c_dtype and d_dtype as IR
types; atom_override derives them from the row through the new
MmaTraits.ir_dtypes(). The driver compares fragment dtypes with them
directly. MmaTraits.accepts takes IR types too, and
MmaTraits.catalog_dtype is the one conversion from an IR type to the
catalog's dtype name (F16 is "fp16"), so selection and a pinned
instruction's check compare dtypes the same way.
When no instruction matches, select names the query and lists what the
target offers in IR type names, e.g. "no MMA instruction found for
a=f16, b=f16, c=f32, shape=(13, 16, 16), family='dense', target='gfx90a'
-- available on 'gfx90a' (a x b -> c, MxNxK): ['f16 x f16 -> f32,
16x16x16', ...]".
MmaTee.from_mma labels its inputs and result from the MMA's IR dtypes;
mma_input_label is the one home of the input label ("f16", or
"fp8e4m3/bf8e5m2" when A and B differ). Labels and widths are
unchanged.
Every call site, example and doc in the tiling package passes IR types.
New tests cover the refusals (including the constant name of every
scalar IR type), the derived and default dtypes, selection by IR types,
catalog_dtype and accepts (the unsigned twin only where signedness is
selectable),
that gfx950 has no tf32 instruction, the fp8 and bf8 wave tiles lowering
to one MFMA per atom on gfx942 in both engines, and the tee's widths and
labels.
Co-Authored-By: Claude <noreply@anthropic.com>
…A/B dtypes TileMma(atom_override=...) is the one way to pin an instruction by name. Its a/b/c/d are now optional checks rather than ignored: each dtype given must be one the instruction takes (for an integer row with selectable signedness that includes the unsigned twin, kept as given), and each one left out is taken from the instruction, the result dtype included. A mismatch names every wrong dtype and the constant to pass, e.g. "MMA instruction dtype mismatch for 'mfma_f32_16x16x16f16' -- a=bf16 but it takes A=f16; c=f16 but it takes C=f32 -- pass a=rocke.core.ir.F16, c=rocke.core.ir.F32, or omit a/b/c to take them from the instruction". An instruction the target does not list is refused as "MMA instruction not available on target". Tiling(atom_shape=...) takes only an (M, N, K) shape; an instruction name there is refused with TypeError naming the atom_override call to make instead. TileMma's signature and docstring name every keyword. Without a pinned instruction, a/b/c are required. A and B may now differ: selection finds a mixed instruction such as the gfx942 fp8 x bf8 MFMA, and a pair with no instruction fails with the usual "no MMA instruction found". The recorder stores the MMA's A and B dtypes on its op as IR types, which an MMA op requires (a missing one is refused when the op is built, a value that is not an IR type raises TypeError, and a transform op takes none), and the compute and epilogue views read operand widths and the input label from there instead of from an LDS read, so a mixed MMA is drawn with each operand at its own width and labelled "fp8e4m3/bf8e5m2". The docs show the atom_override forms and the mixed example, and AGENTS.md names the two ways to choose an atom. Tests cover the derived and checked dtypes, the mismatch, wrong-target and Tiling-name refusals, mixed and integer selection, the recorded dtypes, the op refusing missing or non-IR dtypes, and a recorded fp8 x bf8 MMA's compute tee. Co-Authored-By: Claude <noreply@anthropic.com>
…age says TileMmaPlan.emit_op() now returns the backend op the backend coverage record names for the plan's (instruction, target) pair, fetched from core by that op id instead of searched for by operand dtypes and shape, which several backend ops can share. A pair the record does not class issuable raises NotImplementedError naming its class and what to pin instead. BackendCoverage.mma_op fetches the op and refuses, naming the regeneration, a record that names an op core no longer registers. An integer instruction with selectable operand signedness passes the plan's A and B dtypes, so u8 operands issue the unsigned op; a scaled instruction takes scale_format=, which the record requires when it holds more than one scale format. TileMma.emit_op() forwards scale_format. emit_op reads the committed record through committed_backend_coverage (exported from the traits package), which loads and validates it once per process, and names a selectable integer variant by MmaTraits.catalog_dtype of the plan's A and B; load_backend_coverage still reads a given path on every call. A BackendCoverage holds read-only copies of its mappings, so the record every kernel shares cannot be changed by one caller. The driver checks, before any IR and in this order: operand dtypes; that the instruction's family is one it feeds (dense; a scaled or sparse instruction raises NotImplementedError); that an instruction whose result dtype differs from its accumulator takes one instruction along K per C subtile; that rocke issues the instruction; then the layout validation. Its result Fragment is typed by the instruction's result dtype, and the grid walk assembles a result of another dtype into a zero vector of that dtype's storage type. The plan no longer refuses such an instruction, so TileMma(..., d=BF16) on the WMMA target that has it and atom_override="wmma_bf16f32_16x16x32_bf16" now drive it. Tests: emit_op of every issuable record entry, per variant, returns the recorded op; emit_op loads the committed record once; a loaded record refuses writes; the refusals and their texts; the dense op chosen where scaled ops share its operands; the bf16-result tiles; WMMA, gfx11 16-bit-result, fp8 x bf8 and f32 16x16x4 tiles lower to their intrinsics in both engines, compiled where the arbiter's clang targets the arch; the stale-record refusal. The traits README and the API surface doc say how issue follows the record; the README gains a start-here table and an "Adding a dtype" checklist, and gfx_support_audit.md and the tiling skill name the record as the authority for what is emitted. Co-Authored-By: Claude <noreply@anthropic.com>
The TileMma driver now slices each instruction's A and B operand out of the fragment's storage registers and hands b.mma the vector type the instruction's A or B dtype takes (mma_io.mma_operand_ir_type), bitcasting only when the two differ. A packed sub-byte operand is sliced as the whole 32-bit words that hold its elements (storage_slots); a TF32 operand held as i32 is reinterpreted as tf32. A u8 operand is sliced as the i8 registers of the same bits, since core's zero_vec has no u8 form, and handed over as u8. Every other dtype is sliced and passed exactly as before, so its IR is unchanged. The grid walk writes each subtile's result back into the tile C with register_moves.move_registers_into, the register move the memory verbs use, instead of a private copy of it; the IR is the same. A TF32 instruction also reads F32 fragments, dropping their low mantissa bits (_MMA_TRUNCATING_SOURCES), so F32 tensors feed xf32 MFMAs with no conversion. The operand dtype error names that accepted source. Docs: the TileMma section states the operand feed and TF32 rule and how a hand-written b.mma call passes operands; the gfx942 arch doc's atom catalog gains the note on when to convert to TF32 first (round to nearest even instead of truncation). Tests, lowering in both engines: xf32 tiles on gfx942 from F32 and from TF32 tensors issue one MFMA per atom and compile; an f16 fragment on a TF32 plan is refused; a two-instruction gfx950 fp4 tile feeds words [0, 4) and [4, 8) to the f8f6f4 instruction and compiles; u8 x i8 WMMA tiles of two K steps, four C subtiles and both issue every instruction with a u8 A operand. Co-Authored-By: Claude <noreply@anthropic.com>
The K check compared A and B slot for slot. That is how every same-format MMA instruction pairs its operands, but not a mixed-width f8f6f4 one: its 8-bit operand holds K in two runs of 16 per lane and its 6/4-bit operand in one run of 32. The slot rule therefore refused the canonical fp8 x fp4 pair and accepted an fp4 B loaded with fp8's map, which multiplies the wrong K. mma_pair_k_aligned now takes the instruction's canonical A and B layouts as required arguments. It reads from them which A and B positions hold the same hardware K, moves B's K labels into A's slot order through that coupling and runs the unchanged three-tier comparison. For a same-format instruction the coupling is the identity, so every result and message stays as it was. mma_pair_compatible forwards its canonicals, so the driver's issue check, verify_mma_soundness and the MmaTee diagnostic all pair K through the coupling. The MmaTee diagnostic moves into MmaTee.pairing_diagnostic() with one branch: a tee without canonicals has no coupling and judges no pairing. evaluate_transform and optimize_layout take k_partner_canon and refuse a k_partner without both canonicals. An instruction whose lane holds several M blocks of A multiplies each block with the same B registers. The coupling pairs every M block of A with the lane's one B column through K, and the driver's atom-contiguity check takes such a lane's atom block as that many runs of one M row each, all in one K step. A whole wave tile compared with too few atoms, whose B lane holds several columns, still gets no coupling. The interleaved construction lays A and B out with one K map, so the dense interleave sweep now proves its mixed-width f8f6f4 pairs K-misaligned instead of passing them. Tests: the mixed-width cases on gfx950 16x16x128 canonicals through both checks; a driven gfx950 fp8 x fp4 tile that issues and compiles; the soundness gate on a recorded fp8 x fp4 pipeline; the MmaTee diagnostic; the optimizer refusal; and one test that drives every pair the backend coverage record issues for a dense instruction (each operand-dtype variant, unsigned ones included) and checks the tile lowers to the instruction's intrinsic and returns the instruction's result dtype, the instruction whose lane holds two M blocks of A included. For that instruction: the canonical pair passes both checks, a K mismatch on the second M block alone fails them, and a tile of two atoms along M, N and K issues all eight. Where the current GPU is a gfx950, the driven fp8 x fp4 tile also runs on random K-dependent operands and must equal an independent numpy product exactly; elsewhere it skips. The same run of the two-M-block f4 tile is written but skipped until it has run on its WMMA target's hardware, and until then the three rows whose lane holds two M blocks of A are pin-only (auto_select: false). The tiles lower in both engines, and an atom that does not split into its M blocks is refused. Docs: mma_is_machinery.md condition 3 (with the coupling of several M blocks), its consequences and trap 1; the interleaving design, the layout-viz skill and the API design standard show the new signatures, and where they say the hardware pairs A slot s with B slot s they now say so for same-format instructions only. Co-Authored-By: Claude <noreply@anthropic.com>
… stores Add fp8e4m3, bf8e5m2 and i8 to the storage elements the tiling memory verbs move with vector accesses. A register run of these 8-bit types that is contiguous along the tensor's stride-1 axis is now read from global memory, read from LDS and written to LDS as one <N x i8> access instead of one byte at a time; the comment on the set names those three users. An 8-bit run is read from LDS as i8 bytes, bitcast for fp8, bf8 and u8, so the read claims the alignment of a byte run whatever width the lowering assumes for an 8-bit float, and a u8 fragment can be read from LDS. The built-in LDS swizzle moves blocks of 8 elements, so an access under it is at most one block wide; its block and group sizes have one home. known_gaps.md records the tiling package as a Python-only authoring layer whose kernels are lowered by both engines. Tests, lowering in both engines: an fp8 or bf8 32x32x64 wave tile loads its operands as eight 8-byte global loads claiming align 8, with no single-byte load, and an fp8, bf8 or i8 fragment staged through an LDS buffer uses only 8-byte vector LDS loads and stores claiming align 8 and compiles for gfx942 (compile-verified, not run). Under the built-in swizzle an i8 or fp8 access covers exactly the elements the per-element swizzle names, and on gfx90a an i8 tile stored with 16-element runs and read back through the swizzle round-trips bit-exactly. A u8 tile staged through LDS lowers identically in both engines and, on gfx90a, round-trips every byte value bit-exactly. Co-Authored-By: Claude <noreply@anthropic.com>
BrianHarrisonAMD
left a comment
There was a problem hiding this comment.
Approving for dnn-providers (changes are not related)
|
Rebased onto current develop to clear the merge conflict.
|
bb253bb to
ab34eac
Compare
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #15007 — feat(rocke): one exact MMA catalog row per instruction variant, and IR dtypes through the tiling layer (cgmillette, AICK-2572, AICK-2576) — 🟠 First-round: IR tests skip in CI
Utility: 🟢 Positive — exact MMA rows, IR dtypes and a tiling CI lane; fixes 14667 Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted First-round review: first pass by a local qwen3.6 model, verified by the orchestrator (Claude Opus 5.5). Checks: merge with develop, build + ctest, byte identity and goldens at four flavors, the whole tiling tree (both engines, and installed as CI runs it), a CP1252 locale run, the coverage record at llvm24, CI logs against develop's own runs, code quality. Flags whether a closer look is needed; not an approval or rejection. |
|
GPU results at ab34eac (Python engine; the tiling test tree):
The red |
The C++ zero_vec admits u8, but no C++ test built one, so dropping u8 from it passed every test. Build an 8-lane u8 zero vector next to the f64 and i16 ones and check that it is a vector of u8, not of i8. Co-Authored-By: Claude <noreply@anthropic.com>
Add lower_with_engine, assert_engine_rejects and their HIP forms, which check one engine named by the caller, and ENGINES, the engines a test parametrizes over. A test parametrized over ENGINES checks the Python engine in every lane, including the lanes that do not build the rocke_engine binding; only its C++ cases skip there. The both-engine helpers now compose them and keep their behaviour. Co-Authored-By: Claude <noreply@anthropic.com>
…ice check Two test modules each probed device 0 for a gfx90a and skipped when it was absent, so a lane that sets ROCKE_REQUIRE_GFX90A=1 passed without running them. A gfx90a_gpu fixture in the package's conftest now asks rocke.tests.gpu.current_device_arch_in, which checks the device the runtime launches on and fails instead of skipping under ROCKE_REQUIRE_GFX90A=1. Co-Authored-By: Claude <noreply@anthropic.com>
The IR tests of the tiling MMA and fragment-storage modules lowered every kernel through lower_with_both_engines, which skips the whole test when the rocke_engine binding is not built. CI's installed tiling lane has no binding, so none of their IR checks ran there. They are now parametrized over ENGINES and lower through lower_with_engine: the Python case runs in every lane, the C++ case also proves the C++ engine emits the same IR, and only the C++ case skips without the binding. The comment on the installed tiling CTest entry lists that skip. Compiling a module does not depend on the engine, so each such test's compile and link now runs once, in its own test, through rocke.tests.ir_validity.assert_kernel_target_accepts, which lowers with the Python engine at the arbiter clang's flavor. It replaces a private helper of test_mma_issue that four other modules imported, and require_arbiter_llvm_flavor replaces the backend coverage tests' own copy of the flavor-or-skip step. The issuable-pair tests build each pair's tile once and share it between those tests, as building the TileMma takes most of their time, and zero their u8 fragments with fill_fragment, which now builds u8 zero vectors. Co-Authored-By: Claude <noreply@anthropic.com>
|
Addressed the latest automated review in four commits on top of
Gates:
|
…lane (AICK-2821) Drop the tiling entry (owned by #15007), the heuristics entry (CI images lack pandas, sklearn and lightgbm) and test_wavescope_pmc.py (30 s subprocess limit is unreliable on a loaded host). Report a child pytest that ran no test as a skip. Update TESTING.md G7. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #15007 — feat(rocke): one exact MMA catalog row per instruction variant, and IR dtypes through the tiling layer (cgmillette, AICK-2572, AICK-2576) — 🟠 Follow-up: N1–N2 fixed; re-measurement deferred
Utility: 🟢 Positive — N1 and N2 fixed; the tiling IR checks now run in CI Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted Follow-up review: first pass by a local qwen3.6 model, verified by the orchestrator (Claude Opus 5.5). Checks: merge with today's develop, build + ctest, byte identity and goldens (4 flavors), changed tests and the whole tiling lane (installed, as CI runs it), mutation checks of N1–N3, CP1252 scans, CI logs against develop's runs, code quality. Flags whether a closer look is needed; not an approval or rejection. |
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #15007 — feat(rocke): one exact MMA catalog row per instruction variant, and IR dtypes through the tiling layer (cgmillette, AICK-2572, AICK-2576) — 🟠 Follow-up: clean merge; re-measurement deferred
Utility: 🟢 Positive — clean develop merge; the tiling lane passes on both Linux shards Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted Follow-up review: first pass by a local qwen3.6 model, verified by the orchestrator (Claude Opus 5.5). Checks: merge with today's develop and the head's remerge-diff, build + ctest, byte identity and goldens (4 flavors), changed tests (source and installed), the coverage re-measurement on the merge, CP1252 scans, CI logs and recent Windows shard times, code quality. Flags whether a closer look is needed; not an approval or rejection. |
…est lane (AICK-2821) (#14903) JIRA ID: AICK-2821 ## Summary `rocke_pytest` collects `platform/tests` only. The `benchmark/perf/tests` directory sits inside the `rocke` package and runs in no CI lane. The `gfx942` cases added in #14713 never run in CI. This PR adds one pytest item for each file in that directory. Each item runs the file in a child pytest process. This follows finding N1 in the automated review of #14713. ## What changes - **`platform/tests/in_tree/test_in_tree_suites.py`** (new). It creates one item for each `benchmark/perf/tests` file. `rocke_pytest` uses `--timeout=60` for each item. A whole directory takes longer. - **A child that runs no test reports a skip.** pytest exits with 0 when every test skips. The item then skips and prints the child output. - **`benchmark/perf/harness.py` and `tests/test_cli.py`.** Changes that make the perf tests pass on Windows and in the child process. - **`TESTING.md` G7.** The entry lists the directories that still run in no lane. ## Left out of this PR - **`tests/helpers/tiling`.** #15007 owns it through `rocke_tiling_pytest`. Both lanes together would run 40 test files twice. - **`heuristics/tests`.** The tests need `pandas`, `sklearn` and `lightgbm`. The CI images install none of them, so every file would skip. A comment in `test_in_tree_suites.py` names the line to add after CI installs these packages. - **`test_wavescope_pmc.py`.** Each test starts a profiler subprocess with a 30 s limit. On a loaded host, 7 of 26 tests exceeded it. - **`portable_ir/tests`.** `test_portable_ir.py::test_unit` already runs it. ## How this was tested - `pytest tests/in_tree`: 16 passed. 12 items run the perf files. 4 items test the skip rule. - Pre-commit hooks over the branch range: pass. - CI: the Linux shards and Windows shard 1/2 pass. Windows gfx110X shard 2/2 timed out in `rocke_pytest` on the slow runner pool. The new head reruns it. ## For AI reviewers - The runner resolves the package directory with `importlib.util.find_spec`. No path anchors on `platform/python`. - The child process runs with `-c os.devnull`. It does not read `platform/tests/pytest.ini`. - A child exit code other than 0 fails the item. A child with no passed test skips the item. - No emitter, `core`, `helpers`, `instances` or `library/kernels` file changes. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
The MMA compatibility check can accept layouts that pair different logical K values, allowing incorrect GEMM results.
4 open issues
What changed in this PR
Re-lands rocKE’s instruction-variant MMA catalog work and carries IR dtypes through the tiling authoring layer, supporting mixed-format operands and packed register storage without changing the binary ABI.
Changes:
- Introduces storage-aware fragment shuffles and typed A/B/C/D MMA APIs, recorder metadata, and visualization.
- Checks operand K alignment through canonical instruction coupling and explicitly rejects unsupported layouts.
- Adds portable tests, per-engine instruction checks, and a scoped installed tiling CTest entry; documents source-only validation gaps.
| File | Description |
|---|---|
| dnn-providers/hip-kernel-provider/rocke/TESTING.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/style/API_DESIGN_STANDARD.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tools/gen_arch_domain.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tools/check_ir_validity.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/tools/test_arch_domain_artifact.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/test_engine_linkage.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_zero_vec_and_lds_element_types.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_storage.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_mma_io.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_f64_lowering.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_constant_int.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/ir_validity.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_layout_render.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_block_diagram.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_auto_pipeline.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/traits/layout_descriptor.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_scheduling.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_lds_conflict.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_emit_wide_access.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_coalescing_probe.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/tensor_transforms/test_accesses.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_warp_encoding.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issuable_pairs.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_styles.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_oracle.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/kernels/test_tiling_gemm_interleaved_demo.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/conftest.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/auto_pipeline.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/verb.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/observers.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/_core.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/__init__.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/tiling_recorder.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/register_moves.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/README.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/warp_encoding.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/styles/interleaved.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/styles/base.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/memory.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/lds_conflict.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/layout_optimizer.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_interleaved_demo.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_crc_demo/crc_interleaved_gemm.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/visualization_api_surface.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/mma_is_machinery.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/lds_banks.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/api_proposals/PROPOSAL_c_transpose_crossed_ab.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/descriptors.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/soundness.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/geometry.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/team_members/lds_expert.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/rocke-tiling-api/SKILL.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/__init__.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/mma_io.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/storage.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/passes.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/ir.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/core/arch/wmma_scale.py | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/dtype_support.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx942.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/ir_lowering/ir_model.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/known_gaps.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/extending.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/multi_arch_data_layout.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/architecture/gfx_support_audit.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/ir_arith.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/cpp/core/ir/core_types.cpp | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/platform/AGENTS.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/rocke/KERNEL_AUTHORING.md | Updated as part of this pull request. |
| dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml | Updated as part of this pull request. |
Review effort: Max
| b_free, b_k = b_slots[lane][slot] | ||
| coupled = (lane, slot) if b_k == k else b_where.get((b_free, k)) |
| mma_pair_k_aligned( | ||
| tgt, k_partner, a_canon=canon, b_canon=k_partner_canon | ||
| ).severity |
| from rocke.tests.gpu import current_device_arch | ||
|
|
||
|
|
||
| @pytest.fixture | ||
| def gfx90a_device() -> None: | ||
| """Skip unless the current GPU is a gfx90a: the demo kernels are built for gfx90a.""" | ||
| arch = current_device_arch() | ||
| if arch != "gfx90a": | ||
| pytest.skip(f"the interleaved demo runs on gfx90a, not {arch}") |
| **What is gated:** the kernels it builds are lowered by both engines. Its tests lower each form the | ||
| package emits with `rocke.tests.engines.lower_with_both_engines`, which requires identical IR: 8-bit |



Motivation
This PR re-lands #14329. #14329 was merged into #14328's branch instead of develop and was reverted there by #14880, so develop, which now has #14328, has none of it. The commits here are #14329's, rebased onto develop, with the fixes for the review threads left open on #14329 folded into the commits they belong to (listed under "Review feedback from #14329" below), and with the new tiling CTest entry following develop's installed compiler test policy from #14868.
The tiling layer's MMA catalog (
helpers/tiling/traits/data/mma_traits.json) is the table the tiling layer selects MMA instructions from and lays out their registers with. It had three problems:op_id, so the multi-format f8f6f4 builtins loaded "last one wins".fp8bf8,iu8, ...), some of them wrong, instead of real dtypes.On top of the catalog, the tiling layer named MMA dtypes with strings, held every Fragment in its logical vector type, and found the backend instruction by searching on dtypes and shape, which could pick a different instruction from the catalog row it had selected. As a result:
This PR does both halves:
The binary ABI does not change.
Technical Details
Test lane (commits 1 and 2)
test(rocke): make the tiling package's tests portabletest(rocke): run the tiling package's tests in CTestrocke_tiling_pytest. It runsrocke/tests/helpers/tilingfrom the installed package and is listed in the quick and ffm-quick categories. Every later commit runs under it.rocke_pytest, it gets the compiler tool locations (_rocke_compiler_tool_environment) but not the compiler requirement that only the MLA entry takes (fix(rocke): scope installed compiler test policy #14868). The installed CTest environment contract test (test_installed_ctest_entries_get_scoped_environment_contracts) now checks this entry's environment too, in both layouts, with and withoutROCM_PATH.rocke_enginebinding, so the C++ engine cases of the IR tests skip there (their Python engine cases run; commit 23). Every test that loadsplatform/tools/also skips there and the entry still passes, because the installed tree has notools/. That is the backend coverage re-measurement, the compile of every issued MMA instruction, and the gfx950 fp8 x fp4 hardware run. They run from a source checkout only; the CMake comment and the catalog README say so.Catalog (commits 3 to 9)
fix(rocke): describe one MMA instruction variant per catalog rowa_dtype,b_dtype,c_dtypeandd_dtype, in canonicalcore.dtypesnames;layout_paramsobject;null, see commit 4);auto_select: falseon pin-only rows.core.dtypes.dtype_infoknows, spelled by its canonical name, so an alias, a typo or a retired token (iu8) fails.select()takesa=,b=,c=, an optionald=, andscale_block_k=. The selection key is (target, family, A, B, C, D, shape, scale block size).TileMmarefuses a row whose result dtype differs from its accumulator dtype; commit 16 lifts this.fix(rocke): store one register layout per MMA catalog row, from its sourcelayout_params).16x16x128tile computes C exactly on hardware. Any other ISA-table or calculator row is as right as its transcription; the README says so.index_layoutisnullthere; the README lists it under "Unverified layouts", and a test requires every sparse MFMA row to record one.fix(rocke): list the real gfx targets and LLVM intrinsic of every MMA instructionrefactor(rocke): make the IR-validity tool's compile helpers publictools/check_ir_validity.py'sllvm_tool,compile_code_object,target_unsupportedandfirst_diagnosticare public;tools/gen_arch_domain.py,rocke.tests.ir_validityand the artifact tests use the public names. Behaviour is unchanged.feat(rocke): add a compiler cross-check tool for the MMA catalogcore.arch.wmma_scale.SCALE_FORMATS, now public, the table the scaled WMMA lowering selects them from.tools/check_mma_catalog.pyasks the installed clang four questions about the catalog: on which targets each row's builtin compiles and links, which LLVM intrinsic it lowers to, which scale formats a scaled builtin accepts, and whether every matrix builtin that compiles on a backend target has a row.feat(rocke): add the unscaled gfx950 f8f6f4 MFMA variants to the MMA catalogfeat(rocke): record which MMA catalog instructions rocke's backend issuestraits/data/backend_coverage.jsonrecords one measured class and the matched backend op for every (row, target) pair. It is read throughload_backend_coverage(). ACoverageEntryvalidates its own fields when built.issuablemeans the backend issues the row's instruction there. The other classes say why not, for exampleno_backend_op,other_instruction,multi_block,operand_form_words_onlyorlane_map_differs.helpers/mma_io.mma_operand_ir_typeis the one home of the operand-form rule.core.archfails that test until the record is regenerated;dsl_docs/development/extending.md(section 5.5),KERNEL_AUTHORING.mdand the catalog README give the command and its clang requirement.Core primitives the tiling layer builds on (commits 10 to 12)
feat(rocke): build zero vectors of u8 in both enginesIRBuilder.zero_vectakes u8 in the Python and the C++ builder alike; it lowers to<n x i8> zeroinitializer. Kernels that build no u8 zero vector lower to the same bytes as before.feat(rocke): name the register-packing and storage rules in core.storageREGISTER_BITSandis_register_packed(dtype);register_packingis built on them, with unchanged results and messages.storage_dtype(dtype)is the one rule for what holds a dtype in registers and LDS: i32 words for sub-byte dtypes, i32 for tf32, i8 for the scale encodings e8m0 and e5m3 (SCALE_ENCODINGS, which core.ir's buffer check reads too), the dtype itself otherwise.storage_group(8 elements in 1 word at 4 bits, 16 in 3 at 6, 1 in 1 otherwise),elements_to_storage/storage_to_elementsandelement_bytesare the unit conversions every caller uses.feat(rocke): read a compile-time integer through core.passes.constant_intIR dtypes through the tiling layer (commits 13 to 19)
feat(rocke): hold Fragment registers in the storage type of their dtypecore.storage.storage_dtype;mma_io.storage_element_type/storage_vector_typein IR types): the logical vector for arithmetic dtypes, i32 for tf32, and whole i32 words for packed sub-byte dtypes, packed low bit first.emit.lds_allocation(tensor, lds_swizzle)gives the buffer to allocate, and the verbs refuse any other buffer naming it. A tf32 tensor's LDS buffer is now i32 (the LLVM IR is the same).fragments.storage_slotsis the one rule for which storage registers hold which logical registers;register_moves.move_registersis the one register shuffle.packed_runs_aligned=True, and a packed LDS access is no wider than its start's word alignment.feat(rocke): take IR dtypes for A, B, C and the result in TileMmaa/b/cand a newd=are scalar IR types, e.g.rocke.core.ir.F16; a string is refused withTypeErrornaming the constant to pass.MmaTraits.acceptstakes IR types;MmaTraits.catalog_dtypeis the one conversion from an IR type to the catalog's dtype name.TileMmacall passes IR types.feat(rocke): pin MMA instructions by full signature and accept mixed A/B dtypesatom_override=is the one way to pin a row. Its dtypes are optional checks; any left out are taken from the row.feat(rocke): issue MMA instructions as rocke's measured backend coverage saysemit_op()returns the backend op the coverage record names for the (row, target) pair, using the author's A/B dtypes for selectable-signedness rows andscale_format=for multi-format scaled rows.committed_backend_coverageis exported from the traits package.feat(rocke): feed MMA operands from their fragments' storage registersregister_moves.move_registers_into.fix(rocke): pair A and B K through the instruction's canonical coupling32x16x128_f4rows, whose lane holds two M blocks of A, are pin-only until their driven-tile test has run on their WMMA target's hardware.feat(rocke): move fp8, bf8 and i8 fragments with vector loads and LDS stores<N x i8>access instead of one byte at a time, in global memory and LDS. An LDS read of an 8-bit run is an i8 read (bitcast for fp8 and bf8), so it claims the alignment of the byte run.IR. The only C++ engine change is commit 10 (u8 in
zero_vec), made in both engines. The IR of shipped kernels does not change, so byte identity and the goldens are unaffected. The IR of 8-bit fragment moves the tiling package emits changes on purpose (commit 19); no golden builds a tiling kernel, so the tiling tests pin those accesses,alignincluded.Engines. The tiling package is a Python-only authoring layer (now recorded in
dsl_docs/development/known_gaps.md, section 9). Its IR tests run once per engine (lower_with_engineoverENGINES), and the C++ engine must emit the Python engine's IR: each form it emits (8-bit vectors, packed words, tf32 held as i32) and one single-instruction tile for every dense instruction the coverage record classes issuable. Where therocke_enginebinding is not built, as in the installed lane, only the C++ cases skip.Review feedback on this PR (commits 20 to 23)
test(rocke): check that the C++ builder makes u8 zero vectors: the C++ half of commit 10 now has a test; it builds a u8 zero vector and checks that its type is a vector of u8.test(rocke): lower or refuse a kernel with one named engine in tests:rocke/tests/engines.pyas develop has it:lower_with_engine,assert_engine_rejects, their HIP forms andENGINES.test(rocke): gate the tiling package's gfx90a tests on the shared device check: agfx90a_gpufixture built onrocke.tests.gpu.current_device_arch_inreplaces two local probes, soROCKE_REQUIRE_GFX90A=1fails instead of skipping without a gfx90a.test(rocke): check the tiling MMA tests' IR with each engine: the IR tests run per engine, so their Python cases run without the binding. The engine-independent compile and link runs once, in its own test, throughrocke.tests.ir_validity.assert_kernel_target_accepts.Review feedback from #14329
The open review threads on #14329 are addressed in the commits above:
<N x i8>(commit 19); the tests assertalign 8.packed_runs_aligned=True; the LDS access width follows the origin's word alignment (commit 13).lower_with_engine, commit 23);known_gaps.mdrecords the tiling package (commit 19).alignassertions: reworded above; the regexes includealign N.auto_select: false) until their driven-tile test has run on hardware (commit 18).index_layout: the README lists the sparse WMMA index operand as unverified, and a test requires it on every sparse MFMA row.TileMmasignature and stale docs, the contributor way-in ("Start here", "Adding a dtype") and three misleading messages, the API-standard items (CoverageEntry.__post_init__,TypeError, refusal tests, annotations) and the import/private-helper items (relative imports, public homes for shared rules, one clang arbiter path).tools/in the installed tree, and more layout point facts with more gfx950 hardware tiles, are follow-up work (AICK-2975, AICK-2976).Test Plan
llvm.amdgcn.*name.CoverageEntryrefusal, the reader's errors and the stale-record refusal are pinned.rocke_tiling_pytestin both layouts, with and withoutROCM_PATH.element_bytes,constant_inton every kind of value.storage_slots,lds_allocationfor every kind of dtype, and every refusal; thealigneach packed and 8-bit access claims; bit-exact GPU round trips on gfx90a for tf32 through LDS, for fp4, fp6 and i4 in gfx950's canonical f8f6f4 layout through global memory and LDS (including a clipped packed window with a compile-time and a runtime bound), for u8 through LDS, and for i8 through the built-in LDS swizzle.TypeError) with the constant name of every scalar type, derived and default dtypes, selection by IR type on gfx90a, gfx942 and gfx950, mixed and integer selection.emit_opof every issuable coverage-record entry, per operand-dtype variant, returns the recorded op; every refusal text.Test Result
test_binding_compiler_selection.py::test_python_bindings_use_the_compiling_library, failed the same way on develop and passes after the rebase onto develop's fix (test(rocke): expect segment and reduce from the gfx942 tiled 3D binding #14836).ISSUE ID : AICK-2572
ISSUE ID : AICK-2576
Submission Checklist
🤖 Generated with Claude Code