feat: amd: rank the matrix entries by issues, with a place

 for cycles

Choosing the narrowest entry that fits is a guess, not a rule.  Thirteen
output columns take four 4x4 issues or one 16x16 that wastes three of its
sixteen, and which is faster is a property of the two instructions rather
than of the waste.

So the selection is a ranking.  issues() counts what an entry takes for
the whole contraction across its four axes; rank() orders by that and
multiplies in CYCLES where it knows an entry.  CYCLES is empty and stays
empty until the numbers are read off the ISA and checked like every other
row in the catalogue -- guessing them would make a count read as a cost
model.  Filling it is then a table edit rather than a rewrite.

spare_products() is the axis predicate the packing work needed and
packing.py deliberately does not carry: a term product occupies every
output column, so an entry holds a further one only where it is that much
wider.  Nine columns leave 23 of a 32-wide entry, which is two more
products; 7 of a 16-wide one, which is none.  The contraction axis has
the small quantum and no spare, and the lane axis costs a cross-lane
reduction, so the output axis is the one worth spending.

All three schemes enumerate candidates now, not only the emulated one --
exchange_ops and emu_tiles beside the singular selections, with the
conditions still stated once.  ComputeShape carries the leading dimension
and the depth so a ranking can read them; an extent of 0 counts as one
tile, so a caller that does not know them ranks by what it does.

emu_tiles keeps the emitter's k == m condition and now says what it
costs: the 32-wide entry, where three term products would share an issue
instead of one.

AI-generated. Model: Opus 5

Author: davschneller

src/tensorforge/backend/instructions/compute/multilinear.py

@@ -610,7 +610,9 @@ def _plan(self) -> Tuple[Span, ...]:
         shape = ComputeShape(threads=self._num_threads,
                              accumulator=self._idest.get_fptype(),
                              sparse=self._second_operand_is_sparse(),
-                             explicit_simd=_explicit_simd(self._context))
+                             explicit_simd=_explicit_simd(self._context),
+                             lead=self._ns[0][1] - self._ns[0][0],
+                             depth=math.prod(mx - mi for mi, mx in self._ks))
         chosen = choose_strategy(
             legal_strategies(module.strategies(shape, self._context)),
             self._context.get_vm().get_hw_descr().vendor)

src/tensorforge/backend/instructions/compute/primitives/amd/__init__.py

@@ -42,7 +42,7 @@
 from .arch import amdarch, cdna2, gfx1250, gfx1251, rdna
 from .caps import has_fmacdpp4, has_fmacdpp8, has_fmacdpp16
 from .catalog import (DEFINED_SPLITS, DEFINED_TRANSPOSES, MANTISSA,
-                      MATRIX_OPS, MFMA_TILES, emu_tile_for,
+                      MATRIX_OPS, MFMA_TILES, emu_tile_for, emu_tiles,
                       NOT_MODELLED, Call, Fragment, MatrixOp,
                       MfmaTile, lane_batched_ops, mfma_tile_for, ops_for,
                       usable_mfma_tiles)
@@ -54,9 +54,10 @@
                       accumulator_gathers, broadcast_feeds_a,
                       fragment_cost, fragment_moves)
 from .codegen import hfma, matmul32, matmulemu, matmuldpp
-from .exchange_codegen import exchange_op, matmul_exchange
-from .tiling import (EMULATION, EXCHANGE, Fit, Scheme, boundary,
-                     choose, offers)
+from .exchange_codegen import exchange_op, exchange_ops, matmul_exchange
+from .tiling import (CYCLES, EMULATION, EXCHANGE, Fit, Scheme, boundary,
+                     candidates, choose, issues, offers, rank,
+                     spare_products)
 from .emitters import fmadpp, fmadpp4, fmadpp8, fmadpp16, fmascalar
 from .relayout import (BROADCAST, MOVDPP16, RELAYOUTS, TRANSPOSE4X4, Relayout,
                        find_relayout)
@@ -70,7 +71,9 @@
     'FEATURE_TARGETS', 'has_feature', 'wave_size',
     'Call', 'Fragment', 'MatrixOp', 'MATRIX_OPS', 'MANTISSA',
     'DEFINED_SPLITS', 'emu_tile_for', 'matmulemu', 'EMULATION',
-    'EXCHANGE', 'Scheme', 'Fit', 'exchange_op', 'matmul_exchange',
+    'EXCHANGE', 'Scheme', 'Fit', 'exchange_op', 'exchange_ops',
+    'matmul_exchange', 'emu_tiles', 'CYCLES', 'candidates', 'issues',
+    'rank', 'spare_products',
     'NOT_MODELLED', 'ops_for',
     'MfmaTile', 'DEFINED_TRANSPOSES', 'MFMA_TILES', 'usable_mfma_tiles',
     'lane_batched_ops', 'mfma_tile_for',
@@ -144,7 +147,8 @@ def plan(strategy, shape, n, ctx):
     """
     if strategy is not Strategy.MATRIX:
         return whole(strategy, n)
-    fit = choose(shape.threads, shape.accumulator, ctx)
+    fit = choose(shape.threads, shape.accumulator, ctx,
+                 columns=n, lead=shape.lead, depth=shape.depth)
     # Asked of the scheme that will run rather than of one of them.  Only the
     # lane-batched one draws a boundary at all, and the threshold behind it is
     # a measurement against its own block width.
@@ -168,7 +172,8 @@ def matmul(writer, ops, ctx, span):
     if span.strategy is Strategy.BROADCAST:
         return broadcast.matmul(writer, ops, ctx, span)
     if span.strategy is Strategy.MATRIX:
-        fit = choose(threads, dtype, ctx)
+        fit = choose(threads, dtype, ctx, columns=N,
+                     lead=ops.lead_slots * threads, depth=K + kx)
         if fit is None or (ops.a, ops.b) != (fit.reads, fit.reads):
             # The entry was selected by what it accumulates in; what it
             # multiplies is a separate property of the same entry, and an

src/tensorforge/backend/instructions/compute/primitives/amd/catalog.py

@@ -594,7 +594,27 @@ def emu_tile_for(threads, dtype, ctx):
     already holds, and `layouts.position` puts them at the fragment's slots
     only when the two widths agree.
     """
+    found = emu_tiles(threads, dtype, ctx)
+    return found[0] if found else None
+
+
+def emu_tiles(threads, dtype, ctx):
+    """Every tile an emulated path could run here, largest first.
+
+    Restricted to entries whose k-vector is exactly the block width.  That is
+    the *emitter's* limit and not the catalogue's: `matmulemu` hands the
+    instruction the registers one `kk` step already holds, which are the
+    fragment's slots only when the two widths agree.  The stacking itself
+    holds wider -- `layouts.position` says so for the 16- and 32-wide entries
+    too -- and lifting the condition is what would reach them.
+
+    It is not free to leave.  A term product occupies every output column, so
+    an entry wider than the shape holds several: at nine columns the 32-wide
+    entry takes three products per issue where these take one, which is six
+    issues against two for a three-term split.
+    """
     from ...split import MANTISSA, terms as _terms
+    out = []
     for op in ops_for(dtype, ctx, threads):
         if op.arithmetic is not None or op.a.dtype is dtype:
             continue
@@ -606,12 +626,12 @@ def emu_tile_for(threads, dtype, ctx):
             continue
         if op.a.per_lane != op.k or op.b.per_lane != op.k:
             continue
-        tile = MfmaTile(op=op, transpose='tensorforge::transpose4x4b32',
+        tile = MfmaTile(op=op, transpose=_TILE_TRANSPOSES.get(op.m, (None, False))[0],
                         transpose_has_separate_outputs=True)
         if not tile.fits(threads) or tile.transpose not in DEFINED_TRANSPOSES:
             continue
-        return tile, _terms(MANTISSA[op.a.dtype], dtype)
-    return None
+        out.append((tile, _terms(MANTISSA[op.a.dtype], dtype)))
+    return tuple(out)
 
 
 def mfma_tile_for(threads, dtype, ctx):

src/tensorforge/backend/instructions/compute/primitives/amd/exchange_codegen.py

@@ -50,6 +50,18 @@ def exchange_op(dtype, threads, ctx) -> Optional[object]:
     a plan.  Asking all three here rather than discovering the third halfway
     through an emission is what lets `matmul` gate on one call.
     """
+    found = exchange_ops(dtype, threads, ctx)
+    return found[0] if found else None
+
+
+def exchange_ops(dtype, threads, ctx) -> tuple:
+    """Every instruction this path could emit here, largest first.
+
+    The same conditions, kept in one place: a selection that reads them and a
+    ranking that reads them again are two chances to disagree about what the
+    path can serve.
+    """
+    out = []
     for op in sorted(MATRIX_OPS, key=lambda o: (-o.m * o.n, o.builtin)):
         if op.call is not Call.MFMA or not op.available_for(dtype, ctx):
             continue
@@ -63,8 +75,8 @@ def exchange_op(dtype, threads, ctx) -> Optional[object]:
             continue
         if accumulator_gathers(op, 0) is None:
             continue
-        return op
-    return None
+        out.append(op)
+    return tuple(out)
 
 
 def _swapped(writer, value, blocks, ftype):

src/tensorforge/backend/instructions/compute/primitives/amd/tiling.py

@@ -27,11 +27,11 @@
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import FrozenSet, Optional, Tuple
+from typing import Dict, FrozenSet, Iterable, List, Optional, Tuple
 
 from ... import packing
-from .catalog import emu_tile_for, mfma_tile_for
-from .exchange_codegen import exchange_op
+from .catalog import emu_tile_for, emu_tiles, mfma_tile_for
+from .exchange_codegen import exchange_op, exchange_ops
 
 
 class Scheme(Enum):
@@ -119,23 +119,45 @@ def offers(threads, accumulator, ctx) -> FrozenSet[Scheme]:
                              Scheme.LANE_BATCHED)
 
 
-def choose(threads, accumulator, ctx) -> Optional[Fit]:
-    """The scheme this target prefers among the ones that fit, ready to emit."""
+def candidates(scheme: Scheme, threads, accumulator, ctx) -> Tuple[Fit, ...]:
+    """Every entry this scheme could run here, unranked.
+
+    Each scheme answers in its own shape -- a tile, a bare entry, a tile and a
+    term count -- and this is where that stops mattering to anyone else.
+    """
+    if scheme is Scheme.LANE_BATCHED:
+        tile = mfma_tile_for(threads, accumulator, ctx)
+        # One, and the policy in `mfma_tile_for` is what keeps it one: the
+        # wider tiles need a staging step that is not written.
+        return () if tile is None else (
+            Fit(scheme, tile.op, tile=tile, reads=tile.op.a.dtype),)
+    if scheme is Scheme.EXCHANGE:
+        return tuple(Fit(scheme, op, reads=op.a.dtype)
+                     for op in exchange_ops(accumulator, threads, ctx))
+    return tuple(
+        # What the split reads is the accumulator: the narrow type is what
+        # comes out of it, not what goes in.
+        Fit(scheme, tile.op, tile=tile, terms=terms, reads=accumulator)
+        for tile, terms in emu_tiles(threads, accumulator, ctx))
+
+
+def choose(threads, accumulator, ctx, columns: int = 0, lead: int = 0,
+           depth: int = 0) -> Optional[Fit]:
+    """The scheme this target prefers, and the cheapest entry it can run it on.
+
+    Two rankings, and they are not the same question.  Which scheme is a
+    preference between mappings and reads `ORDER`; which entry within it is a
+    count over the shape, because an entry wider than the output wastes the
+    difference in every issue and one narrower than it needs several.
+    """
     legal = offers(threads, accumulator, ctx)
     for scheme in ORDER:
         if scheme not in legal:
             continue
-        if scheme is Scheme.LANE_BATCHED:
-            tile = mfma_tile_for(threads, accumulator, ctx)
-            return Fit(scheme, tile.op, tile=tile, reads=tile.op.a.dtype)
-        if scheme is Scheme.EXCHANGE:
-            op = exchange_op(accumulator, threads, ctx)
-            return Fit(scheme, op, reads=op.a.dtype)
-        tile, terms = emu_tile_for(threads, accumulator, ctx)
-        # What the split reads, which is the accumulator: the narrow type is
-        # what comes out of it, not what goes in.
-        return Fit(scheme, tile.op, tile=tile, terms=terms,
-                   reads=accumulator)
+        fits = rank(candidates(scheme, threads, accumulator, ctx),
+                    columns, lead, depth)
+        if fits:
+            return fits[0]
     return None
 
 
@@ -155,3 +177,79 @@ def boundary(fit: Fit, n: int) -> int:
         return n
     empty = packing.waste(n, fit.width)
     return ((n // fit.width) * fit.width) if empty in (0, fit.width - 1) else n
+
+
+# -- how much of an entry a shape actually uses ---------------------------- #
+
+def spare_products(op, columns: int) -> int:
+    """Term products one issue holds on the output axis, the real one included.
+
+    A term product is a whole ``C += A_i x B_j`` over the tile, so it occupies
+    every output column the shape has -- which is why partial spare is no use
+    here and the count is a division rather than a remainder.  An entry
+    narrower than the shape holds one and no more.
+
+    This is the axis worth spending, and the only one of the three.  The
+    contraction axis has the small quantum -- one slot per step -- and no
+    spare, since the contraction fills it; the lane axis has spare but a
+    product needs the whole leading dimension there, and reclaiming it is a
+    cross-lane reduction rather than an add chain inside a lane.
+    """
+    if columns <= 0:
+        return 1
+    return max(1, op.m // columns)
+
+
+def issues(op, columns: int, lead: int = 0, depth: int = 0,
+           products: int = 1) -> int:
+    """Instruction issues this entry takes for the whole contraction.
+
+    The four axes multiplied: output columns against `m`, the leading
+    dimension against the lanes one issue spans, the contraction against `k`,
+    and the term products against what one issue can hold beside the real one.
+    An extent given as 0 is one the caller does not know, and counts as a
+    single tile rather than as nothing.
+
+    It counts issues, not time.  Two entries that differ in passes do not cost
+    the same per issue, which is what `CYCLES` is for and why this is not
+    called a cost.
+    """
+    def tiles(demand, capacity):
+        return packing.tiles(demand, capacity) if demand > 0 else 1
+
+    return (tiles(columns, op.m)
+            * tiles(lead, op.n * op.blocks)
+            * tiles(depth, op.k)
+            * packing.tiles(products, spare_products(op, columns)))
+
+
+#: Issue cost per instruction, in cycles.  Empty.
+#:
+#: The numbers are a hardware fact like every other row in the catalogue, and
+#: they have to be read off the ISA guide and checked the same way -- AMD
+#: documents them per instruction; for other vendors this table would need its
+#: own source.  Guessing them is worse than not having them: the ranking below
+#: then reads as a cost model and is a count.
+#:
+#: Until it is filled, `rank` orders by issues alone, which is exactly wrong
+#: where two entries differ in passes.  A 32x32 issue is not one 4x4 issue,
+#: and fewer issues of a longer instruction can be slower.
+CYCLES: Dict[str, int] = {}
+
+
+def rank(fits, columns: int, lead: int = 0, depth: int = 0):
+    """The candidates, cheapest first.
+
+    Fewest issues wins, and where `CYCLES` knows an entry its issues are
+    weighted by it.  Deliberately not "the narrowest that fits": four 4x4
+    issues and one 16x16 issue both serve thirteen columns, the second wastes
+    three of its sixteen, and which is faster is a property of the two
+    instructions rather than of the waste.  Stating the order here is what
+    lets that be answered by filling a table instead of by rewriting a
+    selection.
+    """
+    def key(fit):
+        count = issues(fit.op, columns, lead, depth, fit.terms)
+        return (count * CYCLES.get(fit.op.builtin, 1), count, fit.op.builtin)
+
+    return tuple(sorted(fits, key=key))

src/tensorforge/backend/instructions/compute/strategy.py

@@ -87,6 +87,13 @@ class ComputeShape:
     #: broadcast costs differs between them by more than its spelling.
     explicit_simd: bool
 
+    #: Elements the leading dimension spans, and steps the contraction takes.
+    #: 0 where the caller does not know them -- a count that reads an extent
+    #: of 0 treats it as one tile rather than as nothing, so an unfilled pair
+    #: narrows what a ranking can tell apart without making it wrong.
+    lead: int = 0
+    depth: int = 0
+
 
 def is_contraction(operands: int, lead_width: int) -> bool:
     """Whether anything but the nest could compute this at all.

tests/test_compute_strategy.py

@@ -286,3 +286,56 @@ def test_only_the_lane_batched_scheme_draws_a_boundary():
         assert tiling.boundary(swap, n) == n
     assert tiling.boundary(lane, 9) == 8
     assert tiling.boundary(lane, 10) == 10
+
+
+# -- choosing the entry within a scheme ------------------------------------ #
+
+def test_the_entry_is_ranked_by_issues_not_by_fit():
+    """Thirteen columns take four 4x4 issues or one 16x16 that wastes three
+    of its sixteen.  Which is faster is a property of the two instructions
+    rather than of the waste, so the count is what is stated and `CYCLES` is
+    where the rest would go."""
+    from tensorforge.backend.instructions.compute.primitives.amd import (
+        MATRIX_OPS, tiling)
+    ops = {o.builtin: o for o in MATRIX_OPS}
+    narrow, wide = ops['mfma_f32_4x4x1f32'], ops['mfma_f32_16x16x1f32']
+    assert tiling.issues(narrow, 13, 56, 56) == 4 * tiling.issues(
+        wide, 13, 56, 56)
+    assert tiling.CYCLES == {}, 'a guessed cycle count reads as a measurement'
+
+
+def test_a_term_product_takes_the_whole_output_axis():
+    """Which is why partial spare is no use: a product is a complete tile,
+    not a column of one."""
+    from tensorforge.backend.instructions.compute.primitives.amd import (
+        MATRIX_OPS, tiling)
+    ops = {o.builtin: o for o in MATRIX_OPS}
+    for name, expected in (('mfma_f32_4x4x4bf16_1k', 1),
+                           ('mfma_f32_16x16x4bf16_1k', 1),
+                           ('mfma_f32_32x32x4bf16_1k', 3)):
+        assert tiling.spare_products(ops[name], 9) == expected, name
+
+
+def test_an_unknown_extent_counts_as_one_tile():
+    """A shape that does not carry its leading dimension or its depth still
+    ranks by what it does carry, rather than ranking everything equal."""
+    from tensorforge.backend.instructions.compute.primitives.amd import (
+        MATRIX_OPS, tiling)
+    ops = {o.builtin: o for o in MATRIX_OPS}
+    op = ops['mfma_f32_4x4x1f32']
+    assert tiling.issues(op, 9) == tiling.issues(op, 9, lead=1, depth=1)
+
+
+def test_the_emulated_limit_has_a_number_on_it():
+    """The emitter takes only entries whose k-vector is the block width, and
+    that is what keeps the 32-wide one out -- where three term products would
+    share an issue instead of one."""
+    from tensorforge.backend.instructions.compute.primitives.amd import (
+        MATRIX_OPS, emu_tiles, tiling)
+    ctx = _hip()
+    offered = {tile.op.builtin for tile, _ in emu_tiles(64, Datatype.F32, ctx)}
+    ops = {o.builtin: o for o in MATRIX_OPS}
+    assert offered == {'mfma_f32_4x4x4bf16_1k'}
+    wide = ops['mfma_f32_32x32x4bf16_1k']
+    assert wide.k != wide.m
+    assert tiling.spare_products(wide, 9) == 3