refactor: the epilogue is the solver's answer, not a second

 opinion

accumulator_gathers walked the lanes, asked layouts.position where each
element sat, grouped by region and checked that one XOR carried every
lane of it. bitlayout.moves is that question asked of two layouts, and
patch 15 showed the two agree on 1888 regions -- so keeping both was
keeping a second implementation of one thing, which is how they drift.

It reads the solver now. The plan is identical for every entry and every
column: 756 answers, 692 of them plans, compared field by field against
the implementation it replaces.

moves takes a base on each side. Which one carries the group offset
depends on which of the two the wave is divided over -- the epilogue
gathers into the nest, so it is the destination, while the transpose
check has it on the source -- and stating both is what lets one function
serve both directions. Move.lanes are the destination lanes, because
that is the region a caller has to mask when it merges.

AI-generated. Model: Opus 5

Author: davschneller

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

@@ -211,7 +211,8 @@ class Move:
 
 
 def moves(have: BitLayout, want: BitLayout, indices,
-          base: Position = Position()) -> Optional[Tuple[Move, ...]]:
+          base_have: Position = Position(),
+          base_want: Position = Position()) -> Optional[Tuple[Move, ...]]:
     """How to get from one distribution to the other, or `None`.
 
     The `swap` family does not permute bits: `swap<B>` reads the lane `B` away,
@@ -221,10 +222,14 @@ def moves(have: BitLayout, want: BitLayout, indices,
     lane and a slot, which is a transpose and `relayout.find_relayout`'s
     business.
 
-    `base` is a constant added to the source, which is what a caller holding
-    one of several groups of a wave contributes: the group index is an offset
-    on the lane and not a bit of any index the two layouts share, so it cannot
-    be an axis of either.
+    The two bases are constants added to each side, which is what a caller
+    holding one of several groups of a wave contributes: the group index is an
+    offset on the lane and not a bit of any index the two layouts share, so it
+    cannot be an axis of either.  Which side carries it depends on which of
+    them the wave is divided over, so both are stated rather than one assumed.
+
+    `Move.lanes` are the *destination* lanes -- the region of `want`'s register
+    this writes -- because that is what a caller has to mask when it merges.
 
     `None` where some pair of slots needs more than one toggle.  Not a
     limitation to route around: that is the case where no sequence of reads
@@ -239,7 +244,8 @@ def moves(have: BitLayout, want: BitLayout, indices,
         return None
     regions = {}
     for index in indices:
-        here, there = base + have.locate(*index), want.locate(*index)
+        here = base_have + have.locate(*index)
+        there = base_want + want.locate(*index)
         if here.element or there.element:
             # A vector element is neither a lane nor a slot, so no lane toggle
             # reaches it; a packed operand needs the element bits accounted

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

@@ -87,6 +87,8 @@
 
 from . import layouts
 
+from ... import bitlayout
+
 #: `quad_perm:[0,1,2,3]`.  The DPP control that moves nothing, so that
 #: `dppUpdate` is a masked merge and not also a shuffle.
 IDENTITY_DPP = 0xE4
@@ -400,21 +402,32 @@ def accumulator_gathers(op, column: int) -> Optional[Tuple[Gather, ...]]:
     span = op.n * op.blocks
     if op.wave % span:
         return None
+    fragment = layouts.fragment_layout(op, 'D')
+    if fragment is None:
+        return None
 
-    regions = {}
-    for lane in range(op.wave):
-        group, local = divmod(lane, span)
-        block, index = divmod(local, op.n)
-        slot, source = layouts.position(op, 'D', column, index, block)
-        regions.setdefault((group, slot), []).append((lane, source))
+    # The nest's own register: lane `l` holds leading-dimension element `l`,
+    # and nothing varies with the column -- each one is a separate store.
+    nest = bitlayout.BitLayout((
+        tuple(bitlayout.Bit(bitlayout.Place.LANE, op.n << bit)
+              for bit in range((op.blocks - 1).bit_length())),
+        (),
+        tuple(bitlayout.Bit(bitlayout.Place.LANE, 1 << bit)
+              for bit in range((op.n - 1).bit_length())),
+    ))
+    indices = [(block, column, index)
+               for block in range(op.blocks) for index in range(op.n)]
 
     gathers = []
-    for (group, slot), pairs in sorted(regions.items()):
-        masks = {lane ^ source for lane, source in pairs}
-        if len(masks) != 1:
+    for group in range(op.wave // span):
+        found = bitlayout.moves(
+            fragment, nest, indices,
+            base_want=bitlayout.Position(lane=group * span))
+        if found is None:
             return None
-        gathers.append(Gather(group, slot, _swaps_for(masks.pop()),
-                              Select.of((lane for lane, _ in pairs), op.wave)))
+        gathers += [Gather(group, move.source, _swaps_for(move.xor),
+                           Select.of(move.lanes, op.wave))
+                    for move in found]
     return tuple(gathers)
 
 

tests/test_bitlayout.py

@@ -278,7 +278,7 @@ def test_the_solver_finds_the_swaps_the_plan_states(name):
                    for b in range(op.blocks) for i in range(op.n)]
         for group in range(op.wave // span):
             found = bitlayout.moves(have, fragment, indices,
-                                    base=Position(lane=group * span))
+                                    base_have=Position(lane=group * span))
             assert found is not None, (name, column, group)
             for move in found:
                 assert stated[(group, move.target)] == reorder._swaps_for(
@@ -415,3 +415,31 @@ def test_a_packed_operand_is_not_reached_by_a_transpose():
     ))
     assert relayout.transposes_between(
         packed, relayout.transposed(4, 64), 4) is None
+
+
+def test_the_gathers_are_the_solver_s_answer():
+    """`accumulator_gathers` reads the solver now instead of walking lanes.
+
+    The plan it returns is the same for every entry and every column -- 756
+    answers, 692 of them plans -- which is what makes reading it from two
+    layouts a refactor rather than a second opinion.
+    """
+    from tensorforge.backend.instructions.compute.primitives.amd import reorder
+    planned = sum(1 for op in MATRIX_OPS for column in range(op.m)
+                  if reorder.accumulator_gathers(op, column) is not None)
+    assert planned == 692
+
+
+def test_a_base_on_each_side():
+    """Which side carries the group offset depends on which of the two the
+    wave is divided over: the epilogue gathers *into* the nest, so the offset
+    is on the destination, and the transpose check has it on the source."""
+    layout = BitLayout(((Bit(Place.LANE, 1),),))
+    indices = [(0,), (1,)]
+    forward = bitlayout.moves(layout, layout, indices,
+                              base_have=Position(lane=8))
+    backward = bitlayout.moves(layout, layout, indices,
+                               base_want=Position(lane=8))
+    assert forward is not None and backward is not None
+    assert {m.xor for m in forward} == {8}
+    assert {m.xor for m in backward} == {8}

tools/mutation_check.py

@@ -421,10 +421,14 @@ def make():
          sub(Path('src/tensorforge/backend/instructions/compute/bitlayout.py'),
              '        regions.setdefault((here.slot, there.slot), []).append(',
              '        regions.setdefault((0, there.slot), []).append(')),
-        ('the group offset dropped from the source',
+        ('the group offset dropped from the destination',
          sub(Path('src/tensorforge/backend/instructions/compute/bitlayout.py'),
-             '        here, there = base + have.locate(*index), want.locate(*index)',
-             '        here, there = have.locate(*index), want.locate(*index)')),
+             '        there = base_want + want.locate(*index)',
+             '        there = want.locate(*index)')),
+        ('the destination lanes reported as the source lanes',
+         sub(Path('src/tensorforge/backend/instructions/compute/bitlayout.py'),
+             'tuple(sorted({there for _, there in pairs}))',
+             'tuple(sorted({here for here, _ in pairs}))')),
         ('a slot weight read as a lane weight',
          sub(Path('src/tensorforge/backend/instructions/compute/bitlayout.py'),
              'Bit(Place.LANE if weight > 0 else Place.SLOT, abs(weight))',