From 28816932cd755de61a930978753b794c1bcecd0c Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Tue, 27 Oct 2026 19:11:54 -0400 Subject: [PATCH 1/7] feat(rocke): hold Fragment registers in the storage type of their dtype A Fragment's dtype is the logical number format; its value now holds the registers in Fragment.storage_type, derived from the 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. 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. 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 packed LDS buffer holds i32 words, sized from the descriptor; any other buffer, and an LDS swizzle of packed data, is refused. 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. A u8 fragment is held as vec and stores like any byte dtype. fill_fragment and load_fragment refuse it, because they start from a zero vector of the storage type and core's zero_vec has no u8 form; the message names the route that works: an I8 fragment whose value is bitcast to vec. transform_fragment reorders the storage registers: the elements of a byte-or-wider dtype (i32 for tf32; a u8 fragment through the i8 registers of the same bits) and the words of a packed fragment. A packed reorder must move whole word groups, the fewest elements that fill whole words, 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 (the LDS bank model) refuse a sub-byte dtype instead of mis-sizing it. Tests cover the storage table, every refusal and its message, the word index of the largest whole-group element index i32 holds, in i32 arithmetic, 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 --- .../rocke/helpers/tiling/analysis/geometry.py | 20 +- .../helpers/tiling/docs/tiling_api_surface.md | 16 + .../python/rocke/helpers/tiling/emit.py | 606 ++++++++++---- .../python/rocke/helpers/tiling/fragments.py | 144 +++- .../rocke/helpers/tiling/lds_conflict.py | 8 +- .../python/rocke/helpers/tiling/memory.py | 4 +- .../rocke/helpers/tiling/transforms/verb.py | 53 +- .../helpers/tiling/test_fragment_storage.py | 746 ++++++++++++++++++ 8 files changed, 1378 insertions(+), 219 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment_storage.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/geometry.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/geometry.py index 6c164fd08318..3d432dea0deb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/geometry.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/geometry.py @@ -18,6 +18,8 @@ from typing import Any +from ....core.dtypes import dtype_info + _BINARY = { "arith.add": lambda a, b: a + b, "arith.sub": lambda a, b: a - b, @@ -112,7 +114,19 @@ def _arch_wave(pipeline: Any) -> tuple[str, int]: return pipeline.arch, pipeline.wave_size -def _elem_bytes(dtype_name: str) -> int: - from ..emit import _BYTE_WIDTH +def _elem_bits(dtype_name: str) -> int: + """Bits of one element of `dtype_name` (``core.dtypes.dtype_info``) -- the width home of the + views.""" + return dtype_info(dtype_name).encoded_bits + - return _BYTE_WIDTH[dtype_name] +def _elem_bytes(dtype_name: str) -> int: + """Bytes of one element of `dtype_name`, for the byte-addressed views (the LDS bank model). A + sub-byte element has no byte address of its own and is refused.""" + bits = _elem_bits(dtype_name) + if bits % 8: + raise NotImplementedError( + f"{dtype_name} elements are {bits} bits, not a whole number of bytes -- " + "byte-addressed views (LDS bank model) of packed data -- TODO: not implemented yet" + ) + return bits // 8 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md index 5b91d94877db..9170c614b06f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md @@ -187,6 +187,10 @@ c_td = make_tensor_desc((M_LEN, N_LEN), (N_LEN, 1), F16) # C (M, N), N-contigu a_td = make_tensor_desc((M_LEN, K_LEN), (K_LEN, 1), BF16) # bf16 in c_td = make_tensor_desc((M_LEN, N_LEN), (N_LEN, 1), F32) # f32 store (no narrowing cast) ``` +The descriptor's dtype is the logical format; the verbs move its storage units and never read the +pointer's element type. Pass a TF32 tensor as `PtrType(TF32)` (core's convention) and a packed +sub-byte tensor (fp4, fp6, i4) as `PtrType(I8)`, its bytes packed low bit first; its descriptor is +still in elements. ### LDS buffers: a flat allocation, described by its descriptor @@ -211,6 +215,9 @@ b_frag = load_fragment(b, smem, make_window(b_lds_td, (n0, k0)), mma.b_desc, lan needs 512 swizzled. A swizzle callable must keep every position inside the descriptor's lengths. - **Padding** is a larger outer stride: a `(rows, cols)` buffer padded by `pad` elements per row is `make_tensor_desc((rows, cols), (cols + pad, 1), dtype)`. +- **Packed sub-byte tensors** (fp4, fp6, i4) live in a buffer of i32 words: + `smem_alloc(I32, [ceil(lds_element_span(td) * bits / 32)])`, addressed by the same element index + converted to words. A swizzle is refused for them. - **Vector width.** A wide access claims the alignment of its width, so the verbs pick a width that divides every stride except the stride-1 axis's, leading axes included (in global memory too): a pad or batch stride that is not a multiple of the natural vector narrows the access. @@ -370,6 +377,15 @@ a_frag = load_fragment(b, a_ptr, make_window(a_td, (m, k)), a_desc, lane) # fr ``` - dtype lives on the `Fragment` (D16); `store_fragment` casts fragment dtype -> desc dtype on the honest path only (identity, or f32->{f16,bf16}); anything else fails fast. +- `dtype` is the logical format; the registers in `value` are held in `frag.storage_type`, derived + from it: the logical vector `vec` for every arithmetic dtype, `vec` for tf32 + (it has no arithmetic), and whole i32 words `vec` for a packed sub-byte dtype + (fp4, fp6, i4; `n * bits` must be a multiple of 32). A logical `vec` is refused with the + `b.bitcast` to i32 that fixes it. A u8 fragment is held as `vec` and stores, but + `fill_fragment` and `load_fragment` refuse it (core's `zero_vec` has no u8 form): build an I8 + fragment and `b.bitcast` its value to `vec`. The verbs move storage units: a packed + fragment's registers must form runs of whole 32-bit words along the tensor's stride-1 axis, and a + clipped packed window's bound along that axis must be whole words; both are refused otherwise. - `tile_desc`/`dtype` are read-only; `value` is the one writable field (re-bound each K step: `accumulator.value = b.mma(...)` for a single-atom tile -- wave shape == atom shape, one `b.mma`; a multi-atom tile goes through `TileMma`). The write and read checks are documented on diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/emit.py index 92a20eec305f..08d86d527656 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/emit.py @@ -19,11 +19,21 @@ from __future__ import annotations import math +from collections.abc import Sequence from typing import Any, Callable +from ...core.dtypes import dtype_info +from ...core.ir import I8, I32, U8, VectorType +from ...core.storage import BitPacking from .descriptors import TensorDesc, TensorWindow from .encoding import WarpDistributionEncoding -from .fragments import Fragment, TileDesc +from .fragments import ( + Fragment, + TileDesc, + _fragment_storage_type, + _is_packed, + _storage_element, +) __all__ = [ "emit_tensor_coordinates", @@ -33,31 +43,47 @@ "store_fragment", ] -# Element byte-widths, keyed by the rocke ``ir.Type`` name. Drives the load/store -# alignment so NO dtype is baked into the verb -- the width follows the fragment's type. -_BYTE_WIDTH = { - "i8": 1, - "fp8e4m3": 1, - "bf8e5m2": 1, - "i16": 2, - "f16": 2, - "bf16": 2, - "i32": 4, - "f32": 4, - "i64": 8, - "f64": 8, -} - - -def _align_of(dtype: Any) -> int: - """Natural alignment (bytes) for an ``ir.Type`` -- the element's own width.""" - try: - return _BYTE_WIDTH[dtype.name] - except KeyError as exc: - raise NotImplementedError( - f"no alignment known for dtype -- dtype={dtype.name!r}, " - f"expected one of {sorted(_BYTE_WIDTH)}" - ) from exc +# Sub-byte elements are packed into 32-bit words; one word is the smallest unit a lane accesses. +_WORD_BITS = 32 + + +def _element_bits(dtype: Any) -> int: + """Bits of one logical element of `dtype` (``core.dtypes.dtype_info``).""" + return dtype_info(dtype.name).encoded_bits + + +def _unit_bits(dtype: Any) -> int: + """Bits of one storage unit of `dtype` -- what one memory access moves per element of its + vector: the fragment storage element (``i32`` for tf32 and for packed sub-byte dtypes, the dtype + itself otherwise). Also the access's natural alignment, in bits.""" + return dtype_info(_storage_element(dtype).name).encoded_bits + + +def _unit_count(dtype: Any, elements: int) -> int: + """How many storage units hold `elements` logical elements of `dtype` (whole units).""" + return elements * _element_bits(dtype) // _unit_bits(dtype) + + +def _unit_index(b: Any, element_index: Any, dtype: Any) -> Any: + """Convert a flat element index into the storage-unit index an access is addressed by, + ``element_index * bits / unit_bits``. Emits nothing when one element is one unit (every + byte-or-wider dtype), so their IR is unchanged. A packed index is divided by the elements of one + whole-word group before it is multiplied by that group's words, ``(element_index / 16) * 3`` for + a 6-bit dtype: exact, because every packed access starts on a whole word and so on a whole group + (:func:`_check_packed_runs`), and it never overflows i32 where the element index does not, as + multiplying first would past a third of the i32 range for a 6-bit dtype.""" + bits, unit = _element_bits(dtype), _unit_bits(dtype) + if bits == unit: + return element_index + common = math.gcd(bits, unit) + groups = b.div(element_index, b.const_i32(unit // common)) + return groups if bits == common else b.mul(groups, b.const_i32(bits // common)) + + +def _packed_group(dtype: Any) -> tuple[int, int]: + """The smallest whole-word group of a packed `dtype`: ``(elements, words)`` -- (8, 1) for a + 4-bit dtype, (16, 3) for a 6-bit one.""" + return BitPacking(_element_bits(dtype)).group(_WORD_BITS) def _cast_element(b: Any, value: Any, src: Any, dst: Any) -> Any: @@ -161,19 +187,56 @@ def emit_tensor_coordinates( ) +def _refuse_u8_fragment(verb: str, dtype: Any) -> None: + """Refuse a u8 fragment in a verb that builds the register vector from a zero vector of the + storage type: core's ``zero_vec`` has no u8 form. The i8 zero vector reinterpreted is the same + bits, so the message names that route. + """ + if dtype == U8: + raise NotImplementedError( + f"{verb} cannot build a u8 fragment: core's zero_vec has no u8 form -- TODO: not " + f"implemented yet; build an I8 fragment with {verb} and b.bitcast its value to " + f"vec" + ) + + +def _move_registers( + b: Any, vector: Any, moves: Sequence[tuple[int, int]], count: int +) -> Any: + """A new ``count``-element register vector of ``vector``'s element type holding, for each + ``(src, dst)`` of `moves` in order, element ``src`` of `vector` at element ``dst``: a + compile-time register shuffle (``vec_extract`` / ``vec_insert`` into a zero vector). Elements no + move writes stay zero. + + A u8 vector is moved as the i8 vector of the same bits, since core's ``zero_vec`` has no u8 + form, and the result reinterpreted back to u8. + """ + elem = vector.type.elem + if elem == U8: + as_i8 = b.bitcast(vector, VectorType(I8, vector.type.count)) + return b.bitcast(_move_registers(b, as_i8, moves, count), VectorType(U8, count)) + out = b.zero_vec(elem, count) + for src, dst in moves: + out = b.vec_insert(out, b.vec_extract(vector, src), dst) + return out + + def fill_fragment(b: Any, fragment: Fragment, scalar: Any) -> None: - """Set every register of `fragment` to `scalar`, element-wise (no layout, no addressing). + """Set every register of `fragment` to `scalar`, element-wise (no layout, no addressing). A zero + fill is the zero of the fragment's storage type, so a tf32 or packed fragment is all-zero words. Only a compile-time zero fill is implemented today. A nonzero or runtime `scalar` is REJECTED rather than silently zero-filled, so a caller asking for ones gets a loud failure instead of a - wrong fragment. + wrong fragment. A u8 fragment is refused (:func:`_refuse_u8_fragment`). """ + _refuse_u8_fragment("fill_fragment", fragment.dtype) if not (isinstance(scalar, int) and scalar == 0): raise NotImplementedError( f"fill_fragment supports only a compile-time zero fill -- got scalar={scalar!r}; " "nonzero/runtime scalar splat is not implemented" ) - fragment.value = b.zero_vec(fragment.dtype, fragment.tile_desc.register_count) + storage = fragment.storage_type + fragment.value = b.zero_vec(storage.elem, storage.count) def _as_value(b: Any, x: Any) -> Any: @@ -254,6 +317,12 @@ def _address( return address +def _lds_storage_element(dtype: Any) -> Any: + """The element type of an LDS buffer holding a `dtype` tensor: ``i32`` words for a packed + sub-byte dtype, the tensor dtype itself otherwise.""" + return I32 if _is_packed(dtype) else dtype + + def _check_lds_buffer( verb: str, buffer: Any, @@ -261,24 +330,45 @@ def _check_lds_buffer( lds_swizzle: bool | Callable[[Any, list], list], ) -> None: """Fail fast unless ``buffer`` can hold ``tensor``: an LDS buffer is flat (rank 1), holds the - tensor's dtype, and holds at least :func:`lds_element_span` elements (the descriptor's span, or - the built-in swizzle's reach). The descriptor is the one home of the buffer's layout; the - allocation only gives it room, so a buffer that disagrees with it is refused here rather than - read or written out of range (LDS accesses are unclipped). Window origins are trusted, as for - every LDS access.""" + tensor's LDS storage element (:func:`_lds_storage_element`), and holds the storage units of at + least :func:`lds_element_span` elements (the descriptor's span, or the built-in swizzle's reach). + The descriptor is the one home of the buffer's layout; the allocation only gives it room, so a + buffer that disagrees with it is refused here rather than read or written out of range (LDS + accesses are unclipped). Window origins are trusted, as for every LDS access. A packed tensor is + addressed in whole words, which no swizzle keeps together, so a swizzle is refused for it. + """ buffer_type = buffer.type + dtype = tensor.dtype + element = _lds_storage_element(dtype) span = lds_element_span(tensor, lds_swizzle) + units = -(-span * _element_bits(dtype) // _element_bits(element)) + if _is_packed(dtype) and lds_swizzle: + raise NotImplementedError( + f"{verb}: an LDS swizzle of packed {dtype.name} data -- a swizzle moves elements, and " + f"{dtype.name} elements are addressed in whole 32-bit words -- TODO: not implemented yet" + ) if len(buffer_type.shape) != 1: raise ValueError( f"{verb}: the LDS buffer {buffer_type.name} has rank {len(buffer_type.shape)} -- an LDS " f"buffer is flat; its layout lives in the tensor descriptor. Allocate " - f"smem_alloc({tensor.dtype.name}, [{span}])" + f"smem_alloc({element.name}, [{units}])" ) - if buffer_type.elem.name != tensor.dtype.name: + if buffer_type.elem.name != element.name: + if _is_packed(dtype): + raise ValueError( + f"LDS buffer for {dtype.name} fragments must hold its storage words -- allocate " + f"smem_alloc(I32, [words]) sized from the descriptor; got {buffer_type.name}" + ) raise ValueError( - f"{verb}: LDS buffer element {buffer_type.elem.name} != tensor dtype {tensor.dtype.name}" + f"{verb}: LDS buffer element {buffer_type.elem.name} != tensor dtype {dtype.name}" ) - if span > buffer_type.shape[0]: + if units > buffer_type.shape[0]: + if _is_packed(dtype): + raise ValueError( + f"{verb}: the LDS tensor spans {span} {dtype.name} elements, {units} i32 words, " + f"but the buffer holds {buffer_type.shape[0]} -- size the allocation as " + f"ceil(lds_element_span * {_element_bits(dtype)} / 32) words" + ) raise ValueError( f"{verb}: the LDS tensor spans {span} elements but the buffer holds " f"{buffer_type.shape[0]} -- size the allocation with lds_element_span" @@ -343,16 +433,26 @@ def _register_positions( def _register_index( - b: Any, window: TensorWindow, tile_desc: TileDesc, base_index: Any, register: int + b: Any, + window: TensorWindow, + tile_desc: TileDesc, + base_index: Any, + register: int, + dtype: Any, ) -> Any: - """The register's element index: the lane base index plus the register's compile-time coordinate - offset linearized by the tensor strides. Consecutive registers' indices then differ by a constant, - which the backend vectorizer can merge; it is bit-exact with :func:`_address` of the register's - positions by linearity.""" + """The register's storage-unit index: the lane base index (already in storage units, see + :func:`_unit_index`) plus the register's compile-time coordinate offset linearized by the tensor + strides and converted to units. Consecutive registers' indices then differ by a constant, which + the backend vectorizer can merge; it is bit-exact with :func:`_address` of the register's + positions by linearity. For every byte-or-wider dtype a unit is an element.""" deltas = _register_coord_offsets(tile_desc.layout, register) lead = _leading_rank(window, len(deltas)) - offset = sum( - delta * window.tensor.strides[lead + axis] for axis, delta in enumerate(deltas) + offset = _unit_count( + dtype, + sum( + delta * window.tensor.strides[lead + axis] + for axis, delta in enumerate(deltas) + ), ) return base_index if offset == 0 else b.add(base_index, b.const_i32(offset)) @@ -398,8 +498,8 @@ def _clip_mask( def _zero_scalar(b: Any, dtype: Any) -> Any: - """A scalar 0 of `dtype` -- the zero-pad value handed to ``masked_global_load``.""" - return b.vec_extract(b.zero_vec(dtype, 1), 0) + """A scalar 0 of `dtype`'s storage unit -- the zero-pad value handed to ``masked_global_load``.""" + return b.vec_extract(b.zero_vec(_storage_element(dtype), 1), 0) _WIDE_OK = { @@ -407,78 +507,163 @@ def _zero_scalar(b: Any, dtype: Any) -> Any: "bf16", "f32", "i32", -} # dtypes global_load_vN / smem_load_vN vectorize +} # storage elements global_load_vN / smem_load_vN vectorize (tf32 is held and moved as i32) -def _contiguous_run( - tile_desc: TileDesc, window: TensorWindow, dtype: Any, is_lds: bool -) -> int: - """Width of the innermost stride-1 register run -> a single WIDE load; else 1 (scalar). +def _register_run(tile_desc: TileDesc, window: TensorWindow) -> tuple[int, int] | None: + """``(tensor axis, length)`` of the innermost register bucket when its registers are consecutive + elements in memory, else ``None``. The innermost register bucket varies fastest in the register index. Its registers are consecutive positions on their axis only when the bucket is that axis's INNERMOST hierarchical level (levels of length 1 inside it do not count). A bucket on an outer level steps the position by the product of the levels inside it -- an interleaved layout holds positions 0, 4, 8, 12 in four consecutive registers -- so it gives no run. When the bucket is the innermost level of the - tensor's stride-1 axis, its registers are memory-contiguous and can be one `global_load_vN` / - `smem_load_vN`. Width is the largest supported vector size for the dtype (8/4/2) dividing the run. - - A wide GLOBAL load is UNMASKED, so it is emitted only for a tile that is FULLY IN BOUNDS on the - stride-1 run axis: the effective clip there (`window.bounds[axis]` or the tensor length) must be a - COMPILE-TIME multiple of the tile extent, so no tile overhangs the tensor. A ragged/odd tensor, or - a runtime clip that can't be proven aligned, falls to the SCALAR path (which masks against the - length) -- never an out-of-bounds wide read (the page-fault risk Principle 4 kills). LDS is EXEMPT: - the verbs check that the flat buffer holds the descriptor's whole span (:func:`_check_lds_buffer`) - and LDS accesses are unclipped, so the gate applies to global only. Because both spaces address - an element by ``sum(position * stride)``, consecutive registers of a run on a stride-1 axis are - consecutive elements in LDS exactly as in global memory. - - The width must also divide every OTHER stride, leading axes included: a run then starts at a - multiple of the width whatever the positions on the other axes, so the ``align`` the access claims - holds (an odd batch stride would otherwise start a wide run mid-vector).""" + tensor's stride-1 axis, its registers are memory-contiguous. Because both spaces address an + element by ``sum(position * stride)``, that holds in LDS exactly as in global memory. + """ layout = tile_desc.layout - if dtype.name not in _WIDE_OK: - return 1 majors, minors = layout.register_to_rh_major, layout.register_to_rh_minor if not majors: - return 1 + return None lead = _leading_rank(window, len(tile_desc.shape)) tile_axis = majors[-1] - 1 # innermost register bucket's X-dim -> tile axis - strides = window.tensor.strides run_axis = lead + tile_axis - if tile_axis < 0 or strides[run_axis] != 1: - return 1 + if tile_axis < 0 or window.tensor.strides[run_axis] != 1: + return None inner_levels = layout.hierarchical_lengths[tile_axis][minors[-1] + 1 :] if math.prod(inner_levels) != 1: - return 1 # an outer level of its axis: consecutive registers are not consecutive positions - other_strides = [stride for axis, stride in enumerate(strides) if axis != run_axis] + return None # an outer level of its axis: consecutive registers are not consecutive positions + return run_axis, layout.bucket_length(majors[-1], minors[-1]) + + +def _tile_in_bounds(tile_desc: TileDesc, window: TensorWindow) -> bool: + """True when no tile of this window can overhang an axis it spans, so a GLOBAL access needs no + mask: every tile axis's effective clip (``window.bounds[axis]`` or the tensor length) is a + COMPILE-TIME multiple of the tile extent, and the tile origin is on the tile grid. + + Not just the stride-1 run axis: a wide K-run at a fixed but OOB M row (odd M, aligned K) is just as + out of bounds as an overhanging run. A COMPILE-TIME mid-tile origin is caught (it could straddle + the boundary and overhang); a runtime (SSA) origin is TRUSTED grid-aligned (every caller: + block_id*tile). This is a PARTIAL guard -- an SSA mid-tile origin (attention) is not caught; that + is DEFERRED with attention support (see _origin_misaligned). Only the tile axes are checked: a + leading axis is positioned by its origin and not clipped, exactly as in _clip_mask. + """ + lead = _leading_rank(window, len(tile_desc.shape)) + for tile_axis, extent in enumerate(tile_desc.shape): + ax = lead + tile_axis + clip = ( + window.bounds[ax] + if window.bounds is not None and window.bounds[ax] is not None + else window.tensor.lengths[ax] + ) + if not (isinstance(clip, int) and clip % extent == 0): + return False + if _origin_misaligned(window.origin[ax], extent): + return False + return True + + +def _check_packed_runs( + verb: str, tile_desc: TileDesc, window: TensorWindow, dtype: Any, is_lds: bool +) -> int: + """Fail fast unless every access of a packed `dtype` moves whole 32-bit words, and return the + run length in elements. + + A lane cannot write part of a word another lane owns, so each lane's registers must form runs of + whole words along the tensor's stride-1 axis (:func:`_register_run`; a multiple of the dtype's + whole-word group, 8 elements for 4 bits, 16 for 6 bits), and each run must start on a group + boundary: every other stride and a compile-time origin on the run axis are multiples of the + group. A clipped GLOBAL window masks whole groups, so its compile-time bound along the stride-1 + axis must be a multiple of the group too (a runtime bound is trusted aligned, as runtime origins + are). LDS accesses are unclipped.""" + group, _ = _packed_group(dtype) + run = _register_run(tile_desc, window) + length = 1 if run is None else run[1] + # "an 8-element boundary", "a 16-element boundary" + article = "an" if str(group).startswith("8") else "a" + need = ( + f"{verb} of {dtype.name} needs each lane's registers to form runs of whole 32-bit words " + f"along the tensor's stride-1 axis (multiples of {group} elements, starting on {article} " + f"{group}-element boundary)" + ) + if run is None or length % group: + raise ValueError( + f"{need}; this layout gives runs of {length} " + f"{'element' if length == 1 else 'elements'} -- choose a layout whose innermost " + "register run lies along the stride-1 axis" + ) + run_axis = run[0] + for axis, stride in enumerate(window.tensor.strides): + if axis != run_axis and stride % group: + raise ValueError( + f"{need}; tensor axis {axis} steps by {stride} elements, so its runs start off " + f"that boundary -- pad that stride to a multiple of {group}" + ) + origin = window.origin[run_axis] + if isinstance(origin, int) and origin % group: + raise ValueError( + f"{need}; the window origin on the stride-1 axis is {origin} -- place the window on a " + f"multiple of {group}" + ) if not is_lds: - # Fully-in-bounds gate (global only): a wide load is UNMASKED, so it is safe only if the tile - # overhangs NO axis it spans -- not just the stride-1 run axis. A wide K-run at a fixed but - # OOB M row (odd M, aligned K) is just as out of bounds as an overhanging run. So EVERY axis's - # effective clip must be a compile-time multiple of the tile extent on that axis; otherwise - # some tile overhangs and the unmasked read faults -> scalarize (the masked path is safe). - # Every axis's effective clip must be a compile-time tile multiple AND the tile origin must be on - # the tile grid: a COMPILE-TIME mid-tile origin is caught (it could straddle the boundary and - # overhang); a runtime (SSA) origin is TRUSTED grid-aligned (every caller: block_id*tile). This - # is a PARTIAL guard -- an SSA mid-tile origin (attention) is not caught; that is DEFERRED with - # attention support (see _origin_misaligned). Only the tile axes are checked: a leading axis - # is positioned by its origin and not clipped, exactly as in _clip_mask. - for tile_axis, extent in enumerate(tile_desc.shape): - ax = lead + tile_axis - clip = ( - window.bounds[ax] - if window.bounds is not None and window.bounds[ax] is not None - else window.tensor.lengths[ax] + bound = ( + window.bounds[run_axis] + if window.bounds is not None and window.bounds[run_axis] is not None + else window.tensor.lengths[run_axis] + ) + if isinstance(bound, int) and bound % group: + raise ValueError( + f"{verb} of a clipped {dtype.name} window needs its bound along the stride-1 axis " + f"to be a multiple of {group} elements (whole 32-bit words) -- got {bound}" ) - if not (isinstance(clip, int) and clip % extent == 0): - return 1 - if _origin_misaligned(window.origin[ax], extent): - return 1 - length = layout.bucket_length(majors[-1], minors[-1]) - max_vw = ( - 16 // _BYTE_WIDTH[dtype.name] - ) # one load = dwordx4 = 16 bytes; longer runs -> N loads + return length + + +def _contiguous_run( + tile_desc: TileDesc, + window: TensorWindow, + dtype: Any, + is_lds: bool, + *, + verb: str = "load_fragment / store_fragment", +) -> int: + """Storage units one access moves: the widest vector over the innermost stride-1 register run + (:func:`_register_run`) -> a single WIDE load; else 1 (one element). `dtype` is the logical + dtype; the width is counted in its storage units (:func:`_unit_bits`), which are its elements for + every byte-or-wider dtype. Width is the largest supported vector size (16/8/4/2) of at most 16 + bytes dividing the run. + + A wide GLOBAL load is UNMASKED, so it is emitted only for a tile that is FULLY IN BOUNDS + (:func:`_tile_in_bounds`). A ragged/odd tensor, or a runtime clip that can't be proven aligned, + falls to the SCALAR path (which masks against the length) -- never an out-of-bounds wide read (the + page-fault risk Principle 4 kills). LDS is EXEMPT: the verbs check that the flat buffer holds the + descriptor's whole span (:func:`_check_lds_buffer`) and LDS accesses are unclipped, so the gate + applies to global only. + + The width must also divide every OTHER stride, leading axes included: a run then starts at a + multiple of the width whatever the positions on the other axes, so the ``align`` the access claims + holds (an odd batch stride would otherwise start a wide run mid-vector). + + A packed sub-byte dtype has no one-element access: its runs must be whole words + (:func:`_check_packed_runs`, which raises with `verb` otherwise). A global tile that may overhang + moves one whole-word group per access, masked as a unit; any other packed access is the widest + whole-group vector of words dividing the run.""" + if _is_packed(dtype): + return _packed_access_units(verb, tile_desc, window, dtype, is_lds) + if _storage_element(dtype).name not in _WIDE_OK: + return 1 + run = _register_run(tile_desc, window) + if run is None: + return 1 + run_axis, length = run + if not is_lds and not _tile_in_bounds(tile_desc, window): + return 1 + other_strides = [ + stride for axis, stride in enumerate(window.tensor.strides) if axis != run_axis + ] + # One load = dwordx4 = 16 bytes; longer runs -> N loads. + max_vw = 128 // _unit_bits(dtype) for vw in (16, 8, 4, 2): if ( vw <= max_vw @@ -489,6 +674,42 @@ def _contiguous_run( return 1 +def _packed_access_units( + verb: str, tile_desc: TileDesc, window: TensorWindow, dtype: Any, is_lds: bool +) -> int: + """The words one access of a packed `dtype` moves (see :func:`_contiguous_run`): one whole-word + group for a global tile that may overhang (masked as a unit), else the widest whole-group count + of at most 4 words (16 bytes) that divides the run and keeps every run start on a multiple of + itself.""" + length = _check_packed_runs(verb, tile_desc, window, dtype, is_lds) + _, group_words = _packed_group(dtype) + if not is_lds and not _tile_in_bounds(tile_desc, window): + return group_words + run_axis, _ = _register_run(tile_desc, window) + bits = _element_bits(dtype) + run_words = _unit_count(dtype, length) + other_strides = [ + stride for axis, stride in enumerate(window.tensor.strides) if axis != run_axis + ] + for words in (4, 3, 2, 1): + if ( + words % group_words == 0 + and run_words % words == 0 + and all( + stride * bits % (words * _WORD_BITS) == 0 for stride in other_strides + ) + ): + return words + raise AssertionError("a whole-word group always divides a checked packed run") + + +def _access_align(units: int, dtype: Any) -> int: + """The alignment (bytes) an access of `units` storage units claims: the access size when it is a + power of two, else the largest power of two dividing it (a 3-word access claims one word). + """ + return (_unit_bits(dtype) // 8) * (units & -units) + + def _is_lds(source: Any) -> bool: """True when `source` is an LDS buffer (an ``smem<...>`` value from ``smem_alloc``) rather than a global/constant ``ptr<...>``. Lets one verb serve both spaces -- the load/store op is @@ -590,27 +811,29 @@ def _lds_index( register: int, lds_swizzle: bool | Callable[[Any, list], list], base_index: Any, + dtype: Any, ) -> Any: - """The flat LDS element index of `register`'s run start, ``sum(position * stride)`` over the - window's descriptor -- the same address rule as global memory. + """The flat LDS storage-unit index of `register`'s run start: ``sum(position * stride)`` over + the window's descriptor -- the same address rule as global memory -- in units of `dtype`'s LDS + storage element (:func:`_unit_index`; the element index for every byte-or-wider dtype). - Without a swizzle it is the lane base index (``base_index``, from :func:`_lane_base_index`, - emitted once per access) plus the register's compile-time offset: consecutive runs differ by a - constant, which the backend folds into the instruction offset and pairs into ``ds_read2`` / - ``ds_write2``. A swizzle is not linear, so a swizzled run rebuilds its positions, swizzles them - and linearises the result (``base_index`` is unused then). + Without a swizzle it is the lane base index (``base_index``, from :func:`_lane_base_index` in + storage units, emitted once per access) plus the register's compile-time offset: consecutive runs + differ by a constant, which the backend folds into the instruction offset and pairs into + ``ds_read2`` / ``ds_write2``. A swizzle is not linear, so a swizzled run rebuilds its positions, + swizzles them and linearises the result (``base_index`` is unused then). The swizzle sees the position on EVERY tensor axis, leading axes (their origins) first, so an element's swizzled address does not depend on the rank of the tile that reaches it. """ if not lds_swizzle: - return _register_index(b, window, tile_desc, base_index, register) + return _register_index(b, window, tile_desc, base_index, register, dtype) coords = emit_tensor_coordinates(b, tile_desc.layout, thread, register) lead = _leading_rank(window, len(coords)) leading = [_as_value(b, window.origin[axis]) for axis in range(lead)] swizzle = _swizzle_lds_positions if lds_swizzle is True else lds_swizzle swizzled = swizzle(b, leading + _positions(b, window, coords)) - return _address(b, window, swizzled[lead:], swizzled[:lead]) + return _unit_index(b, _address(b, window, swizzled[lead:], swizzled[:lead]), dtype) def _checked_lds_base_index( @@ -623,14 +846,16 @@ def _checked_lds_base_index( lds_swizzle: bool | Callable[[Any, list], list], ) -> Any: """Check the LDS buffer (and the built-in swizzle's layout rule), then return the lane base - index the unswizzled runs add their offsets to (``None`` when a swizzle is active). + index, in storage units, the unswizzled runs add their offsets to (``None`` when a swizzle is + active). """ if lds_swizzle is True: _check_builtin_swizzle(window.tensor) _check_lds_buffer(verb, buffer, window.tensor, lds_swizzle) if lds_swizzle: return None - return _lane_base_index(b, window, tile_desc, thread)[1] + base_index = _lane_base_index(b, window, tile_desc, thread)[1] + return _unit_index(b, base_index, window.tensor.dtype) def _swizzle_vw(lds_swizzle: Any, dist_vw: int, align: int) -> int: @@ -669,57 +894,68 @@ def load_fragment( TensorDesc (strides) + the tile_desc. `ptr` is the SOURCE the fragment is loaded from: a global ``ptr<...>`` (VMEM) OR an LDS ``smem<...>`` buffer from ``smem_alloc`` -- the load op is chosen from the source (:func:`_is_lds`), so the same verb stages global<->register and LDS<->register. - The fragment's dtype IS ``window.tensor.dtype`` (nothing is baked in). Where the window clips (an + The fragment's dtype IS ``window.tensor.dtype`` (nothing is baked in); its registers are read as + that dtype's storage units (:attr:`Fragment.storage_type`: i32 for tf32, whole i32 words for a + packed sub-byte dtype, whose runs must be whole words). Where the window clips (an element past its bound), OOB elements are filled with `pad` (`masked_global_load`); otherwise the plain `global_load` fast path (byte-identical to no-clip). In both spaces an element's index is ``sum(position * stride)`` over the window's descriptor. An LDS buffer is flat (``b.smem_alloc(td.dtype, [lds_element_span(td, lds_swizzle)])``) and the descriptor is its - layout: the buffer must be rank 1, hold the tensor's dtype and hold that span (refused - otherwise); LDS loads are unclipped. + layout: the buffer must be rank 1, hold the tensor's dtype (i32 words for a packed dtype) and + hold that span (refused otherwise); LDS loads are unclipped. `pad` default `0`; a non-zero `constant(value)` fill is reserved. `coherency` reserved. A rank-t tile covers the window's last t axes (t <= the window's rank, refused otherwise); each - leading axis is read at its origin. + leading axis is read at its origin. A u8 tensor is refused (:func:`_refuse_u8_fragment`). """ + _refuse_u8_fragment("load_fragment", window.tensor.dtype) _check_tile_rank("load_fragment", window, tile_desc) _check_leading_bounds("load_fragment", window, tile_desc) dtype = window.tensor.dtype - align = _align_of(dtype) lds = _is_lds(ptr) - vw = _contiguous_run(tile_desc, window, dtype, lds) # >1 -> one wide load per run + # Storage units per access (>1 -> one wide load per run) and the registers they hold. + vw = _contiguous_run(tile_desc, window, dtype, lds, verb="load_fragment") # A custom swizzle relocates whole units of its granularity (vw_elems), capped + range-checked to # [1, natural run]. The built-in bool swizzle preserves its run. if lds: - vw = _swizzle_vw(lds_swizzle, vw, align) - value = b.zero_vec(dtype, tile_desc.register_count) + vw = _swizzle_vw(lds_swizzle, vw, _unit_bits(dtype) // 8) + step = vw * _unit_bits(dtype) // _element_bits(dtype) + storage = _fragment_storage_type(dtype, tile_desc.register_count) + unit = storage.elem + value = b.zero_vec(unit, storage.count) if lds: - # LDS source: one flat element index per run (wide when contiguous), lane base once. + # LDS source: one flat storage-unit index per run (wide when contiguous), lane base once. base_index = _checked_lds_base_index( "load_fragment", b, ptr, window, tile_desc, thread, lds_swizzle ) - for register in range(0, tile_desc.register_count, vw): + for register in range(0, tile_desc.register_count, step): index = _lds_index( - b, window, tile_desc, thread, register, lds_swizzle, base_index + b, window, tile_desc, thread, register, lds_swizzle, base_index, dtype ) - loaded = b.smem_load_vN(ptr, index, dtype=dtype, n=vw) + loaded = b.smem_load_vN(ptr, index, dtype=unit, n=vw) + slot = _unit_count(dtype, register) for i in range(vw): - value = b.vec_insert(value, b.vec_extract(loaded, i), register + i) + value = b.vec_insert(value, b.vec_extract(loaded, i), slot + i) return Fragment(tile_desc, dtype, value) - zero: Any = None # the pad scalar, emitted lazily on the first clipped element - for register in range(0, tile_desc.register_count, vw): + align = _unit_bits(dtype) // 8 + zero: Any = None # the pad scalar, emitted lazily on the first clipped access + for register in range(0, tile_desc.register_count, step): coords = emit_tensor_coordinates(b, tile_desc.layout, thread, register) positions = _positions(b, window, coords) - address = _address(b, window, positions) - if vw > 1: + address = _unit_index(b, _address(b, window, positions), dtype) + slot = _unit_count(dtype, register) + # A wide access is only chosen for a tile that cannot overhang, so it never needs a mask; a + # packed whole-word group may, and is masked as a unit (its bound is whole groups). + mask = _clip_mask(b, window, positions, tile_desc.shape) + if mask is None and vw > 1: # Contiguous stride-1 run -> single vectorised global load (global_load_dwordxN). - loaded = b.global_load_vN(ptr, address, dtype, vw, align=vw * align) + loaded = b.global_load_vN( + ptr, address, unit, vw, align=_access_align(vw, dtype) + ) for i in range(vw): - value = b.vec_insert(value, b.vec_extract(loaded, i), register + i) + value = b.vec_insert(value, b.vec_extract(loaded, i), slot + i) continue - mask = _clip_mask(b, window, positions, tile_desc.shape) - if mask is None: - loaded = b.global_load(ptr, address, dtype, align=align) - else: + if mask is not None: if pad != 0: # Only a zero clip pad is implemented; a per-stage pad (e.g. -inf for a max) is # reserved for the reduction stages. @@ -729,11 +965,34 @@ def load_fragment( ) if zero is None: zero = _zero_scalar(b, dtype) - loaded = b.masked_global_load(ptr, address, mask, zero, dtype, align=align) - value = b.vec_insert(value, loaded, register) + for i in range(vw): + unit_address = address if i == 0 else b.add(address, b.const_i32(i)) + if mask is None: + loaded = b.global_load(ptr, unit_address, unit, align=align) + else: + loaded = b.masked_global_load( + ptr, unit_address, mask, zero, unit, align=align + ) + value = b.vec_insert(value, loaded, slot + i) return Fragment(tile_desc, dtype, value) +def _stored_unit(b: Any, fragment: Fragment, slot: int, out_dtype: Any) -> Any: + """Storage register `slot` of `fragment`, cast to `out_dtype` (:func:`_cast_element`).""" + return _cast_element( + b, b.vec_extract(fragment.value, slot), fragment.dtype, out_dtype + ) + + +def _global_store_units( + b: Any, ptr: Any, address: Any, units: list[Any], align: int +) -> None: + """Store `units` at consecutive storage-unit addresses from `address`, one scalar store each.""" + for i, value in enumerate(units): + unit_address = address if i == 0 else b.add(address, b.const_i32(i)) + b.global_store(ptr, unit_address, value, align=align) + + def store_fragment( b: Any, ptr: Any, @@ -758,20 +1017,26 @@ def store_fragment( """ _check_tile_rank("store_fragment", window, fragment.tile_desc) _check_leading_bounds("store_fragment", window, fragment.tile_desc) + tile_desc = fragment.tile_desc out_dtype = window.tensor.dtype - align = _align_of(out_dtype) + align = _unit_bits(out_dtype) // 8 lds = _is_lds(ptr) # Wide LDS stores when the innermost register run is contiguous (ds_write_b{32,64,128}); the - # global store stays scalar (the C epilogue is one-shot + may clip/cast). - vw = ( - _contiguous_run(fragment.tile_desc, window, out_dtype, is_lds=True) - if lds - else 1 - ) - # A custom swizzle relocates whole units of its granularity (vw_elems) -> the access is capped and - # range-checked to [1, natural run]. The built-in bool swizzle preserves its run. + # global store stays scalar (the C epilogue is one-shot + may clip/cast): one element, or one + # whole-word group of a packed dtype, per access. if lds: + vw = _contiguous_run(tile_desc, window, out_dtype, True, verb="store_fragment") + # A custom swizzle relocates whole units of its granularity (vw_elems) -> the access is + # capped and range-checked to [1, natural run]. The built-in bool swizzle preserves its run. vw = _swizzle_vw(lds_swizzle, vw, align) + elif _is_packed(out_dtype): + _check_packed_runs("store_fragment", tile_desc, window, out_dtype, False) + vw = _packed_group(out_dtype)[1] + else: + vw = 1 + step = vw * _unit_bits(out_dtype) // _element_bits(out_dtype) + unit = _storage_element(out_dtype) + # Global path: compute the LANE base-address ONCE (register 0 = the thread-only address, no register # offset), then emit each store as base + -- the fixed per-register coordinate # offset linearized by the strides. Consecutive registers' addresses then differ by a constant, which @@ -782,47 +1047,46 @@ def store_fragment( # LDS path uses the same base-plus-constant form per wide run (see _lds_index). if lds: lds_base_index = _checked_lds_base_index( - "store_fragment", b, ptr, window, fragment.tile_desc, thread, lds_swizzle + "store_fragment", b, ptr, window, tile_desc, thread, lds_swizzle ) - for register in range(0, fragment.tile_desc.register_count, vw): + for register in range(0, tile_desc.register_count, step): index = _lds_index( b, window, - fragment.tile_desc, + tile_desc, thread, register, lds_swizzle, lds_base_index, + out_dtype, ) + slot = _unit_count(fragment.dtype, register) if vw == 1: - el = _cast_element( - b, - b.vec_extract(fragment.value, register), - fragment.dtype, - out_dtype, + b.smem_store_vN( + ptr, [index], _stored_unit(b, fragment, slot, out_dtype), 1 ) - b.smem_store_vN(ptr, [index], el, 1) + elif vw & (vw - 1): + # No LDS vector store of this width (a 3-word packed group): one store per word. + for i in range(vw): + el = _stored_unit(b, fragment, slot + i, out_dtype) + b.smem_store_vN(ptr, [b.add(index, b.const_i32(i))], el, 1) else: - vec = b.zero_vec(out_dtype, vw) + vec = b.zero_vec(unit, vw) for i in range(vw): - el = _cast_element( - b, - b.vec_extract(fragment.value, register + i), - fragment.dtype, - out_dtype, - ) + el = _stored_unit(b, fragment, slot + i, out_dtype) vec = b.vec_insert(vec, el, i) b.smem_store_vN(ptr, [index], vec, vw) return - base_positions, base_index = _lane_base_index(b, window, fragment.tile_desc, thread) - for register in range(0, fragment.tile_desc.register_count, vw): - positions = _register_positions(b, fragment.tile_desc, base_positions, register) - address = _register_index(b, window, fragment.tile_desc, base_index, register) - element = b.vec_extract(fragment.value, register) - value = _cast_element(b, element, fragment.dtype, out_dtype) - mask = _clip_mask(b, window, positions, fragment.tile_desc.shape) + base_positions, base_index = _lane_base_index(b, window, tile_desc, thread) + base_index = _unit_index(b, base_index, out_dtype) + for register in range(0, tile_desc.register_count, step): + positions = _register_positions(b, tile_desc, base_positions, register) + address = _register_index(b, window, tile_desc, base_index, register, out_dtype) + slot = _unit_count(fragment.dtype, register) + units = [_stored_unit(b, fragment, slot + i, out_dtype) for i in range(vw)] + mask = _clip_mask(b, window, positions, tile_desc.shape) if mask is None: - b.global_store(ptr, address, value, align=align) + _global_store_units(b, ptr, address, units, align) else: with b.scf_if(mask): - b.global_store(ptr, address, value, align=align) + _global_store_units(b, ptr, address, units, align) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/fragments.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/fragments.py index a37f5c268ebe..1bd34316802d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/fragments.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/fragments.py @@ -12,17 +12,35 @@ field, and every write to it is checked against the layout and dtype. The IR verbs in :mod:`rocke.helpers.tiling.emit` fill / load / store the `Fragment`. -Terminology: a *register* here is one logical vector slot holding one element of the fragment's -`dtype` (``register_count`` slots per lane), not a physical 32-bit VGPR. LLVM packs the slots into -VGPRs: a gfx90a MFMA A operand of ``vec`` occupies two 32-bit registers, and a 64-bit -element occupies two. Sub-byte formats, which LLVM cannot pack, are not supported here. +Terminology: a *register* here is one logical element slot of the fragment's `dtype` +(``register_count`` slots per lane), not a physical 32-bit VGPR. The fragment's `dtype` is the +logical number format; the SSA `value` holds the registers in the fragment's *storage type* +(:attr:`Fragment.storage_type`), derived from the dtype and never supplied: + + dtype storage_type for n registers + f16, bf16, f32, f64, i8, fp8, ... vec (the logical vector) + tf32 vec (tf32 has no arithmetic) + u8 vec (filled and loaded: not yet) + fp4e2m1, fp6e2m3, fp6e3m2, i4 vec (packed 32-bit words) + +LLVM packs a byte-or-wider vector into VGPRs itself: a gfx90a MFMA A operand of ``vec`` +occupies two 32-bit registers, and a 64-bit element occupies two. A sub-byte element has no LLVM +value type, so its registers are packed low-bit-first into whole i32 words +(:func:`rocke.core.storage.register_packing`), and ``n * bits`` must be a multiple of 32. + +A u8 fragment stores like any byte dtype, but ``fill_fragment`` and ``load_fragment`` refuse it: they +start from a zero vector of the storage type, and core's ``zero_vec`` has no u8 form. An I8 fragment's +value, bitcast to ``vec``, holds the same bits. """ from __future__ import annotations from dataclasses import dataclass -from ...core.ir import Type, Value, VectorType +from ...core.dtypes import dtype_info +from ...core.ir import I32, Type, Value, VectorType +from ...core.storage import register_packing +from ..mma_io import storage_ir_type from .encoding import WarpDistributionEncoding from .register_mapper import RegisterMapper @@ -133,18 +151,56 @@ def _fragment_dtype_ok(dtype: object) -> tuple[bool, str]: ) +def _is_packed(dtype: Type) -> bool: + """True for a sub-byte `dtype` (fp4, fp6, i4): it has no LLVM value type, so its registers are + packed into i32 words. Byte-or-wider dtypes are LLVM vectors (the rule of + :func:`rocke.core.storage.register_packing`).""" + return dtype_info(dtype.name).encoded_bits < 8 + + +def _storage_element(dtype: Type) -> Type: + """The scalar type a `dtype` fragment's registers are held in: ``i32`` words for a packed dtype, + else the dtype's storage unit (:func:`rocke.helpers.mma_io.storage_ir_type`: ``i32`` for tf32, + the dtype itself for every arithmetic type).""" + return I32 if _is_packed(dtype) else storage_ir_type(dtype.name) + + +def _fragment_storage_type(dtype: Type, register_count: int) -> VectorType: + """The SSA vector type holding `register_count` logical elements of `dtype`: the + :func:`~rocke.core.storage.register_packing` words (``vec``) for a packed dtype, + which raises unless the elements fill whole 32-bit words, else + ``vec``.""" + packing = register_packing(dtype.name, register_count) + if packing is not None: + return VectorType(I32, packing.carrier_count) + return VectorType(storage_ir_type(dtype.name), register_count) + + +def _element_fix(subject: str, got: Type, dtype: Type, storage: Type) -> str: + """The fix for registers of element type `got` where the fragment holds `storage`: a logical + `dtype` value is reinterpreted to its storage type; any other element type is a different + fragment.""" + if got == dtype: + return ( + f"{dtype.name} registers are held as their storage type {storage.name}; " + "reinterpret them with b.bitcast" + ) + return f"{subject} {got.name}, not {storage.name}; build the Fragment at {got.name}" + + def _fragment_value_type_ok( tile_desc: TileDesc, dtype: Type, value: object ) -> tuple[bool, str]: """Whether `value` can hold the registers of a `tile_desc` fragment at element `dtype`: an IR - `Value` of type ``vec``. + `Value` of the fragment's storage type (:func:`_fragment_storage_type`). - One element per register slot -- packed sub-byte representations are not supported. A - single-register fragment is a ``vec`` vector, never a bare scalar. This checks the - register TYPE only: a value of the right type whose registers are in the wrong order, or an SSA - value used outside the scope that defines it, is not caught here. + A single-register fragment is a ``vec`` vector, never a bare scalar. A value of the + logical vector type of a dtype stored differently (``vec``), or anything but the words of a + packed dtype, is refused with the storage type to use. This checks the register TYPE only: a value + of the right type whose registers are in the wrong order, or an SSA value used outside the scope + that defines it, is not caught here. """ - expected = VectorType(dtype, tile_desc.register_count) + expected = _fragment_storage_type(dtype, tile_desc.register_count) where = f"for TileDesc(shape={tile_desc.shape})" if not isinstance(value, Value): return False, ( @@ -152,6 +208,22 @@ def _fragment_value_type_ok( f"{type(value).__name__} {value!r}" ) got = f"Fragment value expected {expected.name} {where}, got" + if value.type == expected: + return True, "ok" + if _is_packed(dtype): + words = expected.count + return False, ( + f"{got} {value.type.name} -- {dtype.name} registers are held packed, " + f"{tile_desc.register_count} elements of {dtype_info(dtype.name).encoded_bits} bits " + f"in {words} i32 {'word' if words == 1 else 'words'}; build them with load_fragment " + "or fill_fragment" + ) + if expected.elem != dtype and value.type == VectorType(dtype, expected.count): + return False, ( + f"{got} {value.type.name} -- {dtype.name} registers are held as their storage type " + f"{expected.elem.name}; reinterpret them with b.bitcast(value, " + f"VectorType({expected.elem.name.upper()}, {expected.count}))" + ) if not isinstance(value.type, VectorType): if type(value.type) is not Type: return False, ( @@ -159,20 +231,17 @@ def _fragment_value_type_ok( "handle; read memory into one with load_fragment" ) if tile_desc.register_count == 1: - fix = f"wrap the single register as a vec<{dtype.name}x1>" + fix = f"wrap the single register as a vec<{expected.elem.name}x1>" else: fix = ( f"this TileDesc lays out {tile_desc.register_count} registers per lane; pass " "them as one vector" ) - if value.type != dtype: - fix += ( - f"; and the register is {value.type.name}, not {dtype.name}; build the Fragment " - f"at {value.type.name}" + if value.type != expected.elem: + fix += "; and " + _element_fix( + "the register is", value.type, dtype, expected.elem ) return False, f"{got} scalar {value.type.name} -- {fix}" - if value.type == expected: - return True, "ok" fixes: list[str] = [] if value.type.count != expected.count: fixes.append( @@ -181,26 +250,26 @@ def _fragment_value_type_ok( "these registers, or, if they are one MMA atom's result, drive the whole tile " "through TileMma instead of raw b.mma" ) - if value.type.elem != dtype: + if value.type.elem != expected.elem: fixes.append( - f"the registers are {value.type.elem.name}, not {dtype.name}; build the Fragment at " - f"{value.type.elem.name}" + _element_fix("the registers are", value.type.elem, dtype, expected.elem) ) return False, f"{got} {value.type.name} -- " + "; and ".join(fixes) class Fragment: - """Per-lane register data for a tile: the `tile_desc` that lays it out, its element `dtype`, - and the SSA `value` holding the registers. + """Per-lane register data for a tile: the `tile_desc` that lays it out, its logical element + `dtype`, and the SSA `value` holding the registers in the dtype's :attr:`storage_type`. `tile_desc` and `dtype` are fixed at construction (read-only); `dtype` must be a scalar IR - `Type` (``F16``, ``F32``, ...) -- a string token or a vector / pointer / LDS type raises - ``ValueError``; so does a `tile_desc` that is not a :class:`TileDesc`. `value` is the one + `Type` (``F16``, ``F32``, ``TF32``, ``FP4E2M1``, ...) -- a string token or a vector / pointer / + LDS type raises ``ValueError``; so does a `tile_desc` that is not a :class:`TileDesc`, or a + packed dtype whose ``register_count`` elements do not fill whole 32-bit words. `value` is the one mutable field, because a kernel re-binds it every K step (``accumulator.value = b.mma(...)`` for a single-atom tile -- wave shape == atom shape, one ``b.mma``; a multi-atom tile is driven - through `TileMma`). Every write, construction included, must be a - ``VectorType(dtype, tile_desc.register_count)``; anything else raises ``ValueError`` naming the - expected and actual type. A too-short vector would otherwise still lower, and every register + through `TileMma`). Every write, construction included, must be of type :attr:`storage_type` + (``vec`` for every arithmetic dtype); anything else raises + ``ValueError`` naming the expected and actual type. A too-short vector would otherwise still lower, and every register slice past its end would extract out-of-range vector elements as LLVM poison -- a silent wrong answer; a too-long one would silently drop its extra registers. A wrong element type mislabels every register. @@ -211,7 +280,7 @@ class Fragment: filled `Fragment`. Two fragments compare equal only if they are the same object. """ - __slots__ = ("_tile_desc", "_dtype", "_value") + __slots__ = ("_tile_desc", "_dtype", "_storage_type", "_value") def __init__( self, tile_desc: TileDesc, dtype: Type, value: Value | None = None @@ -226,6 +295,7 @@ def __init__( raise ValueError(why) self._tile_desc = tile_desc self._dtype = dtype + self._storage_type = _fragment_storage_type(dtype, tile_desc.register_count) self._value: Value | None = None if value is not None: self.value = value @@ -237,12 +307,19 @@ def tile_desc(self) -> TileDesc: @property def dtype(self) -> Type: - """The element type; fixed at construction.""" + """The logical element type; fixed at construction.""" return self._dtype + @property + def storage_type(self) -> VectorType: + """The register vector type of :attr:`value`, derived from :attr:`dtype` and the tile: + ``vec`` for tf32, ``vec`` for a packed sub-byte dtype, + ``vec`` otherwise.""" + return self._storage_type + @property def value(self) -> Value: - """The SSA ``vec`` vector holding this lane's registers.""" + """The SSA :attr:`storage_type` vector holding this lane's registers.""" if self._value is None: raise ValueError( f"Fragment(shape={self._tile_desc.shape}, dtype={self._dtype.name}) not filled -- " @@ -253,9 +330,8 @@ def value(self) -> Value: @value.setter def value(self, value: Value) -> None: if value is None: - expected = VectorType(self._dtype, self._tile_desc.register_count) raise ValueError( - f"Fragment value cannot be None -- assign a {expected.name}, or build a new " + f"Fragment value cannot be None -- assign a {self._storage_type.name}, or build a new " "unfilled Fragment with make_fragment" ) ok, why = _fragment_value_type_ok(self._tile_desc, self._dtype, value) @@ -278,7 +354,7 @@ def __repr__(self) -> str: def make_fragment( tile_desc: TileDesc, dtype: Type, value: Value | None = None ) -> Fragment: - """Free factory: a `Fragment` for `tile_desc` at element `dtype`. Without `value` it is + """Free factory: a `Fragment` for `tile_desc` at logical element `dtype`. Without `value` it is unfilled until its first write (`fill_fragment`, or a checked ``value =`` write); a given `value` is checked like any other write (see `Fragment`). `load_fragment` builds and returns a filled fragment itself.""" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/lds_conflict.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/lds_conflict.py index 9dc4e3b71e9d..0db4bf43f5d6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/lds_conflict.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/lds_conflict.py @@ -496,12 +496,12 @@ def access_width(tile_desc, tensor, lds_swizzle): granularity exactly as :func:`~rocke.helpers.tiling.emit.load_fragment` and :func:`~rocke.helpers.tiling.emit.store_fragment` cap it.""" from rocke.helpers.tiling.descriptors import make_window - from rocke.helpers.tiling.emit import _align_of, _contiguous_run, _swizzle_vw + from rocke.helpers.tiling.emit import _contiguous_run, _swizzle_vw, _unit_bits window = make_window(tensor, (0,) * tensor.rank) vw = _contiguous_run(tile_desc, window, tensor.dtype, is_lds=True) if lds_swizzle: - vw = _swizzle_vw(lds_swizzle, vw, _align_of(tensor.dtype)) + vw = _swizzle_vw(lds_swizzle, vw, _unit_bits(tensor.dtype) // 8) return vw @@ -519,6 +519,7 @@ def addr_map(tile_desc, tensor, *, origin, n_lanes, lds_swizzle): The positions are carried rather than decoded from the index, which is only possible for one stride order. """ + from rocke.helpers.tiling.analysis.geometry import _elem_bytes from rocke.helpers.tiling.descriptors import make_window from rocke.helpers.tiling.emit import ( _check_builtin_swizzle, @@ -528,6 +529,8 @@ def addr_map(tile_desc, tensor, *, origin, n_lanes, lds_swizzle): _swizzle_lds_positions, ) + # The map is byte-addressed (the bank model), so a packed sub-byte dtype is refused. + _elem_bytes(tensor.dtype.name) if lds_swizzle is True: _check_builtin_swizzle(tensor) window = make_window(tensor, tuple(origin)) @@ -555,6 +558,7 @@ def addr_map(tile_desc, tensor, *, origin, n_lanes, lds_swizzle): reg0, lds_swizzle, None if lds_swizzle else base_index, + tensor.dtype, ) accesses.append( { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/memory.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/memory.py index 990896f9c819..cb17a62b9be7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/memory.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/memory.py @@ -44,8 +44,8 @@ def cooperative_load_width( ) -> int | None: """Widest GEOMETRY-PERMITTED per-lane free-run (in elements) for a cooperative load of this operand, or ``None`` if none works. This is the distribution ceiling only -- the dtype / 128-bit HARDWARE ceiling - is applied separately in ``emit``: ``_contiguous_run`` caps a vector-eligible dtype at 16 B (an f32 - 8-run -> 2x ``dwordx4``) and scalarizes an ineligible one (f64) outright, so a wide result here need not + is applied separately in ``emit``: ``_contiguous_run`` caps a vector-eligible storage unit at 16 B (an + f32 8-run -> 2x ``dwordx4``) and scalarizes an ineligible one (f64) outright, so a wide result here need not be a single global access. DERIVED, and PER OPERAND -- A and B have independent free extents so may want different widths. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/verb.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/verb.py index 47b862d32ed7..9e51fec522d1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/verb.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/verb.py @@ -12,6 +12,8 @@ from typing import Any +from ....core.storage import FragmentPacking, register_packing +from ..emit import _move_registers from ..fragments import Fragment, TileDesc from .observers import classify_transform @@ -22,9 +24,12 @@ def transform_fragment(b: Any, fragment: Fragment, target_desc: TileDesc) -> Fra The delta between the two layouts is solved from their encodings and classified: a `reorder` (same element set per lane, one lane-uniform register permutation) is emitted as a compile-time - register shuffle (`vec_extract`/`vec_insert`, element-granular so it is correct for any dtype - packing). A `cross_lane` delta (an element changes lanes) needs cross-lane movement (DPP / - ds_bpermute / LDS) and raises `NotImplementedError` -- the reserved seam a future cross-lane + register shuffle (`vec_extract`/`vec_insert`) of the fragment's storage registers + (:attr:`Fragment.storage_type`): the logical elements themselves for a byte-or-wider dtype (``i32`` + words for tf32), the 32-bit words for a packed sub-byte dtype. A packed reorder must move whole + word groups (:func:`_packed_word_moves`); one that moves elements within a word raises + `NotImplementedError`. A `cross_lane` delta (an element changes lanes) needs cross-lane movement + (DPP / ds_bpermute / LDS) and raises `NotImplementedError` -- the reserved seam a future cross-lane realisation fills. Same element set required (else `ValueError`). Dtype is carried through unchanged. """ plan = classify_transform(fragment.tile_desc.layout, target_desc.layout) @@ -32,8 +37,42 @@ def transform_fragment(b: Any, fragment: Fragment, target_desc: TileDesc) -> Fra raise NotImplementedError( f"cross-lane fragment transform is not supported yet -- {plan.reason}" ) - n = target_desc.register_count - out = b.zero_vec(fragment.dtype, n) - for src_reg, dst_reg in enumerate(plan.permutation): - out = b.vec_insert(out, b.vec_extract(fragment.value, src_reg), dst_reg) + packing = register_packing(fragment.dtype.name, target_desc.register_count) + if packing is None: + moves = list(enumerate(plan.permutation)) + else: + moves = _packed_word_moves(fragment.dtype.name, packing, plan.permutation) + out = _move_registers(b, fragment.value, moves, fragment.storage_type.count) return Fragment(target_desc, fragment.dtype, out) + + +def _packed_word_moves( + dtype: str, packing: FragmentPacking, permutation: tuple[int, ...] +) -> list[tuple[int, int]]: + """The (source word, destination word) moves that realise `permutation` (source register -> + destination register) on the 32-bit words of a `dtype` fragment packed as `packing`. + + Elements pack low-bit-first into words, so the smallest unit that moves without repacking bits is + a word group: the fewest elements that fill whole words (8 elements in one word at 4 bits, 16 in + three words at 6 bits). The permutation must carry every word group, in order, onto a word group; + otherwise it raises `NotImplementedError`, since repacking the bits of a packed fragment is not + implemented. + """ + group, words = packing.packing.group(packing.carrier_bits) + moves = [] + for source in range(0, len(permutation), group): + destination = permutation[source] + run = tuple(permutation[source : source + group]) + if destination % group or run != tuple(range(destination, destination + group)): + raise NotImplementedError( + f"transform_fragment cannot move {dtype} registers {source}..{source + group - 1} " + f"to {list(run)}: {dtype} registers are held packed, so a reorder must move each " + f"run of {group} registers ({words} i32 {'word' if words == 1 else 'words'}) whole " + "and in order onto another such run -- TODO: not implemented yet; repacking the " + "bits of a packed fragment" + ) + moves.extend( + (source // group * words + w, destination // group * words + w) + for w in range(words) + ) + return moves diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment_storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment_storage.py new file mode 100644 index 000000000000..6b8fb429f971 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment_storage.py @@ -0,0 +1,746 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""A Fragment holds its registers in the storage type of its dtype. + +The fragment's ``dtype`` is the logical number format an author names; its ``value`` holds the +registers in ``storage_type``, derived from the dtype and never supplied: + +- every arithmetic dtype (f16, bf16, f32, f64, i8, i16, i32, fp8, bf8) is its own logical vector; +- tf32 has no arithmetic, so its registers are held as i32; +- a sub-byte dtype (fp4, fp6, i4) has no value type, so its registers are packed into whole i32 + words, low bit first. + +The memory verbs move storage units: a packed fragment moves whole 32-bit words, so a layout or +window that would split a word is refused, and its LDS buffer holds i32 words. These tests check +the storage table, the refusals and their messages on the CPU, and on a gfx90a GPU that tf32 +payloads and packed operand tiles (in gfx950's canonical operand layouts) round-trip bit-exactly +through global memory and LDS, and that the packed words a lane loads are exactly +``FragmentPacking.pack`` of the elements its layout places there. +""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +pytest.importorskip("rocke.core.ir", reason="rocke IR substrate required") + +from rocke.core.ir import ( + BF8E5M2, + BF16, + F16, + F32, + F64, + FP4E2M1, + FP6E2M3, + FP6E3M2, + FP8E4M3, + I4, + I8, + I16, + I32, + TF32, + U8, + IRBuilder, + PtrType, + VectorType, +) +from rocke.core.storage import BitPacking, register_packing +from rocke.helpers.tiling import ( + Fragment, + TileMma, + fill_fragment, + load_fragment, + make_fragment, + make_tensor_desc, + make_tile_desc, + make_window, + store_fragment, + transform_fragment, +) +from rocke.helpers.tiling.analysis.geometry import _elem_bits, _elem_bytes +from rocke.helpers.tiling.emit import emit_tensor_coordinates +from rocke.helpers.tiling.lds_conflict import NumBuilder + +_WAVE = 64 + + +def _flat_tile(registers: int): + """A 1-D tile whose 64 lanes each hold `registers` consecutive elements.""" + return make_tile_desc( + shape=[_WAVE * registers], + thread_tile=[registers], + thread_dist=[_WAVE], + wave_size=_WAVE, + ) + + +def _operand_desc(dtype_name: str): + """The canonical A-operand layout of gfx950's 16x16x128 f8f6f4 instruction for a 4-bit or 6-bit + operand: each lane holds one run of 32 K elements.""" + return TileMma( + (16, 16, 128), a=dtype_name, b=dtype_name, c="f32", target="gfx950" + ).a_desc + + +# ---- the storage type --------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "dtype,registers,storage", + [ + (F16, 8, "vec"), + (BF16, 8, "vec"), + (F32, 4, "vec"), + (F64, 2, "vec"), + (I8, 16, "vec"), + (U8, 16, "vec"), + (I16, 8, "vec"), + (I32, 4, "vec"), + (FP8E4M3, 16, "vec"), + (BF8E5M2, 16, "vec"), + (TF32, 8, "vec"), + (FP4E2M1, 32, "vec"), + (I4, 32, "vec"), + (FP6E2M3, 32, "vec"), + (FP6E3M2, 16, "vec"), + ], +) +def test_storage_type_is_derived_from_the_dtype(dtype, registers, storage) -> None: + fragment = make_fragment(_flat_tile(registers), dtype) + assert fragment.dtype == dtype + assert fragment.tile_desc.register_count == registers + assert fragment.storage_type.name == storage + + +def test_a_packed_fragment_must_fill_whole_words() -> None: + # 4 fp4 elements are 16 bits: half a register, which no lane can own alone. + with pytest.raises(ValueError, match="not a whole number of 32-bit registers"): + make_fragment(_flat_tile(4), FP4E2M1) + + +def test_a_tf32_fragment_takes_its_i32_registers() -> None: + b = IRBuilder("tf32_value") + value = b.zero_vec(I32, 8) + fragment = Fragment(_flat_tile(8), TF32, value) + assert fragment.value is value + + +def test_a_tf32_value_must_be_reinterpreted_to_its_storage_type() -> None: + tile = make_tile_desc( + shape=[64, 8], thread_tile=[1, 8], thread_dist=[64, 1], wave_size=_WAVE + ) + b = IRBuilder("tf32_logical_value") + ptr = b.param("IN", PtrType(TF32, "global")) + logical = b.global_load_vN(ptr, b.const_i32(0), TF32, 8) + expected = ( + "Fragment value expected vec for TileDesc(shape=(64, 8)), got vec -- " + "tf32 registers are held as their storage type i32; reinterpret them with " + "b.bitcast(value, VectorType(I32, 8))" + ) + with pytest.raises(ValueError) as raised: + Fragment(tile, TF32, logical) + assert str(raised.value) == expected + # The fix the message names is accepted. + Fragment(tile, TF32, b.bitcast(logical, VectorType(I32, 8))) + + +def test_a_packed_value_must_be_its_words() -> None: + a_desc = _operand_desc("fp4e2m1") + b = IRBuilder("fp4_bytes_value") + expected = ( + "Fragment value expected vec for TileDesc(shape=(16, 128)), got vec -- " + "fp4e2m1 registers are held packed, 32 elements of 4 bits in 4 i32 words; build them " + "with load_fragment or fill_fragment" + ) + with pytest.raises(ValueError) as raised: + Fragment(a_desc, FP4E2M1, b.zero_vec(I8, 32)) + assert str(raised.value) == expected + + +@pytest.mark.parametrize("dtype,words", [(FP4E2M1, 4), (FP6E2M3, 6), (TF32, 32)]) +def test_a_zero_fill_is_all_zero_storage_words(dtype, words) -> None: + b = IRBuilder("zero_fill") + fragment = make_fragment(_operand_desc("fp4e2m1"), dtype) + fill_fragment(b, fragment, 0) + assert fragment.value.type == VectorType(I32, words) + (fill,) = [op for op in b.kernel.body.ops if op.name == "arith.constant_vec"] + assert fill.attrs["elem"] == "i32" and fill.attrs["fill"] == 0.0 + + +# ---- u8: stored, but neither filled nor loaded ----------------------------------------------------------- + +_U8_TILE = make_tile_desc( + shape=[64, 16], thread_tile=[1, 16], thread_dist=[64, 1], wave_size=_WAVE +) +_U8_TENSOR = make_tensor_desc((64, 16), (16, 1), U8) + + +def _u8_refusal(verb: str) -> str: + return ( + f"{verb} cannot build a u8 fragment: core's zero_vec has no u8 form -- TODO: not " + f"implemented yet; build an I8 fragment with {verb} and b.bitcast its value to " + f"vec" + ) + + +@pytest.mark.parametrize("source", ["global", "lds"]) +def test_a_u8_fragment_is_refused_by_load_fragment(source) -> None: + b = IRBuilder("u8_load") + tid = b.thread_id_x() + if source == "global": + src = b.param("IN", PtrType(U8, "global")) + else: + src = b.smem_alloc(U8, [_U8_TENSOR.element_span], name_hint="lds") + with pytest.raises(NotImplementedError) as raised: + load_fragment(b, src, make_window(_U8_TENSOR, (0, 0)), _U8_TILE, tid) + assert str(raised.value) == _u8_refusal("load_fragment") + + +def test_a_u8_fragment_is_refused_by_fill_fragment() -> None: + b = IRBuilder("u8_fill") + with pytest.raises(NotImplementedError) as raised: + fill_fragment(b, make_fragment(_U8_TILE, U8), 0) + assert str(raised.value) == _u8_refusal("fill_fragment") + + +@pytest.mark.parametrize("dest", ["global", "lds"]) +def test_the_named_u8_route_builds_a_fragment_that_stores(dest) -> None: + # An I8 fragment's value, bitcast to vec, is a u8 fragment with the same bits; it stores + # to global memory and to LDS. + b = IRBuilder("u8_route") + tid = b.thread_id_x() + src = b.param("IN", PtrType(I8, "global")) + i8_tensor = make_tensor_desc((64, 16), (16, 1), I8) + loaded = load_fragment(b, src, make_window(i8_tensor, (0, 0)), _U8_TILE, tid) + fragment = Fragment(_U8_TILE, U8, b.bitcast(loaded.value, VectorType(U8, 16))) + if dest == "global": + dst = b.param("OUT", PtrType(U8, "global")) + else: + dst = b.smem_alloc(U8, [_U8_TENSOR.element_span], name_hint="lds") + store_fragment(b, dst, make_window(_U8_TENSOR, (0, 0)), fragment, tid) + stores = [op for op in b.kernel.body.ops if "store" in op.name] + assert len(stores) == 16 + + +# ---- refusals: a packed access moves whole words -------------------------------------------------------- + + +def _load(td, tile, *, origin=(0, 0), bounds=None, lds_elem=None): + b = IRBuilder("packed_load") + tid = b.thread_id_x() + if lds_elem is None: + src = b.param("IN", PtrType(I8, "global")) + else: + src = b.smem_alloc(lds_elem, [td.element_span], name_hint="lds") + return load_fragment(b, src, make_window(td, origin, bounds), tile, tid) + + +def test_a_k_major_strided_packed_tensor_is_refused() -> None: + # The operand layout's K runs lie along axis 1, but this tensor is M-contiguous, so each lane's + # registers would be single 4-bit elements 16 apart. + td = make_tensor_desc((16, 128), (1, 16), FP4E2M1) + with pytest.raises(ValueError) as raised: + _load(td, _operand_desc("fp4e2m1")) + assert str(raised.value) == ( + "load_fragment of fp4e2m1 needs each lane's registers to form runs of whole 32-bit " + "words along the tensor's stride-1 axis (multiples of 8 elements, starting on an " + "8-element boundary); this layout gives runs of 1 element -- choose a layout whose " + "innermost register run lies along the stride-1 axis" + ) + + +def test_runs_shorter_than_a_word_group_are_refused() -> None: + # Runs of 4 fp6 elements are 24 bits; whole words start every 16 fp6 elements. + tile = make_tile_desc( + shape=[32, 32], thread_tile=[4, 4], thread_dist=[8, 8], wave_size=_WAVE + ) + td = make_tensor_desc((32, 32), (32, 1), FP6E2M3) + with pytest.raises(ValueError) as raised: + _load(td, tile) + assert str(raised.value) == ( + "load_fragment of fp6e2m3 needs each lane's registers to form runs of whole 32-bit " + "words along the tensor's stride-1 axis (multiples of 16 elements, starting on a " + "16-element boundary); this layout gives runs of 4 elements -- choose a layout whose " + "innermost register run lies along the stride-1 axis" + ) + + +def test_a_row_stride_that_splits_a_word_is_refused() -> None: + td = make_tensor_desc((16, 128), (132, 1), FP4E2M1) + with pytest.raises(ValueError, match="tensor axis 0 steps by 132 elements"): + _load(td, _operand_desc("fp4e2m1")) + + +def test_an_origin_inside_a_word_is_refused() -> None: + td = make_tensor_desc((16, 256), (256, 1), FP4E2M1) + with pytest.raises(ValueError, match="the window origin on the stride-1 axis is 4"): + _load(td, _operand_desc("fp4e2m1"), origin=(0, 4)) + + +@pytest.mark.parametrize("verb", ["load", "store"]) +def test_a_clipped_packed_window_needs_a_whole_word_bound(verb) -> None: + a_desc = _operand_desc("fp4e2m1") + td = make_tensor_desc((16, 128), (128, 1), FP4E2M1) + window = make_window(td, (0, 0), (None, 100)) + b = IRBuilder("packed_clip") + tid = b.thread_id_x() + ptr = b.param("P", PtrType(I8, "global")) + with pytest.raises(ValueError) as raised: + if verb == "load": + load_fragment(b, ptr, window, a_desc, tid) + else: + fragment = make_fragment(a_desc, FP4E2M1) + fill_fragment(b, fragment, 0) + store_fragment(b, ptr, window, fragment, tid) + assert str(raised.value) == ( + f"{verb}_fragment of a clipped fp4e2m1 window needs its bound along the stride-1 axis " + "to be a multiple of 8 elements (whole 32-bit words) -- got 100" + ) + + +def test_a_packed_lds_buffer_must_hold_words() -> None: + td = make_tensor_desc((16, 128), (128, 1), FP4E2M1) + with pytest.raises(ValueError) as raised: + _load(td, _operand_desc("fp4e2m1"), lds_elem=F16) + assert str(raised.value) == ( + "LDS buffer for fp4e2m1 fragments must hold its storage words -- allocate " + "smem_alloc(I32, [words]) sized from the descriptor; got smem" + ) + + +def test_a_packed_lds_buffer_is_sized_in_words() -> None: + # 16 x 128 fp4 elements are 256 words: a buffer of 255 words is refused, one of 256 accepted. + td = make_tensor_desc((16, 128), (128, 1), FP4E2M1) + a_desc = _operand_desc("fp4e2m1") + for words, ok in ((255, False), (256, True)): + b = IRBuilder("packed_lds_size") + tid = b.thread_id_x() + lds = b.smem_alloc(I32, [words], name_hint="lds") + if ok: + load_fragment(b, lds, make_window(td, (0, 0)), a_desc, tid) + else: + with pytest.raises( + ValueError, match="256 i32 words, but the buffer holds 255" + ): + load_fragment(b, lds, make_window(td, (0, 0)), a_desc, tid) + + +class _I32Arithmetic: + """Evaluates the index arithmetic the memory verbs emit the way the GPU computes it: in + two's-complement i32, so a product past the i32 range wraps instead of growing.""" + + @staticmethod + def _wrap(value: int) -> int: + return (value + (1 << 31)) % (1 << 32) - (1 << 31) + + def const_i32(self, value: int) -> int: + return self._wrap(value) + + def mul(self, a: int, b: int) -> int: + return self._wrap(a * b) + + def div(self, a: int, b: int) -> int: + return self._wrap(int(a / b)) + + +@pytest.mark.parametrize("dtype,bits", [(FP6E2M3, 6), (FP4E2M1, 4), (I4, 4)]) +def test_a_packed_word_index_is_exact_for_every_i32_element_index(dtype, bits) -> None: + # The word an access starts at, for the largest whole-group element index i32 can hold: the + # element index is in range, so its word index must be too. A 6-bit index multiplied by 3 + # before the division by 16 would wrap here (the product is about 6.4e9). + from rocke.helpers.tiling.emit import _unit_index + + group = 32 // math.gcd(bits, 32) + element_index = (2**31 - 1) // group * group + assert ( + _unit_index(_I32Arithmetic(), element_index, dtype) + == element_index * bits // 32 + ) + + +# ---- the view widths ---------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "name,bits", [("f16", 16), ("tf32", 32), ("fp8e4m3", 8), ("fp4e2m1", 4)] +) +def test_elem_bits(name, bits) -> None: + assert _elem_bits(name) == bits + + +def test_elem_bytes() -> None: + assert _elem_bytes("f64") == 8 + assert _elem_bytes("f16") == 2 + + +def test_byte_views_refuse_a_sub_byte_dtype() -> None: + with pytest.raises(NotImplementedError) as raised: + _elem_bytes("fp4e2m1") + assert str(raised.value) == ( + "fp4e2m1 elements are 4 bits, not a whole number of bytes -- byte-addressed views " + "(LDS bank model) of packed data -- TODO: not implemented yet" + ) + + +# ---- GPU (gfx90a) --------------------------------------------------------------------------------------- + + +def _gfx90a_present() -> bool: + try: + from rocke.runtime.hip_module import get_device_arch + + return get_device_arch(0) == "gfx90a" + except Exception: # noqa: BLE001 + # Any HIP or driver failure means there is no usable GPU here. + return False + + +_needs_gfx90a = pytest.mark.skipif(not _gfx90a_present(), reason="needs a gfx90a GPU") + + +def _run(kernel, signature, inputs: dict, outputs: dict) -> None: + """Compile `kernel` for gfx90a, launch one wave on host arrays and copy `outputs` back.""" + from rocke.helpers.compile import compile_kernel + from rocke.runtime.hip_module import Runtime + from rocke.runtime.host_buffers import as_u8_buffer + from rocke.runtime.launcher import ( + DeviceMem, + KernelLauncher, + LaunchConfig, + synchronize_and_release, + ) + + art = compile_kernel(kernel, arch="gfx90a") + launcher = KernelLauncher( + hsaco=art.hsaco, kernel_name=art.kernel_name, signature=signature + ) + rt = Runtime() + arrays = {**inputs, **outputs} + dev = {name: DeviceMem(a.nbytes) for name, a in arrays.items()} + for name, a in arrays.items(): + rt.memcpy_h2d(dev[name].ptr(), as_u8_buffer(a), a.nbytes) + launcher(dev, config=LaunchConfig(grid=(1, 1, 1), block=(_WAVE, 1, 1))) + synchronize_and_release() + for name, a in outputs.items(): + rt.memcpy_d2h(as_u8_buffer(a), dev[name].ptr(), a.nbytes) + + +def _param(b, name, elem, **flags): + return b.param(name, PtrType(elem, "global"), noalias=True, align=16, **flags) + + +@_needs_gfx90a +def test_tf32_payload_round_trips_through_lds() -> None: + # A tf32 tile is loaded from global memory with 4-element runs, stored into an LDS buffer that + # holds it transposed, read back with runs along the other axis and stored to OUT. tf32 has no + # arithmetic, so every 32-bit pattern, NaN payloads included, must come back unchanged. + from rocke.helpers.spec import SignatureBuilder + + rows, cols = 16, 32 + global_td = make_tensor_desc((rows, cols), (cols, 1), TF32) + lds_td = make_tensor_desc((rows, cols), (1, rows), TF32) + load_tile = make_tile_desc( + shape=[rows, cols], thread_tile=[1, 8], thread_dist=[16, 4], wave_size=_WAVE + ) + read_tile = make_tile_desc( + shape=[rows, cols], thread_tile=[8, 1], thread_dist=[2, 32], wave_size=_WAVE + ) + b = IRBuilder("tf32_round_trip") + b.kernel.attrs["max_workgroup_size"] = _WAVE + src = _param(b, "IN", TF32, readonly=True) + out = _param(b, "OUT", TF32, writeonly=True) + tid = b.thread_id_x() + lds = b.smem_alloc(TF32, [lds_td.element_span], name_hint="lds") + window = make_window(global_td, (0, 0)) + fragment = load_fragment(b, src, window, load_tile, tid) + assert fragment.storage_type == VectorType(I32, 8) + store_fragment(b, lds, make_window(lds_td, (0, 0)), fragment, tid) + b.sync() + staged = load_fragment(b, lds, make_window(lds_td, (0, 0)), read_tile, tid) + store_fragment(b, out, window, staged, tid) + b.ret() + signature = SignatureBuilder().ptr("IN", "tf32").ptr("OUT", "tf32").build() + rng = np.random.default_rng(0) + data = rng.integers(0, 2**32, size=(rows, cols), dtype=np.uint64).astype(np.uint32) + result = np.zeros_like(data) + _run(b.kernel, signature, {"IN": data}, {"OUT": result}) + np.testing.assert_array_equal(result, data) + + +def _packed_bytes(patterns: np.ndarray, bits: int) -> np.ndarray: + """Row-major `patterns` packed low-bit-first at `bits` per element.""" + packed = BitPacking(bits).pack([int(p) for p in patterns.reshape(-1)]) + return np.frombuffer(packed, dtype=np.uint8).copy() + + +def _expected_words(tile, patterns: np.ndarray, dtype_name: str) -> np.ndarray: + """Per lane, the words ``FragmentPacking.pack`` makes of the elements the layout gives it.""" + packing = register_packing(dtype_name, tile.register_count) + words = [] + for lane in range(_WAVE): + elements = [ + int( + patterns[ + tuple( + emit_tensor_coordinates(NumBuilder(lane), tile.layout, lane, r) + ) + ] + ) + for r in range(tile.register_count) + ] + words.append(packing.pack(elements)) + return np.array(words, dtype=np.uint32) + + +_PACKED = [("fp4e2m1", FP4E2M1, 4), ("fp6e2m3", FP6E2M3, 6), ("i4", I4, 4)] + + +@_needs_gfx90a +@pytest.mark.parametrize("name,dtype,bits", _PACKED, ids=[p[0] for p in _PACKED]) +def test_packed_operand_tile_round_trips(name, dtype, bits) -> None: + # An A-operand tile (16 x 128, K contiguous) in gfx950's canonical f8f6f4 layout is loaded from a + # packed tensor. WORDS receives every lane's registers, which must be the packing of the + # elements the layout places there. The fragment is stored back to OUT directly, and to OUT_LDS + # through an LDS buffer of i32 words; both must equal the input bytes. + from rocke.helpers.spec import SignatureBuilder + + a_desc = _operand_desc("fp4e2m1" if bits == 4 else name) + rows, cols = a_desc.shape + td = make_tensor_desc((rows, cols), (cols, 1), dtype) + words = register_packing(name, a_desc.register_count).carrier_count + lds_words = rows * cols * bits // 32 + + b = IRBuilder(f"packed_round_trip_{name}") + b.kernel.attrs["max_workgroup_size"] = _WAVE + src = _param(b, "IN", I8, readonly=True) + out = _param(b, "OUT", I8, writeonly=True) + out_lds = _param(b, "OUT_LDS", I8, writeonly=True) + dump = _param(b, "WORDS", I32, writeonly=True) + tid = b.thread_id_x() + lds = b.smem_alloc(I32, [lds_words], name_hint="lds") + window = make_window(td, (0, 0)) + fragment = load_fragment(b, src, window, a_desc, tid) + assert fragment.storage_type == VectorType(I32, words) + lane_base = b.mul(tid, b.const_i32(words)) + for j in range(words): + b.global_store( + dump, + b.add(lane_base, b.const_i32(j)), + b.vec_extract(fragment.value, j), + align=4, + ) + store_fragment(b, out, window, fragment, tid) + store_fragment(b, lds, window, fragment, tid) + b.sync() + staged = load_fragment(b, lds, window, a_desc, tid) + store_fragment(b, out_lds, window, staged, tid) + b.ret() + signature = ( + SignatureBuilder() + .ptr("IN", "i8") + .ptr("OUT", "i8") + .ptr("OUT_LDS", "i8") + .ptr("WORDS", "i32") + .build() + ) + + rng = np.random.default_rng(bits) + patterns = rng.integers(0, 1 << bits, size=(rows, cols)) + data = _packed_bytes(patterns, bits) + result, result_lds = np.zeros_like(data), np.zeros_like(data) + lane_words = np.zeros((_WAVE, words), np.uint32) + _run( + b.kernel, + signature, + {"IN": data}, + {"OUT": result, "OUT_LDS": result_lds, "WORDS": lane_words}, + ) + np.testing.assert_array_equal(lane_words, _expected_words(a_desc, patterns, name)) + np.testing.assert_array_equal(result, data) + np.testing.assert_array_equal(result_lds, data) + + +@_needs_gfx90a +@pytest.mark.parametrize( + "name,dtype,bits", _PACKED[:2], ids=[p[0] for p in _PACKED[:2]] +) +def test_clipped_packed_window_masks_whole_words(name, dtype, bits) -> None: + # The window is bounded at K = 64. Its load reads zero words past the bound (OUT_FULL stores the + # whole tile unclipped), and its store drops the words past the bound (OUT keeps its 0xff fill + # there); before the bound both equal the input. + from rocke.helpers.spec import SignatureBuilder + + a_desc = _operand_desc(name) + rows, cols = a_desc.shape + bound = 64 + td = make_tensor_desc((rows, cols), (cols, 1), dtype) + b = IRBuilder(f"packed_clip_{name}") + b.kernel.attrs["max_workgroup_size"] = _WAVE + src = _param(b, "IN", I8, readonly=True) + out_full = _param(b, "OUT_FULL", I8, writeonly=True) + out = _param(b, "OUT", I8, writeonly=True) + tid = b.thread_id_x() + clipped = make_window(td, (0, 0), (None, bound)) + fragment = load_fragment(b, src, clipped, a_desc, tid) + store_fragment(b, out_full, make_window(td, (0, 0)), fragment, tid) + whole = load_fragment(b, src, make_window(td, (0, 0)), a_desc, tid) + store_fragment(b, out, clipped, whole, tid) + b.ret() + signature = ( + SignatureBuilder().ptr("IN", "i8").ptr("OUT_FULL", "i8").ptr("OUT", "i8") + ) + signature = signature.build() + + rng = np.random.default_rng(bits + 1) + patterns = rng.integers(1, 1 << bits, size=(rows, cols)) + data = _packed_bytes(patterns, bits) + result_full = np.full_like(data, 0xFF) + result = np.full_like(data, 0xFF) + _run(b.kernel, signature, {"IN": data}, {"OUT_FULL": result_full, "OUT": result}) + row_bytes = cols * bits // 8 + kept = bound * bits // 8 + expected = data.reshape(rows, row_bytes) + result_full = result_full.reshape(rows, row_bytes) + result = result.reshape(rows, row_bytes) + np.testing.assert_array_equal(result_full[:, :kept], expected[:, :kept]) + np.testing.assert_array_equal(result_full[:, kept:], 0) + np.testing.assert_array_equal(result[:, :kept], expected[:, :kept]) + np.testing.assert_array_equal(result[:, kept:], 0xFF) + + +# ---- transform_fragment moves storage registers ------------------------------------------------ + +# A (2, 2, 512) tile whose 64 lanes each hold 2 x 2 x 8 registers: register 16 * i + 8 * j + r holds +# element (i, j, 8 * lane + r). Making j the most significant register axis instead moves registers +# 8-15 to 16-23 and back, whole runs of 8; making the run axis r the middle one splits every run. +_REORDER_TILE = make_tile_desc( + shape=[2, 2, 8 * _WAVE], + thread_tile=[2, 2, 8], + thread_dist=[1, 1, _WAVE], + wave_size=_WAVE, +) +_RUNS_SWAPPED = _REORDER_TILE.reorder_registers((1, 0, 2)) +_RUNS_SPLIT = _REORDER_TILE.reorder_registers((0, 2, 1)) + + +def _filled(b, dtype): + """A zero-filled `_REORDER_TILE` fragment of `dtype`; a u8 one is the i8 zero reinterpreted.""" + if dtype == U8: + zero = b.bitcast(b.zero_vec(I8, 32), VectorType(U8, 32)) + return Fragment(_REORDER_TILE, U8, zero) + fragment = make_fragment(_REORDER_TILE, dtype) + fill_fragment(b, fragment, 0) + return fragment + + +@pytest.mark.parametrize( + "dtype", [TF32, U8, FP4E2M1, FP6E2M3], ids=["tf32", "u8", "fp4e2m1", "fp6e2m3"] +) +def test_transform_fragment_keeps_the_storage_type(dtype) -> None: + # The identity transform and, for every dtype whose runs of registers stay whole, the swap of + # runs return a fragment of the source's storage type: i32 for tf32, u8 for u8, packed words for + # the sub-byte dtypes. + b = IRBuilder("transform_storage") + fragment = _filled(b, dtype) + targets = [_REORDER_TILE] + ([_RUNS_SWAPPED] if dtype != FP6E2M3 else []) + for target in targets: + moved = transform_fragment(b, fragment, target) + assert (moved.dtype, moved.storage_type) == (dtype, fragment.storage_type) + assert moved.value.type == fragment.storage_type + + +@pytest.mark.parametrize( + "dtype,target,message", + [ + ( + FP4E2M1, + _RUNS_SPLIT, + ( + "transform_fragment cannot move fp4e2m1 registers 0..7 to [0, 2, 4, 6, 8, 10, " + "12, 14]: fp4e2m1 registers are held packed, so a reorder must move each run of " + "8 registers (1 i32 word) whole and in order onto another such run -- TODO: not " + "implemented yet; repacking the bits of a packed fragment" + ), + ), + ( + FP6E2M3, + _RUNS_SWAPPED, + ( + "transform_fragment cannot move fp6e2m3 registers 0..15 to [0, 1, 2, 3, 4, 5, 6, " + "7, 16, 17, 18, 19, 20, 21, 22, 23]: fp6e2m3 registers are held packed, so a " + "reorder must move each run of 16 registers (3 i32 words) whole and in order onto " + "another such run -- TODO: not implemented yet; repacking the bits of a packed " + "fragment" + ), + ), + ], + ids=["fp4-split-runs", "fp6-half-runs"], +) +def test_a_packed_reorder_inside_a_word_group_is_refused( + dtype, target, message +) -> None: + b = IRBuilder("transform_packed_refusal") + with pytest.raises(NotImplementedError) as excinfo: + transform_fragment(b, _filled(b, dtype), target) + assert str(excinfo.value) == message + + +_TRANSFORMED = [ + ("tf32", TF32, "tf32"), + ("u8", U8, "i8"), + ("fp4e2m1", FP4E2M1, "i8"), +] + + +@_needs_gfx90a +@pytest.mark.parametrize( + "name,dtype,carrier", _TRANSFORMED, ids=[t[0] for t in _TRANSFORMED] +) +def test_transform_fragment_round_trips_on_the_gpu(name, dtype, carrier) -> None: + # The tile is loaded in `_REORDER_TILE`'s layout, transformed to `_RUNS_SWAPPED` (registers + # 8-15 and 16-23 trade places) and stored in that layout: OUT equals IN only if every storage + # register moved with its elements. A u8 tile is loaded as i8 and reinterpreted. + from rocke.helpers.spec import SignatureBuilder + + shape = tuple(_REORDER_TILE.shape) + load_dtype = I8 if dtype == U8 else dtype + load_td = make_tensor_desc(shape, (2 * 8 * _WAVE, 8 * _WAVE, 1), load_dtype) + store_td = make_tensor_desc(shape, (2 * 8 * _WAVE, 8 * _WAVE, 1), dtype) + elem = {"tf32": TF32, "i8": I8}[carrier] + b = IRBuilder(f"transform_round_trip_{name}") + b.kernel.attrs["max_workgroup_size"] = _WAVE + src = _param(b, "IN", elem, readonly=True) + out = _param(b, "OUT", U8 if dtype == U8 else elem, writeonly=True) + tid = b.thread_id_x() + fragment = load_fragment( + b, src, make_window(load_td, (0, 0, 0)), _REORDER_TILE, tid + ) + if dtype == U8: + fragment = Fragment( + _REORDER_TILE, U8, b.bitcast(fragment.value, VectorType(U8, 32)) + ) + moved = transform_fragment(b, fragment, _RUNS_SWAPPED) + store_fragment(b, out, make_window(store_td, (0, 0, 0)), moved, tid) + b.ret() + signature = ( + SignatureBuilder() + .ptr("IN", carrier) + .ptr("OUT", "u8" if dtype == U8 else carrier) + .build() + ) + + rng = np.random.default_rng(7) + count = 2 * 2 * 8 * _WAVE + if dtype == TF32: + data = rng.integers(0, 2**32, size=count, dtype=np.uint64).astype(np.uint32) + elif dtype == U8: + data = rng.integers(0, 256, size=count, dtype=np.uint16).astype(np.uint8) + else: + data = _packed_bytes(rng.integers(0, 16, size=count), 4) + result = np.zeros_like(data) + _run(b.kernel, signature, {"IN": data}, {"OUT": result}) + np.testing.assert_array_equal(result, data) From b7dbdb4cc3c09fc6c7cd5613153b14733cae8fbe Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Tue, 27 Oct 2026 19:39:45 -0400 Subject: [PATCH 2/7] feat(rocke): take IR dtypes for A, B, C and the result in TileMma 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 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 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. 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, that gfx950 has no tf32 instruction, the fp8 and bf8 wave tiles lowering to one MFMA per atom on gfx942, and the tee's widths and labels. Co-Authored-By: Claude --- .../tiling/.claude/skills/layout-viz/SKILL.md | 2 +- .../python/rocke/helpers/tiling/README.md | 3 +- .../PROPOSAL_c_transpose_crossed_ab.md | 4 +- .../helpers/tiling/docs/tiling_api_surface.md | 46 +++--- .../tiling/docs/visualization_api_surface.md | 4 +- .../crc_interleaved_gemm.py | 12 +- .../kernels/tiling_gemm_interleaved_demo.py | 6 +- .../python/rocke/helpers/tiling/mma/driver.py | 7 +- .../rocke/helpers/tiling/mma/mma_operation.py | 27 +++- .../python/rocke/helpers/tiling/mma/plan.py | 114 +++++++++++--- .../rocke/helpers/tiling/traits/README.md | 5 +- .../rocke/helpers/tiling/traits/mma_traits.py | 79 +++++++--- .../tiling/visualization/layout_render.py | 21 +-- .../helpers/tiling/mma/test_layout_oracle.py | 5 +- .../helpers/tiling/mma/test_layout_styles.py | 15 +- .../helpers/tiling/mma/test_mma_dtypes.py | 106 +++++++++++++ .../helpers/tiling/mma/test_mma_operation.py | 147 ++++++++++++++---- .../tiling/mma/test_mma_validation_recipe.py | 18 +-- .../helpers/tiling/test_coalescing_probe.py | 6 +- .../helpers/tiling/test_emit_wide_access.py | 6 +- .../tests/helpers/tiling/test_fragment.py | 6 +- .../helpers/tiling/test_fragment_storage.py | 26 ++-- .../helpers/tiling/test_layout_optimizer.py | 7 +- .../tiling/test_lds_descriptor_addressing.py | 12 +- .../tests/helpers/tiling/test_scheduling.py | 7 +- .../tests/helpers/tiling/test_transforms.py | 13 +- .../helpers/tiling/traits/test_mma_traits.py | 87 +++++++---- .../visualization/test_layout_render.py | 11 +- 28 files changed, 577 insertions(+), 225 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_dtypes.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md index 735e4f8257ee..35233a4757cc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md @@ -222,7 +222,7 @@ cooperative map; `classify_epilogue(...)` → the branch. Worked example: ## Tool catalog (load-bearing knobs; full list in `visualization_api_surface.md`) -- **`MmaTee`** — `from_mma(mma, **overrides)` (canonical refs + dtypes from the atom) OR explicit args. ✗ +- **`MmaTee`** — `from_mma(mma, **overrides)` (canonical refs from the atom + per-operand dtypes from the TileMma: its A, B and result dtypes) OR explicit args. ✗ `from_mma` with no `a_enc`/`b_enc` renders CANONICAL (it sources canonical from the `canonical_layouts` helper, NOT `mma.a_layout`) — for interleaved you MUST pass explicit `a_enc`/`b_enc` and VERIFY the encoding is `!= canonical_layouts(mma.traits, mma.subtiles)` + sound (`mma.a_layout` is the CONFIGURED style now, NOT a diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/README.md index e1d1f06b3df9..7b2d27410193 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/README.md @@ -6,12 +6,13 @@ tiled GEMM/MMA kernel -- or any custom thread/register distribution -- without h bound target so authorship stays target-agnostic. ```python +from rocke.core.ir import F16, F32 from rocke.helpers.tiling import ( TileMma, make_tensor_desc, make_window, make_fragment, fill_fragment, load_fragment, store_fragment, ) -mma = TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") # IR-free resolve +mma = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") # IR-free resolve acc = make_fragment(mma.c_desc, F32) fill_fragment(b, acc, 0) for k0 in range(0, K_LEN, 16): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/api_proposals/PROPOSAL_c_transpose_crossed_ab.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/api_proposals/PROPOSAL_c_transpose_crossed_ab.md index 939a4c7ad9a0..c46fe98bebbe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/api_proposals/PROPOSAL_c_transpose_crossed_ab.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/api_proposals/PROPOSAL_c_transpose_crossed_ab.md @@ -17,7 +17,7 @@ hand-plumb three coupled changes and get them all consistent: ```python # today (awkward) — the "ab_swap" crossed path, hand-wired in build_crc_gemm: -mma = TileMma((warp_n, warp_m, tile_k), a="f16", b="f16", c="f32", target=arch) # 1. swap M<->N in the shape +mma = TileMma((warp_n, warp_m, tile_k), a=F16, b=F16, c=F32, target=arch) # 1. swap M<->N in the shape ... acc = mma(b_frag, a_frag, acc) # 2. feed B into the A-slot, A into the B-slot (M<->N relabel) ... @@ -39,7 +39,7 @@ transform (register-identity), not cross-lane. ```python # proposed API -mma = TileMma((M, N, K), a="f16", b="f16", c="f32", target=arch, c_transpose=True) +mma = TileMma((M, N, K), a=F16, b=F16, c=F32, target=arch, c_transpose=True) acc = mma(a_frag, b_frag, acc) # normal call; object handles the crossed feed store_fragment(b, c_ptr, make_window(mma.c_desc, (m0, n0)), acc, lane) # c_desc already (N,M)-oriented ``` diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md index 9170c614b06f..a5d0e04a8f61 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md @@ -63,7 +63,9 @@ descriptor). The smallest complete kernel body. Everything not stated is resolved from the MMA atom. ```python -mma = TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") +from rocke.core.ir import F16, F32 # MMA dtypes are IR types + +mma = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") acc = make_fragment(mma.c_desc, F32) # desc + dtype -- both defaulted by the MMA object fill_fragment(b, acc, 0) @@ -491,7 +493,9 @@ The demos (`tiling_gemm_{interwave,pipelined,interleaved}_demo.py`) show real, n ## 6. `TileMma` -- the resolver + subtile driver (where the defaults live) -`TileMma(shape, *, a, b, c, target, tiling=None)`. `shape` is the **wave tile**. Everything +`TileMma(shape, *, a, b, c, d=None, target, tiling=None)`. `shape` is the **wave tile**; `a`, `b`, `c` +are IR types from `rocke.core.ir` (`F16`, `F32`, `FP8E4M3`, ...; a dtype string is refused with the constant +to pass), and `d` is the result dtype, omitted when it equals the accumulator `c`. Everything below is resolved for you and is inspectable. The **instruction** catalog is **`traits/data/mma_traits.json`** (one row per MMA instruction variant: atom shape + per-operand dtypes; hand-maintained, schema and edit procedure in `traits/README.md`; resolved automatically). Its **canonical** @@ -500,7 +504,9 @@ hand-roll valid non-canonical (interleaved / custom) layouts (canonical is never ### Default: single MMA, everything resolved from traits ```python -mma = TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") +from rocke.core.ir import BF16, F16, F32 + +mma = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") mma.op_id # 'mfma_f32_16x16x16f16' (resolved intrinsic) mma.atom_shape # (16, 16, 16) mma.subtiles # (1, 1, 1) (single atom) @@ -512,7 +518,7 @@ mma.c_desc # TileDesc((16,16), c_layout) -- use for the accumulator Fragmen ### Override: wave macro-tile -- the object owns the subtile grid ```python -mma = TileMma((32, 32, 32), a="f16", b="f16", c="f32", target="gfx90a", +mma = TileMma((32, 32, 32), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16))) mma.subtiles # (2, 2, 2) -> the object walks a 2x2x2 atom grid inside one mma() call # author STILL writes one load/mma/store per K-tile; no atom loop. @@ -520,9 +526,9 @@ mma.subtiles # (2, 2, 2) -> the object walks a 2x2x2 atom grid inside one mm ### Validation is automatic and fail-fast ```python -TileMma((32, 32, 16), a="f16", b="bf16", c="f32", target="gfx90a") # ValueError: A/B dtypes -TileMma((13, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") # ValueError: no intrinsic -TileMma((32, 32, 24), a="f16", b="f16", c="f32", target="gfx90a", +TileMma((32, 32, 16), a=F16, b=BF16, c=F32, target="gfx90a") # ValueError: A/B dtypes +TileMma((13, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") # ValueError: no MMA instruction found +TileMma((32, 32, 24), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16,16,16))) # ValueError: wave K not a multiple of atom K ``` @@ -534,24 +540,26 @@ TileMma((32, 32, 24), a="f16", b="f16", c="f32", target="gfx90a", **Default (None): single MMA -- atom == wave shape.** ```python -TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") # tiling=None +from rocke.core.ir import F16, F32 + +TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") # tiling=None # -> atom_shape == (16,16,16), subtiles == (1,1,1) ``` **Override by SHAPE tuple: target-agnostic.** The atom is resolved for the bound target + dtypes -- MFMA on CDNA, WMMA on RDNA -- with no author-side gfx branch. ```python -TileMma((32, 32, 16), a="f16", b="f16", c="f32", target="gfx90a", +TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16))) # op_id -> mfma_f32_16x16x16f16 # same author code, different target (verified): -TileMma((32, 32, 16), a="f16", b="f16", c="f32", target="gfx1201", +TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx1201", tiling=Tiling(atom_shape=(16, 16, 16))) # op_id -> wmma_f32_16x16x16_f16_w32_gfx12 (wave 32) ``` **Override by NAME string: the escape hatch (target-specific by construction).** Pin the exact backend intrinsic; validated (exists / runs on target / dtypes match) fail-fast. ```python -TileMma((32, 32, 16), a="f16", b="f16", c="f32", target="gfx90a", +TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape="mfma_f32_16x16x16f16")) # exact op resolved in the backend ``` @@ -560,7 +568,7 @@ resolve the *native* 32x32 atom (a single MMA), NOT a 2x2 grid of 16x16 atoms. T lets you force the small atom either way: ```python -TileMma((32, 32, 8), a="f16", b="f16", c="f32", target="gfx90a") # -> 32x32x8 native, subtiles (1,1,1) +TileMma((32, 32, 8), a=F16, b=F16, c=F32, target="gfx90a") # -> 32x32x8 native, subtiles (1,1,1) TileMma((32, 32, 16), ..., tiling=Tiling(atom_shape=(16,16,16))) # -> 16x16 atom, subtiles (2,2,1) ``` @@ -634,27 +642,29 @@ The same GEMM, adding exactly one override at a time -- the loop body never chan `TileMma` construction. Knobs travel as constructor arguments (`shape`, `tiling=`, `style=`). ```python +from rocke.core.ir import F16, F32 + # (0) DEFAULT: single 16x16x16 atom, f16->f32, MNK order. Say almost nothing. -mma = TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") +mma = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") # (1) BIGGER K per tile: 4 K-subtiles, MMA owns the inner K loop. -mma = TileMma((16, 16, 64), a="f16", b="f16", c="f32", target="gfx90a", +mma = TileMma((16, 16, 64), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16))) # (2) COOPERATIVE M/N: a 2x2 output grid, MMA walks it internally. -mma = TileMma((32, 32, 16), a="f16", b="f16", c="f32", target="gfx90a", +mma = TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16))) # (3) FULL GRID + iteration order: 2x2x2, K outermost. -mma = TileMma((32, 32, 32), a="f16", b="f16", c="f32", target="gfx90a", +mma = TileMma((32, 32, 32), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16), order="KNM")) # (4) PIN THE EXACT INTRINSIC by name (target-specific escape hatch). -mma = TileMma((32, 32, 16), a="f16", b="f16", c="f32", target="gfx90a", +mma = TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape="mfma_f32_16x16x16f16")) # (5) SWITCH THE OPERAND-LAYOUT PROFILE: wide free-contiguous loads. -mma = TileMma((32, 32, 16), a="f16", b="f16", c="f32", target="gfx90a", +mma = TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), style=InterleavedStyle()) ``` diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/visualization_api_surface.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/visualization_api_surface.md index 345f389686d3..5ef4d12e76c1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/visualization_api_surface.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/visualization_api_surface.md @@ -100,8 +100,8 @@ label per cell (the C-body trick). `atom`, `addr_fn`, `detail_first` refine bord ## 4. Composite views ### `MmaTee.from_mma(mma, **overrides)` (or the dataclass with `a_enc`/`b_enc`/`c_enc`/`atom_shape`) -The A (left wing) / B (top wing) -> **derived C** tee in one figure. `from_mma` fills canonical refs + -dtypes from the atom; override `a_enc=`/`b_enc=` to render a non-canonical (interleaved) operand. Inside the +The A (left wing) / B (top wing) -> **derived C** tee in one figure. `from_mma` fills canonical refs from the +atom and per-operand dtypes from the TileMma (its A, B and result dtypes); override `a_enc=`/`b_enc=` to render a non-canonical (interleaved) operand. Inside the tee the **B and C wings are transposed** (tid horizontal) to align K/N -- this is the ONLY place a register file is transposed. Knobs: `issue_order` (C register grouping), `show_logical_inputs`/`show_static`, `trace_a`/`trace_b`, `show_diagnostics` (soundness/K-match banner). `.render(out_dir, name)` -> path; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_crc_demo/crc_interleaved_gemm.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_crc_demo/crc_interleaved_gemm.py index df5dadeaf653..2e28e44973b9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_crc_demo/crc_interleaved_gemm.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_crc_demo/crc_interleaved_gemm.py @@ -194,18 +194,18 @@ def build_crc_gemm( if ab_swap: mma = TileMma( (warp_n, warp_m, tile_k), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target=arch, tiling=Tiling(atom_shape=(_ATOM, _ATOM, _ATOM)), ) else: mma = TileMma( (warp_m, warp_n, tile_k), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target=arch, tiling=Tiling(atom_shape=(_ATOM, _ATOM, _ATOM)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_interleaved_demo.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_interleaved_demo.py index 5ec66ad4934b..6e79b3cf0df6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_interleaved_demo.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/kernels/tiling_gemm_interleaved_demo.py @@ -340,9 +340,9 @@ def build_interleaved_gemm( # The driver is configured with the WARP tile -- it owns the per-wave m/n/k atom grid + the op. mma = TileMma( (warp_m, warp_n, tile_k), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target=arch, tiling=Tiling(atom_shape=(_ATOM, _ATOM, _ATOM)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py index 3fa0b2f54c43..7bdae4356acd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py @@ -12,7 +12,6 @@ from __future__ import annotations -from ....core.ir import dtype_to_ir_type from ..fragments import Fragment, TileDesc, fragment_length from ..register_mapper import RegisterMapper from ..transforms import ( @@ -208,10 +207,8 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): ("B", b_fragment), ("C", accumulator), ): - want = dtype_to_ir_type( - {"A": plan._a_dtype, "B": plan._b_dtype, "C": plan._c_dtype}[name] - ) - if fragment.dtype.name != want.name: + want = {"A": plan.a_dtype, "B": plan.b_dtype, "C": plan.c_dtype}[name] + if fragment.dtype != want: raise ValueError( f"MMA operand dtype mismatch -- operand={name}, " f"fragment={fragment.dtype.name!r}, expected {want.name!r}" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py index 1dee03aef3d3..3e35216253af 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py @@ -25,9 +25,10 @@ class TileMma: """A target-resolved wave-tile MMA operation + subtile driver (dense; the M1 surface). - Author call: ``TileMma((16, 16, 64), a="f16", b="f16", c="f32", target="gfx90a", - tiling=Tiling(atom_shape=(16, 16, 16)))``. Resolution + validation happen at construction, - fail-fast, with no IR. Calling the instance (``mma(b, a_frag, b_frag, acc)``) walks the atom grid + Author call: ``TileMma((16, 16, 64), a=F16, b=F16, c=F32, target="gfx90a", + tiling=Tiling(atom_shape=(16, 16, 16)))`` with the dtypes imported from ``rocke.core.ir`` + (``d=`` names a result dtype other than ``c``; see :class:`TileMmaPlan`). Resolution + + validation happen at construction, fail-fast, with no IR. Calling the instance (``mma(b, a_frag, b_frag, acc)``) walks the atom grid internally and returns the updated accumulator ``Fragment``. Reach ``.plan`` / ``.driver`` for the composed pieces. """ @@ -39,6 +40,7 @@ def __init__( a=None, b=None, c=None, + d=None, target, tiling=None, style=None, @@ -50,6 +52,7 @@ def __init__( a=a, b=b, c=c, + d=d, target=target, tiling=tiling, style=style, @@ -90,6 +93,22 @@ def tiling(self): def target(self): return self._plan.target + @property + def a_dtype(self): + return self._plan.a_dtype + + @property + def b_dtype(self): + return self._plan.b_dtype + + @property + def c_dtype(self): + return self._plan.c_dtype + + @property + def d_dtype(self): + return self._plan.d_dtype + @property def op_id(self): return self._plan.op_id @@ -142,5 +161,5 @@ def __repr__(self) -> str: return ( f"TileMma(shape={p.shape}, atom={p.atom_shape}, " f"a={p.a_dtype!r}, b={p.b_dtype!r}, c={p.c_dtype!r}, " - f"target={p.target!r}, op_id={p.op_id!r})" + f"d={p.d_dtype!r}, target={p.target!r}, op_id={p.op_id!r})" ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py index d3b68e8fbd7b..0a74f445b421 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py @@ -21,8 +21,9 @@ from typing import Optional from ....core.dtypes import normalize_dtype +from ....core.ir import Type, dtype_to_ir_type from ..encoding import WarpDistributionEncoding -from ..fragments import TileDesc +from ..fragments import TileDesc, _fragment_dtype_ok from ..traits import MmaTraits, MmaTraitsCatalog, load_mma_traits from .styles import CanonicalStyle, LayoutStyle from .warp_encoding import canonical_layouts @@ -37,6 +38,39 @@ _SUBTILE_ORDERS = ("MNK", "MKN", "NMK", "NKM", "KMN", "KNM") +def _ir_constant_name(dtype: Type) -> str: + """The ``rocke.core.ir`` constant an author writes for ``dtype`` (``F16`` for ``f16``): every + scalar IR type is exported under its name in upper case.""" + return f"rocke.core.ir.{dtype.name.upper()}" + + +def _require_scalar_ir_type(role: str, value: object) -> Type: + """Return ``value`` when it is a scalar IR type; otherwise raise naming the constant to pass. + + ``role`` is the keyword the value came in (``"a"``, ``"b"``, ``"c"`` or ``"d"``). A dtype string + is refused rather than converted, so every MMA dtype the plan holds is the IR type the author + wrote. + """ + ok, _ = _fragment_dtype_ok(value) + if ok: + return value # type: ignore[return-value] + head = f"{role}= must be a scalar IR type such as rocke.core.ir.F16" + if isinstance(value, str): + try: + dtype = dtype_to_ir_type(value) + except ValueError: + raise ValueError( + f"{head}, not the string {value!r} -- rocke.core.ir has no type for " + f"{value!r}" + ) from None + raise ValueError( # noqa: TRY004 -- every refused dtype argument is a ValueError + f"{head}, not the string {value!r} -- pass {_ir_constant_name(dtype)}" + ) + # A composite IR type (vector, pointer, LDS) is shown by its IR spelling, anything else by repr. + shown = value.name if isinstance(value, Type) else repr(value) + raise ValueError(f"{head} -- got {shown} ({type(value).__name__})") + + @dataclass(frozen=True) class Tiling: """The wave-tile iteration policy -- the MMA object's knobs. @@ -90,24 +124,34 @@ def __post_init__(self) -> None: class TileMmaPlan: """A target-resolved wave-tile MMA design (dense; the M1 surface). IR-free. - ``TileMmaPlan((16, 16, 64), a="f16", b="f16", c="f32", target="gfx90a", + ``TileMmaPlan((16, 16, 64), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)))`` -- resolution + validation happen here, fail-fast, with no IR. Feed the resolved layouts / ``emit_op`` to a :class:`TileMmaDriver` to walk the atom grid. + + ``a``, ``b``, ``c`` and ``d`` are scalar IR types from ``rocke.core.ir`` (``F16``, ``F32``, + ``FP8E4M3``, ...): A and B are the operand dtypes, ``c`` the accumulator taken in and ``d`` the + result returned. ``d`` omitted means the result has the accumulator's dtype (``d = c``), which + is what every accumulate-in-place instruction does. A dtype string is refused with the constant + to pass instead. """ def __init__( self, shape: tuple[int, int, int] | None = None, *, - a: str | None = None, - b: str | None = None, - c: str | None = None, + a: Type | None = None, + b: Type | None = None, + c: Type | None = None, + d: Type | None = None, target: str, tiling: Tiling | None = None, style: LayoutStyle | None = None, catalog: MmaTraitsCatalog | None = None, atom_override: str | None = None, ) -> None: + for role, value in (("a", a), ("b", b), ("c", c), ("d", d)): + if value is not None: + _require_scalar_ir_type(role, value) tiling = tiling if tiling is not None else Tiling() style = style if style is not None else CanonicalStyle() catalog = catalog if catalog is not None else load_mma_traits() @@ -116,14 +160,14 @@ def __init__( if atom_override is not None: # Name a specific intrinsic and DERIVE its M/N/K + dtypes from the catalog row (target # still required, to validate/resolve). `shape`, if given, is the wave tile (a multiple of - # the atom); otherwise it IS the atom -- a single MMA. Any a/b/c passed are overridden. + # the atom); otherwise it IS the atom -- a single MMA. Any a/b/c/d passed are overridden. if tiling.atom_shape is not None: raise ValueError( "specify the atom via atom_override OR tiling.atom_shape, not both" ) self._traits = self._resolve_atom(catalog, atom_override, target) atom_shape = (self._traits.m, self._traits.n, self._traits.k) - a, b, c = self._traits.a_dtype, self._traits.b_dtype, self._traits.c_dtype + a, b, c, d = self._traits.ir_dtypes() shape = tuple(shape) if shape is not None else atom_shape else: if shape is None: @@ -140,27 +184,29 @@ def __init__( f"MFMA requires matching A/B dtypes -- a={a!r}, b={b!r}" ) # Resolve the atom + its traits from the knob: explicit intrinsic NAME, atom SHAPE, - # or (None) the wave shape itself. Both paths end with a resolved traits row. The - # author's dtype strings are kept as given; they are normalised to canonical names - # only here, where they meet the catalog. - dtypes = tuple(normalize_dtype(name) for name in (a, b, c)) + # or (None) the wave shape itself. Both paths end with a resolved traits row. atom_knob = tiling.atom_shape if isinstance(atom_knob, str): - self._traits = self._resolve_by_name(catalog, atom_knob, target, dtypes) + self._traits = self._resolve_by_name( + catalog, atom_knob, target, a, b, c, d + ) atom_shape = (self._traits.m, self._traits.n, self._traits.k) + d = d if d is not None else self._traits.ir_dtypes()[3] else: atom_shape = atom_knob if atom_knob is not None else shape m, n, k = atom_shape self._traits = catalog.select( target=target, - a=dtypes[0], - b=dtypes[1], - c=dtypes[2], + a=a, + b=b, + c=c, + d=d, m=m, n=n, k=k, family="dense", ) + d = d if d is not None else c self._refuse_result_dtype_change(self._traits) if len(shape) != 3: @@ -187,6 +233,7 @@ def __init__( self._a_dtype = a self._b_dtype = b self._c_dtype = c + self._d_dtype = d self._target = target # C-oracle (rev-7 Principle 8): the DERIVED native C the style produces must equal the machine's @@ -223,23 +270,31 @@ def _resolve_by_name( catalog: MmaTraitsCatalog, op_id: str, target: str, - dtypes: tuple[str, str, str], + a: Type, + b: Type, + c: Type, + d: Type | None, ) -> MmaTraits: """Resolve an EXPLICIT intrinsic name (escape hatch) to its traits, validating that it - exists, runs on ``target``, and takes the requested canonical A/B/C ``dtypes``. - ``catalog.get`` already fails fast for an unknown op_id.""" + exists, runs on ``target``, and takes the requested A/B/C dtypes (and returns ``d`` when + given). ``catalog.get`` already fails fast for an unknown op_id.""" traits = catalog.get(op_id) if not traits.supports(target): raise ValueError( f"MMA intrinsic not available on target -- op_id={op_id!r}, " f"target={target!r}, supported={list(traits.supported_targets)}" ) - a, b, c = dtypes - if not traits.accepts(a=a, b=b, c=c, d=traits.d_dtype): + # The catalog spells dtypes by their canonical core.dtypes names. + asked = { + role: normalize_dtype(dtype.name) + for role, dtype in (("a", a), ("b", b), ("c", c)) + } + asked["d"] = traits.d_dtype if d is None else normalize_dtype(d.name) + if not traits.accepts(**asked): raise ValueError( f"MMA intrinsic dtype mismatch -- op_id={op_id!r}, intrinsic " f"a/b/c=({traits.a_dtype},{traits.b_dtype},{traits.c_dtype}), " - f"requested ({a},{b},{c})" + f"requested ({asked['a']},{asked['b']},{asked['c']})" ) return traits @@ -287,17 +342,26 @@ def target(self) -> str: return self._target @property - def a_dtype(self) -> str: + def a_dtype(self) -> Type: + """The A operand dtype (an IR type).""" return self._a_dtype @property - def b_dtype(self) -> str: + def b_dtype(self) -> Type: + """The B operand dtype (an IR type).""" return self._b_dtype @property - def c_dtype(self) -> str: + def c_dtype(self) -> Type: + """The accumulator dtype the instruction takes in (an IR type).""" return self._c_dtype + @property + def d_dtype(self) -> Type: + """The result dtype the instruction returns (an IR type); equal to :attr:`c_dtype` for + every accumulate-in-place instruction.""" + return self._d_dtype + @property def op_id(self) -> str: """The concrete intrinsic the target resolved to (e.g. mfma_f32_16x16x16f16).""" @@ -434,5 +498,5 @@ def __repr__(self) -> str: return ( f"TileMmaPlan(shape={self._shape}, atom={self._atom_shape}, " f"a={self._a_dtype!r}, b={self._b_dtype!r}, c={self._c_dtype!r}, " - f"target={self._target!r}, op_id={self.op_id!r})" + f"d={self._d_dtype!r}, target={self._target!r}, op_id={self.op_id!r})" ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md index c1f3a008ce5c..eddfec111c52 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md @@ -183,8 +183,9 @@ compile-verified only, until a hardware lane-map run confirms them. ## Selection `load_mma_traits().select(target=, a=, b=, c=, m=, n=, k=, d=None, family="dense", -scale_block_k=None)` returns the one row for an intent. Inputs are canonical dtype names. `d=None` -means the result has the accumulator's dtype. `scale_block_k` is required for family `scaled` and +scale_block_k=None)` returns the one row for an intent. The dtypes are IR types (`rocke.core.ir.F16`, +...), matched against the row's canonical names; `MmaTraits.ir_dtypes()` gives a row's (A, B, C, D) as IR +types. `d=None` means the result has the accumulator's dtype. `scale_block_k` is required for family `scaled` and refused otherwise: block-16 and block-32 scaled instructions decode scales differently, so a kernel names the one it holds. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/mma_traits.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/mma_traits.py index 4ce742b354bd..a91f34e4824d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/mma_traits.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/mma_traits.py @@ -12,7 +12,9 @@ - :class:`MmaTraitsCatalog` -- the loaded table, with lookup by name and selection by intent. - :func:`load_mma_traits` -- read the JSON into a catalog; any malformed row fails the load. -Operand dtypes are rocke's canonical ``core.dtypes`` names (``fp16``, ``fp8e4m3``, ``i8``, ...). +A row stores its operand dtypes as rocke's canonical ``core.dtypes`` names (``fp16``, ``fp8e4m3``, +``i8``, ...); :meth:`MmaTraits.ir_dtypes` gives them as IR types, and +:meth:`MmaTraitsCatalog.select` takes IR types (``rocke.core.ir.F16``, ...). """ from __future__ import annotations @@ -24,7 +26,8 @@ from pathlib import Path from ....core.arch.target import MmaScaleBlockK, MmaScaleDType -from ....core.dtypes import dtype_info +from ....core.dtypes import dtype_info, normalize_dtype +from ....core.ir import Type, dtype_to_ir_type __all__ = [ "MmaLayoutParams", @@ -288,6 +291,13 @@ def accepts(self, *, a: str, b: str, c: str, d: str) -> bool: and (c, d) == (self.c_dtype, self.d_dtype) ) + def ir_dtypes(self) -> tuple[Type, Type, Type, Type]: + """The row's (A, B, C, D) dtypes as IR types (``rocke.core.ir.F16``, ...).""" + return tuple( # type: ignore[return-value] + dtype_to_ir_type(name) + for name in (self.a_dtype, self.b_dtype, self.c_dtype, self.d_dtype) + ) + @classmethod def from_record(cls, record: Mapping[str, object]) -> MmaTraits: """Build a validated ``MmaTraits`` from one ``mma_traits.json`` row. @@ -451,22 +461,22 @@ def select( self, *, target: str, - a: str, - b: str, - c: str, + a: Type, + b: Type, + c: Type, m: int, n: int, k: int, - d: str | None = None, + d: Type | None = None, family: str = "dense", scale_block_k: int | None = None, ) -> MmaTraits: """Select the one instruction for this intent on ``target``. - ``a``, ``b``, ``c`` and ``d`` are canonical ``core.dtypes`` names; ``d=None`` means the - result has the accumulator's dtype (``d = c``). ``scale_block_k`` (the K elements sharing - one scale) is required for family ``scaled`` and refused for every other family. Selection - keys are unique by a load-time invariant, so at most one row matches. + ``a``, ``b``, ``c`` and ``d`` are scalar IR types (``rocke.core.ir.F16``, ...); ``d=None`` + means the result has the accumulator's dtype (``d = c``). ``scale_block_k`` (the K elements + sharing one scale) is required for family ``scaled`` and refused for every other family. + Selection keys are unique by a load-time invariant, so at most one row matches. """ if family == "scaled" and scale_block_k is None: raise ValueError( @@ -478,28 +488,57 @@ def select( f"scale_block_k={scale_block_k!r} applies to family 'scaled' only " f"-- family={family!r}" ) - d = c if d is None else d + # The catalog spells dtypes by their canonical core.dtypes names: that is the match key. + asked = { + role: normalize_dtype(dtype.name) + for role, dtype in ( + ("a", a), + ("b", b), + ("c", c), + ("d", c if d is None else d), + ) + } for traits in self.by_op_id.values(): if ( self._is_selectable(traits, target, family) - and traits.accepts(a=a, b=b, c=c, d=d) + and traits.accepts(**asked) and (traits.m, traits.n, traits.k) == (m, n, k) and _scale_block_k(traits) == scale_block_k ): return traits - selectable = sorted( - t.op_id - for t in self.by_op_id.values() - if self._is_selectable(t, target, family) + query = f"a={a.name}, b={b.name}, c={c.name}" + if d is not None: + query += f", d={d.name}" + query += f", shape={(m, n, k)}, family={family!r}" + if scale_block_k is not None: + query += f", scale_block_k={scale_block_k}" + available = sorted( + { + _describe_row(t) + for t in self.by_op_id.values() + if self._is_selectable(t, target, family) + } ) - block = "" if scale_block_k is None else f" scale_block_k={scale_block_k}" + listed = f" (a x b -> c, MxNxK): {available}" if available else ": none" raise ValueError( - f"no MMA instruction for target={target!r} a={a!r} b={b!r} c={c!r} d={d!r} " - f"shape=({m},{n},{k}) family={family!r}{block} " - f"-- selectable on {target!r}: {selectable}" + f"no MMA instruction found for {query}, target={target!r} " + f"-- available on {target!r}{listed}" ) +def _describe_row(traits: MmaTraits) -> str: + """Spell a row's intent for an author: ``'f16 x f16 -> f32, 16x16x16'`` in IR type names, with + the result dtype when it differs from the accumulator's and the scale block size of a scaled + row.""" + a, b, c, d = (dtype.name for dtype in traits.ir_dtypes()) + text = f"{a} x {b} -> {c}" + if d != c: + text += f" (d={d})" + text += f", {traits.m}x{traits.n}x{traits.k}" + block_k = _scale_block_k(traits) + return text if block_k is None else f"{text}, scale_block_k={block_k}" + + def _check_unique_selection(rows: list[MmaTraits]) -> None: """Load-time invariant: no two selectable rows answer the same selection key.""" answers: dict[tuple, list[str]] = {} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py index 621a4c7564d8..c0d68d92cc1b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py @@ -28,7 +28,7 @@ from types import SimpleNamespace from ....core.dtypes import dtype_info -from ....core.ir import dtype_to_ir_type +from ....core.ir import Type from ..encoding import WarpDistributionEncoding from ..register_mapper import RegisterMapper from ..transforms import ( @@ -2870,6 +2870,12 @@ def flow_kloop_operand( ) +def mma_input_label(a: Type, b: Type) -> str: + """Spell an MMA's input dtypes for a label: the IR type name (``"f16"``) when A and B agree, + else ``"/"`` (``"fp8e4m3/bf8e5m2"``).""" + return a.name if a == b else f"{a.name}/{b.name}" + + def flow_wave_mma(mma, **overrides): """Workflow 3 -- 'show the wave-tile MMA': the A/B register-file wings feeding the machine-derived C (logical body + C register file), with each operand's boxed static distribution (A at A's top-left, @@ -2963,12 +2969,9 @@ def from_mma(cls, mma, **overrides): # `a_enc`/`b_enc` override to render a styled operand. C derives by flowing the supplied labels # through the refs. a_canon, b_canon, c_canon = canonical_layouts(t, mma.subtiles) - # Widths and labels come from the row's per-operand dtypes. The C/D fragment is drawn with - # the result's (D) width; C and D differ only where the instruction returns another dtype. - # Labels are IR type names ("f16", "fp8e4m3"); A and B are spelled "/" when they differ. - a_name, b_name, d_name = ( - dtype_to_ir_type(name).name for name in (t.a_dtype, t.b_dtype, t.d_dtype) - ) + # Widths come from the row's per-operand dtypes; labels from the MMA's IR dtypes. The C/D + # fragment is drawn and labelled with the result's (D) dtype; C and D differ only where the + # instruction returns another dtype. base = dict( a_enc=a_canon, b_enc=b_canon, @@ -2980,8 +2983,8 @@ def from_mma(cls, mma, **overrides): a_dtype_bits=dtype_info(t.a_dtype).encoded_bits, b_dtype_bits=dtype_info(t.b_dtype).encoded_bits, c_dtype_bits=dtype_info(t.d_dtype).encoded_bits, - in_dtype=a_name if a_name == b_name else f"{a_name}/{b_name}", - out_dtype=d_name, + in_dtype=mma_input_label(mma.a_dtype, mma.b_dtype), + out_dtype=mma.d_dtype.name, op_id=mma.op_id, wave_shape=mma.shape, ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_oracle.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_oracle.py index 175e3fb8808c..73dcc4c01623 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_oracle.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_oracle.py @@ -19,6 +19,7 @@ "is available (e.g. the gfx90a host)", ) +from rocke.core.ir import F16, F32 from rocke.helpers.tiling.mma.warp_encoding import c_warp_encoding # noqa: E402 from rocke.helpers.tiling.traits import load_mma_traits # noqa: E402 @@ -35,7 +36,7 @@ def _assert_field_equal(ours, oracle) -> None: def test_c_encoding_field_equals_rocke_oracle_16x16x16_f16() -> None: catalog = load_mma_traits() ours = c_warp_encoding( - catalog.select(target="gfx90a", a="fp16", b="fp16", c="fp32", m=16, n=16, k=16) + catalog.select(target="gfx90a", a=F16, b=F16, c=F32, m=16, n=16, k=16) ) oracle = rocke_atoms.make_c_warp_dstr_encoding(rocke_atoms.MfmaAtom.f16_16x16x16()) _assert_field_equal(ours, oracle) @@ -47,7 +48,7 @@ def test_c_encoding_field_equals_rocke_oracle_32x32x8_f16() -> None: pytest.skip("rocke MfmaAtom has no f16_32x32x8 factory in this build") catalog = load_mma_traits() ours = c_warp_encoding( - catalog.select(target="gfx90a", a="fp16", b="fp16", c="fp32", m=32, n=32, k=8) + catalog.select(target="gfx90a", a=F16, b=F16, c=F32, m=32, n=32, k=8) ) oracle = rocke_atoms.make_c_warp_dstr_encoding(factory()) _assert_field_equal(ours, oracle) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_styles.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_styles.py index 3f22bb51e2f4..e095103a7912 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_styles.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_layout_styles.py @@ -10,6 +10,7 @@ import pytest +from rocke.core.ir import F16, F32, FP4E2M1 from rocke.helpers.tiling.mma.styles import CanonicalStyle, InterleavedStyle from rocke.helpers.tiling.traits import load_mma_traits @@ -62,9 +63,9 @@ def test_shipped_styles_build_a_plan_without_raising() -> None: for shape, tiling in configs: TileMmaPlan( shape, - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", style=style, tiling=tiling, @@ -82,9 +83,9 @@ def test_a_operand_with_several_m_blocks_builds_canonical_and_refuses_interleave def plan(style): return TileMmaPlan( (32, 16, 128), - a="fp4e2m1", - b="fp4e2m1", - c="f32", + a=FP4E2M1, + b=FP4E2M1, + c=F32, target="gfx1250", style=style, atom_override="wmma_f32_32x16x128_f4", @@ -113,4 +114,4 @@ def test_unsound_operand_rejected_at_plan_construction(monkeypatch) -> None: lambda *a, **k: Diagnostic("error", "forced-unsound (test)"), ) with pytest.raises(ValueError, match="not sound"): - TileMmaPlan((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") + TileMmaPlan((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_dtypes.py new file mode 100644 index 000000000000..14b144a4bafe --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_dtypes.py @@ -0,0 +1,106 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""TileMma with IR-typed 8-bit dtypes, from selection to LLVM IR, and the visualizer's view of an +MMA's dtypes. + +A TileMma names its dtypes as IR types (``FP8E4M3``, ``BF8E5M2``); these tests check that an fp8 or +bf8 wave tile lowers to one MFMA call per atom, and that ``MmaTee.from_mma`` sizes and labels each +operand from those dtypes. +""" + +from __future__ import annotations + +import pytest + +from rocke.core.ir import ( + BF8E5M2, + F16, + F32, + FP8E4M3, + IRBuilder, + PtrType, +) +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.helpers.tiling import ( + fill_fragment, + load_fragment, + make_fragment, + make_tensor_desc, + make_window, + store_fragment, +) +from rocke.helpers.tiling.analysis.geometry import _elem_bits +from rocke.helpers.tiling.mma import TileMma, Tiling +from rocke.helpers.tiling.visualization import layout_render as lr + + +def _wave_tile_kernel(dtype, mma: TileMma): + """One wave computes C = A * B^T for the MMA's wave tile, loading A and B from global memory.""" + wave_m, wave_n, wave_k = mma.shape + b = IRBuilder(f"mma_dtypes_{dtype.name}") + b.kernel.attrs["max_workgroup_size"] = mma.wave_size + a_ptr = b.param("A", PtrType(dtype, "global"), noalias=True, readonly=True) + b_ptr = b.param("B", PtrType(dtype, "global"), noalias=True, readonly=True) + c_ptr = b.param("C", PtrType(F32, "global"), noalias=True, writeonly=True) + lane = b.thread_id_x() + zero = b.const_i32(0) + a_td = make_tensor_desc((wave_m, wave_k), (wave_k, 1), dtype) + b_td = make_tensor_desc((wave_n, wave_k), (wave_k, 1), dtype) + c_td = make_tensor_desc((wave_m, wave_n), (wave_n, 1), F32) + acc = make_fragment(mma.c_desc, F32) + fill_fragment(b, acc, 0) + fa = load_fragment(b, a_ptr, make_window(a_td, (zero, zero)), mma.a_desc, lane) + fb = load_fragment(b, b_ptr, make_window(b_td, (zero, zero)), mma.b_desc, lane) + acc = mma(b, fa, fb, acc) + store_fragment(b, c_ptr, make_window(c_td, (zero, zero)), acc, lane) + b.ret() + return b.kernel + + +@pytest.mark.parametrize( + "dtype,intrinsic", + [ + (FP8E4M3, "@llvm.amdgcn.mfma.f32.16x16x32.fp8.fp8"), + (BF8E5M2, "@llvm.amdgcn.mfma.f32.16x16x32.bf8.bf8"), + ], +) +def test_8bit_wave_tile_issues_one_mfma_per_atom(dtype, intrinsic) -> None: + # A 32x32x64 wave tile over the gfx942 16x16x32 atom is a 2 x 2 x 2 atom grid: eight calls. + mma = TileMma( + (32, 32, 64), + a=dtype, + b=dtype, + c=F32, + target="gfx942", + tiling=Tiling(atom_shape=(16, 16, 32)), + ) + assert mma.a_dtype is dtype and mma.b_dtype is dtype and mma.c_dtype is F32 + # Pinned to the Python engine: these tests describe its output; the C++ engine is checked by + # the byte-identity gate. + ir = _lower_kernel_to_llvm_python(_wave_tile_kernel(dtype, mma), arch="gfx942") + calls = [line for line in ir.splitlines() if "call" in line and intrinsic in line] + assert len(calls) == 8 + + +def test_mma_tee_sizes_and_labels_operands_from_the_mma_dtypes() -> None: + fp8 = lr.MmaTee.from_mma( + TileMma((16, 16, 32), a=FP8E4M3, b=FP8E4M3, c=F32, target="gfx942") + ) + assert (fp8.a_dtype_bits, fp8.b_dtype_bits, fp8.c_dtype_bits) == (8, 8, 32) + assert (fp8.in_dtype, fp8.out_dtype) == ("fp8e4m3", "f32") + f16 = lr.MmaTee.from_mma( + TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") + ) + assert (f16.a_dtype_bits, f16.c_dtype_bits) == (16, 32) + assert (f16.in_dtype, f16.out_dtype) == ("f16", "f32") + + +def test_mma_input_label_spells_both_inputs_only_when_they_differ() -> None: + assert lr.mma_input_label(F16, F16) == "f16" + assert lr.mma_input_label(FP8E4M3, BF8E5M2) == "fp8e4m3/bf8e5m2" + + +@pytest.mark.parametrize("name,bits", [("f16", 16), ("tf32", 32), ("fp8e4m3", 8)]) +def test_elem_bits_reads_the_dtype_width(name, bits) -> None: + assert _elem_bits(name) == bits diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py index a0f4cc3fdf1c..3c2a892e7942 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py @@ -7,11 +7,13 @@ import pytest -from rocke.helpers.tiling.mma import TileMma, Tiling +from rocke.core import ir +from rocke.core.ir import BF16, F16, F32, Type, VectorType +from rocke.helpers.tiling.mma import TileMma, TileMmaPlan, Tiling def test_mma_resolves_gfx90a_16x16x16_f16() -> None: - mma_op = TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") + mma_op = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") assert mma_op.op_id == "mfma_f32_16x16x16f16" assert mma_op.wave_size == 64 # Layouts are the calculator's encodings (IR-free). @@ -19,10 +21,92 @@ def test_mma_resolves_gfx90a_16x16x16_f16() -> None: assert mma_op.a_layout.hierarchical_lengths == ((16,), (1, 4, 4)) +# ---- dtypes are IR types -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "token,fix", + [ + ("f16", "-- pass rocke.core.ir.F16"), + ("f128", "-- rocke.core.ir has no type for 'f128'"), + ], +) +def test_string_dtype_is_refused_with_the_constant_to_pass(token, fix) -> None: + # The string goes through a variable so that a search for string dtypes in TileMma calls finds + # none in this package. + with pytest.raises(ValueError) as excinfo: + TileMma((16, 16, 16), a=token, b=F16, c=F32, target="gfx90a") + assert str(excinfo.value) == ( + f"a= must be a scalar IR type such as rocke.core.ir.F16, not the string {token!r} {fix}" + ) + + +def test_every_dtype_keyword_is_checked() -> None: + token = "bf16" + with pytest.raises( + ValueError, match=r"^d= must be a scalar IR type .* rocke.core.ir.BF16$" + ): + TileMma((16, 16, 16), a=F16, b=F16, c=F32, d=token, target="gfx90a") + + +def test_composite_ir_type_is_refused() -> None: + with pytest.raises(ValueError) as excinfo: + TileMma((16, 16, 16), a=VectorType(F16, 4), b=F16, c=F32, target="gfx90a") + assert str(excinfo.value) == ( + "a= must be a scalar IR type such as rocke.core.ir.F16 -- got vec (VectorType)" + ) + + +_SCALAR_IR_TYPES = sorted( + (value for value in vars(ir).values() if type(value) is Type), + key=lambda t: t.name, +) + + +@pytest.mark.parametrize("dtype", _SCALAR_IR_TYPES, ids=lambda t: t.name) +def test_refusal_names_the_constant_of_every_scalar_type(dtype) -> None: + # The refusal spells the constant to pass as rocke.core.ir.; that constant must + # exist and be the very type the string names. + with pytest.raises(ValueError) as excinfo: + TileMma((16, 16, 16), a=dtype.name, b=F16, c=F32, target="gfx90a") + constant = dtype.name.upper() + assert str(excinfo.value).endswith(f"-- pass rocke.core.ir.{constant}") + assert getattr(ir, constant) is dtype + + +def test_dtypes_are_the_ir_types_given() -> None: + mma_op = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") + assert mma_op.a_dtype is F16 and mma_op.b_dtype is F16 + assert mma_op.plan.c_dtype is F32 + # An omitted result dtype is the accumulator's. + assert mma_op.d_dtype is F32 + + +def test_atom_override_derives_ir_dtypes() -> None: + plan = TileMmaPlan(target="gfx90a", atom_override="mfma_f32_16x16x16f16") + assert (plan.a_dtype, plan.b_dtype, plan.c_dtype, plan.d_dtype) == ( + F16, + F16, + F32, + F32, + ) + + +def test_tf32_has_no_instruction_on_gfx950() -> None: + # gfx950 has no tf32 MMA instruction; the request fails like any other dtype with no row. + with pytest.raises(ValueError) as excinfo: + TileMma((16, 16, 8), a=ir.TF32, b=ir.TF32, c=F32, target="gfx950") + message = str(excinfo.value) + assert message.startswith( + "no MMA instruction found for a=tf32, b=tf32, c=f32, shape=(16, 16, 8), " + "family='dense', target='gfx950' -- available on 'gfx950' (a x b -> c, MxNxK): [" + ) + + def test_mma_is_target_agnostic_authorship_swap_target() -> None: # Same authored intent, different bound target -> distinct valid resolution. - on_gfx90a = TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") - on_gfx942 = TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx942") + on_gfx90a = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") + on_gfx942 = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx942") assert on_gfx90a.op_id == on_gfx942.op_id == "mfma_f32_16x16x16f16" # (both CDNA MFMA here; the point is authorship carries no gfx branch) assert on_gfx90a.wave_size == 64 @@ -31,7 +115,7 @@ def test_mma_is_target_agnostic_authorship_swap_target() -> None: def test_default_atom_picks_native_32x32_single_atom() -> None: # With no atom knob, a 32x32x8 wave tile resolves the NATIVE 32x32 atom -> single atom, # NOT a grid of 16x16 atoms. This is the behaviour we override below. - mma_op = TileMma((32, 32, 8), a="f16", b="f16", c="f32", target="gfx90a") + mma_op = TileMma((32, 32, 8), a=F16, b=F16, c=F32, target="gfx90a") assert mma_op.op_id == "mfma_f32_32x32x8f16" assert mma_op.atom_shape == (32, 32, 8) assert mma_op.subtiles == (1, 1, 1) @@ -70,9 +154,9 @@ def test_forcing_16x16_atom_on_32x32_tile() -> None: # a 2x2x1 subtile grid, NOT the native 32x32 atom. mma_op = TileMma( (32, 32, 16), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), ) @@ -91,9 +175,9 @@ def test_atom_selected_by_explicit_intrinsic_name() -> None: # shape tuple, and the wave tile subtiles over it. mma_op = TileMma( (32, 32, 16), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape="mfma_f32_16x16x16f16"), ) @@ -106,9 +190,9 @@ def test_atom_name_unknown_fails_fast() -> None: with pytest.raises(ValueError) as excinfo: TileMma( (16, 16, 16), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape="mfma_f32_16x16x16_does_not_exist"), ) @@ -120,9 +204,9 @@ def test_atom_name_wrong_target_fails_fast() -> None: with pytest.raises(ValueError) as excinfo: TileMma( (16, 16, 32), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape="mfma_f32_16x16x32_f16"), ) @@ -133,9 +217,9 @@ def test_atom_name_dtype_mismatch_fails_fast() -> None: with pytest.raises(ValueError) as excinfo: TileMma( (16, 16, 16), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape="mfma_i32_16x16x16i8"), ) @@ -148,9 +232,9 @@ def test_mma_exposes_per_operand_tiledescs() -> None: # The MMA object hands back ready-to-use TileDescs (wave shape + resolved layout). mma_op = TileMma( (32, 32, 16), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), ) @@ -163,15 +247,16 @@ def test_mma_exposes_per_operand_tiledescs() -> None: def test_mma_mismatched_ab_dtype_fails_fast() -> None: with pytest.raises(ValueError) as excinfo: - TileMma((16, 16, 16), a="f16", b="bf16", c="f32", target="gfx90a") + TileMma((16, 16, 16), a=F16, b=BF16, c=F32, target="gfx90a") assert "matching A/B dtypes" in str(excinfo.value) def test_mma_unsupported_shape_fails_fast_with_candidates() -> None: with pytest.raises(ValueError) as excinfo: - TileMma((13, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") - assert "no MMA instruction" in str(excinfo.value) - assert "selectable on 'gfx90a'" in str(excinfo.value) + TileMma((13, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") + assert "no MMA instruction found" in str(excinfo.value) + assert "(13, 16, 16)" in str(excinfo.value) + assert "available on 'gfx90a'" in str(excinfo.value) # The gfx1250 'wmma_bf16f32_16x16x32_bf16' instruction accumulates in fp32 but returns bf16. The @@ -185,7 +270,7 @@ def test_mma_unsupported_shape_fails_fast_with_candidates() -> None: def test_selection_by_intent_never_picks_a_result_dtype_change() -> None: - mma_op = TileMma((16, 16, 32), a="bf16", b="bf16", c="fp32", target="gfx1250") + mma_op = TileMma((16, 16, 32), a=BF16, b=BF16, c=F32, target="gfx1250") assert mma_op.op_id == "wmma_f32_16x16x32_bf16" assert mma_op.traits.d_dtype == mma_op.traits.c_dtype @@ -200,9 +285,9 @@ def test_atom_name_refuses_a_result_dtype_change() -> None: with pytest.raises(ValueError) as excinfo: TileMma( (16, 16, 32), - a="bf16", - b="bf16", - c="fp32", + a=BF16, + b=BF16, + c=F32, target="gfx1250", tiling=Tiling(atom_shape="wmma_bf16f32_16x16x32_bf16"), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py index ef6e180e3f9b..484d13a9dbb5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py @@ -32,9 +32,9 @@ def _mma(shape=(32, 32, 16), atom=(16, 16, 16), **kw): return TileMma( shape, - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=atom), **kw, @@ -60,9 +60,9 @@ def test_canonical_layouts_byte_identical_to_plan_accessors(shape, atom): # accessors differ from canonical_layouts by design. p = TileMmaPlan( shape, - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=atom), ) @@ -78,9 +78,9 @@ def test_canonical_layouts_matches_canonical_style_operands(): # soundness gate would check default-style operands against a diverged reference and never notice. p = TileMmaPlan( (32, 32, 16), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_coalescing_probe.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_coalescing_probe.py index 75f26abd8871..0f3952e194dc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_coalescing_probe.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_coalescing_probe.py @@ -118,9 +118,9 @@ def _gemm_operand_loads(): """ mma = TileMma( (32, 64, 8), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(32, 32, 8)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_emit_wide_access.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_emit_wide_access.py index 190bcc1280cd..b2f81516caef 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_emit_wide_access.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_emit_wide_access.py @@ -42,9 +42,9 @@ def _mma(interleaved: bool) -> TileMma: style = {"style": InterleavedStyle()} if interleaved else {} return TileMma( (64, 64, 32), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), **style, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment.py index 735bf9cc6326..9d0169d62956 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment.py @@ -42,9 +42,9 @@ # 64-lane wave. _MMA = TileMma( (32, 32, 16), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment_storage.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment_storage.py index 6b8fb429f971..cc717eafd016 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment_storage.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_fragment_storage.py @@ -78,12 +78,10 @@ def _flat_tile(registers: int): ) -def _operand_desc(dtype_name: str): +def _operand_desc(dtype): """The canonical A-operand layout of gfx950's 16x16x128 f8f6f4 instruction for a 4-bit or 6-bit operand: each lane holds one run of 32 K elements.""" - return TileMma( - (16, 16, 128), a=dtype_name, b=dtype_name, c="f32", target="gfx950" - ).a_desc + return TileMma((16, 16, 128), a=dtype, b=dtype, c=F32, target="gfx950").a_desc # ---- the storage type --------------------------------------------------------------------------------- @@ -149,7 +147,7 @@ def test_a_tf32_value_must_be_reinterpreted_to_its_storage_type() -> None: def test_a_packed_value_must_be_its_words() -> None: - a_desc = _operand_desc("fp4e2m1") + a_desc = _operand_desc(FP4E2M1) b = IRBuilder("fp4_bytes_value") expected = ( "Fragment value expected vec for TileDesc(shape=(16, 128)), got vec -- " @@ -164,7 +162,7 @@ def test_a_packed_value_must_be_its_words() -> None: @pytest.mark.parametrize("dtype,words", [(FP4E2M1, 4), (FP6E2M3, 6), (TF32, 32)]) def test_a_zero_fill_is_all_zero_storage_words(dtype, words) -> None: b = IRBuilder("zero_fill") - fragment = make_fragment(_operand_desc("fp4e2m1"), dtype) + fragment = make_fragment(_operand_desc(FP4E2M1), dtype) fill_fragment(b, fragment, 0) assert fragment.value.type == VectorType(I32, words) (fill,) = [op for op in b.kernel.body.ops if op.name == "arith.constant_vec"] @@ -244,7 +242,7 @@ def test_a_k_major_strided_packed_tensor_is_refused() -> None: # registers would be single 4-bit elements 16 apart. td = make_tensor_desc((16, 128), (1, 16), FP4E2M1) with pytest.raises(ValueError) as raised: - _load(td, _operand_desc("fp4e2m1")) + _load(td, _operand_desc(FP4E2M1)) assert str(raised.value) == ( "load_fragment of fp4e2m1 needs each lane's registers to form runs of whole 32-bit " "words along the tensor's stride-1 axis (multiples of 8 elements, starting on an " @@ -272,18 +270,18 @@ def test_runs_shorter_than_a_word_group_are_refused() -> None: def test_a_row_stride_that_splits_a_word_is_refused() -> None: td = make_tensor_desc((16, 128), (132, 1), FP4E2M1) with pytest.raises(ValueError, match="tensor axis 0 steps by 132 elements"): - _load(td, _operand_desc("fp4e2m1")) + _load(td, _operand_desc(FP4E2M1)) def test_an_origin_inside_a_word_is_refused() -> None: td = make_tensor_desc((16, 256), (256, 1), FP4E2M1) with pytest.raises(ValueError, match="the window origin on the stride-1 axis is 4"): - _load(td, _operand_desc("fp4e2m1"), origin=(0, 4)) + _load(td, _operand_desc(FP4E2M1), origin=(0, 4)) @pytest.mark.parametrize("verb", ["load", "store"]) def test_a_clipped_packed_window_needs_a_whole_word_bound(verb) -> None: - a_desc = _operand_desc("fp4e2m1") + a_desc = _operand_desc(FP4E2M1) td = make_tensor_desc((16, 128), (128, 1), FP4E2M1) window = make_window(td, (0, 0), (None, 100)) b = IRBuilder("packed_clip") @@ -305,7 +303,7 @@ def test_a_clipped_packed_window_needs_a_whole_word_bound(verb) -> None: def test_a_packed_lds_buffer_must_hold_words() -> None: td = make_tensor_desc((16, 128), (128, 1), FP4E2M1) with pytest.raises(ValueError) as raised: - _load(td, _operand_desc("fp4e2m1"), lds_elem=F16) + _load(td, _operand_desc(FP4E2M1), lds_elem=F16) assert str(raised.value) == ( "LDS buffer for fp4e2m1 fragments must hold its storage words -- allocate " "smem_alloc(I32, [words]) sized from the descriptor; got smem" @@ -315,7 +313,7 @@ def test_a_packed_lds_buffer_must_hold_words() -> None: def test_a_packed_lds_buffer_is_sized_in_words() -> None: # 16 x 128 fp4 elements are 256 words: a buffer of 255 words is refused, one of 256 accepted. td = make_tensor_desc((16, 128), (128, 1), FP4E2M1) - a_desc = _operand_desc("fp4e2m1") + a_desc = _operand_desc(FP4E2M1) for words, ok in ((255, False), (256, True)): b = IRBuilder("packed_lds_size") tid = b.thread_id_x() @@ -508,7 +506,7 @@ def test_packed_operand_tile_round_trips(name, dtype, bits) -> None: # through an LDS buffer of i32 words; both must equal the input bytes. from rocke.helpers.spec import SignatureBuilder - a_desc = _operand_desc("fp4e2m1" if bits == 4 else name) + a_desc = _operand_desc(FP4E2M1 if bits == 4 else dtype) rows, cols = a_desc.shape td = make_tensor_desc((rows, cols), (cols, 1), dtype) words = register_packing(name, a_desc.register_count).carrier_count @@ -574,7 +572,7 @@ def test_clipped_packed_window_masks_whole_words(name, dtype, bits) -> None: # there); before the bound both equal the input. from rocke.helpers.spec import SignatureBuilder - a_desc = _operand_desc(name) + a_desc = _operand_desc(dtype) rows, cols = a_desc.shape bound = 64 td = make_tensor_desc((rows, cols), (cols, 1), dtype) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py index 7fdab43308f1..51f77bef682a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py @@ -7,6 +7,7 @@ ladder; it never consults a table. These lock the ladder ordering + the validity gates on real layouts. """ +from rocke.core.ir import F16, F32 from rocke.helpers.tiling.layouts.tile_distribution import make_tile_desc from rocke.helpers.tiling.mma.mma_operation import TileMma, Tiling from rocke.helpers.tiling.transforms import TransformPlan, as_forward_map @@ -22,9 +23,9 @@ def _mma(): return TileMma( (64, 64, 32), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_lds_descriptor_addressing.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_lds_descriptor_addressing.py index c83ba7b1d221..04699b528616 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_lds_descriptor_addressing.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_lds_descriptor_addressing.py @@ -720,9 +720,9 @@ def test_interleaved_accumulator_lds_round_trip() -> None: # would undo it. c_desc = TileMma( (64, 64, 32), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), style=InterleavedStyle(), @@ -756,9 +756,9 @@ def _convention_gemm(wave): m_len = n_len = k_len = 64 mma = TileMma( wave, - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_scheduling.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_scheduling.py index f5859746b1fa..9714a6b264a5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_scheduling.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_scheduling.py @@ -16,6 +16,7 @@ "rocke.core.ir", reason="rocke IR substrate required (build-only; no GPU)" ) +from rocke.core.ir import F16, F32 from rocke.helpers.tiling import ( # noqa: E402 InstrClass, TileMma, @@ -44,9 +45,9 @@ def test_instrclass_masks_are_the_hardware_bits(): def _tile_mma(warp_m, warp_n, tile_k): return TileMma( (warp_m, warp_n, tile_k), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(_ATOM, _ATOM, _ATOM)), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py index c02379ba10bf..8e0da152cd9b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py @@ -14,6 +14,7 @@ import pytest +from rocke.core.ir import F16, F32 from rocke.helpers.tiling.encoding import WarpDistributionEncoding from rocke.helpers.tiling.mma.warp_encoding import a_warp_encoding, b_warp_encoding from rocke.helpers.tiling.register_mapper import RegisterMapper @@ -224,17 +225,17 @@ def test_interleaved_style_operand_desc_is_distinct() -> None: canon = TileMma( (32, 32, 32), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), ) inter = TileMma( (32, 32, 32), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16, 16, 16)), style=InterleavedStyle(), diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/traits/test_mma_traits.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/traits/test_mma_traits.py index 68c3161af30f..7beb5b2cf33f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/traits/test_mma_traits.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/traits/test_mma_traits.py @@ -15,6 +15,7 @@ import pytest +from rocke.core.ir import BF16, F16, F32, FP4E2M1, FP8E4M3, I8, I32, TF32, U8 from rocke.helpers.tiling.mma import TileMmaPlan from rocke.helpers.tiling.traits import MmaLayoutParams, load_mma_traits @@ -22,7 +23,7 @@ def _m1_atom(): catalog = load_mma_traits() return catalog, catalog.select( - target="gfx90a", a="fp16", b="fp16", c="fp32", m=16, n=16, k=16 + target="gfx90a", a=F16, b=F16, c=F32, m=16, n=16, k=16 ) @@ -69,13 +70,15 @@ def test_gfx90a_16x16x16_f16_matches_the_catalog() -> None: def test_select_unknown_combo_raises_valueerror_with_query() -> None: catalog = load_mma_traits() with pytest.raises(ValueError) as excinfo: - catalog.select(target="gfx90a", a="fp16", b="fp16", c="fp32", m=13, n=16, k=16) + catalog.select(target="gfx90a", a=F16, b=F16, c=F32, m=13, n=16, k=16) message = str(excinfo.value) assert message.startswith( - "no MMA instruction for target='gfx90a' a='fp16' b='fp16' c='fp32' d='fp32' " - "shape=(13,16,16) family='dense' -- selectable on 'gfx90a': [" + "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): [" ) - assert "'mfma_f32_16x16x16f16'" in message # lists the candidates + # The candidates are spelled by their dtypes in IR type names, not by instruction name. + assert "'f16 x f16 -> f32, 16x16x16'" in message + assert "fp32" not in message @pytest.mark.parametrize( @@ -90,7 +93,7 @@ def test_select_unknown_combo_raises_valueerror_with_query() -> None: ) def test_select_on_real_target_ids(target, op_id) -> None: chosen = load_mma_traits().select( - target=target, a="fp16", b="fp16", c="fp32", m=16, n=16, k=16 + target=target, a=F16, b=F16, c=F32, m=16, n=16, k=16 ) assert chosen.op_id == op_id @@ -98,12 +101,10 @@ def test_select_on_real_target_ids(target, op_id) -> None: def test_family_name_is_not_a_target() -> None: # Rows list real target ids; a family name such as 'gfx11' selects nothing. with pytest.raises(ValueError) as excinfo: - load_mma_traits().select( - target="gfx11", a="fp16", b="fp16", c="fp32", m=16, n=16, k=16 - ) + load_mma_traits().select(target="gfx11", a=F16, b=F16, c=F32, m=16, n=16, k=16) assert str(excinfo.value) == ( - "no MMA instruction for target='gfx11' a='fp16' b='fp16' c='fp32' d='fp32' " - "shape=(16,16,16) family='dense' -- selectable on 'gfx11': []" + "no MMA instruction found for a=f16, b=f16, c=f32, shape=(16, 16, 16), " + "family='dense', target='gfx11' -- available on 'gfx11': none" ) @@ -153,28 +154,50 @@ def test_scaled_rows_carry_their_scale_block_size() -> None: def test_result_dtype_is_part_of_the_selection() -> None: # gfx1250 has two dense bf16 x bf16, f32-accumulator 16x16x32 instructions that differ only in - # the result: d defaults to c, and naming d="bf16" picks the bf16-result instruction. + # the result: d defaults to c, and naming d=BF16 picks the bf16-result instruction. catalog = load_mma_traits() query = { "target": "gfx1250", - "a": "bf16", - "b": "bf16", - "c": "fp32", + "a": BF16, + "b": BF16, + "c": F32, "m": 16, "n": 16, "k": 32, } assert catalog.select(**query).op_id == "wmma_f32_16x16x32_bf16" - assert catalog.select(**query, d="bf16").op_id == "wmma_bf16f32_16x16x32_bf16" + assert catalog.select(**query, d=BF16).op_id == "wmma_bf16f32_16x16x32_bf16" + + +@pytest.mark.parametrize( + "dtype,k,op_id", + [(TF32, 8, "mfma_f32_16x16x8_xf32"), (FP8E4M3, 32, "mfma_f32_16x16x32_fp8_fp8")], +) +def test_select_takes_ir_types(dtype, k, op_id) -> None: + chosen = load_mma_traits().select( + target="gfx942", a=dtype, b=dtype, c=F32, m=16, n=16, k=k + ) + assert chosen.op_id == op_id + + +def test_ir_dtypes_gives_a_row_its_ir_types() -> None: + catalog = load_mma_traits() + assert catalog.get("mfma_f32_16x16x8_xf32").ir_dtypes() == (TF32, TF32, F32, F32) + assert catalog.get("wmma_bf16f32_16x16x32_bf16").ir_dtypes() == ( + BF16, + BF16, + F32, + BF16, + ) def test_scaled_selection_names_the_scale_block_size() -> None: catalog = load_mma_traits() query = { "target": "gfx1250", - "a": "fp8e4m3", - "b": "fp4e2m1", - "c": "fp32", + "a": FP8E4M3, + "b": FP4E2M1, + "c": F32, "m": 16, "n": 16, "k": 128, @@ -197,15 +220,15 @@ def test_unscaled_gfx950_f8f6f4_rows_are_selected_by_their_formats() -> None: # is its own row, selected by the two formats, and the scaled instruction of the same formats # stays a separate row of family 'scaled'. catalog = load_mma_traits() - query = {"target": "gfx950", "c": "fp32", "m": 16, "n": 16, "k": 128} - same = catalog.select(**query, a="fp8e4m3", b="fp8e4m3") + query = {"target": "gfx950", "c": F32, "m": 16, "n": 16, "k": 128} + same = catalog.select(**query, a=FP8E4M3, b=FP8E4M3) assert same.op_id == "mfma_f32_16x16x128_f8f6f4_fp8e4m3_fp8e4m3" assert (same.family, same.scale) == ("dense", None) - mixed = catalog.select(**query, a="fp8e4m3", b="fp4e2m1") + mixed = catalog.select(**query, a=FP8E4M3, b=FP4E2M1) assert mixed.op_id == "mfma_f32_16x16x128_f8f6f4_fp8e4m3_fp4e2m1" assert (mixed.a_dtype, mixed.b_dtype) == ("fp8e4m3", "fp4e2m1") scaled = catalog.select( - **query, a="fp8e4m3", b="fp8e4m3", family="scaled", scale_block_k=32 + **query, a=FP8E4M3, b=FP8E4M3, family="scaled", scale_block_k=32 ) assert scaled.op_id == "mfma_scale_f32_16x16x128_f8f6f4_fp8e4m3_fp8e4m3" @@ -215,9 +238,9 @@ def test_scale_block_size_is_refused_outside_family_scaled() -> None: with pytest.raises(ValueError) as excinfo: catalog.select( target="gfx90a", - a="fp16", - b="fp16", - c="fp32", + a=F16, + b=F16, + c=F32, m=16, n=16, k=16, @@ -235,7 +258,7 @@ def test_pin_only_rows_are_never_selected() -> None: pinned = catalog.get("wmma_f32_16x16x128_f8f6f4_fp8e4m3_fp8e4m3") assert not pinned.auto_select chosen = catalog.select( - target="gfx1250", a="fp8e4m3", b="fp8e4m3", c="fp32", m=16, n=16, k=128 + target="gfx1250", a=FP8E4M3, b=FP8E4M3, c=F32, m=16, n=16, k=128 ) assert chosen.op_id == "wmma_f32_16x16x128_fp8_fp8" @@ -244,14 +267,14 @@ def test_selectable_signedness_also_takes_unsigned_operands() -> None: # Integer WMMA has a signedness flag per operand, so u8 matches its i8 row, independently for A # and B; integer MFMA is signed only. catalog = load_mma_traits() - for a, b in (("u8", "i8"), ("i8", "u8"), ("u8", "u8")): - chosen = catalog.select(target="gfx1250", a=a, b=b, c="i32", m=16, n=16, k=64) + for a, b in ((U8, I8), (I8, U8), (U8, U8)): + chosen = catalog.select(target="gfx1250", a=a, b=b, c=I32, m=16, n=16, k=64) assert chosen.op_id == "wmma_i32_16x16x64_iu8" assert chosen.operand_signedness == "selectable" - signed = catalog.select(target="gfx942", a="i8", b="i8", c="i32", m=16, n=16, k=32) + signed = catalog.select(target="gfx942", a=I8, b=I8, c=I32, m=16, n=16, k=32) assert signed.operand_signedness == "signed" with pytest.raises(ValueError, match="no MMA instruction"): - catalog.select(target="gfx942", a="u8", b="i8", c="i32", m=16, n=16, k=32) + catalog.select(target="gfx942", a=U8, b=I8, c=I32, m=16, n=16, k=32) # ---- multi-block rows ----------------------------------------------------------------------------- @@ -274,7 +297,7 @@ def test_multi_block_row_loads_and_plan_build_refuses_it() -> None: def test_multi_block_rows_are_never_selected() -> None: catalog = load_mma_traits() with pytest.raises(ValueError, match="no MMA instruction"): - catalog.select(target="gfx90a", a="fp32", b="fp32", c="fp32", m=16, n=16, k=1) + catalog.select(target="gfx90a", a=F32, b=F32, c=F32, m=16, n=16, k=1) def test_get_unknown_op_raises_valueerror() -> None: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_layout_render.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_layout_render.py index 55bec7b047e0..f4044ae98acc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_layout_render.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_layout_render.py @@ -14,6 +14,7 @@ import pytest +from rocke.core.ir import F16, F32 from rocke.helpers.tiling.mma.mma_operation import TileMma from rocke.helpers.tiling.visualization import _canvas as cv from rocke.helpers.tiling.visualization import layout_render as lr @@ -21,7 +22,7 @@ def _mma(): - return TileMma((16, 16, 16), a="f16", b="f16", c="f32", target="gfx90a") + return TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") # ---------------------------------------------------------------- primitive-layer invariants @@ -109,15 +110,15 @@ def test_from_mma_default_renders_canonical_not_the_configured_style(): atom = Tiling(atom_shape=(16, 16, 16)) styled = TileMma( (32, 32, 16), - a="f16", - b="f16", - c="f32", + a=F16, + b=F16, + c=F32, target="gfx90a", tiling=atom, style=InterleavedStyle(), ) canon = TileMma( - (32, 32, 16), a="f16", b="f16", c="f32", target="gfx90a", tiling=atom + (32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", tiling=atom ) # canonical sibling # self-check: the configured a_layout really is NON-canonical, or this test proves nothing. assert styled.a_layout != canon.a_desc.layout From 38c1b7098218a09cac236114ab11bcf552a91ae3 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Tue, 27 Oct 2026 20:02:43 -0400 Subject: [PATCH 3/7] feat(rocke): pin MMA instructions by full signature and accept mixed 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 the atom_override call to make instead. 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, 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. Tests cover the derived and checked dtypes, the mismatch, wrong-target and Tiling-name refusals, mixed and integer selection, the recorded dtypes, the views refusing an op without them, and a recorded fp8 x bf8 MMA's compute tee. Co-Authored-By: Claude --- .../helpers/tiling/docs/tiling_api_surface.md | 39 ++-- .../python/rocke/helpers/tiling/mma/plan.py | 172 +++++++++--------- .../rocke/helpers/tiling/tiling_recorder.py | 5 + .../rocke/helpers/tiling/traits/README.md | 4 +- .../tiling/visualization/auto_pipeline.py | 68 ++++--- .../helpers/tiling/mma/test_mma_operation.py | 157 +++++++++++----- .../visualization/test_auto_pipeline.py | 74 ++++++++ 7 files changed, 353 insertions(+), 166 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md index a5d0e04a8f61..6629eae9675e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md @@ -23,8 +23,8 @@ loop-nest order" -- with the *same* dataflow (`load -> mma -> store`) the whole DEFAULT (MMA drives everything) FULL OVERRIDE (author pins everything) |-----------------------------------------------------------------------------------------| TileMma((16,16,16), a,b,c, target) TileMma((64,64,32), a,b,c, target, - -> atom = shape, single MMA, layouts, tiling=Tiling(atom_shape="mfma_f32_16x16x16f16", - wave_size, op_id, frag sizes all order="KNM")) + -> atom = shape, single MMA, layouts, atom_override="mfma_f32_16x16x16f16", + wave_size, op_id, frag sizes all tiling=Tiling(order="KNM")) resolved from traits -> pinned intrinsic by name, pinned 4x4x2 subtile grid + iteration order ``` @@ -526,7 +526,7 @@ mma.subtiles # (2, 2, 2) -> the object walks a 2x2x2 atom grid inside one mm ### Validation is automatic and fail-fast ```python -TileMma((32, 32, 16), a=F16, b=BF16, c=F32, target="gfx90a") # ValueError: A/B dtypes +TileMma((32, 32, 16), a=F16, b=BF16, c=F32, target="gfx90a") # ValueError: no MMA instruction found TileMma((13, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") # ValueError: no MMA instruction found TileMma((32, 32, 24), a=F16, b=F16, c=F32, target="gfx90a", tiling=Tiling(atom_shape=(16,16,16))) # ValueError: wave K not a multiple of atom K @@ -536,11 +536,11 @@ TileMma((32, 32, 24), a=F16, b=F16, c=F32, target="gfx90a", ## 7. `Tiling` -- the two knobs (atom selection + iteration order) -### 7a. `atom_shape` -- THREE ways to pick the atom +### 7a. The atom -- THREE ways to pick it **Default (None): single MMA -- atom == wave shape.** ```python -from rocke.core.ir import F16, F32 +from rocke.core.ir import BF16, F16, F32 TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") # tiling=None # -> atom_shape == (16,16,16), subtiles == (1,1,1) @@ -556,11 +556,24 @@ TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx1201", tiling=Tiling(atom_shape=(16, 16, 16))) # op_id -> wmma_f32_16x16x16_f16_w32_gfx12 (wave 32) ``` -**Override by NAME string: the escape hatch (target-specific by construction).** Pin the -exact backend intrinsic; validated (exists / runs on target / dtypes match) fail-fast. +**Pin one instruction by NAME with `atom_override`: the escape hatch (target-specific by +construction).** Validated fail-fast: the instruction exists, runs on the target, and takes every +dtype you pass; each dtype you omit is taken from the instruction. `Tiling(atom_shape=...)` takes +only an `(M, N, K)` shape; a name there is refused with the `atom_override` call to make. ```python +TileMma((32, 32, 16), target="gfx90a", + atom_override="mfma_f32_16x16x16f16") # a/b/c taken from it: f16, f16, f32 TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", - tiling=Tiling(atom_shape="mfma_f32_16x16x16f16")) # exact op resolved in the backend + atom_override="mfma_f32_16x16x16f16") # the dtypes given match it: the same plan +TileMma((32, 32, 16), a=BF16, c=F16, target="gfx90a", + atom_override="mfma_f32_16x16x16f16") # ValueError: MMA instruction dtype mismatch +``` + +A and B may differ wherever the target has such an instruction, by intent or by name: +```python +from rocke.core.ir import BF8E5M2, FP8E4M3 + +TileMma((16, 16, 32), a=FP8E4M3, b=BF8E5M2, c=F32, target="gfx942") # op_id -> mfma_f32_16x16x32_fp8_bf8 ``` **Why it matters (the concern this solves):** a bare 32x32 tile with `atom=None` would @@ -621,7 +634,7 @@ designed-in seam, not yet implemented. | operand dtype | -- | f16, bf16, f32, f8/bf8 (via desc + `a/b/c=`) | BUILT (f16->f32 proven; casts honest/fail-fast) | | data layout | -- | any via `make_tensor_desc(..., strides)` + `.permute([...])`, axes in the one convention (sec 3) | BUILT | | tile distribution | `mma.a/b/c_desc` | author any via `make_tile_desc(...)` (quantity-major) | BUILT | -| atom selection | single MMA (`None`) | shape tuple (agnostic) / intrinsic name (pinned) | BUILT | +| atom selection | single MMA (`None`) | `Tiling(atom_shape=)` shape tuple (agnostic) / `atom_override=` name (pinned) | BUILT | | wave tile size | == atom | any integer multiple of atom (M and/or N and/or K) | BUILT | | subtile order | `"MNK"` | any of 6 `MNK` permutations | BUILT | | target | -- | any gfx via `target=` / `arch=` (agnostic authorship) | BUILT (gfx90a proven; RDNA resolves) | @@ -661,7 +674,7 @@ mma = TileMma((32, 32, 32), a=F16, b=F16, c=F32, target="gfx90a", # (4) PIN THE EXACT INTRINSIC by name (target-specific escape hatch). mma = TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", - tiling=Tiling(atom_shape="mfma_f32_16x16x16f16")) + atom_override="mfma_f32_16x16x16f16") # (5) SWITCH THE OPERAND-LAYOUT PROFILE: wide free-contiguous loads. mma = TileMma((32, 32, 16), a=F16, b=F16, c=F32, target="gfx90a", @@ -677,8 +690,8 @@ see the C-oracle in `TileMmaPlan`). That is the composability claim, made concre ## 11. Customising -- your own atom, layout, or style -- **Pin the atom** with `Tiling(atom_shape=...)` -- a shape tuple, an intrinsic name (escape hatch), - or `None` for a single-MMA wave. +- **Pin the atom** with `Tiling(atom_shape=...)` -- a shape tuple, or `None` for a single-MMA wave + -- or pin the exact instruction by name with `TileMma(atom_override=...)` (escape hatch). - **Bring your own operand layout** by holding `mma.plan` and building a `Fragment` from your own `TileDesc`; the driver validates it identically to a derived one (`mma_operand_layout_sound` + `mma_pair_k_aligned` + the SOA atom-contiguity slice guard). Overrides are OPERANDS only -- the @@ -710,7 +723,7 @@ and you can always see the one the MMA object produced. | Knob | Default (MMA-driven) | Manual override | |---|---|---| -| intrinsic | resolved from atom+dtypes+target | `Tiling(atom_shape="mfma_...")` | +| intrinsic | resolved from atom+dtypes+target | `atom_override="mfma_..."` | | atom shape | == wave shape (single MMA) | `Tiling(atom_shape=(16,16,16))` | | wave tile | 1 atom | `tile=(32,32,32)` etc. (multiple of atom) | | subtile order | `"MNK"` (K innermost) | `order="KNM"` (any permutation) | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py index 0a74f445b421..5ffacaa6d8ea 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py @@ -18,7 +18,6 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Optional from ....core.dtypes import normalize_dtype from ....core.ir import Type, dtype_to_ir_type @@ -75,13 +74,11 @@ def _require_scalar_ir_type(role: str, value: object) -> Type: class Tiling: """The wave-tile iteration policy -- the MMA object's knobs. - * ``atom_shape`` -- how to pick the hardware MMA atom. Two flavours (plus a default): - - a shape ``(M, N, K)`` tuple -> TARGET-AGNOSTIC: resolved to the atom for the bound - target + dtypes (MFMA on CDNA, WMMA on RDNA -- authorship carries no gfx branch); - - an explicit intrinsic name ``str`` (e.g. ``"mfma_f32_16x16x16f16"``) -> the escape - hatch: that exact backend op is used (target-specific by construction); - - ``None`` -> the wave shape itself (a single MMA). - The inner-K count the object runs is ``wave_K / atom_K``. + * ``atom_shape`` -- the hardware MMA atom's ``(M, N, K)``, a shape hint only: it is resolved + to the instruction for the bound target + dtypes (MFMA on CDNA, WMMA on RDNA -- authorship + carries no gfx branch). ``None`` (the default) means the wave shape itself (a single MMA). + The inner-K count the object runs is ``wave_K / atom_K``. To pin one exact instruction, name + it with ``TileMma(atom_override=...)`` instead. * ``order`` -- the M/N/K subtile loop-nest order (a permutation of ``"MNK"``). Stride convention: the RIGHT-MOST axis varies fastest (innermost); ``"MNK"`` runs K innermost. Iterated inside one ``TileMmaDriver`` call. @@ -94,18 +91,23 @@ class Tiling: regardless -- there is no dense body to protect. """ - atom_shape: Optional[tuple[int, int, int] | str] = None + atom_shape: tuple[int, int, int] | None = None order: str = "MNK" mac_prio: int = 0 def __post_init__(self) -> None: atom = self.atom_shape - if atom is not None and not isinstance(atom, str): - if len(atom) != 3 or any(not isinstance(d, int) or d <= 0 for d in atom): - raise ValueError( - f"atom_shape must be 3 positive ints or an intrinsic name -- " - f"atom_shape={atom!r}" - ) + if isinstance(atom, str): + raise ValueError( # noqa: TRY004 -- an atom_shape the author mistyped is a ValueError + f"atom_shape is the instruction's (M, N, K) -- got the string {atom!r}; pin an " + f"instruction with TileMma(atom_override={atom!r})" + ) + if atom is not None and ( + len(atom) != 3 or any(not isinstance(d, int) or d <= 0 for d in atom) + ): + raise ValueError( + f"atom_shape must be 3 positive ints -- atom_shape={atom!r}" + ) if self.order not in _SUBTILE_ORDERS: raise ValueError( f"unknown subtile order -- order={self.order!r}, " @@ -132,7 +134,11 @@ class TileMmaPlan: ``FP8E4M3``, ...): A and B are the operand dtypes, ``c`` the accumulator taken in and ``d`` the result returned. ``d`` omitted means the result has the accumulator's dtype (``d = c``), which is what every accumulate-in-place instruction does. A dtype string is refused with the constant - to pass instead. + to pass instead. A and B may differ (``a=FP8E4M3, b=BF8E5M2``) wherever the target has such an + instruction. + + ``atom_override=""`` pins one exact instruction by name; ``a``/``b``/``c``/``d`` are then + optional: each one given must be one the instruction takes, and the rest are taken from it. """ def __init__( @@ -158,16 +164,17 @@ def __init__( self._traits: MmaTraits if atom_override is not None: - # Name a specific intrinsic and DERIVE its M/N/K + dtypes from the catalog row (target - # still required, to validate/resolve). `shape`, if given, is the wave tile (a multiple of - # the atom); otherwise it IS the atom -- a single MMA. Any a/b/c/d passed are overridden. + # Pin one instruction by name. Its M/N/K come from the row; so do any of a/b/c/d the + # author leaves out, and each one given must be one the row takes. `shape`, if given, is + # the wave tile (a multiple of the atom); otherwise it IS the atom -- a single MMA. if tiling.atom_shape is not None: raise ValueError( "specify the atom via atom_override OR tiling.atom_shape, not both" ) - self._traits = self._resolve_atom(catalog, atom_override, target) + self._traits, (a, b, c, d) = self._resolve_named( + catalog, atom_override, target, a, b, c, d + ) atom_shape = (self._traits.m, self._traits.n, self._traits.k) - a, b, c, d = self._traits.ir_dtypes() shape = tuple(shape) if shape is not None else atom_shape else: if shape is None: @@ -176,37 +183,25 @@ def __init__( ) if a is None or b is None or c is None: raise ValueError( - "a/b/c dtypes are required unless atom_override is given" + "a/b/c dtypes are required unless the instruction is pinned with " + "atom_override=..." ) shape = tuple(shape) - if a != b: - raise ValueError( - f"MFMA requires matching A/B dtypes -- a={a!r}, b={b!r}" - ) - # Resolve the atom + its traits from the knob: explicit intrinsic NAME, atom SHAPE, - # or (None) the wave shape itself. Both paths end with a resolved traits row. - atom_knob = tiling.atom_shape - if isinstance(atom_knob, str): - self._traits = self._resolve_by_name( - catalog, atom_knob, target, a, b, c, d - ) - atom_shape = (self._traits.m, self._traits.n, self._traits.k) - d = d if d is not None else self._traits.ir_dtypes()[3] - else: - atom_shape = atom_knob if atom_knob is not None else shape - m, n, k = atom_shape - self._traits = catalog.select( - target=target, - a=a, - b=b, - c=c, - d=d, - m=m, - n=n, - k=k, - family="dense", - ) - d = d if d is not None else c + # The atom is the tiling knob's (M, N, K), or (None) the wave shape itself. + atom_shape = tiling.atom_shape if tiling.atom_shape is not None else shape + m, n, k = atom_shape + self._traits = catalog.select( + target=target, + a=a, + b=b, + c=c, + d=d, + m=m, + n=n, + k=k, + family="dense", + ) + d = d if d is not None else c self._refuse_result_dtype_change(self._traits) if len(shape) != 3: @@ -254,8 +249,8 @@ def _refuse_result_dtype_change(traits: MmaTraits) -> None: The plan and the driver give the backend only the accumulator dtype, and each instruction's result feeds the next one as its accumulator, so an instruction that returns - a different dtype cannot be driven here. Every path that resolves a row (selection by - intent, an explicit name, ``atom_override``) ends in this check. + a different dtype cannot be driven here. Both paths that resolve a row (selection by + intent and ``atom_override``) end in this check. """ if traits.d_dtype != traits.c_dtype: raise ValueError( @@ -266,50 +261,59 @@ def _refuse_result_dtype_change(traits: MmaTraits) -> None: ) @staticmethod - def _resolve_by_name( + def _resolve_named( catalog: MmaTraitsCatalog, op_id: str, target: str, - a: Type, - b: Type, - c: Type, + a: Type | None, + b: Type | None, + c: Type | None, d: Type | None, - ) -> MmaTraits: - """Resolve an EXPLICIT intrinsic name (escape hatch) to its traits, validating that it - exists, runs on ``target``, and takes the requested A/B/C dtypes (and returns ``d`` when - given). ``catalog.get`` already fails fast for an unknown op_id.""" + ) -> tuple[MmaTraits, tuple[Type, Type, Type, Type]]: + """Resolve an ``atom_override`` instruction name to its row and the plan's (A, B, C, D). + + The row must exist (``catalog.get`` fails fast listing the valid names) and run on + ``target``. Each dtype the author gives must be one the row takes -- for an integer row + whose signedness is selectable, that includes the unsigned twin, which is kept as given -- + and every dtype left out is taken from the row. An omitted ``d`` is the row's result + dtype, never defaulted to ``c``. + """ traits = catalog.get(op_id) if not traits.supports(target): raise ValueError( - f"MMA intrinsic not available on target -- op_id={op_id!r}, " + f"MMA instruction not available on target -- op_id={op_id!r}, " f"target={target!r}, supported={list(traits.supported_targets)}" ) - # The catalog spells dtypes by their canonical core.dtypes names. - asked = { - role: normalize_dtype(dtype.name) - for role, dtype in (("a", a), ("b", b), ("c", c)) + row = dict(zip("abcd", traits.ir_dtypes())) + row_names = { + "a": traits.a_dtype, + "b": traits.b_dtype, + "c": traits.c_dtype, + "d": traits.d_dtype, } - asked["d"] = traits.d_dtype if d is None else normalize_dtype(d.name) - if not traits.accepts(**asked): - raise ValueError( - f"MMA intrinsic dtype mismatch -- op_id={op_id!r}, intrinsic " - f"a/b/c=({traits.a_dtype},{traits.b_dtype},{traits.c_dtype}), " - f"requested ({asked['a']},{asked['b']},{asked['c']})" - ) - return traits - - @staticmethod - def _resolve_atom(catalog: MmaTraitsCatalog, op_id: str, target: str) -> MmaTraits: - """Resolve an ``atom_override`` intrinsic NAME to its traits: validate it exists and runs - on ``target``. M/N/K + dtypes are DERIVED from it, so there is no dtype-match check (unlike - ``_resolve_by_name``, which validates against author-supplied dtypes).""" - traits = catalog.get(op_id) - if not traits.supports(target): + given = {"a": a, "b": b, "c": c, "d": d} + clauses, fixes = [], [] + for role, dtype in given.items(): + if dtype is None: + continue + # The catalog spells dtypes by their canonical core.dtypes names: that is the match key. + asked = {**row_names, role: normalize_dtype(dtype.name)} + if not traits.accepts(**asked): + verb = "returns" if role == "d" else "takes" + clauses.append( + f"{role}={dtype.name} but it {verb} {role.upper()}={row[role].name}" + ) + fixes.append(f"{role}={_ir_constant_name(row[role])}") + if clauses: + omit = "a/b/c/d" if given["d"] is not None else "a/b/c" raise ValueError( - f"MMA intrinsic not available on target -- op_id={op_id!r}, " - f"target={target!r}, supported={list(traits.supported_targets)}" + f"MMA instruction dtype mismatch for {op_id!r} -- {'; '.join(clauses)} -- " + f"pass {', '.join(fixes)}, or omit {omit} to take them from the instruction" ) - return traits + resolved = tuple( + given[role] if given[role] is not None else row[role] for role in "abcd" + ) + return traits, resolved # type: ignore[return-value] @property def shape(self) -> tuple[int, int, int]: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/tiling_recorder.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/tiling_recorder.py index dd30d08fd1d6..4669dbfa85ff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/tiling_recorder.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/tiling_recorder.py @@ -110,6 +110,9 @@ class PipelineOp: 1 # free-dim atom counts (m_sub / n_sub) -- for the per-atom K-match soundness gate ) b_free_atoms: int = 1 + # The MMA's A and B math dtypes (IR types): the one home views read operand widths and labels from. + a_dtype: Any = None + b_dtype: Any = None note: str = "" produces: int | None = ( None # id() of the SSA Value this op yields (transform / mma output) @@ -305,6 +308,8 @@ def add_mma( atom_count=atom_count, a_free_atoms=m_sub, b_free_atoms=n_sub, + a_dtype=mma.a_dtype, + b_dtype=mma.b_dtype, note="TileMma", produces=produces, consumes=tuple(consumes), diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md index eddfec111c52..fae1fc68a2c6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md @@ -198,8 +198,8 @@ variant beside the dedicated fp8 instruction), keep the dedicated one selectable Multi-block rows are never selected; they are reachable by name. `TileMma` drives only rows whose result dtype is the accumulator dtype (`d_dtype == c_dtype`): its -driver keeps the result in the accumulator. It refuses any other row on every path, by intent, by -name and by `atom_override` (today only `wmma_bf16f32_16x16x32_bf16`). `select(d=...)` still returns +driver keeps the result in the accumulator. It refuses any other row on both paths, by intent and +by `atom_override` (today only `wmma_bf16f32_16x16x32_bf16`). `select(d=...)` still returns such a row to a caller that asks for it. A pin-only row lays out A, B and C exactly as its selectable twin does, with one exception: the diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/auto_pipeline.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/auto_pipeline.py index e86223aaabea..4765504e76c7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/auto_pipeline.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/auto_pipeline.py @@ -21,7 +21,13 @@ # Geometry helpers the renderers USE, imported from the analysis package (machinery: pure calc, one # home next to coalescing/vectorization). No re-export shim -- callers import the gates from # `rocke.helpers.tiling.analysis` directly, so the renderers no longer re-couple to them. -from ..analysis.geometry import _arch_wave, _elem_bytes, _lane_span, resolve_origin +from ..analysis.geometry import ( + _arch_wave, + _elem_bits, + _elem_bytes, + _lane_span, + resolve_origin, +) # -------------------------------------------------------------------------------------------------- @@ -904,19 +910,42 @@ def _view_compute_flow(pipeline, *, lds_view, operand, wave, out_path): import os - from .layout_render import MmaTee + tee = _compute_flow_tee(pipeline) + in_label, out_label = tee.in_dtype, tee.out_dtype + root = os.path.splitext(_with_dtype(out_path, in_label))[0] + return tee.render( + os.path.dirname(root) or ".", + name=os.path.basename(root), + title=f"COMPUTE flow [{in_label}->{out_label}]: LDS read A/B -> MMA (A x B) -> C", + ) + + +def _recorded_operand_dtypes(op): + """The (A, B) math dtypes of a recorded MMA op -- the one home the views read operand widths and + labels from (A and B may differ, so neither is taken from the other or from a memory access). + """ + if op.a_dtype is None or op.b_dtype is None: + raise ValueError( + "the recorded MMA op carries no A/B dtypes -- record the build with record_build, " + "whose TileMma hook stores them on the op" + ) + return op.a_dtype, op.b_dtype + + +def _compute_flow_tee(pipeline): + """The MMA tee of the COMPUTE flow, built from the recording alone: the recorded operands and + canonical refs, A and B widths and labels from the recorded MMA op, C from the global C store. + """ + from .layout_render import MmaTee, mma_input_label mma_op = next(o for o in pipeline.ops if o.kind == "mma") - reads = [t for t in pipeline.transactions if t.space == "lds" and t.kind == "load"] + a_dtype, b_dtype = _recorded_operand_dtypes(mma_op) cst = next( t for t in pipeline.transactions if t.space == "global" and t.kind == "store" ) - dtype = reads[0].dtype_name # A/B input dtype (f16) - ab_bits, c_bits = _elem_bytes(dtype) * 8, _elem_bytes(cst.dtype_name) * 8 - # FLOW the shade from the operand's real PRODUCER (a direct LDS read) via the shared axis-agnostic helper -- # the tee's A/B colours are the true load-transaction groups, matching the LDS diagram. - tee = MmaTee( + return MmaTee( a_enc=mma_op.a_enc, b_enc=mma_op.b_enc, c_enc=mma_op.c_enc, @@ -924,9 +953,9 @@ def _view_compute_flow(pipeline, *, lds_view, operand, wave, out_path): a_canon=mma_op.a_canon, b_canon=mma_op.b_canon, c_canon=mma_op.c_canon, - a_dtype_bits=ab_bits, - b_dtype_bits=ab_bits, - c_dtype_bits=c_bits, + a_dtype_bits=_elem_bits(a_dtype.name), + b_dtype_bits=_elem_bits(b_dtype.name), + c_dtype_bits=_elem_bits(cst.dtype_name), dims_a=("M", "K"), dims_b=("N", "K"), dims_c=("M", "N"), @@ -934,15 +963,9 @@ def _view_compute_flow(pipeline, *, lds_view, operand, wave, out_path): shade_b=_lds_load_shade(pipeline, mma_op, "B", mma_op.b_enc), show_static=True, show_logical_inputs=False, - in_dtype=dtype, + in_dtype=mma_input_label(a_dtype, b_dtype), out_dtype=cst.dtype_name, ) - root = os.path.splitext(_with_dtype(out_path, dtype))[0] - return tee.render( - os.path.dirname(root) or ".", - name=os.path.basename(root), - title=f"COMPUTE flow [{dtype}->{cst.dtype_name}]: LDS read A/B -> MMA (A x B) -> C", - ) def _read_feeding_combine(pipeline, mma_op, space_id): @@ -1467,14 +1490,11 @@ def _view_epilogue_flow(pipeline, *, out_path): # REACHABILITY, not position: the first recorded reorder is an A-operand bridge on any kernel that # bridges its operands. Raises rather than falling through to a degenerate panel. reo = select_c_shuffle(pipeline) - reads = [t for t in pipeline.transactions if t.space == "lds" and t.kind == "load"] + a_dtype, b_dtype = _recorded_operand_dtypes(mma_op) cst = next( t for t in pipeline.transactions if t.space == "global" and t.kind == "store" ) - ab_bits, c_bits = ( - _elem_bytes(reads[0].dtype_name) * 8, - _elem_bytes(cst.dtype_name) * 8, - ) + c_bits = _elem_bits(cst.dtype_name) tee = MmaTee( a_enc=mma_op.a_enc, b_enc=mma_op.b_enc, @@ -1483,8 +1503,8 @@ def _view_epilogue_flow(pipeline, *, out_path): a_canon=mma_op.a_canon, b_canon=mma_op.b_canon, c_canon=mma_op.c_canon, - a_dtype_bits=ab_bits, - b_dtype_bits=ab_bits, + a_dtype_bits=_elem_bits(a_dtype.name), + b_dtype_bits=_elem_bits(b_dtype.name), c_dtype_bits=c_bits, dims_a=("M", "K"), dims_b=("N", "K"), diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py index 3c2a892e7942..9df57dc07526 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py @@ -8,7 +8,18 @@ import pytest from rocke.core import ir -from rocke.core.ir import BF16, F16, F32, Type, VectorType +from rocke.core.ir import ( + BF8E5M2, + BF16, + F16, + F32, + FP8E4M3, + I8, + I32, + U8, + Type, + VectorType, +) from rocke.helpers.tiling.mma import TileMma, TileMmaPlan, Tiling @@ -145,7 +156,7 @@ def test_atom_override_conflicts_with_tiling_atom_shape() -> None: TileMma( target="gfx90a", atom_override="mfma_f32_16x16x16f16", - tiling=Tiling(atom_shape="mfma_f32_16x16x16f16"), + tiling=Tiling(atom_shape=(16, 16, 16)), ) @@ -179,41 +190,82 @@ def test_atom_selected_by_explicit_intrinsic_name() -> None: b=F16, c=F32, target="gfx90a", - tiling=Tiling(atom_shape="mfma_f32_16x16x16f16"), + atom_override="mfma_f32_16x16x16f16", ) assert mma_op.op_id == "mfma_f32_16x16x16f16" assert mma_op.atom_shape == (16, 16, 16) assert mma_op.subtiles == (2, 2, 1) -def test_atom_name_unknown_fails_fast() -> None: +def test_atom_override_with_matching_dtypes_is_the_same_plan() -> None: + # Dtypes given next to a pinned instruction are checked, not used to choose: when they are the + # ones it takes, the plan is the one built without them. + given = TileMma( + (32, 32, 16), + a=F16, + b=F16, + c=F32, + target="gfx90a", + atom_override="mfma_f32_16x16x16f16", + ) + derived = TileMma( + (32, 32, 16), target="gfx90a", atom_override="mfma_f32_16x16x16f16" + ) + assert repr(given.plan) == repr(derived.plan) + assert given.c_layout == derived.c_layout + + +def test_atom_override_with_some_dtypes_takes_the_rest_from_the_instruction() -> None: + mma_op = TileMma(c=F32, target="gfx90a", atom_override="mfma_f32_16x16x16f16") + assert (mma_op.a_dtype, mma_op.b_dtype, mma_op.d_dtype) == (F16, F16, F32) + + +def test_atom_override_dtype_mismatch_names_each_dtype_and_the_fix() -> None: with pytest.raises(ValueError) as excinfo: - TileMma( - (16, 16, 16), - a=F16, - b=F16, - c=F32, - target="gfx90a", - tiling=Tiling(atom_shape="mfma_f32_16x16x16_does_not_exist"), - ) + TileMma(a=BF16, c=F16, target="gfx90a", atom_override="mfma_f32_16x16x16f16") + assert str(excinfo.value) == ( + "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" + ) + + +def test_atom_override_result_dtype_mismatch_names_the_result() -> None: + # The result dtype is checked like the others; an omitted d is the instruction's, not c. + with pytest.raises(ValueError) as excinfo: + TileMma(d=BF16, target="gfx90a", atom_override="mfma_f32_16x16x16f16") + assert str(excinfo.value) == ( + "MMA instruction dtype mismatch for 'mfma_f32_16x16x16f16' -- d=bf16 but it returns " + "D=f32 -- pass d=rocke.core.ir.F32, or omit a/b/c/d to take them from the instruction" + ) + + +def test_atom_override_keeps_an_unsigned_operand_the_instruction_takes() -> None: + # An integer instruction whose operand signedness is selectable takes u8 as well as i8; the + # author's u8 is kept (it decides the signedness the instruction is issued with) and the + # omitted B is the instruction's i8. + plan = TileMmaPlan(a=U8, target="gfx1250", atom_override="wmma_i32_16x16x64_iu8") + assert (plan.a_dtype, plan.b_dtype, plan.c_dtype) == (U8, I8, I32) + + +def test_atom_override_unknown_fails_fast() -> None: + with pytest.raises(ValueError) as excinfo: + TileMma(target="gfx90a", atom_override="mfma_f32_16x16x16_does_not_exist") assert "mfma_f32_16x16x16_does_not_exist" in str(excinfo.value) -def test_atom_name_wrong_target_fails_fast() -> None: - # A real gfx950-only wide atom named while bound to gfx90a must fail fast. +def test_atom_override_wrong_target_fails_fast() -> None: + # A real wide atom of another target named while bound to gfx90a must fail fast. with pytest.raises(ValueError) as excinfo: - TileMma( - (16, 16, 32), - a=F16, - b=F16, - c=F32, - target="gfx90a", - tiling=Tiling(atom_shape="mfma_f32_16x16x32_f16"), - ) - assert "mfma_f32_16x16x32_f16" in str(excinfo.value) + TileMma(target="gfx90a", atom_override="mfma_f32_16x16x32_f16") + message = str(excinfo.value) + assert message.startswith( + "MMA instruction not available on target -- op_id='mfma_f32_16x16x32_f16', " + "target='gfx90a', supported=[" + ) -def test_atom_name_dtype_mismatch_fails_fast() -> None: +def test_atom_override_operand_dtype_mismatch_fails_fast() -> None: with pytest.raises(ValueError) as excinfo: TileMma( (16, 16, 16), @@ -221,11 +273,29 @@ def test_atom_name_dtype_mismatch_fails_fast() -> None: b=F16, c=F32, target="gfx90a", - tiling=Tiling(atom_shape="mfma_i32_16x16x16i8"), + atom_override="mfma_i32_16x16x16i8", ) - msg = str(excinfo.value) - assert "mfma_i32_16x16x16i8" in msg - assert "dtype mismatch" in msg or "not available" in msg + message = str(excinfo.value) + assert message.startswith( + "MMA instruction dtype mismatch for 'mfma_i32_16x16x16i8' -- a=f16 but it takes A=i8" + ) + + +def test_tiling_refuses_an_instruction_name_as_atom_shape() -> None: + with pytest.raises(ValueError) as excinfo: + Tiling(atom_shape="mfma_f32_16x16x16f16") + assert str(excinfo.value) == ( + "atom_shape is the instruction's (M, N, K) -- got the string 'mfma_f32_16x16x16f16'; " + "pin an instruction with TileMma(atom_override='mfma_f32_16x16x16f16')" + ) + + +def test_dtypes_are_required_without_a_pinned_instruction() -> None: + with pytest.raises(ValueError) as excinfo: + TileMma((16, 16, 16), a=F16, b=F16, target="gfx90a") + assert str(excinfo.value) == ( + "a/b/c dtypes are required unless the instruction is pinned with atom_override=..." + ) def test_mma_exposes_per_operand_tiledescs() -> None: @@ -245,10 +315,24 @@ def test_mma_exposes_per_operand_tiledescs() -> None: assert mma_op.c_desc.layout == mma_op.c_layout -def test_mma_mismatched_ab_dtype_fails_fast() -> None: +def test_mma_mixed_ab_without_instruction_fails_fast() -> None: + # A and B may differ, but gfx90a has no f16 x bf16 instruction. with pytest.raises(ValueError) as excinfo: TileMma((16, 16, 16), a=F16, b=BF16, c=F32, target="gfx90a") - assert "matching A/B dtypes" in str(excinfo.value) + assert str(excinfo.value).startswith( + "no MMA instruction found for a=f16, b=bf16, c=f32, shape=(16, 16, 16)" + ) + + +def test_mma_mixed_fp8_bf8_selects_the_mixed_instruction() -> None: + mma_op = TileMma((16, 16, 32), a=FP8E4M3, b=BF8E5M2, c=F32, target="gfx942") + assert mma_op.op_id == "mfma_f32_16x16x32_fp8_bf8" + assert (mma_op.a_dtype, mma_op.b_dtype) == (FP8E4M3, BF8E5M2) + + +def test_mma_selects_an_integer_instruction() -> None: + mma_op = TileMma((16, 16, 32), a=I8, b=I8, c=I32, target="gfx942") + assert mma_op.op_id == "mfma_i32_16x16x32_i8" def test_mma_unsupported_shape_fails_fast_with_candidates() -> None: @@ -279,16 +363,3 @@ def test_atom_override_refuses_a_result_dtype_change() -> None: with pytest.raises(ValueError) as excinfo: TileMma(target="gfx1250", atom_override="wmma_bf16f32_16x16x32_bf16") assert str(excinfo.value) == _RESULT_DTYPE_CHANGE_MESSAGE - - -def test_atom_name_refuses_a_result_dtype_change() -> None: - with pytest.raises(ValueError) as excinfo: - TileMma( - (16, 16, 32), - a=BF16, - b=BF16, - c=F32, - target="gfx1250", - tiling=Tiling(atom_shape="wmma_bf16f32_16x16x32_bf16"), - ) - assert str(excinfo.value) == _RESULT_DTYPE_CHANGE_MESSAGE diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_auto_pipeline.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_auto_pipeline.py index 1c938ff0cd30..458781849a1a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_auto_pipeline.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/visualization/test_auto_pipeline.py @@ -12,12 +12,22 @@ import pytest +from rocke.core.ir import BF8E5M2, F16, F32, FP8E4M3, IRBuilder, PtrType +from rocke.helpers.tiling import ( + fill_fragment, + load_fragment, + make_fragment, + make_tensor_desc, + make_window, + store_fragment, +) from rocke.helpers.tiling import tiling_recorder as tr from rocke.helpers.tiling.kernels.tiling_gemm_crc_demo.crc_interleaved_gemm import ( b32_swizzle, b64_swizzle, build_crc_gemm, ) +from rocke.helpers.tiling.mma import TileMma from rocke.helpers.tiling.visualization import auto_pipeline as ap from rocke.helpers.tiling import ( analysis as an, @@ -231,6 +241,70 @@ def test_view_epilogue_flow_classifies_and_renders(tmp_path): ) +def test_recorded_mma_carries_its_operand_dtypes(): + """The recorder stores the MMA's A and B dtypes on the op: the one home the views read operand + widths and labels from.""" + _k, _m, pipe = _record() + mma_ops = [o for o in pipe.ops if o.kind == "mma"] + assert mma_ops + assert all(o.a_dtype is F16 and o.b_dtype is F16 for o in mma_ops) + + +@pytest.mark.parametrize("flow", ["compute", "epilogue"]) +def test_views_refuse_an_mma_op_without_dtypes(tmp_path, flow): + """A recorded MMA op without A/B dtypes cannot be drawn: neither view guesses them from a memory + access.""" + _k, _m, pipe = _record() + pipe.nodes = [ + ( + dataclasses.replace(n, a_dtype=None, b_dtype=None) + if getattr(n, "kind", None) == "mma" + else n + ) + for n in pipe.nodes + ] + with pytest.raises(ValueError, match="the recorded MMA op carries no A/B dtypes"): + ap.view(pipe, flow=flow, out_path=str(tmp_path / f"{flow}.png")) + + +def _mixed_fp8_bf8_wave_tile(): + """One wave computes C = A * B^T with fp8 A and bf8 B on gfx942, from global memory.""" + mma = TileMma((16, 16, 32), a=FP8E4M3, b=BF8E5M2, c=F32, target="gfx942") + wave_m, wave_n, wave_k = mma.shape + b = IRBuilder("mixed_fp8_bf8_wave_tile") + b.kernel.attrs["max_workgroup_size"] = mma.wave_size + a_ptr = b.param("A", PtrType(FP8E4M3, "global"), noalias=True, readonly=True) + b_ptr = b.param("B", PtrType(BF8E5M2, "global"), noalias=True, readonly=True) + c_ptr = b.param("C", PtrType(F32, "global"), noalias=True, writeonly=True) + lane = b.thread_id_x() + zero = b.const_i32(0) + a_td = make_tensor_desc((wave_m, wave_k), (wave_k, 1), FP8E4M3) + b_td = make_tensor_desc((wave_n, wave_k), (wave_k, 1), BF8E5M2) + c_td = make_tensor_desc((wave_m, wave_n), (wave_n, 1), F32) + acc = make_fragment(mma.c_desc, F32) + fill_fragment(b, acc, 0) + fa = load_fragment(b, a_ptr, make_window(a_td, (zero, zero)), mma.a_desc, lane) + fb = load_fragment(b, b_ptr, make_window(b_td, (zero, zero)), mma.b_desc, lane) + acc = mma(b, fa, fb, acc) + store_fragment(b, c_ptr, make_window(c_td, (zero, zero)), acc, lane) + b.ret() + return b.kernel, mma + + +@pytest.mark.usefixtures("requires_matplotlib") +def test_mixed_recording_draws_each_operand_at_its_own_dtype(tmp_path): + """A recorded fp8 x bf8 MMA: the compute tee sizes A and B from their own 8-bit dtypes and labels + the inputs with both, and the view renders.""" + import pathlib + + (_kernel, _mma), pipe = tr.record_build(_mixed_fp8_bf8_wave_tile) + tee = ap._compute_flow_tee(pipe) + assert (tee.a_dtype_bits, tee.b_dtype_bits, tee.c_dtype_bits) == (8, 8, 32) + assert (tee.in_dtype, tee.out_dtype) == ("fp8e4m3/bf8e5m2", "f32") + got = ap.view(pipe, flow="compute", out_path=str(tmp_path / "compute.png")) + assert pathlib.Path(got).exists() and pathlib.Path(got).stat().st_size > 0 + + @pytest.mark.parametrize("operand,free", [("A", "M"), ("B", "N")]) @pytest.mark.usefixtures("requires_matplotlib") def test_lds_store_is_a_reposition_with_invariant_labels(tmp_path, operand, free): From ab9d8ad4c43cea2d18563e170be9d29d8ef9e619 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Tue, 27 Oct 2026 20:25:12 -0400 Subject: [PATCH 4/7] feat(rocke): issue MMA instructions as rocke's measured backend coverage 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. 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, which loads and validates it once per process; load_backend_coverage still reads a given path on every call. 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 gfx1250 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; 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, compiled where the host clang targets the arch. The traits README and the API surface doc say how issue follows the record. Co-Authored-By: Claude --- .../helpers/tiling/docs/tiling_api_surface.md | 2 + .../python/rocke/helpers/tiling/mma/driver.py | 74 +++- .../rocke/helpers/tiling/mma/mma_operation.py | 4 +- .../python/rocke/helpers/tiling/mma/plan.py | 86 ++-- .../rocke/helpers/tiling/traits/README.md | 12 +- .../helpers/tiling/traits/backend_coverage.py | 12 + .../helpers/tiling/mma/test_mma_issue.py | 412 ++++++++++++++++++ .../helpers/tiling/mma/test_mma_operation.py | 22 - 8 files changed, 539 insertions(+), 85 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issue.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md index 6629eae9675e..afd5ceb9a038 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md @@ -501,6 +501,8 @@ below is resolved for you and is inspectable. The **instruction** catalog is **` procedure in `traits/README.md`; resolved automatically). Its **canonical** A/B/C encodings are the **baseline/fallback** layout — a starting place, NOT a requirement: you are free to hand-roll valid non-canonical (interleaved / custom) layouts (canonical is never required — `mma_is_machinery.md`). +Whether the resolved instruction is issued on the target, and with which backend op (`mma.emit_op()`), follows the +measured backend coverage record; its classes are described in `traits/README.md` "Backend coverage". ### Default: single MMA, everything resolved from traits ```python diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py index 7bdae4356acd..6907f90f3b30 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py @@ -12,7 +12,7 @@ from __future__ import annotations -from ..fragments import Fragment, TileDesc, fragment_length +from ..fragments import Fragment, TileDesc, _fragment_storage_type, fragment_length from ..register_mapper import RegisterMapper from ..transforms import ( mma_accumulator_flow_consistent, @@ -22,6 +22,10 @@ from .plan import TileMmaPlan from .warp_encoding import canonical_layouts +# The catalog row families whose operands the driver supplies: it feeds A, B and C, so a scaled row +# (which also takes scale operands) or a sparse row (which also takes a sparsity index) is not driven. +_DRIVEN_FAMILIES = frozenset({"dense"}) + def _assert_atom_contiguous( tile_desc: TileDesc, @@ -189,17 +193,35 @@ def _subtile_triples(self): triples.append((axis["M"], axis["N"], axis["K"])) return triples + @staticmethod + def _result_base(b, accumulator: Fragment, d_dtype): + """The vector the grid walk assembles the per-subtile results into: the accumulator's own + vector when the result has its dtype (each result replaces the subtile it was computed + from), else a zero vector of the result dtype's storage type, which every subtile then + fills.""" + if d_dtype == accumulator.dtype: + return accumulator.value + storage = _fragment_storage_type(d_dtype, accumulator.tile_desc.register_count) + return b.zero_vec(storage.elem, storage.count) + def __call__(self, b, a_fragment, b_fragment, accumulator): """Walk the M x N x K atom grid for the wave tile (in ``tiling.order``), issuing one ``b.mma`` per atom and accumulating each C subtile. The fragments are subtile-contiguous (from the wave layouts), so every atom is a register slice. - Checks operand dtypes, then runs the single validation (:func:`_validate_mma_issue` -- - operand soundness, pairwise K-match, accumulator consistency, atom contiguity) before - issuing. Each fragment's register vector type is guaranteed by :class:`Fragment` itself, - not re-checked here. - Returns the updated accumulator as a new :class:`Fragment`; the `accumulator` argument is - not mutated, so the result must be assigned (``acc = driver(b, a_frag, b_frag, acc)``). + Before any IR it checks, in order: the operand dtypes; that the instruction's family is + one the driver feeds (A, B and C only); that a result of another dtype than the + accumulator is never fed back as the next K step's accumulator; that rocke issues the + instruction on the target (:meth:`TileMmaPlan.emit_op`); then the single layout validation + (:func:`_validate_mma_issue` -- operand soundness, pairwise K-match, accumulator + consistency, atom contiguity). A reason the instruction cannot be issued is thus never + hidden behind a layout error. Each fragment's register vector type is guaranteed by + :class:`Fragment` itself, not re-checked here. + + Returns the result as a new :class:`Fragment` on the accumulator's tile, of the + instruction's result dtype (``plan.d_dtype``, the accumulator dtype for every + accumulate-in-place instruction); the `accumulator` argument is not mutated, so the result + must be assigned (``acc = driver(b, a_frag, b_frag, acc)``). """ plan = self._plan for name, fragment in ( @@ -214,15 +236,31 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): f"fragment={fragment.dtype.name!r}, expected {want.name!r}" ) - # One validation call before issuing (operand soundness, pairwise K-match, accumulator - # consistency, atom contiguity), with atom counts from plan.subtiles and the canonical reference - # from the canonical_layouts helper. Kept as ONE call so the issue point can't check a subset -- a - # missing accumulator check is exactly what let a mislabeled C through before. dtype agreement - # (above) and emit_op resolution (below) are the other issue-time checks. - _validate_mma_issue(a_fragment, b_fragment, accumulator, plan) + traits = plan.traits + if traits.family not in _DRIVEN_FAMILIES: + raise NotImplementedError( + f"{plan.op_id!r} is a {traits.family} instruction; the TileMma driver supplies " + f"only A, B and C operands -- the plan can be visualised, and emit_op() returns " + f"the backend op for a hand-written b.mma call wherever rocke issues it -- TODO: " + f"not implemented yet" + ) + m_sub, n_sub, k_sub = plan.subtiles + if plan.d_dtype != plan.c_dtype and k_sub > 1: + raise ValueError( + f"{plan.op_id!r} returns {plan.d_dtype.name} from an {plan.c_dtype.name} " + f"accumulator, so its result cannot be the next K step's accumulator -- this " + f"tile issues {k_sub} instructions along K per C subtile; use a tile whose K is " + f"the instruction's K ({plan.atom_shape[2]}), or the {plan.c_dtype.name}-result " + f"instruction (omit d) and convert the finished tile" + ) + # Issue is decided before the layouts are validated (emit_op emits no IR), then ONE + # validation call (operand soundness, pairwise K-match, accumulator consistency, atom + # contiguity), with atom counts from plan.subtiles and the canonical reference from the + # canonical_layouts helper. Kept as ONE call so the issue point can't check a subset -- a + # missing accumulator check is exactly what let a mislabeled C through before. op = plan.emit_op() - m_sub, n_sub, k_sub = plan._m_subtiles, plan._n_subtiles, plan._k_subtiles + _validate_mma_issue(a_fragment, b_fragment, accumulator, plan) # Single C subtile: accumulate in-register over K (byte-identical to the atom path). if m_sub == 1 and n_sub == 1: @@ -230,7 +268,7 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): if k_sub == 1: return Fragment( accumulator.tile_desc, - accumulator.dtype, + plan.d_dtype, b.mma(op, a_fragment.value, b_fragment.value, acc_value), ) a_atom = fragment_length(a_fragment.tile_desc.layout) // k_sub @@ -250,7 +288,7 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): b.s_setprio(mac_prio) if mac_prio: b.s_setprio(0) - return Fragment(accumulator.tile_desc, accumulator.dtype, acc_value) + return Fragment(accumulator.tile_desc, plan.d_dtype, acc_value) # Subtiled M/N grid. Carry a PER-ATOM accumulator SSA for each (mi, nj) C subtile so # that across K every C subtile is touched ONLY by `b.mma` (an MFMA def->use chain) -- @@ -296,10 +334,10 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): b.s_setprio(mac_prio) if mac_prio: b.s_setprio(0) - result = accumulator.value + result = self._result_base(b, accumulator, plan.d_dtype) for idx in range(m_sub * n_sub): result = self._write_subvector(b, result, accs[idx], idx * c_atom, c_atom) - return Fragment(accumulator.tile_desc, accumulator.dtype, result) + return Fragment(accumulator.tile_desc, plan.d_dtype, result) def __repr__(self) -> str: return f"TileMmaDriver(plan={self._plan!r})" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py index 3e35216253af..6cecfa9e9196 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/mma_operation.py @@ -149,8 +149,8 @@ def c_desc(self): def style(self): return self._plan.style - def emit_op(self): - return self._plan.emit_op() + def emit_op(self, *, scale_format=None): + return self._plan.emit_op(scale_format=scale_format) # ---- iteration, delegated to the driver ----------------------------------------------------- def __call__(self, b, a_fragment, b_fragment, accumulator): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py index 5ffacaa6d8ea..4943b1b724c4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/plan.py @@ -23,7 +23,12 @@ from ....core.ir import Type, dtype_to_ir_type from ..encoding import WarpDistributionEncoding from ..fragments import TileDesc, _fragment_dtype_ok -from ..traits import MmaTraits, MmaTraitsCatalog, load_mma_traits +from ..traits import ( + MmaTraits, + MmaTraitsCatalog, + load_mma_traits, +) +from ..traits.backend_coverage import committed_backend_coverage from .styles import CanonicalStyle, LayoutStyle from .warp_encoding import canonical_layouts @@ -202,7 +207,6 @@ def __init__( family="dense", ) d = d if d is not None else c - self._refuse_result_dtype_change(self._traits) if len(shape) != 3: raise ValueError(f"shape must be (M, N, K) -- got {shape!r}") @@ -243,23 +247,6 @@ def __init__( self._assert_operands_sound() self._assert_accumulator_matches_oracle() - @staticmethod - def _refuse_result_dtype_change(traits: MmaTraits) -> None: - """Refuse a row whose result dtype (D) differs from its accumulator dtype (C). - - The plan and the driver give the backend only the accumulator dtype, and each - instruction's result feeds the next one as its accumulator, so an instruction that returns - a different dtype cannot be driven here. Both paths that resolve a row (selection by - intent and ``atom_override``) end in this check. - """ - if traits.d_dtype != traits.c_dtype: - raise ValueError( - f"MMA instruction {traits.op_id!r} accumulates in {traits.c_dtype} but returns " - f"{traits.d_dtype} -- the TileMma driver does not drive an " - f"instruction whose result dtype differs from its accumulator dtype; choose one " - f"whose result dtype is {traits.c_dtype}" - ) - @staticmethod def _resolve_named( catalog: MmaTraitsCatalog, @@ -471,30 +458,53 @@ def _assert_accumulator_matches_oracle(self) -> None: f"atom-derived descriptor cannot express is not supported." ) - def emit_op(self): - """Resolve the backend ``MmaOp`` for ``b.mma`` from the arch op-registry -- keyed by - the ATOM shape (one ``b.mma`` consumes one atom). - - Lazy import: uses ``core.arch`` -- the op registry, NOT the mfma emission helpers -- - so this layer stays independent of ``mfma_gemm_inner``. + def emit_op(self, *, scale_format: tuple[str, str] | None = None): + """Return the backend ``MmaOp`` that issues this plan's instruction on its target. + + Whether rocke issues an instruction, and with which backend op, is what the measured + backend coverage record says (``traits/data/backend_coverage.json``, loaded once per process + by :func:`~rocke.helpers.tiling.traits.backend_coverage.committed_backend_coverage`; + ``traits/README.md`` "Backend coverage" describes its classes). The op is fetched by the op + id the record holds for this (instruction, target) pair, never searched for by operand + dtypes and shape: several backend ops can share those. A pair whose class is not ``issuable`` raises + ``NotImplementedError`` naming the class. One ``b.mma`` with the op consumes one atom. + + The record holds one op per variant for two kinds of instruction. An integer instruction + whose operand signedness is selectable has one op per (A, B) dtype pair; the plan's own A + and B pick it, so a ``u8`` operand issues the unsigned instruction. A scaled instruction has + one op per scale format, the (A, B) scale dtypes it decodes: ``scale_format=("e8m0", + "e8m0")`` names it, and is required when the record holds more than one. The record refuses + ``scale_format`` for an instruction it does not record per scale format. + + Emits no IR. Lazy import: uses ``core.arch`` -- the op registry, NOT the mfma emission + helpers -- so this layer stays independent of ``mfma_gemm_inner``. """ from rocke.core.arch import ArchTarget - target = ArchTarget.from_gfx(self._target) - m, n, k = self._atom_shape - op = target.mma.op_for_shape( - family="mma", - a_dtype=self._traits.a_dtype, - b_dtype=self._traits.b_dtype, - c_dtype=self._traits.c_dtype, - m=m, - n=n, - k=k, - ) + variant: dict[str, tuple[str, str]] = {} + if scale_format is not None: + variant["scale_format"] = scale_format + if self._traits.operand_signedness == "selectable": + variant["operand_dtypes"] = ( + normalize_dtype(self._a_dtype.name), + normalize_dtype(self._b_dtype.name), + ) + coverage = committed_backend_coverage() + status = coverage.status(self.op_id, self._target, **variant) + if status != "issuable": + raise NotImplementedError( + f"{self.op_id!r} on {self._target!r} is not issued by rocke yet (backend " + f"coverage: {status}) -- the plan can be visualised but not issued; the classes " + f"are described in helpers/tiling/traits/README.md -- TODO: not implemented yet" + ) + backend_op = coverage.backend_op(self.op_id, self._target, **variant) + op = ArchTarget.from_gfx(self._target).mma.by_op_id(backend_op) if op is None: raise ValueError( - f"no backend MMA op for atom={self._atom_shape} " - f"{self._a_dtype}->{self._c_dtype} on target={self._target!r}" + f"backend coverage names backend op {backend_op!r} for {self.op_id!r} on " + f"{self._target!r}, but {self._target!r} has no such op -- the record no longer " + f"matches the ops core registers; regenerate it as traits/README.md describes " + f"under Backend coverage" ) return op diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md index fae1fc68a2c6..a9e32f38ce06 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/README.md @@ -197,10 +197,11 @@ variant beside the dedicated fp8 instruction), keep the dedicated one selectable `auto_select: false` on the other: it stays reachable by name (`atom_override`), never by selection. Multi-block rows are never selected; they are reachable by name. -`TileMma` drives only rows whose result dtype is the accumulator dtype (`d_dtype == c_dtype`): its -driver keeps the result in the accumulator. It refuses any other row on both paths, by intent and -by `atom_override` (today only `wmma_bf16f32_16x16x32_bf16`). `select(d=...)` still returns -such a row to a caller that asks for it. +A row whose result dtype differs from its accumulator dtype (`d_dtype != c_dtype`; today only +`wmma_bf16f32_16x16x32_bf16`) is selected by naming its result, `select(d=...)` or +`TileMma(..., d=...)`, or pinned by `atom_override`. The `TileMma` driver issues it once per C +subtile along K and returns the result in `d_dtype`: such a result cannot be the next +instruction's accumulator, so a tile that needs more than one instruction along K is refused. A pin-only row lays out A, B and C exactly as its selectable twin does, with one exception: the gfx1250 `f8f6f4` 8-bit x 8-bit rows are another instruction than the dedicated fp8 rows beside @@ -246,7 +247,8 @@ backend issues. For every (row, target) pair it holds one class and the backend `MmaOp.op_id`) the row was matched to, read through `load_backend_coverage()` (`backend_coverage.py`): `status(op_id, target)` and `backend_op(op_id, target)`. Only `issuable` means the backend issues the row's instruction on that target; it is the one fact a layer that -issues MMA instructions reads, and `backend_op` is the op it emits. +issues MMA instructions reads, and `backend_op` is the op it emits. `TileMmaPlan.emit_op()` reads +it exactly so: it returns the recorded op of an `issuable` pair and refuses any other class. `tests/helpers/tiling/traits/test_mma_backend_coverage.py` measures the record and requires it to equal the committed file exactly, so a backend op that is added, removed or rewired fails the test diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/backend_coverage.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/backend_coverage.py index 1130a47dd177..3e0102ec4a1e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/backend_coverage.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/traits/backend_coverage.py @@ -24,6 +24,7 @@ from __future__ import annotations +import functools import json from collections.abc import Mapping from dataclasses import dataclass @@ -328,3 +329,14 @@ def load_backend_coverage(path: Path = DEFAULT_COVERAGE_PATH) -> BackendCoverage return BackendCoverage( llvm_flavor=flavor, entries=entries, ops_without_row=ops_without_row ) + + +@functools.cache +def committed_backend_coverage() -> BackendCoverage: + """The committed record (``DEFAULT_COVERAGE_PATH``), loaded and validated once per process. + + It is package data, fixed while rocke runs, and :meth:`TileMmaPlan.emit_op` reads it on every MMA + a kernel issues. Treat the result as read-only: every caller shares it. + :func:`load_backend_coverage` reads its path afresh on every call. + """ + return load_backend_coverage() diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issue.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issue.py new file mode 100644 index 000000000000..bb896365b3d0 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issue.py @@ -0,0 +1,412 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""How TileMma decides what to issue: the measured backend coverage record, and the driver's checks. + +``TileMmaPlan.emit_op()`` returns the backend op the record (``traits/data/backend_coverage.json``) +names for the plan's (instruction, target) pair, fetched by that op id; a pair the record does not +class ``issuable`` is refused with its class. The driver then refuses, before any IR, an instruction +family it cannot feed (scaled, sparse) and a result dtype other than the accumulator's that would be +fed back as the next K step's accumulator, and types its result by the instruction's result dtype. + +Lowering tests pin the Python engine (``_lower_kernel_to_llvm_python``); the C++ engine is checked by +the byte-identity gate. Where the host clang targets the arch, the module is also compiled and linked +(compile-verified, never run). +""" + +from __future__ import annotations + +import sys +import tempfile +from pathlib import Path + +import pytest + +from rocke.assets import platform_root +from rocke.core.arch import ArchTarget +from rocke.core.ir import ( + BF8E5M2, + BF16, + F16, + F32, + FP8E4M3, + I8, + I32, + U8, + IRBuilder, + PtrType, + dtype_to_ir_type, +) +from rocke.core.lower_llvm import LLVM_FLAVORS, _lower_kernel_to_llvm_python +from rocke.helpers.tiling import ( + fill_fragment, + load_fragment, + make_fragment, + make_tensor_desc, + make_window, + store_fragment, +) +from rocke.helpers.tiling.mma import TileMma, TileMmaPlan, Tiling +from rocke.helpers.tiling.traits import load_backend_coverage, load_mma_traits +from rocke.helpers.tiling.traits.backend_coverage import committed_backend_coverage + +_CATALOG = load_mma_traits() +_COVERAGE = load_backend_coverage() + + +def _wave_tile_kernel(mma: TileMma, name: str): + """One wave computes C = A * B^T for the MMA's wave tile: A (M, K) and B (N, K) loaded from + global memory, the result stored to a tensor of the result dtype. Returns the kernel and the + result Fragment the driver gave back.""" + wave_m, wave_n, wave_k = mma.shape + b = IRBuilder(name) + b.kernel.attrs["max_workgroup_size"] = mma.wave_size + a_ptr = b.param("A", PtrType(mma.a_dtype, "global"), noalias=True, readonly=True) + b_ptr = b.param("B", PtrType(mma.b_dtype, "global"), noalias=True, readonly=True) + c_ptr = b.param("C", PtrType(mma.d_dtype, "global"), noalias=True, writeonly=True) + lane = b.thread_id_x() + zero = b.const_i32(0) + a_td = make_tensor_desc((wave_m, wave_k), (wave_k, 1), mma.a_dtype) + b_td = make_tensor_desc((wave_n, wave_k), (wave_k, 1), mma.b_dtype) + c_td = make_tensor_desc((wave_m, wave_n), (wave_n, 1), mma.d_dtype) + acc = make_fragment(mma.c_desc, mma.c_dtype) + fill_fragment(b, acc, 0) + fa = load_fragment(b, a_ptr, make_window(a_td, (zero, zero)), mma.a_desc, lane) + fb = load_fragment(b, b_ptr, make_window(b_td, (zero, zero)), mma.b_desc, lane) + result = mma(b, fa, fb, acc) + store_fragment(b, c_ptr, make_window(c_td, (zero, zero)), result, lane) + b.ret() + return b.kernel, result + + +def _calls(ir: str, intrinsic: str) -> list[str]: + return [line for line in ir.splitlines() if "call" in line and intrinsic in line] + + +def _compile_or_skip(kernel, target: str) -> None: + """Lower ``kernel`` at the flavor of the host clang and compile and link it for ``target`` with + the arbiter of ``tools/check_ir_validity.py``; skip when there is no such clang, or it cannot + target ``target`` (source tree only).""" + tools = platform_root() / "tools" + if not (tools / "check_ir_validity.py").is_file(): + pytest.skip("platform/tools is not installed beside rocke (source tree only)") + if str(tools) not in sys.path: + sys.path.insert(0, str(tools)) + import check_ir_validity + from _llvm_identity import flavor_of_clang, tool_identity + + clang = check_ir_validity._llvm_tool("clang") + if clang is None: + pytest.skip("no clang to compile lowered IR") + flavor = flavor_of_clang(tool_identity(clang)) + if flavor not in LLVM_FLAVORS: + pytest.skip(f"clang {clang} reports no LLVM flavor rocke lowers for") + ir = _lower_kernel_to_llvm_python(kernel, arch=target, llvm_flavor=flavor) + with tempfile.TemporaryDirectory() as work: + source = Path(work) / "module.ll" + source.write_text(ir) + ok, diagnostic = check_ir_validity._compile( + clang, source, target, Path(work) / "module.hsaco" + ) + if not ok and check_ir_validity._target_unsupported(diagnostic): + pytest.skip(f"host clang cannot target {target}") + assert ok, diagnostic + + +# ---- emit_op reads the record ------------------------------------------------------------------ + + +def _issuable_variants(): + """Every issuable (row, target) pair of the record, once per recorded variant.""" + for (op_id, target), entries in sorted(_COVERAGE.entries.items()): + for entry in entries: + if entry.status != "issuable": + continue + variant = entry.operand_dtypes or entry.scale_format + label = f"{op_id}-{target}" + (f"-{'-'.join(variant)}" if variant else "") + yield pytest.param(op_id, target, entry, id=label) + + +@pytest.mark.parametrize("op_id,target,entry", list(_issuable_variants())) +def test_emit_op_returns_the_recorded_backend_op(op_id, target, entry) -> None: + # Catches any drift between the record and core's ops. A selectable integer row is pinned with + # the operand dtypes of the variant (u8 picks the unsigned op); a scaled row names its format. + dtypes = {} + if entry.operand_dtypes is not None: + dtypes = dict(zip("ab", map(dtype_to_ir_type, entry.operand_dtypes))) + plan = TileMmaPlan(target=target, atom_override=op_id, catalog=_CATALOG, **dtypes) + op = plan.emit_op(scale_format=entry.scale_format) + assert op.op_id == entry.backend_op + + +def test_emit_op_loads_the_committed_record_once() -> None: + # emit_op runs for every MMA a kernel issues, so it reads the committed record from the + # once-per-process cache: repeated calls neither read nor validate the file again. + plan = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a").plan + plan.emit_op() + before = committed_backend_coverage.cache_info() + for _ in range(3): + plan.emit_op() + after = committed_backend_coverage.cache_info() + assert (after.misses, after.hits) == (before.misses, before.hits + 3) + + +def test_emit_op_refuses_a_pair_rocke_does_not_issue_naming_its_class() -> None: + # The gfx950 sparse instruction's operand key is shared by a dense op, which an operand-key + # search used to return; the record classes the pair other_instruction. + mma = TileMma(target="gfx950", atom_override="smfmac_f32_16x16x32_f16") + with pytest.raises(NotImplementedError) as excinfo: + mma.emit_op() + assert str(excinfo.value) == ( + "'smfmac_f32_16x16x32_f16' on 'gfx950' is not issued by rocke yet (backend coverage: " + "other_instruction) -- the plan can be visualised but not issued; the classes are " + "described in helpers/tiling/traits/README.md -- TODO: not implemented yet" + ) + + +def test_emit_op_takes_the_dense_op_where_scaled_ops_share_its_operands() -> None: + # gfx1250 has a dense fp8 x fp8 16x16x128 op and two scaled ops with the same operand dtypes + # and shape; the dense instruction issues the dense op. + mma = TileMma((16, 16, 128), a=FP8E4M3, b=FP8E4M3, c=F32, target="gfx1250") + assert mma.op_id == "wmma_f32_16x16x128_fp8_fp8" + op = mma.emit_op() + same_key = [ + other.op_id + for other in ArchTarget.from_gfx("gfx1250").mma.ops + if (other.a_dtype, other.b_dtype, other.c_dtype, other.m, other.n, other.k) + == (op.a_dtype, op.b_dtype, op.c_dtype, op.m, op.n, op.k) + ] + assert sorted(same_key) == [ + "wmma_gfx1250_f32_16x16x128_fp8_fp8", + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k16", + "wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32", + ] + assert op.op_id == "wmma_gfx1250_f32_16x16x128_fp8_fp8" + + +def test_emit_op_issues_the_unsigned_op_for_unsigned_operands() -> None: + # One integer instruction takes signed or unsigned 8-bit operands; each pair has its own op. + signed = TileMma((16, 16, 64), a=I8, b=I8, c=I32, target="gfx1250") + unsigned_a = TileMma((16, 16, 64), a=U8, b=I8, c=I32, target="gfx1250") + assert signed.op_id == unsigned_a.op_id == "wmma_i32_16x16x64_iu8" + assert signed.emit_op().op_id == "wmma_gfx1250_i32_16x16x64_iu8" + assert unsigned_a.emit_op().op_id == "wmma_gfx1250_i32_16x16x64_iu8_u8_i8" + + +_SCALED = "wmma_scale16_f32_16x16x128_f8f6f4_bf8e5m2_fp4e2m1" + + +def test_emit_op_of_a_scaled_instruction_takes_its_scale_format() -> None: + mma = TileMma(target="gfx1250", atom_override=_SCALED) + op = mma.emit_op(scale_format=("e8m0", "e5m3")) + assert op.op_id == "wmma_gfx1250_f32_16x16x128_bf8_fp4_scale_e8m0_e5m3_k16" + + +def test_emit_op_of_a_scaled_instruction_needs_the_format_when_it_has_several() -> None: + mma = TileMma(target="gfx1250", atom_override=_SCALED) + with pytest.raises(ValueError) as excinfo: + mma.emit_op() + assert str(excinfo.value) == ( + f"({_SCALED!r}, 'gfx1250') is recorded per scale_format ([('e8m0', 'e8m0'), " + "('e8m0', 'e4m3'), ('e8m0', 'e5m3')]) -- pass scale_format=" + ) + + +def test_emit_op_refuses_a_scale_format_for_an_unscaled_instruction() -> None: + mma = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") + with pytest.raises( + ValueError, match="recorded without variants -- pass no scale_format" + ): + mma.emit_op(scale_format=("e8m0", "e8m0")) + + +# ---- the driver's checks ----------------------------------------------------------------------- + + +def _operand_fragments(b: IRBuilder, mma: TileMma): + """Zero A, B and C fragments of the MMA's own descriptors and dtypes.""" + fragments = [] + for desc, dtype in ( + (mma.a_desc, mma.a_dtype), + (mma.b_desc, mma.b_dtype), + (mma.c_desc, mma.c_dtype), + ): + fragment = make_fragment(desc, dtype) + fill_fragment(b, fragment, 0) + fragments.append(fragment) + return fragments + + +def test_driver_refuses_a_scaled_instruction_before_validating_layouts() -> None: + # The record issues this scaled instruction, but the driver feeds no scale operands. The refusal + # comes first, so the author is never sent to fix a layout of a tile that cannot be issued. + mma = TileMma(target="gfx1250", atom_override=_SCALED) + b = IRBuilder("driver_scaled") + a, bf, acc = _operand_fragments(b, mma) + with pytest.raises(NotImplementedError) as excinfo: + mma(b, a, bf, acc) + assert str(excinfo.value) == ( + f"{_SCALED!r} is a scaled instruction; the TileMma driver supplies only A, B and C " + "operands -- the plan can be visualised, and emit_op() returns the backend op for a " + "hand-written b.mma call wherever rocke issues it -- TODO: not implemented yet" + ) + + +def test_driver_refuses_a_sparse_instruction() -> None: + # Swapped operands would fail the K-alignment check; the family refusal runs before it. + mma = TileMma(target="gfx950", atom_override="smfmac_f32_16x16x32_f16") + b = IRBuilder("driver_sparse") + a, bf, acc = _operand_fragments(b, mma) + with pytest.raises( + NotImplementedError, match="is a sparse instruction; the TileMma" + ): + mma(b, bf, a, acc) + + +def test_driver_refuses_an_instruction_rocke_does_not_issue_before_validating_layouts() -> ( + None +): + # A dense pair the record does not class issuable: the tied gfx11 instruction's backend op lowers + # to the untied intrinsic. The A fragment passed as the accumulator would fail the accumulator + # check; issue is decided first. + pair = ("wmma_f16_16x16x16_f16_tied_w32", "gfx1100") + assert _COVERAGE.status(*pair) == "other_instruction" + mma = TileMma(target=pair[1], atom_override=pair[0]) + b = IRBuilder("driver_not_issued") + a, bf, _ = _operand_fragments(b, mma) + with pytest.raises( + NotImplementedError, match=r"backend coverage: other_instruction" + ): + mma(b, a, bf, a) + + +# ---- a result dtype other than the accumulator's ---------------------------------------------- + + +_BF16_RESULT = "wmma_bf16f32_16x16x32_bf16" + + +def test_result_dtype_selects_the_instruction() -> None: + f32_result = TileMma((16, 16, 32), a=BF16, b=BF16, c=F32, target="gfx1250") + bf16_result = TileMma((16, 16, 32), a=BF16, b=BF16, c=F32, d=BF16, target="gfx1250") + assert (f32_result.op_id, f32_result.d_dtype) == ("wmma_f32_16x16x32_bf16", F32) + assert (bf16_result.op_id, bf16_result.d_dtype) == (_BF16_RESULT, BF16) + assert bf16_result.c_dtype == F32 + + +def test_atom_override_takes_the_result_dtype_from_the_instruction() -> None: + mma = TileMma(target="gfx1250", atom_override=_BF16_RESULT) + assert (mma.a_dtype, mma.b_dtype, mma.c_dtype, mma.d_dtype) == ( + BF16, + BF16, + F32, + BF16, + ) + + +def test_atom_override_refuses_another_result_dtype() -> None: + with pytest.raises(ValueError) as excinfo: + TileMma(target="gfx1250", atom_override=_BF16_RESULT, d=F32) + assert str(excinfo.value) == ( + f"MMA instruction dtype mismatch for {_BF16_RESULT!r} -- d=f32 but it returns D=bf16 " + "-- pass d=rocke.core.ir.BF16, or omit a/b/c/d to take them from the instruction" + ) + + +@pytest.mark.parametrize("shape,calls", [((16, 16, 32), 1), ((32, 32, 32), 4)]) +def test_bf16_result_tile_returns_a_bf16_fragment(shape, calls) -> None: + # One C subtile, then four, each taking one instruction along K: the f32 accumulator goes in, + # a bf16 result comes out, and it is never fed back as an accumulator. + mma = TileMma( + shape, + a=BF16, + b=BF16, + c=F32, + d=BF16, + target="gfx1250", + tiling=Tiling(atom_shape=(16, 16, 32)), + ) + kernel, result = _wave_tile_kernel(mma, "bf16_result") + assert result.dtype == BF16 + assert result.value.type == result.storage_type + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx1250") + assert len(_calls(ir, "@llvm.amdgcn.wmma.bf16f32.16x16x32.bf16")) == calls + + +def test_bf16_result_tile_with_two_instructions_along_k_is_refused() -> None: + mma = TileMma( + (16, 16, 64), + a=BF16, + b=BF16, + c=F32, + d=BF16, + target="gfx1250", + tiling=Tiling(atom_shape=(16, 16, 32)), + ) + assert mma.subtiles == (1, 1, 2) + b = IRBuilder("bf16_result_k_walk") + a, bf, acc = _operand_fragments(b, mma) + with pytest.raises(ValueError) as excinfo: + mma(b, a, bf, acc) + assert str(excinfo.value) == ( + f"{_BF16_RESULT!r} returns bf16 from an f32 accumulator, so its result cannot be the " + "next K step's accumulator -- this tile issues 2 instructions along K per C subtile; use " + "a tile whose K is the instruction's K (32), or the f32-result instruction (omit d) and " + "convert the finished tile" + ) + + +# ---- WMMA and MFMA tiles issued from the record ----------------------------------------------- + + +@pytest.mark.parametrize( + "target,atom,intrinsic", + [ + ("gfx1250", (16, 16, 32), "@llvm.amdgcn.wmma.f32.16x16x32.f16"), + ("gfx1151", (16, 16, 16), "@llvm.amdgcn.wmma.f32.16x16x16.f16"), + ], +) +def test_f16_wmma_tile_lowers_to_wmma(target, atom, intrinsic) -> None: + # A 32x32x32 wave tile: a grid of WMMA atoms, one call each. + mma = TileMma( + (32, 32, 32), + a=F16, + b=F16, + c=F32, + target=target, + tiling=Tiling(atom_shape=atom), + ) + kernel, _ = _wave_tile_kernel(mma, f"wmma_f16_{target}") + ir = _lower_kernel_to_llvm_python(kernel, arch=target) + assert len(_calls(ir, intrinsic)) == mma.plan.mfma_count + _compile_or_skip(kernel, target) + + +def test_gfx11_f16_result_wmma_tile_lowers_and_compiles() -> None: + # The 16-bit-result instruction keeps an F16 accumulator in its 8 logical slots per lane. + mma = TileMma((16, 16, 16), a=F16, b=F16, c=F16, target="gfx1100") + assert mma.op_id == "wmma_f16_16x16x16_f16_w32" + assert mma.c_desc.register_count == 8 + kernel, result = _wave_tile_kernel(mma, "wmma_f16_result") + assert result.dtype == F16 + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx1100") + assert len(_calls(ir, "@llvm.amdgcn.wmma.f16.16x16x16.f16")) == 1 + _compile_or_skip(kernel, "gfx1100") + + +def test_pinned_fp8_by_bf8_wmma_lowers_to_the_mixed_intrinsic() -> None: + # The pinned mixed instruction issues fp8 x bf8, never the fp8 x fp8 one of the same shape. + mma = TileMma(target="gfx1250", atom_override="wmma_f32_16x16x64_fp8_bf8") + assert (mma.a_dtype, mma.b_dtype) == (FP8E4M3, BF8E5M2) + kernel, _ = _wave_tile_kernel(mma, "wmma_fp8_bf8") + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx1250") + assert len(_calls(ir, "@llvm.amdgcn.wmma.f32.16x16x64.fp8.bf8")) == 1 + assert not _calls(ir, "@llvm.amdgcn.wmma.f32.16x16x64.fp8.fp8") + + +def test_f32_mfma_16x16x4_issues_and_compiles() -> None: + mma = TileMma((16, 16, 4), a=F32, b=F32, c=F32, target="gfx942") + kernel, _ = _wave_tile_kernel(mma, "mfma_f32_16x16x4") + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx942") + assert len(_calls(ir, "@llvm.amdgcn.mfma.f32.16x16x4f32")) == 1 + _compile_or_skip(kernel, "gfx942") diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py index 9df57dc07526..0abd350c539b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operation.py @@ -341,25 +341,3 @@ def test_mma_unsupported_shape_fails_fast_with_candidates() -> None: assert "no MMA instruction found" in str(excinfo.value) assert "(13, 16, 16)" in str(excinfo.value) assert "available on 'gfx90a'" in str(excinfo.value) - - -# The gfx1250 'wmma_bf16f32_16x16x32_bf16' instruction accumulates in fp32 but returns bf16. The -# TileMma driver keeps the result in the accumulator dtype, so it must never pick that instruction, -# and naming it must fail rather than silently issue the fp32-result instruction of the same shape. -_RESULT_DTYPE_CHANGE_MESSAGE = ( - "MMA instruction 'wmma_bf16f32_16x16x32_bf16' accumulates in fp32 but returns bf16 -- the " - "TileMma driver does not drive an instruction whose result dtype differs from its accumulator " - "dtype; choose one whose result dtype is fp32" -) - - -def test_selection_by_intent_never_picks_a_result_dtype_change() -> None: - mma_op = TileMma((16, 16, 32), a=BF16, b=BF16, c=F32, target="gfx1250") - assert mma_op.op_id == "wmma_f32_16x16x32_bf16" - assert mma_op.traits.d_dtype == mma_op.traits.c_dtype - - -def test_atom_override_refuses_a_result_dtype_change() -> None: - with pytest.raises(ValueError) as excinfo: - TileMma(target="gfx1250", atom_override="wmma_bf16f32_16x16x32_bf16") - assert str(excinfo.value) == _RESULT_DTYPE_CHANGE_MESSAGE From 7d2414cebe77d280ad37730aa820bdf91727e0f3 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Tue, 27 Oct 2026 20:46:51 -0400 Subject: [PATCH 5/7] feat(rocke): feed MMA operands from their fragments' storage registers 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 (packed_word_range); 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. 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: 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 gfx1250 tiles of two K steps, four C subtiles and both issue every instruction with a u8 A operand. Co-Authored-By: Claude --- .../dsl_docs/optimization/arch/gfx942.md | 8 + .../helpers/tiling/docs/tiling_api_surface.md | 5 + .../python/rocke/helpers/tiling/mma/driver.py | 123 +++++++--- .../tiling/mma/test_mma_operand_feed.py | 222 ++++++++++++++++++ 4 files changed, 322 insertions(+), 36 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operand_feed.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx942.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx942.md index a4ab2d480222..4ccdcc6bd19d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx942.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/arch/gfx942.md @@ -51,6 +51,14 @@ enforced by `instances/common/gemm_universal.py::is_valid_spec` / **f16 / bf16** atoms `16x16x32` and `32x32x16`. Wide-K on CDNA3 exists for **fp8 / bf8 only** — there is no dense wide-K f16/bf16 path. +**xf32 (TF32) MFMA: why pre-convert?** The xf32 instructions +(`16x16x8`, `32x32x4`) take f32 registers and use only their TF32 bits: +the low 13 mantissa bits of each input are dropped (truncation). An f32 +operand can therefore feed the instruction directly, with no conversion +(the tiling driver accepts an F32 fragment for a TF32 `TileMma`). Convert the operand to TF32 first only when the inputs must be +rounded to nearest even instead of truncated; the converted TF32 fragment +then feeds the instruction unchanged. + **Not present on CDNA3 at all:** block-scale / microscaling (MX) MFMA and FP6 / FP4 — these are CDNA4 (gfx950) additions ([ROCm Matrix-Cores blog](https://rocm.blogs.amd.com/software-tools-optimization/matrix-cores-cdna/README.html)). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md index afd5ceb9a038..ce05e40b396b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_api_surface.md @@ -503,6 +503,11 @@ A/B/C encodings are the **baseline/fallback** layout — a starting place, NOT a hand-roll valid non-canonical (interleaved / custom) layouts (canonical is never required — `mma_is_machinery.md`). Whether the resolved instruction is issued on the target, and with which backend op (`mma.emit_op()`), follows the measured backend coverage record; its classes are described in `traits/README.md` "Backend coverage". +The driver feeds each instruction its A and B operands from the fragments' storage registers (whole +32-bit words for a packed dtype); a TF32 instruction reads TF32 fragments or F32 fragments, whose low +mantissa bits it drops (when to convert to TF32 first: `dsl_docs/optimization/arch/gfx942.md` sec 21.1). +A hand-written `b.mma(mma.emit_op(), ...)` call passes operands of the instruction's operand type +itself, e.g. `b.bitcast(value, VectorType(TF32, 2))` for TF32 registers held as `i32`. ### Default: single MMA, everything resolved from traits ```python diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py index 6907f90f3b30..46e41f5dd44a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py @@ -12,6 +12,10 @@ from __future__ import annotations +from ....core.ir import F32, TF32 +from ....core.storage import packed_word_range, register_packing +from ...mma_io import mma_operand_ir_type +from ..emit import _move_registers from ..fragments import Fragment, TileDesc, _fragment_storage_type, fragment_length from ..register_mapper import RegisterMapper from ..transforms import ( @@ -26,6 +30,11 @@ # (which also takes scale operands) or a sparse row (which also takes a sparsity index) is not driven. _DRIVEN_FAMILIES = frozenset({"dense"}) +# Math dtypes whose instruction also reads a wider fragment dtype, dropping its low mantissa bits: a +# TF32 MMA reads F32 registers as TF32 (truncation). An author who wants round-to-nearest-even +# converts the fragment to TF32 first (the "why pre-convert?" note in the gfx942 arch doc, 21.1). +_MMA_TRUNCATING_SOURCES = {TF32: F32} + def _assert_atom_contiguous( tile_desc: TileDesc, @@ -158,13 +167,41 @@ def plan(self) -> TileMmaPlan: return self._plan @staticmethod - def _read_subvector(b, vec, start: int, length: int, dtype): - """Extract one atom's contiguous register slice ``[start:start+length]`` into a fresh - ```` vector for ``b.mma``.""" - out = b.zero_vec(dtype, length) - for i in range(length): - out = b.vec_insert(out, b.vec_extract(vec, start + i), i) - return out + def _read_subvector(b, vec, start: int, length: int): + """Extract one atom's contiguous register slice ``[start:start+length]`` of ``vec`` into a + fresh ``length``-element vector of its element type (u8 included, through + :func:`~rocke.helpers.tiling.emit._move_registers`) for ``b.mma``.""" + return _move_registers(b, vec, [(start + i, i) for i in range(length)], length) + + @staticmethod + def _as_operand(b, value, math_dtype, count: int): + """Return ``value`` (``count`` logical elements' registers) as the vector type ``b.mma`` + takes for an operand of ``math_dtype`` (:func:`~rocke.helpers.mma_io.mma_operand_ir_type`), + bitcasting only when the types differ: TF32 storage ``i32`` and truncated F32 registers + become ``vec``; every other operand, packed words included, is already that type. + """ + operand_type = mma_operand_ir_type(math_dtype.name, count) + if value.type == operand_type: + return value + return b.bitcast(value, operand_type) + + def _read_operand_atom( + self, b, fragment: Fragment, start: int, length: int, math_dtype + ): + """Return one instruction's A or B operand: logical elements ``[start, start+length)`` of + ``fragment``, sliced from its storage registers -- the 32-bit words that hold them + (:func:`~rocke.core.storage.packed_word_range`) for a packed dtype, the elements themselves + otherwise -- then given the operand type of ``math_dtype`` (:meth:`_as_operand`). + """ + packing = register_packing( + fragment.dtype.name, fragment.tile_desc.register_count + ) + if packing is not None: + words = packed_word_range(packing, start, length) + sub = self._read_subvector(b, fragment.value, words.start, len(words)) + else: + sub = self._read_subvector(b, fragment.value, start, length) + return self._as_operand(b, sub, math_dtype, length) @staticmethod def _write_subvector(b, vec, sub, start: int, length: int): @@ -209,6 +246,11 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): ``b.mma`` per atom and accumulating each C subtile. The fragments are subtile-contiguous (from the wave layouts), so every atom is a register slice. + Each A and B operand is fed from the fragment's storage registers: an atom's slice is the + 32-bit words that hold it for a packed dtype, its logical elements otherwise, then given + the vector type ``b.mma`` takes for the plan's A or B dtype. A TF32 instruction reads a + TF32 fragment (held as ``i32``) or an F32 fragment, whose low mantissa bits it drops. + Before any IR it checks, in order: the operand dtypes; that the instruction's family is one the driver feeds (A, B and C only); that a result of another dtype than the accumulator is never fed back as the next K step's accumulator; that rocke issues the @@ -224,17 +266,24 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): must be assigned (``acc = driver(b, a_frag, b_frag, acc)``). """ plan = self._plan - for name, fragment in ( - ("A", a_fragment), - ("B", b_fragment), - ("C", accumulator), + for name, fragment, want in ( + ("A", a_fragment, plan.a_dtype), + ("B", b_fragment, plan.b_dtype), + ("C", accumulator, plan.c_dtype), ): - want = {"A": plan.a_dtype, "B": plan.b_dtype, "C": plan.c_dtype}[name] - if fragment.dtype != want: - raise ValueError( - f"MMA operand dtype mismatch -- operand={name}, " - f"fragment={fragment.dtype.name!r}, expected {want.name!r}" + source = _MMA_TRUNCATING_SOURCES.get(want) if name != "C" else None + if fragment.dtype == want or fragment.dtype == source: + continue + expected = want.name + if source is not None: + expected += ( + f" or {source.name} ({want.name} math also reads {source.name} " + f"fragments, dropping their low mantissa bits)" ) + raise ValueError( + f"MMA operand dtype mismatch -- operand={name}, " + f"fragment={fragment.dtype.name}, expected {expected}" + ) traits = plan.traits if traits.family not in _DRIVEN_FAMILIES: @@ -266,10 +315,22 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): if m_sub == 1 and n_sub == 1: acc_value = accumulator.value if k_sub == 1: + a_op = self._as_operand( + b, + a_fragment.value, + plan.a_dtype, + a_fragment.tile_desc.register_count, + ) + b_op = self._as_operand( + b, + b_fragment.value, + plan.b_dtype, + b_fragment.tile_desc.register_count, + ) return Fragment( accumulator.tile_desc, plan.d_dtype, - b.mma(op, a_fragment.value, b_fragment.value, acc_value), + b.mma(op, a_op, b_op, acc_value), ) a_atom = fragment_length(a_fragment.tile_desc.layout) // k_sub b_atom = fragment_length(b_fragment.tile_desc.layout) // k_sub @@ -277,11 +338,11 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): plan.tiling.mac_prio ) # raise AFTER the first atom (see grid-branch note below) for ki in range(k_sub): - a_sub = self._read_subvector( - b, a_fragment.value, ki * a_atom, a_atom, a_fragment.dtype + a_sub = self._read_operand_atom( + b, a_fragment, ki * a_atom, a_atom, plan.a_dtype ) - b_sub = self._read_subvector( - b, b_fragment.value, ki * b_atom, b_atom, b_fragment.dtype + b_sub = self._read_operand_atom( + b, b_fragment, ki * b_atom, b_atom, plan.b_dtype ) acc_value = b.mma(op, a_sub, b_sub, acc_value) if mac_prio and ki == 0: @@ -301,9 +362,7 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): b_atom = fragment_length(b_fragment.tile_desc.layout) // (n_sub * k_sub) c_atom = fragment_length(accumulator.tile_desc.layout) // (m_sub * n_sub) accs = [ - self._read_subvector( - b, accumulator.value, idx * c_atom, c_atom, accumulator.dtype - ) + self._read_subvector(b, accumulator.value, idx * c_atom, c_atom) for idx in range(m_sub * n_sub) ] # `mac_prio` raises wave issue priority for the matrix-dense body (MFMA or WMMA -- the driver is @@ -315,19 +374,11 @@ def __call__(self, b, a_fragment, b_fragment, accumulator): mac_prio = plan.tiling.mac_prio for i, (mi, nj, ki) in enumerate(self._subtile_triples()): idx = mi * n_sub + nj - a_sub = self._read_subvector( - b, - a_fragment.value, - (mi * k_sub + ki) * a_atom, - a_atom, - a_fragment.dtype, + a_sub = self._read_operand_atom( + b, a_fragment, (mi * k_sub + ki) * a_atom, a_atom, plan.a_dtype ) - b_sub = self._read_subvector( - b, - b_fragment.value, - (nj * k_sub + ki) * b_atom, - b_atom, - b_fragment.dtype, + b_sub = self._read_operand_atom( + b, b_fragment, (nj * k_sub + ki) * b_atom, b_atom, plan.b_dtype ) accs[idx] = b.mma(op, a_sub, b_sub, accs[idx]) if mac_prio and i == 0: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operand_feed.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operand_feed.py new file mode 100644 index 000000000000..a536f9d53bdd --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_operand_feed.py @@ -0,0 +1,222 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""How the TileMma driver feeds A and B operands from their fragments' storage registers. + +A fragment holds its registers in its storage type: TF32 as ``i32``, packed sub-byte dtypes (fp4, +fp6, i4) as 32-bit words, every other dtype as its logical vector. The driver slices each +instruction's operand out of those registers (whole words for a packed dtype) and hands ``b.mma`` +the vector type the instruction's A or B dtype takes. A TF32 instruction also reads F32 fragments, +dropping their low mantissa bits, so an F32 tensor feeds it with no conversion. A u8 operand is sliced +as the i8 registers of the same bits and handed over as u8. + +Lowering tests pin the Python engine (``_lower_kernel_to_llvm_python``); the C++ engine is checked by +the byte-identity gate. Where the host clang targets the arch, the module is also compiled and linked +(compile-verified, never run). +""" + +from __future__ import annotations + +import pytest + +from rocke.core.ir import ( + F16, + F32, + FP4E2M1, + I8, + I32, + TF32, + U8, + IRBuilder, + PtrType, + VectorType, +) +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.helpers.tiling import ( + Fragment, + fill_fragment, + load_fragment, + make_fragment, + make_tensor_desc, + make_window, + store_fragment, +) +from rocke.helpers.tiling.mma import TileMma, Tiling + +from .test_mma_issue import _calls, _compile_or_skip + + +def _wave_tile_kernel(mma: TileMma, tensor_dtype, pointer_dtype): + """One wave computes C = A * B^T for the MMA's wave tile, with A (M, K) and B (N, K) read from + global tensors of ``tensor_dtype`` behind ``pointer_dtype`` pointers and the f32 result stored. + Returns the kernel and the A fragment the driver was given.""" + wave_m, wave_n, wave_k = mma.shape + b = IRBuilder("operand_feed") + b.kernel.attrs["max_workgroup_size"] = mma.wave_size + a_ptr = b.param("A", PtrType(pointer_dtype, "global"), noalias=True, readonly=True) + b_ptr = b.param("B", PtrType(pointer_dtype, "global"), noalias=True, readonly=True) + c_ptr = b.param("C", PtrType(F32, "global"), noalias=True, writeonly=True) + lane = b.thread_id_x() + zero = b.const_i32(0) + a_td = make_tensor_desc((wave_m, wave_k), (wave_k, 1), tensor_dtype) + b_td = make_tensor_desc((wave_n, wave_k), (wave_k, 1), tensor_dtype) + c_td = make_tensor_desc((wave_m, wave_n), (wave_n, 1), F32) + acc = make_fragment(mma.c_desc, F32) + fill_fragment(b, acc, 0) + fa = load_fragment(b, a_ptr, make_window(a_td, (zero, zero)), mma.a_desc, lane) + fb = load_fragment(b, b_ptr, make_window(b_td, (zero, zero)), mma.b_desc, lane) + result = mma(b, fa, fb, acc) + store_fragment(b, c_ptr, make_window(c_td, (zero, zero)), result, lane) + b.ret() + return b.kernel, fa + + +def _tf32_mma(shape, atom): + tiling = Tiling(atom_shape=atom) if atom is not None else None + return TileMma(shape, a=TF32, b=TF32, c=F32, target="gfx942", tiling=tiling) + + +# The xf32 tiles: (wave tile, atom hint, instructions issued). One atom, a 2 x 2 grid of C +# subtiles, a 4-step K walk, and a 2 x 2 x 2 grid on the 32x32x4 instruction. +_XF32_TILES = [ + pytest.param((16, 16, 8), None, 1, id="16x16x8"), + pytest.param((32, 32, 8), (16, 16, 8), 4, id="32x32x8-atom16x16x8"), + pytest.param((16, 16, 32), (16, 16, 8), 4, id="16x16x32-atom16x16x8"), + pytest.param((64, 64, 8), (32, 32, 4), 8, id="64x64x8-atom32x32x4"), +] + + +@pytest.mark.parametrize("shape,atom,calls", _XF32_TILES) +@pytest.mark.parametrize( + "tensor_dtype", + [ + pytest.param(F32, id="from-f32-tensors"), + pytest.param(TF32, id="from-tf32-tensors"), + ], +) +def test_tf32_tile_issues_one_xf32_mfma_per_atom( + shape, atom, calls, tensor_dtype +) -> None: + # F32 fragments feed the TF32 instruction directly (it truncates); TF32 fragments, held as i32, + # are reinterpreted as tf32. Either way every atom is one xf32 MFMA. + mma = _tf32_mma(shape, atom) + kernel, fa = _wave_tile_kernel(mma, tensor_dtype, tensor_dtype) + assert fa.dtype == tensor_dtype + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx942") + assert len(_calls(ir, mma.traits.llvm_intrinsic)) == calls + _compile_or_skip(kernel, "gfx942") + + +def test_tf32_tile_refuses_an_f16_fragment() -> None: + mma = _tf32_mma((16, 16, 8), None) + b = IRBuilder("tf32_f16_operand") + fa = make_fragment(mma.a_desc, F16) + fb = make_fragment(mma.b_desc, TF32) + acc = make_fragment(mma.c_desc, F32) + for fragment in (fa, fb, acc): + fill_fragment(b, fragment, 0) + with pytest.raises(ValueError) as err: + mma(b, fa, fb, acc) + assert str(err.value) == ( + "MMA operand dtype mismatch -- operand=A, fragment=f16, expected tf32 or f32 (tf32 math " + "also reads f32 fragments, dropping their low mantissa bits)" + ) + + +def test_operand_without_a_truncating_source_names_only_its_dtype() -> None: + mma = TileMma((16, 16, 16), a=F16, b=F16, c=F32, target="gfx90a") + b = IRBuilder("f16_f32_operand") + fa = make_fragment(mma.a_desc, F16) + fb = make_fragment(mma.b_desc, F32) + acc = make_fragment(mma.c_desc, F32) + for fragment in (fa, fb, acc): + fill_fragment(b, fragment, 0) + with pytest.raises(ValueError) as err: + mma(b, fa, fb, acc) + assert str(err.value) == ( + "MMA operand dtype mismatch -- operand=B, fragment=f32, expected f16" + ) + + +def _extract_indices(kernel, vector) -> list[int]: + """The register indices read out of ``vector``, in program order.""" + return [ + op.attrs["index"] + for op in kernel.body.ops + if op.name == "vector.extract" and op.operands[0] is vector + ] + + +def test_packed_tile_feeds_each_instruction_its_own_words() -> None: + # 256 fp4 K elements over two 16x16x128 instructions: each lane's A fragment is 64 elements in 8 + # words, and the instructions take words [0, 4) and [4, 8), whole, with no repacking. + mma = TileMma( + (16, 16, 256), + a=FP4E2M1, + b=FP4E2M1, + c=F32, + target="gfx950", + tiling=Tiling(atom_shape=(16, 16, 128)), + ) + kernel, fa = _wave_tile_kernel(mma, FP4E2M1, I8) + assert fa.storage_type == VectorType(I32, 8) + assert _extract_indices(kernel, fa.value) == list(range(8)) + mma_ops = [op for op in kernel.body.ops if op.name == "tile.mma"] + assert len(mma_ops) == 2 + assert all(op.operands[0].type == VectorType(I32, 4) for op in mma_ops) + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx950") + assert len(_calls(ir, mma.traits.llvm_intrinsic)) == 2 + _compile_or_skip(kernel, "gfx950") + + +def _unsigned_tile_kernel(mma: TileMma): + """One wave computes C = A * B^T for a u8 x i8 wave tile with an i32 result. A and B are read + from i8 tensors: load_fragment refuses u8, so A's registers are loaded as i8 and reinterpreted as + u8, the same bits.""" + wave_m, wave_n, wave_k = mma.shape + b = IRBuilder("unsigned_operand_feed") + b.kernel.attrs["max_workgroup_size"] = mma.wave_size + a_ptr = b.param("A", PtrType(I8, "global"), noalias=True, readonly=True) + b_ptr = b.param("B", PtrType(I8, "global"), noalias=True, readonly=True) + c_ptr = b.param("C", PtrType(I32, "global"), noalias=True, writeonly=True) + lane = b.thread_id_x() + zero = b.const_i32(0) + a_td = make_tensor_desc((wave_m, wave_k), (wave_k, 1), I8) + b_td = make_tensor_desc((wave_n, wave_k), (wave_k, 1), I8) + c_td = make_tensor_desc((wave_m, wave_n), (wave_n, 1), I32) + acc = make_fragment(mma.c_desc, I32) + fill_fragment(b, acc, 0) + a_i8 = load_fragment(b, a_ptr, make_window(a_td, (zero, zero)), mma.a_desc, lane) + fa = Fragment( + mma.a_desc, U8, b.bitcast(a_i8.value, VectorType(U8, a_i8.storage_type.count)) + ) + fb = load_fragment(b, b_ptr, make_window(b_td, (zero, zero)), mma.b_desc, lane) + result = mma(b, fa, fb, acc) + store_fragment(b, c_ptr, make_window(c_td, (zero, zero)), result, lane) + b.ret() + return b.kernel + + +# u8 x i8 tiles of the gfx1250 16x16x64 instruction: (wave tile, instructions issued). Two steps +# along K, then a 2 x 2 grid of C subtiles, then both. +_UNSIGNED_TILES = [ + pytest.param((16, 16, 128), 2, id="16x16x128-two-k-steps"), + pytest.param((32, 32, 64), 4, id="32x32x64-four-c-subtiles"), + pytest.param((32, 32, 128), 8, id="32x32x128-both"), +] + + +@pytest.mark.parametrize("shape,calls", _UNSIGNED_TILES) +def test_unsigned_tile_of_several_instructions_issues_each_one(shape, calls) -> None: + # Every instruction takes a u8 A operand sliced from the u8 fragment. + mma = TileMma( + shape, a=U8, b=I8, target="gfx1250", atom_override="wmma_i32_16x16x64_iu8" + ) + kernel = _unsigned_tile_kernel(mma) + mma_ops = [op for op in kernel.body.ops if op.name == "tile.mma"] + assert len(mma_ops) == calls + assert all(op.operands[0].type == VectorType(U8, 32) for op in mma_ops) + assert all(op.operands[1].type == VectorType(I8, 32) for op in mma_ops) + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx1250") + assert len(_calls(ir, mma.traits.llvm_intrinsic)) == calls + _compile_or_skip(kernel, "gfx1250") From 4476e07402570c916ac8f530484b77698adb18bc Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Tue, 27 Oct 2026 21:33:38 -0400 Subject: [PATCH 6/7] fix(rocke): pair A and B K through the instruction's canonical coupling 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 the target, the driven gfx950 fp8 x fp4 tile and that gfx1250 tile also run on random K-dependent operands and must equal an independent numpy product exactly; elsewhere they skip. 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. Co-Authored-By: Claude --- .../tiling/.claude/skills/layout-viz/SKILL.md | 3 +- .../helpers/tiling/analysis/soundness.py | 7 +- .../helpers/tiling/docs/mma_is_machinery.md | 28 +- .../tiling/docs/tiling_interleaving_design.md | 10 +- .../rocke/helpers/tiling/layout_optimizer.py | 24 +- .../python/rocke/helpers/tiling/mma/driver.py | 32 +- .../rocke/helpers/tiling/transforms/_core.py | 120 +++++- .../helpers/tiling/transforms/observers.py | 62 +++- .../tiling/visualization/layout_render.py | 83 ++--- .../tiling/mma/test_mma_issuable_pairs.py | 128 +++++++ .../helpers/tiling/mma/test_mma_k_coupling.py | 351 ++++++++++++++++++ .../tiling/mma/test_mma_validation_recipe.py | 4 +- .../tiling/test_interleave_all_dense.py | 32 +- .../helpers/tiling/test_layout_optimizer.py | 4 +- .../tests/helpers/tiling/test_transforms.py | 32 +- .../rocke/style/API_DESIGN_STANDARD.md | 2 +- 16 files changed, 828 insertions(+), 94 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issuable_pairs.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_k_coupling.py diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md index 35233a4757cc..41724092bb37 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/.claude/skills/layout-viz/SKILL.md @@ -280,7 +280,8 @@ cooperative map; `classify_epilogue(...)` → the branch. Worked example: - `mma_operand_layout_sound(layout, canon, role=)` → `ok` / `error`. `mma_operand_repair_hint(layout, canon)` → `ok` / `warning` (an in-register reorder makes-it-so — name the permutation) / `error` (cross-lane or none). `mma_pair_compatible(a, b, a_canon=, b_canon=)` → both sound AND K-dists match; drives the tee - banner. `mma_pair_k_aligned(a, b)` → `ok`/`warning`/`error` (pairwise K-match, per atom). + banner. `mma_pair_k_aligned(a, b, a_canon=, b_canon=)` → `ok`/`warning`/`error` (pairwise K-match, per + atom, at the positions the canonicals couple). `classify_transform(src, tgt)` → `reorder` (cheap, dtype-graded) vs `cross_lane` (expensive/deferred). - Tee banner red ERROR / orange WARNING = the pair diagnostic — say what's wrong, which rule, cheapest remedy (free symmetry [reposition/source-swap] → dword reorder → sub-dword → cross-lane last). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/soundness.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/soundness.py index 725fbd544507..a09c2172f9b0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/soundness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/analysis/soundness.py @@ -18,10 +18,11 @@ class MmaSoundnessError(RuntimeError): def verify_mma_soundness(pipeline: Any) -> int: """Gate 2: every recorded MMA's CONSUMED operand encodings must be sound MMA operands against the - canonical machine (``mma_operand_layout_sound``) AND share a K-distribution (``mma_pair_k_aligned``) -- bundled - as ``mma_pair_compatible``. This is the operand-correctness the addressing round-trip is BLIND to: a + canonical machine (``mma_operand_layout_sound``) AND hold the same K at every pair of positions the + instruction couples (``mma_pair_k_aligned``; slot for slot on a same-format instruction, through the + canonical layouts on a mixed-width one) -- bundled as ``mma_pair_compatible``. This is the operand-correctness the addressing round-trip is BLIND to: a scrambled/duplicated operand K passes the address gate but is caught here. ``a_canon``/``b_canon`` are - the trusted canonical layouts from the MMA definition; ``a_enc``/``b_enc`` are the kernel's own + the trusted canonical layouts from the MMA definition (they also define the A/B coupling); ``a_enc``/``b_enc`` are the kernel's own (interleaved) operands. Returns the count verified; raises on the first unsound MMA. Correctness SOT: ``docs/mma_is_machinery.md`` (the three-condition sound MAC). """ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/mma_is_machinery.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/mma_is_machinery.md index 19367f09e786..0859760b5455 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/mma_is_machinery.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/mma_is_machinery.md @@ -42,7 +42,25 @@ Everything else follows from that one physical fact. Because the atom *physicall 1. **Atom-fixed accumulation.** Which physical slots sum into each output is fixed by the atom — not a choice. 2. **One $M$ per output on A, one $N$ per output on B.** Every A-slot feeding a given output carries a single fixed $M$; every B-slot a single fixed $N$. If $M$ drifts across those slots, the accumulator holds a sum whose row wanders term to term — not $C_{ij}$ for any $(i,j)$. -3. **$K$-match, position for position.** $A.K\text{-dist} = B.K\text{-dist}$: $A$ and $B$ agree on which contraction index $K$ sits in each paired slot, and $K$ is well-formed (each index present once). +3. **$K$-match at coupled positions.** $A$ and $B$ agree on which contraction index $K$ sits at each pair of positions the atom multiplies together, and $K$ is well-formed (each index present once). The atom's canonical A and B layouts define that coupling: an A position and a B position are coupled when they hold the same hardware $K$. For every same-format atom this is slot $s$ with slot $s$, so the rule reads $A.K\text{-dist} = B.K\text{-dist}$. A mixed-width `f8f6f4` atom is the exception: its 8-bit operand holds $K$ in two runs of 16 per lane, its 6/4-bit operand in one run of 32, so the coupled positions differ: + + ``` + gfx950 16x16x128, fp8 A x fp4 B, lane 0 + A slots 0-15 = K 0-15, slots 16-31 = K 64-79 (two runs of 16) + B slots 0-31 = K 0-31 (one run of 32) + A (lane 0, slot 16) holds K 64 <-> B (lane 32, slot 0) holds K 64 + ``` + + Comparing such a pair slot for slot is wrong both ways: it rejects the canonical fp8 A x fp4 B pair, and it accepts an fp4 B loaded with fp8's map, which multiplies the wrong $K$. + + An atom whose lane holds several $M$ rows of A (`a_m_register_blocks > 1`) multiplies each row with the same B registers, so every $M$ block of A is coupled with B slot for slot: + + ``` + gfx1250 32x16x128 fp4, lane 0 + A slots 0-63 = (M 0, K 0-31 and 64-95), slots 64-127 = (M 16, the same K) + B slots 0-63 = (N 0, K 0-31 and 64-95) + A (lane 0, slot 64) holds K 0 <-> B (lane 0, slot 0) holds K 0 + ``` > The slots feeding one accumulator are fixed by silicon; a sound MAC labels them with **one $M$, one $N$, and a matched $K$** — nothing else is constrained. @@ -89,18 +107,18 @@ Because A's $M$-placement and B's $N$-placement are free, *they* fix C's native Two checks hide under "is this layout OK?", and conflating them is what sends people reaching for canonical: - **Per-operand soundness** (conditions 1–2) is checkable on *one* fragment: does every slot feeding an output carry a single $M$ (A) or $N$ (B)? It asks nothing about the partner, and nothing about canonical. -- **Pairwise acceptability** (condition 3) is a property of the **pair**, never of one fragment: do A and B agree on $K$, slot for slot? Neither operand answers it alone — and neither is checked against a canonical reference. +- **Pairwise acceptability** (condition 3) is a property of the **pair**, never of one fragment: do A and B agree on $K$ at the positions the atom couples (slot for slot on a same-format atom)? Neither operand answers it alone. The canonical layouts enter only to say *which* positions are coupled; neither operand's labels are compared against them. Consequences: - "Make this fragment acceptable" always means "reorder it until its $K$-row matches its **partner's**" — never "until it looks canonical." -- `mma_pair_k_aligned(A, B)` compares the two fragments **to each other**; it never checks either against a canonical reference. +- `mma_pair_k_aligned(A, B, a_canon=..., b_canon=...)` compares the two fragments **to each other** at coupled positions. The canonicals are required because they define the coupling; neither fragment's labels are checked against them. ## The traps agents fall into (name them, then don't) Wrong-turns that come from pattern-matching on training data instead of reasoning from the machinery: -1. **"The MFMA has a fixed layout, so the fragment must be canonical."** — No. The hardware reads slots and multiply-accumulates; the layout is bookkeeping *we* maintain. $K$-pairing is the only **cross-operand** constraint — not canonical resemblance. +1. **"The MFMA has a fixed layout, so the fragment must be canonical."** — No. The hardware reads slots and multiply-accumulates the positions its canonical coupling pairs (slot $s$ with slot $s$, except on mixed-width `f8f6f4` atoms); the labels are bookkeeping *we* maintain. $K$-pairing at coupled positions is the only **cross-operand** constraint — not canonical resemblance. 2. **"A lane can only own these $N$ strided elements, so we're capped at dwordx2."** — Reasoning from *one* canonical distribution. $M/N$ placement is free; choose a distribution where the lane owns a contiguous run and the wide store exists. "Canonical caps us" is never a valid conclusion. 3. **"Just shuffle the canonical C into the store order."** — C's ownership is derived from A/B. The store-friendly C is *produced* by choosing A/B placement, not recovered by post-hoc shuffling a canonical accumulator (often cross-lane and expensive). 4. **Treating `mma_pair_k_aligned` passing as "canonical-valid."** — It only means A and B agree on $K$. A hand-rolled distribution's correctness is confirmed by **bit-exact**, not by resemblance to canonical. @@ -110,7 +128,7 @@ Wrong-turns that come from pattern-matching on training data instead of reasonin For any layout/direction question about an MMA operand or accumulator: -1. **Is this a correctness question?** Correctness = the sound MAC: per-operand $M$/$N$ fixed (1–2) **and** $A.K\text{-dist} = B.K\text{-dist}$ (3). If those hold, the MMA is correct — full stop. +1. **Is this a correctness question?** Correctness = the sound MAC: per-operand $M$/$N$ fixed (1–2) **and** A and B hold the same $K$ at coupled positions (3). If those hold, the MMA is correct — full stop. 2. **Place $M/N$ for the goal, not for canon.** Choose A's $M$-placement and B's $N$-placement for the memory behavior you want (coalesced wide load, store-friendly C, feed a downstream MMA). 3. **Let C be derived.** C emerges at the $M/N$ you chose; shape the *store* around that (or a cheap register reorder *within* the already-derived ownership), not by reshaping a canonical C. 4. **Verify by bit-exact.** The GPU is the arbiter. "It doesn't look canonical" is not evidence of anything. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_interleaving_design.md b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_interleaving_design.md index 31d89e2a9c10..70b31c97db6c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_interleaving_design.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/docs/tiling_interleaving_design.md @@ -216,8 +216,9 @@ is valid provided A holds `A[m, σ(s)]` and B holds `B[σ(s), n]` in the same sl - A's M-register order and B's N-register order are free; together they **determine the C output layout** (§7). -`mma_pair_k_aligned(a_layout, b_layout)` enforces exactly this: `k_distribution(a) == k_distribution(b)` -positionally, **compared per atom** — not across the concatenated whole-fragment K-list (§8 for the +`mma_pair_k_aligned(a_layout, b_layout, a_canon=, b_canon=)` enforces exactly this: `k_distribution(a) == +k_distribution(b)` positionally (at the positions the canonicals couple -- slot for slot on every same-format +atom; see `mma_is_machinery.md` condition 3 for mixed-width `f8f6f4`), **compared per atom** — not across the concatenated whole-fragment K-list (§8 for the rectangular-tile rationale). Lane ownership is guaranteed valid by construction (fragments come from atom-derived descs), so a positional per-atom K-match is sufficient. @@ -401,7 +402,8 @@ There is no "this fragment is MMA-ready" in isolation. The hardware's only requi (M on A / N on B are free) — worked in full at §2's slot tables. So "make this fragment acceptable" always means "reorder it until its K-row matches its **partner's**". -`mma_pair_k_aligned(a, b)` compares the two fragments; it never checks either against a canonical. +`mma_pair_k_aligned(a, b, a_canon=, b_canon=)` compares the two fragments; the canonicals only say which +positions are coupled, and neither fragment is checked against them. ### A, B, C have different roles @@ -734,7 +736,7 @@ multi-patch (`P>1`) and novel-atom cases are discovered.) These change with ever error-prone to transcribe; construct the two states and ask the solver. The optimizer (`helpers/tiling/layout_optimizer.py`) wraps this into the two questions a designer actually asks: -- **`evaluate_transform(source, target, *, canon=…, k_partner=…, dtype_bits=…)`** — is there a VALID path +- **`evaluate_transform(source, target, *, canon=…, k_partner=…, k_partner_canon=…, dtype_bits=…)`** — is there a VALID path (an MMA-sound target via `mma_operand_layout_sound` + a transform via `classify_transform`), and what does the cheapest edge cost? - **`optimize_layout(source, candidates: dict, …)`** — the MINIMIZATION: sweep candidate distributions diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/layout_optimizer.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/layout_optimizer.py index fb2f1b21c66b..04c067c89c68 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/layout_optimizer.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/layout_optimizer.py @@ -36,7 +36,8 @@ edges (relabel / reorder / identity) that beat both are trustworthy winners and DO decide outright. Validity is rocKE's own gates: :func:`mma_operand_layout_sound` (per-operand sound MMA operand), optionally -the pairwise :func:`mma_pair_k_aligned` (``A.K-dist == B.K-dist`` -- pass ``k_partner``), and the tier from +the pairwise :func:`mma_pair_k_aligned` (A and B hold the same K at the positions the instruction couples -- +pass ``k_partner`` with both canonical layouts), and the tier from :func:`classify_transform`. No reference tables, no hand formulas -- read off the actual layouts, so it is correct for any context (atom / wave / #atoms / tile shape / dtype). """ @@ -199,6 +200,7 @@ def evaluate_transform( *, canon=None, k_partner=None, + k_partner_canon=None, dtype_bits: int = 16, through_lds: bool = False, lds_conflict_cost=None, @@ -209,18 +211,30 @@ def evaluate_transform( ``source``/``target`` are ``WarpDistributionEncoding`` or forward maps. Validity gates (all optional, all rocKE's own solvers): - ``canon`` -- the atom's canonical operand ref -> **per-operand** soundness (:func:`mma_operand_layout_sound`). - - ``k_partner`` -- the OTHER operand's target layout -> the **pairwise** ``A.K-dist == B.K-dist`` half of - the sound MAC (:func:`mma_pair_k_aligned`). Without it, that half is the caller's responsibility. + - ``k_partner`` -- the OTHER operand's target layout -> the **pairwise** half of the sound MAC + (:func:`mma_pair_k_aligned`): the target and ``k_partner`` hold the same K at every pair of positions + the instruction couples. The coupling comes from the two canonical layouts, so ``k_partner`` needs + ``canon`` (the target's) and ``k_partner_canon`` (the partner's); without either it raises + ``ValueError``. Without ``k_partner``, that half is the caller's responsibility. Cost knobs: ``through_lds`` allows an LDS reposition (pass ``lds_conflict_cost=(store_bc, read_bc)`` for its true cost; unset = flagged lower bound). With NO ``canon`` and NO ``k_partner`` the result is transform-cost-ONLY and does not imply a valid MMA. """ + if k_partner is not None and (canon is None or k_partner_canon is None): + raise ValueError( + "k_partner needs canon= and k_partner_canon= -- A and B K positions pair through the " + "instruction's canonical layouts, which define the coupling" + ) src, tgt = as_forward_map(source), as_forward_map(target) sound = ( mma_operand_layout_sound(tgt, canon).severity if canon is not None else "n/a" ) k_match = ( - mma_pair_k_aligned(tgt, k_partner).severity if k_partner is not None else "n/a" + mma_pair_k_aligned( + tgt, k_partner, a_canon=canon, b_canon=k_partner_canon + ).severity + if k_partner is not None + else "n/a" ) works = (canon is None or sound == "ok") and (k_partner is None or k_match == "ok") valid_note = f"sound={sound}, K-match={k_match}" @@ -301,6 +315,7 @@ def optimize_layout( *, canon=None, k_partner=None, + k_partner_canon=None, dtype_bits: int = 16, through_lds: bool = False, lds_conflict_cost=None, @@ -322,6 +337,7 @@ def optimize_layout( tgt, canon=canon, k_partner=k_partner, + k_partner_canon=k_partner_canon, dtype_bits=dtype_bits, through_lds=through_lds, lds_conflict_cost=lds_conflict_cost, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py index 46e41f5dd44a..ae6be15d7b8e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/mma/driver.py @@ -42,6 +42,7 @@ def _assert_atom_contiguous( atom_k: int, free_sub: int, k_sub: int, + free_blocks: int, role: str, op_id: str, ) -> None: @@ -58,6 +59,13 @@ def _assert_atom_contiguous( stride (canonical strides atoms by ``atom_free``, interleaved by 1) and does not assume lane 0 distinguishes atoms (some interleavings share a lane-0 coordinate across atoms), so it must sweep all lanes -- AOS (atom-iteration inner) then shows a block whose atom identity varies at some lane. + + ``free_blocks`` is the number of free rows one lane holds per atom (``a_m_register_blocks`` for A, 1 + for B). Such a lane's atom block is ``free_blocks`` equal runs, one free row each, all in one K-atom:: + + gfx1250 32x16x128 fp4 A, lane 0, one atom of 128 registers: + registers 0-63 (M 0, K 0-31 and 64-95) + registers 64-127 (M 16, K 0-31 and 64-95) """ rm = RegisterMapper(tile_desc.layout) nregs = rm.num_vector_items @@ -68,12 +76,20 @@ def _assert_atom_contiguous( f"free_sub*k_sub = {free_sub}*{k_sub} for {op_id!r}" ) atom_len = nregs // denom + if atom_len % free_blocks: + raise ValueError( + f"MMA {role} operand atom of {atom_len} registers does not split into the " + f"{free_blocks} free rows a lane holds per atom for {op_id!r}" + ) + run_len = atom_len // free_blocks for blk in range(denom): for lane in range(rm.num_lanes): key: tuple[int, int] | None = None for j in range(atom_len): free_c, k_c = rm.matrix_coordinates(lane, blk * atom_len + j) here = (free_c, k_c // atom_k) + if j % run_len == 0 and key is not None and here[1] == key[1]: + key = here # the next free row of this atom, in the same K-atom if key is None: key = here elif here != key: @@ -93,8 +109,9 @@ def _validate_mma_issue(a_fragment, b_fragment, accumulator, plan: TileMmaPlan) - ``mma_operand_layout_sound`` on A and B -- each a well-formed MMA operand (one M/N per output row, well-formed K). Run FIRST and short-circuited, because the accumulator check below flows the operands through the machine and is undefined on an unsound operand. - - ``mma_pair_k_aligned`` -- A and B agree on which K sits in each paired slot, compared per atom (the - free-dim atom counts come from ``plan.subtiles``). + - ``mma_pair_k_aligned`` -- A and B agree on which K sits at each pair of positions the instruction + couples (read off the canonical layouts; slot for slot on a same-format instruction), compared per + atom (the free-dim atom counts come from ``plan.subtiles``). - ``mma_accumulator_flow_consistent`` -- the passed accumulator C carries the exact labels the machine produces from these operands, so the store writes the right coordinates. The absence of this check is what let a mislabeled C store to the wrong place. @@ -117,7 +134,14 @@ def _validate_mma_issue(a_fragment, b_fragment, accumulator, plan: TileMmaPlan) f"MMA {role} operand not sound for {plan.op_id!r} -- {d.message}" ) - d = mma_pair_k_aligned(a_lay, b_lay, a_free_atoms=m_sub, b_free_atoms=n_sub) + d = mma_pair_k_aligned( + a_lay, + b_lay, + a_canon=a_canon, + b_canon=b_canon, + a_free_atoms=m_sub, + b_free_atoms=n_sub, + ) if d.severity != "ok": raise ValueError( f"MMA operands not K-aligned for {plan.op_id!r} -- {d.message}" @@ -142,6 +166,7 @@ def _validate_mma_issue(a_fragment, b_fragment, accumulator, plan: TileMmaPlan) atom_k=atom_k, free_sub=m_sub, k_sub=k_sub, + free_blocks=plan.traits.layout_params.a_m_register_blocks, role="A", op_id=plan.op_id, ) @@ -150,6 +175,7 @@ def _validate_mma_issue(a_fragment, b_fragment, accumulator, plan: TileMmaPlan) atom_k=atom_k, free_sub=n_sub, k_sub=k_sub, + free_blocks=1, role="B", op_id=plan.op_id, ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/_core.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/_core.py index 1c131215e2f0..b409482f98d8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/_core.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/_core.py @@ -91,8 +91,9 @@ def k_distribution( ) -> tuple[tuple[int, ...], ...]: """Project the encoding onto its K axis: per lane, the tuple of K coordinates by register slot. - ``k_axis`` is the contraction axis index (1 for both A=(M,K) and B=(N,K)). Two operands are - K-aligned iff their ``k_distribution`` are equal position-for-position. + ``k_axis`` is the contraction axis index (1 for both A=(M,K) and B=(N,K)). For a same-format + instruction two operands are K-aligned iff their ``k_distribution`` are equal position-for-position; + a mixed-width instruction pairs them through its canonical coupling (``mma_pair_k_aligned``). """ m = RegisterMapper(enc) return tuple( @@ -271,6 +272,121 @@ def _atom_k_signature( return tuple(sig), "" +def _atom_slot_coords( + canon, atoms: int, k_axis: int +) -> tuple[tuple[tuple[tuple[int, ...], int], ...], ...] | None: + """Per lane, the canonical ``(free coordinate, K)`` of each slot of one atom's K signature -- the + same slots :func:`_atom_k_signature` keeps (the first ``1/atoms`` of the lane's registers, in + ascending register order). ``free coordinate`` is the coordinate with its K entry removed. Returns + ``None`` when a lane's register count is not a whole number of atoms.""" + fwd = as_forward_map(canon) + slots = [] + for lane in sorted({l for l, _ in fwd}): + regs = sorted(r for (l, r) in fwd if l == lane) + if atoms <= 0 or len(regs) % atoms: + return None + coords = [fwd[(lane, r)] for r in regs[: len(regs) // atoms]] + slots.append(tuple((c[:k_axis] + c[k_axis + 1 :], c[k_axis]) for c in coords)) + return tuple(slots) + + +def _m_blocks_of_b(a_lane_slots, b_lane_slots) -> bool: + """True when one lane's A atom slots are M blocks that each face all of the lane's B atom slots: + the same slot count (one block), or B's slots hold one free row and A's hold whole copies of that + count, one free row per copy.""" + if len(a_lane_slots) == len(b_lane_slots): + return True + width = len(b_lane_slots) + if not width or len({f for f, _ in b_lane_slots}) != 1: + return False + per_row: dict[tuple[int, ...], int] = {} + for free, _ in a_lane_slots: + per_row[free] = per_row.get(free, 0) + 1 + return len(per_row) > 1 and all(count == width for count in per_row.values()) + + +def _k_coupling( + a_canon, + b_canon, + *, + a_free_atoms: int, + b_free_atoms: int, + k_axis: int, +) -> tuple[tuple[tuple[int, int], ...], ...] | str | None: + """Which B slot the instruction multiplies with each A slot, read off its canonical layouts. + + Positions are those of one atom's K signature (:func:`_atom_k_signature`): ``(lane, slot)``. An A + slot and a B slot are coupled when they hold the same hardware K. An A slot keeps its own position + when B's canonical holds the same K there -- every slot of a same-format instruction, so there the + coupling is the identity. Otherwise (a mixed-width instruction: its 8-bit operand holds K in two + runs of 16 per lane, its 6/4-bit operand in one run of 32) it is the B slot that holds A's hardware + K in the column of B's own slot at that position:: + + gfx950 16x16x128, fp8 A x fp4 B, lane 0: + A slots 0-15 hold K 0-15, slots 16-31 hold K 64-79 (two runs of 16) + B slots 0-31 hold K 0-31 (one run of 32) + A (lane 0, slot 16), K 64 <-> B (lane 32, slot 0), K 64 + + An instruction whose lane holds several M blocks of A (``a_m_register_blocks > 1``) multiplies each + block with the same B registers: A's per-atom signature is then that many runs of B's length, and + each A slot faces the B slot that holds its K in the lane's one B column (in its own lane when B + holds that K there):: + + gfx1250 32x16x128 fp4, lane 0: + A slots 0-63 hold (M 0, K 0-31 and 64-95), slots 64-127 hold (M 16, the same K) + B slots 0-63 hold (N 0, K 0-31 and 64-95) + A (lane 0, slot 64), K 0 <-> B (lane 0, slot 0), K 0 + + That case is read off the canonicals as: on every lane, B's atom slots hold one free row (column) + and A's hold several free rows, each with as many slots as B (:func:`_m_blocks_of_b`). A lane + holding several free rows of B is not an instruction's atom but a whole wave tile compared with + the wrong atom counts, and gets no coupling. + + Returns the coupling (per A lane, the coupled ``(lane, slot)`` of each A slot); ``None`` when the + atom counts give A and B per-atom signatures of different shapes that are not one instruction's M + blocks, so there is no slot pairing to read and the caller compares position by position; or an + error message when the two canonicals do not belong to one instruction. + """ + a_slots = _atom_slot_coords(a_canon, a_free_atoms, k_axis) + b_slots = _atom_slot_coords(b_canon, b_free_atoms, k_axis) + if ( + a_slots is None + or b_slots is None + or len(a_slots) != len(b_slots) + or not all(map(_m_blocks_of_b, a_slots, b_slots)) + ): + return None + b_where: dict[tuple[tuple[int, ...], int], tuple[int, int]] = {} + for lane, slots in enumerate(b_slots): + for slot, coord in enumerate(slots): + b_where.setdefault(coord, (lane, slot)) + coupling = [] + for lane, slots in enumerate(a_slots): + row = [] + m_blocks = len(slots) != len(b_slots[lane]) + b_in_lane = {coord: slot for slot, coord in enumerate(b_slots[lane])} + for slot, (_, k) in enumerate(slots): + if m_blocks: + # Every M block of A faces the lane's one B column. + b_free = b_slots[lane][0][0] + b_slot = b_in_lane.get((b_free, k)) + coupled = ( + (lane, b_slot) if b_slot is not None else b_where.get((b_free, k)) + ) + else: + b_free, b_k = b_slots[lane][slot] + coupled = (lane, slot) if b_k == k else b_where.get((b_free, k)) + if coupled is None: + return ( + f"A's canonical slot (lane {lane}, slot {slot}) holds K {k}, which B's canonical " + f"layout does not hold in column {b_free} -- a_canon and b_canon are not the " + "canonical layouts of one instruction" + ) + row.append(coupled) + coupling.append(tuple(row)) + return tuple(coupling) + + def _kdist_from_fwd( fwd: dict[tuple[int, int], tuple[int, ...]], k_axis: int = 1 ) -> tuple[tuple[int, ...], ...]: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/observers.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/observers.py index 3f1603acb7a9..1df7fd7ac483 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/observers.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/transforms/observers.py @@ -10,9 +10,10 @@ carry no "machinery" import. MMA safety (:func:`mma_pair_k_aligned`, :func:`mma_operand_layout_sound`, :func:`mma_pair_compatible`): the -MFMA/WMMA hardware multiply-accumulates by pairing A-slot-s with B-slot-s and summing over K. The sum -is order-independent, so the K-slot ordering is FREE -- the sole CROSS-OPERAND constraint is that A and -B share the SAME positional K-distribution. Per-operand soundness (one M per output on A, one N on B) is +MFMA/WMMA hardware multiply-accumulates by pairing A and B positions through the instruction's canonical +coupling (slot s with slot s for every same-format instruction) and summing over K. The sum is +order-independent, so the K-slot ordering is FREE -- the sole CROSS-OPERAND constraint is that A and B +hold the SAME logical K at every coupled pair of positions. Per-operand soundness (one M per output on A, one N on B) is the OTHER half of the sound MAC. :func:`mma_accumulator_flow_consistent` is the separate C store-coordinate gate (passed C == machine-derived C), NOT a sound-MAC condition. Correctness SOT: ``docs/mma_is_machinery.md`` (the three-condition sound MAC); edge-kind SOT: ``docs/label_flow_and_transforms.md``. @@ -31,6 +32,7 @@ _classify_maps, _dword_aligned, _free_relabel, + _k_coupling, _kdist_from_fwd, as_forward_map, name_permutation, @@ -159,30 +161,41 @@ def mma_pair_k_aligned( a, b, *, + a_canon, + b_canon, a_free_atoms: int = 1, b_free_atoms: int = 1, k_axis: int = 1, ) -> Diagnostic: - """DIAGNOSTIC (observer): do A and B share the SAME positional K PER ATOM -- the PAIRWISE half of the - sound MAC (correctness SOT: ``docs/mma_is_machinery.md``, condition 3)? Accepts an encoding OR a forward - map ``{(lane,reg)->coord}`` for each. + """DIAGNOSTIC (observer): do A and B hold the SAME logical K at the positions the instruction + multiplies together, PER ATOM -- the PAIRWISE half of the sound MAC (correctness SOT: + ``docs/mma_is_machinery.md``, condition 3)? Accepts an encoding OR a forward map + ``{(lane,reg)->coord}`` for each operand and each canonical. - The MFMA/WMMA hardware pairs A-slot-s with B-slot-s and sums over K; the sum is order-independent, so the + The MFMA/WMMA hardware pairs A and B positions through the instruction's canonical coupling (slot s + with slot s for every same-format instruction) and sums over K; the sum is order-independent, so the K ORDER is FREE -- the sole cross-operand constraint is that A and B agree on which logical K sits in - each paired slot. (Per-operand soundness -- one M/N per output -- is the OTHER half, + each coupled pair of positions. (Per-operand soundness -- one M/N per output -- is the OTHER half, :func:`mma_operand_layout_sound`.) M/N register order is unconstrained, and K need NOT match any "canonical" atom order (interleaved-A x interleaved-B is valid iff their K-dists match). + ``a_canon``/``b_canon`` are the instruction's canonical A and B layouts for the same wave tile + (``canonical_layouts(traits, subtiles)``); they define the coupling and are required, because a + mixed-width ``f8f6f4`` instruction does not couple slot s with slot s: its 8-bit operand holds K in two + runs of 16 per lane, its 6/4-bit operand in one run of 32. B's K labels are moved into A's slot order + through the coupling, then compared. + Comparison is PER ATOM: ``a_free_atoms``/``b_free_atoms`` are the free-dim atom counts (M-atoms for A, N-atoms for B) the driver walks. A rectangular wave tile (``a_free_atoms != b_free_atoms``) has whole-fragment K-lists of different lengths even though every issued atom pairs the SAME K, so each operand is reduced to its per-atom K signature first (defaults of 1 = whole-fragment compare, for a caller with no plan/atom-counts in hand). Three-tier: - - ``ok`` -- atom K-signatures match position-for-position. + - ``ok`` -- atom K-signatures match at coupled positions. - ``warning`` -- order differs but every lane holds the SAME K set per atom: reconcilable by an in-register reorder (named, not performed). - - ``error`` -- a malformed atom tiling, a lane-count mismatch, or lanes holding different K sets. + - ``error`` -- a malformed atom tiling, a lane-count mismatch, an operand that is not shaped like its + canonical, or lanes holding different K sets. """ a_sig, a_reason = _atom_k_signature( _kdist_from_fwd(as_forward_map(a), k_axis), a_free_atoms, "A" @@ -199,6 +212,30 @@ def mma_pair_k_aligned( "error", f"A spans {len(a_sig)} lanes but B spans {len(b_sig)} -- operands not MMA-compatible", ) + coupling = _k_coupling( + a_canon, + b_canon, + a_free_atoms=a_free_atoms, + b_free_atoms=b_free_atoms, + k_axis=k_axis, + ) + if isinstance(coupling, str): + return Diagnostic("error", coupling) + if coupling is not None: + b_canon_sig, _ = _atom_k_signature( + _kdist_from_fwd(as_forward_map(b_canon), k_axis), b_free_atoms, "B" + ) + for role, sig, canon_sig in (("A", a_sig, coupling), ("B", b_sig, b_canon_sig)): + if [len(s) for s in sig] != [len(s) for s in canon_sig]: + return Diagnostic( + "error", + f"{role} holds {[len(s) for s in sig]} K-slots per atom per lane but its " + f"canonical layout holds {[len(s) for s in canon_sig]} -- the canonical layouts " + "are not this wave tile's", + ) + b_sig = tuple( + tuple(b_sig[lane][slot] for lane, slot in row) for row in coupling + ) mism = [lane for lane, (ak, bk) in enumerate(zip(a_sig, b_sig)) if ak != bk] if not mism: return Diagnostic("ok", "A.K == B.K per atom (labels K-aligned; valid MMA)") @@ -404,7 +441,8 @@ def mma_pair_compatible( PER ATOM (the relationship, :func:`mma_pair_k_aligned`). Observer only. ``ok`` iff the pair is a valid, meaningful MMA; else the first failing operand's soundness error, or the K-match diagnostic. The full sound MAC = per-operand soundness (conditions 1-2) + pairwise K-match (3); correctness SOT: - ``docs/mma_is_machinery.md``. + ``docs/mma_is_machinery.md``. ``a_canon``/``b_canon`` serve both halves: each operand is judged + against its own canonical, and the pair is K-matched through the coupling the two define. Pass ``a_free_atoms``/``b_free_atoms`` (the free-dim atom counts m_sub/n_sub) so the K-match is per atom -- **REQUIRED for a rectangular wave tile** (``m_sub != n_sub``), whose whole-wave K-lists differ @@ -418,6 +456,8 @@ def mma_pair_compatible( return mma_pair_k_aligned( a_enc, b_enc, + a_canon=a_canon, + b_canon=b_canon, a_free_atoms=a_free_atoms, b_free_atoms=b_free_atoms, k_axis=k_axis, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py index c0d68d92cc1b..d90d28c52329 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/visualization/layout_render.py @@ -34,7 +34,6 @@ from ..transforms import ( as_forward_map, derive_c_distribution, - mma_pair_k_aligned, mma_pair_compatible, classify_transform, describe_edge, @@ -2941,7 +2940,7 @@ class MmaTee: None # (n, k): matched-K B element -> lights the chain -> C[m, n] ) show_diagnostics: bool = ( - True # run the K-match DIAGNOSTIC (observer) and annotate warn/error + True # run the pairing DIAGNOSTIC (observer; needs a_canon/b_canon) and annotate warn/error ) show_legend: bool = ( True # draw the shade legend (hue=thread, shade=vectorized-load order) @@ -3155,6 +3154,30 @@ def c_mapping(self): c_canon=self.c_canon, ) + def pairing_diagnostic(self): + """The A x B diagnostic the tee banner shows: both operands sound against the canonical machine + and K-matched at the positions it couples (:func:`mma_pair_compatible`). ``None`` when the tee has + no ``a_canon``/``b_canon``: the canonicals define which A and B positions the instruction pairs, + so without them there is no pairing to judge. Observes labels only, never mutates C. + + The per-atom K-match needs the free-dim atom counts (m_sub/n_sub); they come from the wave shape + when present, so a rectangular wave tile is not false-flagged.""" + if self.a_canon is None or self.b_canon is None: + return None + if self.wave_shape is None: + am = bn = 1 + else: + am = self.wave_shape[0] // self.atom_shape[0] + bn = self.wave_shape[1] // self.atom_shape[1] + return mma_pair_compatible( + self.a_enc, + self.b_enc, + a_canon=self.a_canon, + b_canon=self.b_canon, + a_free_atoms=am, + b_free_atoms=bn, + ) + def _build_core(self): """Build the tee's core components (A/B register-file wings, derived C body, C register file) at their tee origins, plus the maps/metadata both ``render`` and the exploded view need. Extracted @@ -3508,50 +3531,20 @@ def render(self, out_dir=".", name="mma_tee", title=None, dpi=200): ) for tx, ty, lab in titles: ax.text(tx, ty, lab, ha="center", va="bottom", fontsize=pt, weight="bold") - if self.show_diagnostics: # DIAGNOSTIC only -- observes labels, never mutates C - # full soundness + K-match when the machine refs are present; else K-match alone. - # per-atom K-match needs the free-dim atom counts (m_sub/n_sub) -- derive from the wave shape - # when present, so a rectangular wave tile isn't false-flagged on EITHER branch. - am = ( - self.wave_shape[0] // self.atom_shape[0] - if self.wave_shape is not None - else 1 - ) - bn = ( - self.wave_shape[1] // self.atom_shape[1] - if self.wave_shape is not None - else 1 + diag = self.pairing_diagnostic() if self.show_diagnostics else None + if diag is not None and diag.severity != "ok": + col = "#c00000" if diag.severity == "error" else "#c07000" + ax.text( + x0 + cw / 2, + ymax + 2.2, + f"{diag.severity.upper()} (MMA-compatible): {diag.message}", + ha="center", + va="top", + fontsize=pt, + weight="bold", + color=col, + bbox=dict(boxstyle="round", facecolor="white", edgecolor=col, lw=1.5), ) - if self.a_canon is not None and self.b_canon is not None: - diag = mma_pair_compatible( - self.a_enc, - self.b_enc, - a_canon=self.a_canon, - b_canon=self.b_canon, - a_free_atoms=am, - b_free_atoms=bn, - ) - tag = "MMA-compatible" - else: - diag = mma_pair_k_aligned( - self.a_enc, self.b_enc, a_free_atoms=am, b_free_atoms=bn - ) - tag = "K-match" - if diag.severity != "ok": - col = "#c00000" if diag.severity == "error" else "#c07000" - ax.text( - x0 + cw / 2, - ymax + 2.2, - f"{diag.severity.upper()} ({tag}): {diag.message}", - ha="center", - va="top", - fontsize=pt, - weight="bold", - color=col, - bbox=dict( - boxstyle="round", facecolor="white", edgecolor=col, lw=1.5 - ), - ) # INFO PANEL (top-left corner): atom/wave/tile/dtype/issue/vec/C-out kept IN VIEW when zoomed into the # data flow. Box is sized to the ACTUAL rendered text extent (measured, not estimated) so it always # fits; the shade legend sits below it. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issuable_pairs.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issuable_pairs.py new file mode 100644 index 000000000000..6564e842aa57 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_issuable_pairs.py @@ -0,0 +1,128 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Every instruction the backend coverage record says rocke issues, driven through TileMma. + +The pairs are read from the record (``traits/data/backend_coverage.json``), never written by hand: each +``issuable`` (instruction, target) pair of a family the driver feeds, once per recorded operand-dtype +variant, drives a single-instruction TileMma whose A, B, C and result dtypes are the instruction's own. +The tile must lower to that instruction's LLVM intrinsic, and its result Fragment must carry the +instruction's result dtype (which differs from the accumulator's for one row). + +Lowering pins the Python engine (``_lower_kernel_to_llvm_python``); the C++ engine is checked by the +byte-identity gate. Where the host clang targets the arch, the module is also compiled and linked +(compile-verified, never run). +""" + +from __future__ import annotations + +import re + +import pytest + +from rocke.core.ir import I8, U8, IRBuilder, PtrType, dtype_to_ir_type +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.helpers.tiling import ( + fill_fragment, + make_fragment, + make_tensor_desc, + make_window, + store_fragment, +) +from rocke.helpers.tiling.mma import TileMma +from rocke.helpers.tiling.mma.driver import _DRIVEN_FAMILIES +from rocke.helpers.tiling.traits import load_backend_coverage, load_mma_traits + +from .test_mma_issue import _compile_or_skip + +_CATALOG = load_mma_traits() +_COVERAGE = load_backend_coverage() + + +def _intrinsic_calls(ir: str, intrinsic: str) -> list[str]: + """The call lines of ``intrinsic``, with or without LLVM's overload type suffixes (``.v8i32``, + ``.f16``): another intrinsic that merely starts with the same name does not count. + """ + pattern = re.compile(rf"@{re.escape(intrinsic)}(\.v?\d*(?:i|f|bf)\d+)*\(") + return [line for line in ir.splitlines() if "call" in line and pattern.search(line)] + + +def _driven_issuable_pairs(): + """Every issuable pair of a driven family, once per recorded operand-dtype variant, as + ``(op_id, target, operand_dtypes)``.""" + for (op_id, target), entries in sorted(_COVERAGE.entries.items()): + if _CATALOG.get(op_id).family not in _DRIVEN_FAMILIES: + continue + for entry in entries: + if entry.status == "issuable": + yield op_id, target, entry.operand_dtypes + + +def _params(pairs): + for op_id, target, variant in pairs: + label = f"{op_id}-{target}" + (f"-{'-'.join(variant)}" if variant else "") + yield pytest.param(op_id, target, variant, id=label) + + +_PAIRS = list(_driven_issuable_pairs()) + + +def _plan(op_id: str, target: str, operand_dtypes) -> TileMma: + """A single-instruction TileMma pinned to ``op_id``; a selectable-signedness integer instruction + takes the operand dtypes of its recorded variant (u8 picks the unsigned form).""" + dtypes = {} + if operand_dtypes is not None: + dtypes = dict(zip("ab", map(dtype_to_ir_type, operand_dtypes))) + mma = TileMma(target=target, atom_override=op_id, **dtypes) + assert mma.plan.mfma_count == 1 + return mma + + +def _zero_fragment(b: IRBuilder, desc, dtype): + """A zero Fragment of ``dtype``. fill_fragment zeroes every dtype but u8, which core's zero_vec + has no form for; a u8 zero is the i8 zero vector reinterpreted, bit for bit.""" + fragment = make_fragment(desc, dtype) + if dtype == U8: + count = fragment.storage_type.count + fragment.value = b.bitcast(b.zero_vec(I8, count), fragment.storage_type) + else: + fill_fragment(b, fragment, 0) + return fragment + + +def _single_instruction_tile(b: IRBuilder, mma: TileMma): + """Zero A, B and C fragments of the instruction's own dtypes and descriptors, one MMA, the result + stored to a tensor of the result dtype. Returns the result Fragment.""" + b.kernel.attrs["max_workgroup_size"] = mma.wave_size + d_ptr = b.param("D", PtrType(mma.d_dtype, "global"), noalias=True, writeonly=True) + lane = b.thread_id_x() + zero = b.const_i32(0) + a = _zero_fragment(b, mma.a_desc, mma.a_dtype) + bf = _zero_fragment(b, mma.b_desc, mma.b_dtype) + acc = _zero_fragment(b, mma.c_desc, mma.c_dtype) + result = mma(b, a, bf, acc) + traits = mma.traits + d_td = make_tensor_desc((traits.m, traits.n), (traits.n, 1), mma.d_dtype) + store_fragment(b, d_ptr, make_window(d_td, (zero, zero)), result, lane) + b.ret() + return result + + +@pytest.mark.parametrize("op_id,target,operand_dtypes", list(_params(_PAIRS))) +def test_issuable_instruction_drives_a_tile_that_lowers_to_it( + op_id, target, operand_dtypes +) -> None: + mma = _plan(op_id, target, operand_dtypes) + b = IRBuilder("issuable_pair") + result = _single_instruction_tile(b, mma) + # The result carries the instruction's result dtype, the accumulator's for every row but the + # bf16-from-f32 one. + assert result.dtype == mma.d_dtype == dtype_to_ir_type(mma.traits.d_dtype) + ir = _lower_kernel_to_llvm_python(b.kernel, arch=target) + assert len(_intrinsic_calls(ir, mma.traits.llvm_intrinsic)) == 1 + _compile_or_skip(b.kernel, target) + + +def test_issuable_population() -> None: + # Locks the pair list read from the record, so a record or catalog change is seen here. + assert len(_PAIRS) == 339 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_k_coupling.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_k_coupling.py new file mode 100644 index 000000000000..13dc5ac6d267 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_k_coupling.py @@ -0,0 +1,351 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""A and B K are paired through the instruction's canonical coupling. + +The hardware multiplies the A position and the B position that hold the same hardware K. For a +same-format instruction that is slot s with slot s. A mixed-width ``f8f6f4`` instruction differs: its +8-bit operand holds K in two runs of 16 per lane, its 6/4-bit operand in one run of 32, so comparing A +and B slot for slot rejects the correct canonical pair and accepts a wrong one. The K check reads the +coupling off the two canonical layouts, which every caller must pass. + +Lowering tests pin the Python engine (``_lower_kernel_to_llvm_python``); the C++ engine is checked by +the byte-identity gate. Where the host clang targets the arch, the module is also compiled and linked. +Where the current GPU is the instruction's target, the driven tiles also run on random K-dependent +operands and must equal an independent numpy product exactly; on any other GPU, or none, those tests +skip. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from rocke.core.ir import F32, FP4E2M1, FP8E4M3, I32, IRBuilder, PtrType, VectorType +from rocke.core.lower_llvm import _lower_kernel_to_llvm_python +from rocke.examples.common.scaled_gemm_reference import decode_fp4, pack_fp4_codes +from rocke.helpers.mma_io import storage_ir_type +from rocke.helpers.tiling import ( + fill_fragment, + load_fragment, + make_fragment, + make_tensor_desc, + make_window, + store_fragment, +) +from rocke.helpers.tiling import tiling_recorder as tr +from rocke.helpers.tiling.analysis.soundness import verify_mma_soundness +from rocke.helpers.tiling.layout_optimizer import evaluate_transform, optimize_layout +from rocke.helpers.tiling.mma import TileMma +from rocke.helpers.tiling.mma.warp_encoding import canonical_layouts +from rocke.helpers.tiling.traits import load_mma_traits +from rocke.helpers.tiling.transforms import ( + as_forward_map, + mma_pair_compatible, + mma_pair_k_aligned, +) +from rocke.helpers.tiling.visualization import layout_render as lr +from rocke.tests.gpu import run_on_device + +from .test_mma_issue import _calls, _compile_or_skip + +_CATALOG = load_mma_traits() +_FP8_FP4 = "mfma_f32_16x16x128_f8f6f4_fp8e4m3_fp4e2m1" +_FP8_FP8 = "mfma_f32_16x16x128_f8f6f4_fp8e4m3_fp8e4m3" +_FP4_FP4 = "mfma_f32_16x16x128_f8f6f4_fp4e2m1_fp4e2m1" + + +def _canonicals(op_id: str): + """The single-atom canonical A and B layouts of a gfx950 16x16x128 instruction.""" + a_canon, b_canon, _ = canonical_layouts(_CATALOG.get(op_id), (1, 1, 1)) + return a_canon, b_canon + + +def _reverse_k(layout): + """The same layout with every K label k replaced by 127 - k: a K relabelling.""" + return { + slot: (coord[0], 127 - coord[1]) + for slot, coord in as_forward_map(layout).items() + } + + +def _pairs(): + """The gfx950 16x16x128 pairs of the mixed-width table: (A, B, A canonical, B canonical).""" + a8, b4 = _canonicals(_FP8_FP4) + _, b8 = _canonicals(_FP8_FP8) + a44, b44 = _canonicals(_FP4_FP4) + return { + "canonical-fp8-A-x-fp4-B": (a8, b4, a8, b4), + "same-K-reversal-on-both": (_reverse_k(a8), _reverse_k(b4), a8, b4), + "K-reversal-on-A-only": (_reverse_k(a8), b4, a8, b4), + "fp4-B-loaded-with-fp8-map": (a8, b8, a8, b4), + "canonical-fp4-x-fp4": (a44, b44, a44, b44), + } + + +# Each case and its verdict. The canonical fp8 x fp4 pair and the jointly reversed one multiply the +# same K, so they pass (both were refused slot for slot); the fp4 B loaded with fp8's map multiplies +# A's K 64 with B's K 32 at the coupled positions, so it fails (it passed slot for slot). +_TABLE = [ + ("canonical-fp8-A-x-fp4-B", "ok"), + ("same-K-reversal-on-both", "ok"), + ("K-reversal-on-A-only", "error"), + ("fp4-B-loaded-with-fp8-map", "error"), + ("canonical-fp4-x-fp4", "ok"), +] + + +@pytest.mark.parametrize("case,severity", _TABLE, ids=[c for c, _ in _TABLE]) +@pytest.mark.parametrize( + "check", + [ + pytest.param(mma_pair_k_aligned, id="k-aligned"), + pytest.param(mma_pair_compatible, id="compatible"), + ], +) +def test_mixed_width_pairs_are_judged_through_the_coupling( + check, case, severity +) -> None: + a, b, a_canon, b_canon = _pairs()[case] + d = check(a, b, a_canon=a_canon, b_canon=b_canon) + assert d.severity == severity, d.message + + +def test_the_coupling_pairs_a_second_run_with_another_lane() -> None: + # fp8 A lane 0 holds K 64-79 in slots 16-31; fp4 B holds K 64-95 in lane 32. Moving one B label + # at the coupled position (lane 32, slot 0) breaks the pair, though B's lane 0 is untouched. + a8, b4 = _canonicals(_FP8_FP4) + a_map, b_map = as_forward_map(a8), as_forward_map(b4) + assert (a_map[(0, 16)][1], b_map[(32, 0)][1]) == (64, 64) + moved = dict(b_map) + moved[(32, 0)], moved[(32, 1)] = b_map[(32, 1)], b_map[(32, 0)] + d = mma_pair_k_aligned(a8, moved, a_canon=a8, b_canon=b4) + assert d.severity == "warning", d.message + + +def test_canonicals_of_two_instructions_are_refused() -> None: + # An fp8 A canonical with an fp4 B canonical of another shape describes no instruction: no B slot + # in A's column holds K 64 of the 16x16x128 A. + a8, _ = _canonicals(_FP8_FP4) + _, b_small = canonical_layouts( + _CATALOG.get("mfma_f32_32x32x64_f8f6f4_fp4e2m1_fp4e2m1"), (1, 1, 1) + )[:2] + d = mma_pair_k_aligned(a8, b_small, a_canon=a8, b_canon=b_small) + assert d.severity == "error" + + +# ---- an instruction whose lane holds two M blocks of A ----------------------------------------- + +# Each lane of this instruction holds two M rows of A (M 0 and M 16 on lane 0), each with the same 64 K, +# against one N row of B with those 64 K: the hardware multiplies both A rows with the same B registers. +_F4_TWO_M_BLOCKS = "wmma_f32_32x16x128_f4" + + +def _second_m_block_k_reversed(layout): + """The layout with K labels k -> 127 - k on the second M block of every lane only.""" + fwd = as_forward_map(layout) + half = max(reg for _, reg in fwd) // 2 + 1 + return { + (lane, reg): (coord[0], 127 - coord[1]) if reg >= half else coord + for (lane, reg), coord in fwd.items() + } + + +@pytest.mark.parametrize( + "second_block_reversed,severity", + [ + pytest.param(False, "ok", id="canonical"), + pytest.param(True, "error", id="block-2-reversed"), + ], +) +@pytest.mark.parametrize( + "check", + [ + pytest.param(mma_pair_k_aligned, id="k-aligned"), + pytest.param(mma_pair_compatible, id="compatible"), + ], +) +def test_each_m_block_of_a_is_paired_with_the_same_b( + check, second_block_reversed, severity +) -> None: + # The canonical pair passes; relabelling K on A's second M block alone breaks that block's pairing + # with B while the first block still matches, so the check must compare every block. + a_canon, b_canon, _ = canonical_layouts(_CATALOG.get(_F4_TWO_M_BLOCKS), (1, 1, 1)) + a = _second_m_block_k_reversed(a_canon) if second_block_reversed else a_canon + d = check(a, b_canon, a_canon=a_canon, b_canon=b_canon) + assert d.severity == severity, d.message + + +# ---- the callers: driver, soundness gate, MmaTee, layout optimizer ----------------------------- + + +def _wave_tile(mma: TileMma): + """One wave computes C = A * B^T for the MMA's wave tile, A (M, K) and B (N, K) loaded from global + tensors behind pointers of each dtype's storage type, the f32 result stored.""" + wave_m, wave_n, wave_k = mma.shape + b = IRBuilder("k_coupling") + b.kernel.attrs["max_workgroup_size"] = mma.wave_size + a_ptr = b.param( + "A", + PtrType(storage_ir_type(mma.a_dtype.name), "global"), + noalias=True, + readonly=True, + ) + b_ptr = b.param( + "B", + PtrType(storage_ir_type(mma.b_dtype.name), "global"), + noalias=True, + readonly=True, + ) + c_ptr = b.param("C", PtrType(F32, "global"), noalias=True, writeonly=True) + lane = b.thread_id_x() + zero = b.const_i32(0) + a_td = make_tensor_desc((wave_m, wave_k), (wave_k, 1), mma.a_dtype) + b_td = make_tensor_desc((wave_n, wave_k), (wave_k, 1), mma.b_dtype) + c_td = make_tensor_desc((wave_m, wave_n), (wave_n, 1), F32) + acc = make_fragment(mma.c_desc, F32) + fill_fragment(b, acc, 0) + fa = load_fragment(b, a_ptr, make_window(a_td, (zero, zero)), mma.a_desc, lane) + fb = load_fragment(b, b_ptr, make_window(b_td, (zero, zero)), mma.b_desc, lane) + result = mma(b, fa, fb, acc) + store_fragment(b, c_ptr, make_window(c_td, (zero, zero)), result, lane) + b.ret() + return b.kernel + + +def _fp8_fp4_mma() -> TileMma: + return TileMma((16, 16, 128), a=FP8E4M3, b=FP4E2M1, c=F32, target="gfx950") + + +def test_driven_fp8_by_fp4_tile_issues_and_compiles() -> None: + # The canonical fragments pass the K check. A reaches the instruction as 32 fp8 elements, B as + # its 4 packed words; core widens each to the intrinsic's 8-word operand. + mma = _fp8_fp4_mma() + assert mma.op_id == _FP8_FP4 + kernel = _wave_tile(mma) + (op,) = [op for op in kernel.body.ops if op.name == "tile.mma"] + assert op.operands[0].type == VectorType(FP8E4M3, 32) + assert op.operands[1].type == VectorType(I32, 4) + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx950") + assert len(_calls(ir, mma.traits.llvm_intrinsic)) == 1 + _compile_or_skip(kernel, "gfx950") + + +def test_driven_tile_of_an_instruction_with_two_m_blocks_per_lane_issues_every_atom() -> ( + None +): + # Two atoms along each of M, N and K: every A atom is one contiguous register block holding both + # of a lane's M rows, so the driver slices and issues all eight. + mma = TileMma((64, 32, 256), target="gfx1250", atom_override=_F4_TWO_M_BLOCKS) + assert mma.plan.subtiles == (2, 2, 2) + kernel = _wave_tile(mma) + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx1250") + assert len(_calls(ir, mma.traits.llvm_intrinsic)) == 8 + _compile_or_skip(kernel, "gfx1250") + + +def _fp8e4m3_codes(rng, shape) -> np.ndarray: + """Random fp8 e4m3 codes of magnitude 0.125 to 15 (biased exponent 4 to 10, any mantissa, either + sign): no NaN, no subnormal.""" + exponent = rng.integers(4, 11, size=shape) + mantissa = rng.integers(0, 8, size=shape) + sign = rng.integers(0, 2, size=shape) + return ((sign << 7) | (exponent << 3) | mantissa).astype(np.uint8) + + +def _decode_fp8e4m3(codes: np.ndarray) -> np.ndarray: + """The value of normal fp8 e4m3 codes: (-1)^s * 2^(e - 7) * (1 + m / 8).""" + exponent = ((codes >> 3) & 15).astype(np.int32) + magnitude = np.ldexp(1.0 + (codes & 7) / 8.0, exponent - 7) + return np.where(codes & 0x80, -magnitude, magnitude) + + +def _run_k_coupled_tile(mma: TileMma, a_fp8: bool, target: str) -> None: + """Run `_wave_tile` for `mma` on random K-dependent operands where the current GPU is `target`, + and require C = A * B^T exactly. A is fp8 e4m3 when `a_fp8`, else fp4 like B. Every product is a + multiple of 2^-7 below 2^7 and K is at most 256, so every partial sum is exact in f32 and the + order the hardware adds them in cannot change the result.""" + wave_m, wave_n, wave_k = mma.shape + rng = np.random.default_rng(wave_k + wave_m) + if a_fp8: + a_codes = _fp8e4m3_codes(rng, (wave_m, wave_k)) + a_bytes, a_values = a_codes, _decode_fp8e4m3(a_codes) + else: + a_bytes = pack_fp4_codes( + rng.integers(0, 16, size=(wave_m, wave_k), dtype=np.uint8) + ) + a_values = decode_fp4(a_bytes) + b_bytes = pack_fp4_codes(rng.integers(0, 16, size=(wave_n, wave_k), dtype=np.uint8)) + expected = (a_values @ decode_fp4(b_bytes).T).astype(np.float32) + c = np.full((wave_m, wave_n), np.nan, dtype=np.float32) + run_on_device( + _wave_tile(mma), {"A": a_bytes, "B": b_bytes, "C": c}, mma.wave_size, {target} + ) + np.testing.assert_array_equal(c, expected) + + +def test_driven_fp8_by_fp4_tile_computes_the_product() -> None: + # A wrong K pairing multiplies A's K with another K of B, which random K-dependent operands turn + # into a wrong C. Runs where the current GPU is a gfx950. + _run_k_coupled_tile(_fp8_fp4_mma(), a_fp8=True, target="gfx950") + + +@pytest.mark.skip(reason="has not run on gfx1250 hardware yet; enable once it has") +def test_driven_tile_with_two_m_blocks_per_lane_computes_the_product() -> None: + # Both of a lane's M blocks of A must meet the same B registers, on every one of the eight atoms. + # Runs where the current GPU is a gfx1250. + mma = TileMma((64, 32, 256), target="gfx1250", atom_override=_F4_TWO_M_BLOCKS) + _run_k_coupled_tile(mma, a_fp8=False, target="gfx1250") + + +def test_soundness_gate_accepts_a_recorded_fp8_by_fp4_pipeline() -> None: + _kernel, pipe = tr.record_build(lambda: _wave_tile(_fp8_fp4_mma())) + assert verify_mma_soundness(pipe) == 1 + + +def test_mma_tee_of_a_canonical_fp8_by_fp4_plan_shows_no_error() -> None: + tee = lr.MmaTee.from_mma(TileMma(target="gfx950", atom_override=_FP8_FP4)) + d = tee.pairing_diagnostic() + assert d.severity == "ok", d.message + + +def test_mma_tee_without_canonicals_judges_no_pairing() -> None: + a8, b4 = _canonicals(_FP8_FP4) + mma = TileMma(target="gfx950", atom_override=_FP8_FP4) + tee = lr.MmaTee( + a_enc=a8, b_enc=b4, c_enc=mma.c_desc.layout, atom_shape=mma.atom_shape + ) + assert tee.pairing_diagnostic() is None + + +def test_layout_optimizer_pairs_k_through_both_canonicals() -> None: + a8, b4 = _canonicals(_FP8_FP4) + _, b8 = _canonicals(_FP8_FP8) + ok = evaluate_transform(a8, a8, canon=a8, k_partner=b4, k_partner_canon=b4) + assert ok.k_match == "ok" and ok.works + wrong = evaluate_transform(a8, a8, canon=a8, k_partner=b8, k_partner_canon=b4) + assert wrong.k_match == "error" and not wrong.works + + +_EK = ( + "k_partner needs canon= and k_partner_canon= -- A and B K positions pair through the " + "instruction's canonical layouts, which define the coupling" +) + + +@pytest.mark.parametrize( + "canons", + [ + pytest.param({"canon": "a"}, id="no-k_partner_canon"), + pytest.param({"k_partner_canon": "b"}, id="no-canon"), + ], +) +def test_layout_optimizer_refuses_a_k_partner_without_both_canonicals(canons) -> None: + a8, b4 = _canonicals(_FP8_FP4) + given = {name: {"a": a8, "b": b4}[which] for name, which in canons.items()} + with pytest.raises(ValueError) as excinfo: + evaluate_transform(a8, a8, k_partner=b4, **given) + assert str(excinfo.value) == _EK + with pytest.raises(ValueError) as excinfo: + optimize_layout(a8, {"self": a8}, k_partner=b4, **given) + assert str(excinfo.value) == _EK diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py index 484d13a9dbb5..bdb6bfb93c7e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_validation_recipe.py @@ -158,7 +158,7 @@ def test_rejects_reorderable_k_misorder(): mma = _mma(shape=(16, 16, 32)) # k_sub = 2 gives K register buckets to permute a_reordered = mma.a_desc.reorder_registers((1, 2, 0)) m_sub, n_sub, _ = mma.subtiles - a_canon = mma.a_desc.layout + a_canon, b_canon, _ = canonical_layouts(mma.plan.traits, mma.subtiles) # Self-check the construction: A is still SOUND, and the pair is genuinely a reorder-fixable WARNING -- # so this test cannot pass for the wrong reason if the layout math shifts. assert ( @@ -168,6 +168,8 @@ def test_rejects_reorderable_k_misorder(): mma_pair_k_aligned( a_reordered.layout, mma.b_desc.layout, + a_canon=a_canon, + b_canon=b_canon, a_free_atoms=m_sub, b_free_atoms=n_sub, ).severity diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_interleave_all_dense.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_interleave_all_dense.py index 3be788c951a0..972cb04b0fe7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_interleave_all_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_interleave_all_dense.py @@ -5,7 +5,8 @@ * the operand K-contiguous <-> free-dim-contiguous transpose is cross-lane-free (``reorder`` tier), * the C-shuffle (native accumulator -> store order) is cross-lane-free, and - * the MMA soundness (positional per-atom K-match) holds, + * the MMA soundness (per-atom K-match at the positions the instruction couples) holds -- except on the + mixed-width ``f8f6f4`` rows, where it is PROVEN to fail (see ``_operand_k_maps_differ``), for a single atom AND a 2x2-atom tile (``free_sub>1`` / multi-patch, e.g. the 32x32 atom's ``P=4``). @@ -27,7 +28,10 @@ from rocke.helpers.tiling.traits import load_mma_traits from rocke.helpers.tiling.layouts import make_tile_desc from rocke.helpers.tiling.mma.styles import AtomNumbers, InterleavedStyle -from rocke.helpers.tiling.mma.warp_encoding import encoder_supports_layout +from rocke.helpers.tiling.mma.warp_encoding import ( + canonical_layouts, + encoder_supports_layout, +) from rocke.helpers.tiling.transforms import ( as_forward_map, classify_transform, @@ -71,6 +75,15 @@ def _operand_gate(t) -> bool: ) +def _operand_k_maps_differ(t) -> bool: + """A mixed-width ``f8f6f4`` row: its 8-bit operand holds K in two runs of 16 per lane, its 6/4-bit + operand in one run of 32 (per-operand ``k_num_access`` differ). The interleaved construction lays A and + B out with ONE K map, so on such a row A and B hold different K at the positions the instruction + couples -- the pair would multiply the wrong K, and the K check rejects it.""" + lp = t.layout_params + return lp.a_k_num_access != lp.b_k_num_access + + def _accum_gate(t) -> bool: """§9 accumulator precondition: ``c_lane_rows * atom.n == wave``.""" return (t.m // t.layout_params.c_m_per_lane) * t.n == t.wave_size @@ -159,12 +172,16 @@ def test_dense_mma_interleavable(t): ok ), f"{t.op_id}: operand K<->free is {tier} at (free_atoms={free_atoms}, k_sub={k_sub})" - # MMA soundness (positional per-atom K-match) for the base single-atom operands. Per-operand free - # lanes: A's free axis is M (free_lanes=a.m), B's is N (free_lanes=a.n) -- equal for a square atom. + # MMA soundness (per-atom K-match at coupled positions) for the base single-atom operands. Per-operand + # free lanes: A's free axis is M (free_lanes=a.m), B's is N (free_lanes=a.n) -- equal for a square atom. _, a_mma = _STYLE.operand_descs(t, free_sub=1, k_sub=1, free_lanes=a.m) _, b_mma = _STYLE.operand_descs(t, free_sub=1, k_sub=1, free_lanes=a.n) - d = mma_pair_k_aligned(a_mma.layout, b_mma.layout) - assert d.severity == "ok", f"{t.op_id}: MMA soundness failed -- {d.message}" + a_canon, b_canon, _ = canonical_layouts(t, (1, 1, 1)) + d = mma_pair_k_aligned(a_mma.layout, b_mma.layout, a_canon=a_canon, b_canon=b_canon) + if _operand_k_maps_differ(t): + assert d.severity == "error", f"{t.op_id}: one K map passed a mixed-width pair" + else: + assert d.severity == "ok", f"{t.op_id}: MMA soundness failed -- {d.message}" # C-shuffle (native accumulator -> store), single atom AND 2x2-atom tile (free_sub>1 / multi-patch). for m_sub, n_sub in [(1, 1), (2, 2)]: @@ -197,3 +214,6 @@ def test_proof_population(): if not (_operand_gate(t) and _accum_gate(t) and t.m == t.n) } assert measured_cannot == _CANNOT_INTERLEAVE, measured_cannot ^ _CANNOT_INTERLEAVE + # The mixed-width f8f6f4 rows whose interleaved pair is K-misaligned: gfx950 16x16x128 and 32x32x64 + # (12 each) and gfx1250 16x16x128 (12). + assert sum(map(_operand_k_maps_differ, _DENSE)) == 36 diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py index 51f77bef682a..34d7caad22d6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_layout_optimizer.py @@ -128,12 +128,12 @@ def test_pairwise_k_match_gate(): # per-operand soundness alone can pass while the PAIR is K-mismatched -> k_partner must catch it mma = _mma() a, b = mma.a_layout, mma.b_layout - ok = evaluate_transform(a, a, canon=a, k_partner=b) + ok = evaluate_transform(a, a, canon=a, k_partner=b, k_partner_canon=b) assert ok.k_match == "ok" and ok.works b_bad = { k: (c[0], 0) for k, c in as_forward_map(b).items() } # scramble B's K -> mismatch - bad = evaluate_transform(a, a, canon=a, k_partner=b_bad) + bad = evaluate_transform(a, a, canon=a, k_partner=b_bad, k_partner_canon=b) assert bad.k_match != "ok" and not bad.works diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py index 8e0da152cd9b..36e427a2dfe6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/test_transforms.py @@ -16,7 +16,11 @@ from rocke.core.ir import F16, F32 from rocke.helpers.tiling.encoding import WarpDistributionEncoding -from rocke.helpers.tiling.mma.warp_encoding import a_warp_encoding, b_warp_encoding +from rocke.helpers.tiling.mma.warp_encoding import ( + a_warp_encoding, + b_warp_encoding, + canonical_layouts, +) from rocke.helpers.tiling.register_mapper import RegisterMapper from rocke.helpers.tiling.traits import load_mma_traits from rocke.helpers.tiling.transforms import ( @@ -158,7 +162,7 @@ def test_describe_edge_declared_relabel_must_be_consistent() -> None: def test_mma_pair_k_aligned_accepts_canonical() -> None: tr = _traits() a, b = a_warp_encoding(tr), b_warp_encoding(tr) - d = mma_pair_k_aligned(a, b) + d = mma_pair_k_aligned(a, b, a_canon=a, b_canon=b) assert d.severity == "ok", d.message @@ -172,7 +176,7 @@ def test_mma_pair_k_aligned_rejects_k_misaligned() -> None: ) # reorders which K sits in which register slot if k_distribution(a_k_scrambled) == k_distribution(a): pytest.skip("swap did not perturb the K projection for this atom") - d = mma_pair_k_aligned(a_k_scrambled, b) + d = mma_pair_k_aligned(a_k_scrambled, b, a_canon=a, b_canon=b) assert d.severity in ("warning", "error") # the driver rejects either # a register-axis swap keeps the K SET per lane but changes its order -> reorder-fixable warning # (names transform_fragment); if it also changed the set it would be a hard "different K sets" error. @@ -193,10 +197,18 @@ def test_mma_pair_k_aligned_accepts_rectangular_per_atom() -> None: ) a, b, _ = _wave_descs_interleaved(4, 2, 2) - d_atom = mma_pair_k_aligned(a.layout, b.layout, a_free_atoms=4, b_free_atoms=2) + a_canon, b_canon, _ = canonical_layouts(_traits(), (4, 2, 2)) + d_atom = mma_pair_k_aligned( + a.layout, + b.layout, + a_canon=a_canon, + b_canon=b_canon, + a_free_atoms=4, + b_free_atoms=2, + ) assert d_atom.severity == "ok", d_atom.message d_whole = mma_pair_k_aligned( - a.layout, b.layout + a.layout, b.layout, a_canon=a_canon, b_canon=b_canon ) # default 1,1 -> whole-fragment, rejects rectangular assert d_whole.severity != "ok" @@ -210,7 +222,15 @@ def test_mma_pair_k_aligned_per_atom_still_rejects_k_mismatch() -> None: ) a, b, _ = _wave_descs_interleaved(4, 2, 2) - d = mma_pair_k_aligned(a.layout, b.layout, a_free_atoms=4, b_free_atoms=1) + a_canon, b_canon, _ = canonical_layouts(_traits(), (4, 2, 2)) + d = mma_pair_k_aligned( + a.layout, + b.layout, + a_canon=a_canon, + b_canon=b_canon, + a_free_atoms=4, + b_free_atoms=1, + ) assert d.severity != "ok" diff --git a/dnn-providers/hip-kernel-provider/rocke/style/API_DESIGN_STANDARD.md b/dnn-providers/hip-kernel-provider/rocke/style/API_DESIGN_STANDARD.md index bd66427c50f3..2c54195ec5b5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/style/API_DESIGN_STANDARD.md +++ b/dnn-providers/hip-kernel-provider/rocke/style/API_DESIGN_STANDARD.md @@ -223,7 +223,7 @@ The MMA issue path is the reference implementation of Parts 1–5. | +--> mma_operand_layout_sound(layout, canon, role) one operand vs the machine +--> mma_accumulator_flow_consistent(c, a, b, a_canon,b_canon,c_canon) C from the PASSED a, b - mma_pair_k_aligned(a, b) A vs B agree on K (takes no canon) + mma_pair_k_aligned(a, b, a_canon, b_canon) A vs B agree on K (canon = coupling) _assert_atom_contiguous(a) / (b) atom-slice contract (A and B only) | v From ec941cded35c00fb3cc34615b51263b1bd6e7b20 Mon Sep 17 00:00:00 2001 From: Chris Millette Date: Tue, 27 Oct 2026 22:06:39 -0400 Subject: [PATCH 7/7] feat(rocke): move fp8, bf8 and i8 fragments with vector loads and LDS 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 access instead of one byte at a time; the comment on the set names those three users. Tests: an fp8 or bf8 32x32x64 wave tile loads its operands as eight 8-byte global loads with no single-byte load, and an fp8, bf8 or i8 fragment staged through an LDS buffer uses only vector LDS loads and stores and compiles for gfx942 (compile-verified, not run). Co-Authored-By: Claude --- .../python/rocke/helpers/tiling/emit.py | 7 +- .../helpers/tiling/mma/test_mma_dtypes.py | 75 ++++++++++++++++++- 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/emit.py index 08d86d527656..efa42a7a4f9a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/tiling/emit.py @@ -502,12 +502,17 @@ def _zero_scalar(b: Any, dtype: Any) -> Any: return b.vec_extract(b.zero_vec(_storage_element(dtype), 1), 0) +# Storage elements moved with vector accesses: global loads, LDS loads and LDS stores +# (global_load_vN, smem_load_vN, smem_store_vN). tf32 is held and moved as i32. _WIDE_OK = { "f16", "bf16", "f32", "i32", -} # storage elements global_load_vN / smem_load_vN vectorize (tf32 is held and moved as i32) + "fp8e4m3", + "bf8e5m2", + "i8", +} def _register_run(tile_desc: TileDesc, window: TensorWindow) -> tuple[int, int] | None: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_dtypes.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_dtypes.py index 14b144a4bafe..cb1de3175c68 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_dtypes.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/tests/helpers/tiling/mma/test_mma_dtypes.py @@ -5,12 +5,15 @@ MMA's dtypes. A TileMma names its dtypes as IR types (``FP8E4M3``, ``BF8E5M2``); these tests check that an fp8 or -bf8 wave tile lowers to one MFMA call per atom, and that ``MmaTee.from_mma`` sizes and labels each -operand from those dtypes. +bf8 wave tile lowers to one MFMA call per atom, that its fragments move through global memory and +LDS with vector accesses rather than one byte at a time, and that ``MmaTee.from_mma`` sizes and +labels each operand from those dtypes. """ from __future__ import annotations +import re + import pytest from rocke.core.ir import ( @@ -18,6 +21,7 @@ F16, F32, FP8E4M3, + I8, IRBuilder, PtrType, ) @@ -34,6 +38,8 @@ from rocke.helpers.tiling.mma import TileMma, Tiling from rocke.helpers.tiling.visualization import layout_render as lr +from .test_mma_issue import _compile_or_skip + def _wave_tile_kernel(dtype, mma: TileMma): """One wave computes C = A * B^T for the MMA's wave tile, loading A and B from global memory.""" @@ -83,6 +89,71 @@ def test_8bit_wave_tile_issues_one_mfma_per_atom(dtype, intrinsic) -> None: assert len(calls) == 8 +def _global_loads(ir: str, loaded: str) -> list[str]: + return re.findall(rf"= load {re.escape(loaded)}, ptr addrspace\(1\)", ir) + + +@pytest.mark.parametrize("dtype", [FP8E4M3, BF8E5M2]) +def test_8bit_wave_tile_loads_its_operands_with_vector_loads(dtype) -> None: + # Each lane holds 8 consecutive K elements of A and of B per atom, so the 32x32x64 tile reads + # its operands as 8 loads of 8 bytes (A and B, 2 atoms along M or N, 2 along K) instead of 64 + # one-byte loads. + mma = TileMma( + (32, 32, 64), + a=dtype, + b=dtype, + c=F32, + target="gfx942", + tiling=Tiling(atom_shape=(16, 16, 32)), + ) + ir = _lower_kernel_to_llvm_python(_wave_tile_kernel(dtype, mma), arch="gfx942") + assert len(_global_loads(ir, "<8 x i8>")) == 8 + assert _global_loads(ir, "i8") == [] + + +def _lds_round_trip_kernel(dtype): + """One wave loads the A operand of an fp8 32x32x64 tile from global memory as ``dtype``, stages + it through an LDS buffer (store_fragment, then load_fragment) and writes it back to global. + """ + mma = TileMma( + (32, 32, 64), + a=FP8E4M3, + b=FP8E4M3, + c=F32, + target="gfx942", + tiling=Tiling(atom_shape=(16, 16, 32)), + ) + wave_m, _, wave_k = mma.shape + b = IRBuilder(f"lds_stage_{dtype.name}") + b.kernel.attrs["max_workgroup_size"] = mma.wave_size + src = b.param("A", PtrType(dtype, "global"), noalias=True, readonly=True) + dst = b.param("Out", PtrType(dtype, "global"), noalias=True, writeonly=True) + lane = b.thread_id_x() + zero = b.const_i32(0) + td = make_tensor_desc((wave_m, wave_k), (wave_k, 1), dtype) + window = make_window(td, (zero, zero)) + lds = b.smem_alloc(dtype, [td.element_span]) + staged = load_fragment(b, src, window, mma.a_desc, lane) + store_fragment(b, lds, window, staged, lane) + loaded = load_fragment(b, lds, window, mma.a_desc, lane) + store_fragment(b, dst, window, loaded, lane) + b.ret() + return b.kernel + + +@pytest.mark.parametrize("dtype", [FP8E4M3, BF8E5M2, I8]) +def test_8bit_fragment_moves_through_lds_with_vector_accesses(dtype) -> None: + # The 8 consecutive bytes a lane holds per atom are one LDS store and one LDS load each; no LDS + # access moves a single byte. + kernel = _lds_round_trip_kernel(dtype) + ir = _lower_kernel_to_llvm_python(kernel, arch="gfx942") + lds_loads = re.findall(r"= load ([^,]+), ptr addrspace\(3\)", ir) + lds_stores = re.findall(r"store (<[^>]+>|\S+) \S+, ptr addrspace\(3\)", ir) + assert lds_loads == ["<8 x i8>"] * 4 + assert lds_stores == ["<8 x i8>"] * 4 + _compile_or_skip(kernel, "gfx942") + + def test_mma_tee_sizes_and_labels_operands_from_the_mma_dtypes() -> None: fp8 = lr.MmaTee.from_mma( TileMma((16, 16, 32), a=FP8E4M3, b=FP8E4M3, c=F32, target="gfx942")