Give readers that index a tensor along unrelated dims their own transfer, so self-attention keeps its query and key axes apart

stream/allocation/lowering.py

@@ -17,6 +17,7 @@
 from stream.mapping.mapping import Mapping
 from stream.opt.allocation.constraint_optimization.hardware import build_hardware_facts
 from stream.profiling import span
+from stream.workload.affine_access import map_dim_positions
 from stream.workload.iterator_type import is_state_operand, streamed_operands
 from stream.workload.node import (
     ComputationNode,
@@ -117,45 +118,77 @@ def build_transfer_graph(self) -> Workload:
             srcs = [n for n in self.workload.nodes if isinstance(n, HasOutputs) and tensor in n.outputs]
             assert len(srcs) == 1, f"Expected exactly one source for tensor {tensor}, found {len(srcs)}"
             src = new_nodes[srcs[0].name]
-            dsts = [new_nodes[n.name] for n in self.workload.nodes if isinstance(n, HasInputs) and tensor in n.inputs]
+            readers = [n for n in self.workload.nodes if isinstance(n, HasInputs) and tensor in n.inputs]
+            dsts = [new_nodes[n.name] for n in readers]
             is_constant_o_transfer = any(isinstance(dst, OutEdge) for dst in dsts)
             if is_constant_o_transfer and not isinstance(src, InEdge):
                 self.add_two_transfer_nodes_for_constant_output_transfer(tensor, src, dsts, new_nodes)
             else:
-                self.add_transfer_nodes(tensor, src, dsts, new_nodes)
+                groups = [[new_nodes[n.name] for n in group] for group in self._reader_walks(tensor, readers)]
+                self.add_transfer_nodes(tensor, src, groups, new_nodes)
         new_workload = Workload(new_nodes.values())
         return new_workload
 
-    def add_transfer_nodes(self, tensor: Tensor, src: HasOutputs, dsts: list[HasInputs], new_nodes: dict[str, Node]):
+    def _reader_walks(self, tensor: Tensor, readers: list[HasInputs]) -> list[list[HasInputs]]:
+        """The readers of ``tensor`` grouped so readers in a group share a unique dim at each of its indices.
+
+        One copy is walked along one dim per index, so readers that index it along unrelated dims (the queries and
+        the keys of a self-attention both read its input rows) each get their own transfer; output edges join the
+        first group."""
+        n_unique = len(self.workload.unique_dimensions()[0])
+        groups: list[tuple[list[frozenset[int]], list[HasInputs]]] = []
+        for reader in readers:
+            if not isinstance(reader, HasIterationSpace):
+                continue
+            dims, access = self.workload.get_dims(reader), reader.get_mapping(tensor)
+            walk = [
+                frozenset().union(
+                    *(
+                        map_dim_positions(AffineMap(n_unique, 0, (dims[p],)))
+                        for p in map_dim_positions(AffineMap(access.num_dims, 0, (r,)))
+                    )
+                )
+                for r in access.results
+            ]
+            group = next((g for g in groups if all(a & b for a, b in zip(g[0], walk, strict=True))), None)
+            if group is None:
+                groups.append((walk, [reader]))
+            else:
+                group[1].append(reader)
+        if len(groups) <= 1:
+            return [readers]
+        groups[0][1].extend(r for r in readers if not isinstance(r, HasIterationSpace))
+        return [members for _, members in groups]
+
+    def add_transfer_nodes(
+        self, tensor: Tensor, src: HasOutputs, groups: list[list[HasInputs]], new_nodes: dict[str, Node]
+    ):
         """
-        Move ``tensor`` from its source to its destinations, either directly or staged on a memory tile.
+        Move ``tensor`` from its source to each group of destinations, either directly or staged on a memory tile.
 
         Staging splits the move in two -- source to the on-chip buffer, buffer to every destination --
         so the tile can hold the tensor across reads and re-lay it out on the way through. See
-        :meth:`_stages_on_mem_tile`.
+        :meth:`_stages_on_mem_tile`. Every group after the first is named after its first reader.
         """
-        if not self._stages_on_mem_tile(tensor, src, dsts):
-            transfer_type = self.determine_transfer_type(src, dsts)
-            out_name = f"{tensor.name}_1"
-            transfer_node, updated_tensors = self.generate_transfer_node(dsts, tensor, transfer_type, out_name)
+        dsts = [dst for group in groups for dst in group]
+        staged = self._stages_on_mem_tile(tensor, src, dsts)
+        source = tensor
+        if staged:
+            transfer_type_1 = self.determine_transfer_type(src, dsts, dst_type="memory")
+            hop, (source,) = self.generate_transfer_node([src], tensor, transfer_type_1, f"{tensor.name}_1")
+            new_nodes[hop.name] = hop
+        for i, group in enumerate(groups):
+            tag = f" for {group[0].name}" if i else ""
+            out_name = f"{tensor.name}_{group[0].name}" if i else f"{tensor.name}_{2 if staged else 1}"
+            transfer_type = (
+                self.determine_transfer_type(src, group, src_type="memory")
+                if staged
+                else self.determine_transfer_type(src, group)
+            )
+            transfer_node, updated_tensors = self.generate_transfer_node(group, source, transfer_type, out_name, tag)
             new_nodes[transfer_node.name] = transfer_node
-            for dst, updated_tensor in zip(dsts, updated_tensors, strict=True):
+            for dst, updated_tensor in zip(group, updated_tensors, strict=True):
                 self.update_destination_node_inputs(tensor, src, new_nodes, dst, updated_tensor)
-            return
-        # First transfer node from source to on-chip buffer
-        transfer_type_1 = self.determine_transfer_type(src, dsts, dst_type="memory")
-        out_name_1 = f"{tensor.name}_1"
-        transfer_node_1, updated_tensors_1 = self.generate_transfer_node([src], tensor, transfer_type_1, out_name_1)
-        new_nodes[transfer_node_1.name] = transfer_node_1
-        # Second transfer node from on-chip buffer to destinations
-        out_name_2 = f"{tensor.name}_2"
-        transfer_type_2 = self.determine_transfer_type(src, dsts, src_type="memory")
-        transfer_node_2, updated_tensors_2 = self.generate_transfer_node(
-            dsts, updated_tensors_1[0], transfer_type_2, out_name_2
-        )
-        new_nodes[transfer_node_2.name] = transfer_node_2
-        for dst, updated_tensor in zip(dsts, updated_tensors_2, strict=True):
-            self.update_destination_node_inputs(tensor, src, new_nodes, dst, updated_tensor)
 
     def _stages_on_mem_tile(self, tensor: Tensor, src: HasOutputs, dsts: list[HasInputs]) -> bool:
         """Whether a transfer is staged in a memory tile instead of landing straight on the cores.
@@ -300,12 +333,12 @@ def update_mapping(self):
         return self.mapping.with_updated_workload(self.ssw, self.workload)  # updates FusedGroups
 
     def generate_transfer_node(
-        self, dsts: list[HasInputs], tensor: Tensor, transfer_type: TransferType, out_name: str = ""
+        self, dsts: list[HasInputs], tensor: Tensor, transfer_type: TransferType, out_name: str = "", tag: str = ""
     ) -> tuple[TransferNode, list[Tensor]]:
         transfer_outputs = self.generate_transfer_output_tensors(tensor, dsts, out_name)
         operand_mapping = tuple(AffineMap.identity(len(tensor.shape)) for _ in range(1 + len(dsts)))
         transfer_node = TransferNode(
-            name=f"Transfer({tensor.name})",
+            name=f"Transfer({tensor.name}{tag})",
             inputs=(tensor,),
             outputs=tuple(transfer_outputs),
             transfer_type=transfer_type,

stream/opt/allocation/constraint_optimization/space.py

@@ -351,10 +351,13 @@ def moved_bits(
         )
 
     def copied_bits(self, t: Tensor) -> float:
-        """Bits per firing the transfer of ``t`` delivers, in place or not: a halo once per target holding it."""
+        """Bits per firing the transfers of ``t`` deliver, in place or not: a halo once per target holding it."""
+        if len(readers := [tr for tr in self.transfer_nodes if t in tr.inputs]) > 1:
+            return sum(self._delivered_bits(tr) for tr in readers)
         tr = next((tr for tr in self.transfer_nodes if t in tr.tensors), None)
-        if tr is None:
-            return t.size_bits()
+        return t.size_bits() if tr is None else self._delivered_bits(tr)
+
+    def _delivered_bits(self, tr: TransferNode) -> float:
         if overlaps := self.overlaps(tr):
             return sum(overlaps.values()) * tr.inputs[0].operand_type.bitwidth
         return self.moved_bits(tr, self.path_choices[tr][0])

tests/unit/test_model_catalog.py

@@ -99,3 +99,20 @@ def test_a_recurrence_with_an_initial_state_input_solves_end_to_end(tmp_path, bl
     lowered = allocation.problem.workload
     assert not {t.name for tr in lowered.get_transfer_nodes() for t in tr.inputs} & states
     assert not {t.name for e in lowered.get_in_edges() for t in e.outputs} & states
+
+
+def test_self_attention_keeps_its_query_and_key_axes_apart_in_the_steady_state(tmp_path):
+    """The queries and the keys both read the input rows: each reader gets its own copy, so the scores still span
+    every query by every key."""
+    workload = build_block("attention", batch=1, heads=2, seq=16, d_head=16)
+    allocation = evaluate_mapping(
+        "stream/inputs/examples/hardware/tpu_like_quad_core.yaml", workload, str(tmp_path)
+    ).context.get("allocation")
+    lowered = allocation.problem.workload
+    scores = next(n for n in lowered.get_computation_nodes() if n.name == "scores")
+    query, key = (lowered.get_dims(scores)[p] for p in (2, 3))
+    assert query != key
+    assert {tr.name for tr in lowered.get_transfer_nodes() if tr.inputs[0].name == "x"} == {
+        "Transfer(x)",
+        "Transfer(x for proj_q)",
+    }