From aa845a6c86decfdd1f37932418515edefa798acc Mon Sep 17 00:00:00 2001 From: Arne Symons Date: Tue, 6 Oct 2026 10:08:55 +0200 Subject: [PATCH] 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 | 87 +++++++++++++------ .../constraint_optimization/space.py | 9 +- tests/unit/test_model_catalog.py | 17 ++++ 3 files changed, 83 insertions(+), 30 deletions(-) diff --git a/stream/allocation/lowering.py b/stream/allocation/lowering.py index 80452999..7a28e7f0 100644 --- a/stream/allocation/lowering.py +++ b/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, diff --git a/stream/opt/allocation/constraint_optimization/space.py b/stream/opt/allocation/constraint_optimization/space.py index 9c4cdd04..23d90e92 100644 --- a/stream/opt/allocation/constraint_optimization/space.py +++ b/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]) diff --git a/tests/unit/test_model_catalog.py b/tests/unit/test_model_catalog.py index 63a729a6..11d326fc 100644 --- a/tests/unit/test_model_catalog.py +++ b/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)", + }