From f7fd64b3fb2a63744b79ba1cc34691a1d56228a6 Mon Sep 17 00:00:00 2001 From: Michael Gilbert Date: Fri, 26 Jun 2026 10:52:41 -0400 Subject: [PATCH 1/5] [FFM] Add binding to compatibility; needs cleanup --- .../FFM/_join_pmappings/compatibility.py | 88 ++++++++++- .../make_pmapping_templates.py | 1 + accelforge/model/main.py | 1 + tests/test_compatibility_physical_spatial.py | 137 ++++++++++++++++++ 4 files changed, 224 insertions(+), 3 deletions(-) create mode 100644 tests/test_compatibility_physical_spatial.py diff --git a/accelforge/mapper/FFM/_join_pmappings/compatibility.py b/accelforge/mapper/FFM/_join_pmappings/compatibility.py index b3b97464..89ee6453 100755 --- a/accelforge/mapper/FFM/_join_pmappings/compatibility.py +++ b/accelforge/mapper/FFM/_join_pmappings/compatibility.py @@ -31,6 +31,8 @@ # 1. Each tensor is stored above some loop index. 0 is the outermost loop, 1 the # next-innermost... # 2. All loops above any shared tensor are co-tiled and must match between PmappingGroups. +# 3. Spatial loops *below* a physically-distributed storage (i.e., the data binding) +# must match. These are in TensorReservations.physical_spatial_loops. T = TypeVar("T", bound="Updatable") @@ -58,11 +60,14 @@ class Loop(Updatable): rank_name: Rank tile_pattern: TilePattern | None is_spatial: bool + spatial_dim: str | None = None + # The architecture spatial dimension (e.g. "X", "Y") this loop fans out over. def __post_init__(self): assert isinstance(self.rank_name, Rank) assert isinstance(self.tile_pattern, Number | TilePattern | str | None) assert isinstance(self.is_spatial, bool) + assert isinstance(self.spatial_dim, str | None) assert isinstance( self.tile_pattern.initial_tile_shape, Number | str | None, @@ -162,10 +167,16 @@ class TensorReservation(Updatable): name: TensorName resource_name: str persistent: bool = False + # Spatial loops *below* this storage that distribute the tensor across physical + # instances + physical_spatial_loops: tuple[Loop, ...] = () def __post_init__(self): if self.persistent: assert len(self.loops) == 0, "Persistent tensors be above all loops" + assert all( + isinstance(l, Loop) and l.is_spatial for l in self.physical_spatial_loops + ), "physical_spatial_loops must all be spatial Loops" @property def above_loop_index(self) -> int: @@ -175,7 +186,12 @@ def __str__(self): return f"[{self.resource_name}] {self.name} below {self.loops}" def __repr__(self): - return f"Reservation({repr(self.name)}, {repr(self.loops)}, {repr(self.resource_name)})" + phys = ( + f", physical_spatial_loops={repr(self.physical_spatial_loops)}" + if self.physical_spatial_loops + else "" + ) + return f"Reservation({repr(self.name)}, {repr(self.loops)}, {repr(self.resource_name)}{phys})" def pydot_str(self): return f"[{self.resource_name}] {self.name}" @@ -216,6 +232,9 @@ def _prepend_symbols(self, prepend: str) -> "TensorReservation": def clear_symbolic_tile_patterns(self) -> "TensorReservation": return self.update( loops=tuple(l.clear_symbolic_tile_patterns() for l in self.loops), + physical_spatial_loops=tuple( + l.clear_symbolic_tile_patterns() for l in self.physical_spatial_loops + ), ) def make_fused_loop_symbols( @@ -243,7 +262,20 @@ def _rename_to_match( l_mine, new_renames = l_mine._rename_to_match(l_other) _update_rename_dict(renames, new_renames) new_loops.append(l_mine) - return self.update(loops=tuple(new_loops)), renames + new_physical = [] + for l_mine, l_other in zip( + self.physical_spatial_loops, other.physical_spatial_loops + ): + l_mine, new_renames = l_mine._rename_to_match(l_other) + _update_rename_dict(renames, new_renames) + new_physical.append(l_mine) + return ( + self.update( + loops=tuple(new_loops), + physical_spatial_loops=tuple(new_physical), + ), + renames, + ) class SplitKind(Enum): @@ -533,10 +565,16 @@ def from_mapping( mapping: Mapping, tensors: set[TensorName], einsum: Einsum, + flattened_arch=None, ) -> "Compatibility": """ Create Compatibility from a mapping, a set of fusable tensors, and the workload. + + If ``flattened_arch`` is given, spatial loops below a physically-distributed + storage are recorded in each ``TensorReservation.physical_spatial_loops`` so that + fused Einsums sharing a tensor must agree on its physical placement. If it is + ``None`` (the default), no such loops are recorded and behavior is unchanged. """ if not isinstance(einsum, Einsum): raise TypeError(f"einsum should be an Einsum, but {type(einsum)} instead") @@ -544,7 +582,6 @@ def from_mapping( t.name: t.rank_variable2ranks for t in einsum.tensor_accesses } - # TODO: update compatibility to handle spatial-for loop per-tensor update tensor_indices = [] split_above_loop_indices = [] reservation_indices = [] @@ -597,6 +634,48 @@ def make_loops(above_index: int, tensor_name: TensorName) -> list[MappingLoop]: ] return tuple(loops) + def make_physical_spatial_loops( + above_index: int, tensor_name: TensorName, resource: str + ) -> tuple[Loop, ...]: + # Spatial loops below a distributed storage fix which physical instance holds + # which slice of the tensor; fused Einsums must agree on them. Only relevant + # when the backing storage has a physical stride. + if flattened_arch is None: + return () + try: + storage = flattened_arch[resource] + except (KeyError, ValueError): + return () + if not getattr(storage, "_is_distributed", lambda: False)(): + return () + out = [] + for n in mapping.nodes[above_index + 1 :]: + # Stop at the next storage level: loops below it belong to that storage. + if isinstance(n, (MappingReservation, TensorHolder)): + break + if not isinstance(n, Spatial): + continue + # Only dimensions the storage is physically distributed along matter. + if not storage._has_physical_dim(n.name): + continue + rank = get_rank(n.rank_variable, tensor_name) + # Skip loops that don't partition this tensor (don't index it). + if rank == Rank("NO RANK. RECOMPUTED."): + continue + # Clear symbolic tile patterns: below-storage spatial loops never get + # DataFrame columns, so they're matched structurally (rank + dim, plus a + # concrete fanout when known; symbolic fanouts are pinned by the fixed + # hardware fanout along the dimension). + out.append( + Loop( + rank_name=rank, + tile_pattern=n.tile_pattern._symbol2str(), + is_spatial=True, + spatial_dim=n.name, + ).clear_symbolic_tile_patterns() + ) + return tuple(out) + return cls( tensors=fzs( TensorReservation( @@ -604,6 +683,9 @@ def make_loops(above_index: int, tensor_name: TensorName) -> list[MappingLoop]: loops=make_loops(i, mapping.nodes[i].purpose), resource_name=mapping.nodes[i].resource, persistent=mapping.nodes[i].persistent, + physical_spatial_loops=make_physical_spatial_loops( + i, mapping.nodes[i].purpose, mapping.nodes[i].resource + ), ) for i in tensor_indices ), diff --git a/accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py b/accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py index 6adbfe7f..afba294d 100755 --- a/accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py +++ b/accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py @@ -455,6 +455,7 @@ def make_pmapping_templates(job: Job, print_progress: bool = True) -> SameEinsum mapping_with_reservations, new_job.fusable_tensors, new_job.einsum, + flattened_arch=new_job.flattened_arch, ) jobs.append(new_job) diff --git a/accelforge/model/main.py b/accelforge/model/main.py index 7f9cb5e5..c51cea75 100644 --- a/accelforge/model/main.py +++ b/accelforge/model/main.py @@ -183,6 +183,7 @@ def evaluate_mapping( job.mapping, job.fusable_tensors, einsum, + flattened_arch=job.flattened_arch, ) symbol_renames, compatibility = compatibility.make_fused_loop_symbols( einsum_name diff --git a/tests/test_compatibility_physical_spatial.py b/tests/test_compatibility_physical_spatial.py new file mode 100644 index 00000000..0dbf4ae2 --- /dev/null +++ b/tests/test_compatibility_physical_spatial.py @@ -0,0 +1,137 @@ +""" +Tests for recording spatial-for loops below a physically-distributed storage in +``Compatibility`` (see ``TensorReservation.physical_spatial_loops``). + + pytest tests/test_compatibility_physical_spatial.py -v + python -m unittest tests.test_compatibility_physical_spatial +""" + +import types +import unittest + +import accelforge as af +from accelforge.frontend.spec import Spec +from accelforge.frontend.mapping import Reservation, Spatial +from accelforge.mapper.FFM._join_pmappings.compatibility import Compatibility, Loop + + +class _StubStorage: + """A storage that is distributed along a fixed set of physical dimensions.""" + + def __init__(self, distributed, physical_dims): + self._distributed = distributed + self._physical_dims = set(physical_dims) + + def _is_distributed(self): + return self._distributed + + def _has_physical_dim(self, dim_name): + return dim_name in self._physical_dims + + +class _StubArch: + def __init__(self, name_to_storage): + self._n2s = name_to_storage + + def __getitem__(self, name): + return self._n2s[name] # raises KeyError if absent, like FlattenedArch + + +def _matmul_einsum(): + spec = Spec.from_yaml( + af.examples.arches.simple, + af.examples.workloads.matmuls, + jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64}, + ) + name = list(spec.workload.einsum_names)[0] + return spec.workload.einsums[name] + + +def _mapping(nodes): + return types.SimpleNamespace(nodes=list(nodes)) + + +class TestPhysicalSpatialLoops(unittest.TestCase): + def setUp(self): + self.einsum = _matmul_einsum() + # T1 is indexed by rank variables m -> M and n1 -> N1; n0 -> N0 does NOT index it. + self.res = Reservation(purposes=["T1"], resource="DistBuf") + self.sp_m = Spatial(rank_variable="m", name="X", component="DistBuf", tile_shape=4) + self.sp_n1 = Spatial(rank_variable="n1", name="Y", component="DistBuf", tile_shape=4) + self.sp_n0 = Spatial(rank_variable="n0", name="X", component="DistBuf", tile_shape=4) + self.mapping = _mapping([self.res, self.sp_m, self.sp_n1, self.sp_n0]) + + def _from(self, arch): + return Compatibility.from_mapping( + self.mapping, {"T1"}, self.einsum, flattened_arch=arch + ) + + def _phys(self, compat): + (t,) = compat.tensors + return t.physical_spatial_loops + + def test_no_arch_records_nothing(self): + self.assertEqual(self._phys(self._from(None)), ()) + + def test_non_distributed_records_nothing(self): + arch = _StubArch({"DistBuf": _StubStorage(False, {"X", "Y"})}) + self.assertEqual(self._phys(self._from(arch)), ()) + + def test_dim_filter(self): + # Only X is a physical dim: only the 'm' loop (over X, indexing T1) is recorded. + arch = _StubArch({"DistBuf": _StubStorage(True, {"X"})}) + phys = self._phys(self._from(arch)) + self.assertEqual(len(phys), 1) + loop = phys[0] + self.assertTrue(loop.is_spatial) + self.assertEqual(str(loop.rank_name), "M") + self.assertEqual(loop.spatial_dim, "X") + # Concrete fanout is preserved through clearing. + self.assertEqual(loop.tile_pattern.tile_shape, 4) + + def test_excludes_loop_not_indexing_tensor(self): + # X and Y physical: 'm'->X and 'n1'->Y recorded; 'n0' (over X) doesn't index T1. + arch = _StubArch({"DistBuf": _StubStorage(True, {"X", "Y"})}) + phys = self._phys(self._from(arch)) + self.assertEqual(len(phys), 2) + self.assertEqual({str(l.rank_name) for l in phys}, {"M", "N1"}) + self.assertEqual({l.spatial_dim for l in phys}, {"X", "Y"}) + + def test_stops_at_next_storage_level(self): + # A lower reservation ends the region; spatial loops below it are not recorded. + lower = Reservation(purposes=["T1"], resource="DistBuf") + mapping = _mapping([self.res, lower, self.sp_m]) + compat = Compatibility.from_mapping( + mapping, {"T1"}, self.einsum, + flattened_arch=_StubArch({"DistBuf": _StubStorage(True, {"X"})}), + ) + # Backing reservation (first) has the loop below the lower one excluded. + backing = compat.get_reservation_of_tensor("T1") + self.assertEqual(backing.physical_spatial_loops, ()) + + def test_distribution_difference_breaks_equality(self): + arch_x = _StubArch({"DistBuf": _StubStorage(True, {"X"})}) + arch_xy = _StubArch({"DistBuf": _StubStorage(True, {"X", "Y"})}) + a = self._from(arch_x) + b = self._from(arch_xy) + c = self._from(arch_x) + self.assertEqual(a, c) + self.assertEqual(hash(a), hash(c)) + self.assertNotEqual(a, b) + # Structural (cleared) compatibilities also differ. + self.assertNotEqual( + a.clear_symbolic_tile_patterns(), b.clear_symbolic_tile_patterns() + ) + + def test_regular_loops_unaffected(self): + # A reservation with above-storage loops keeps physical_spatial_loops empty + # when the storage is not distributed, and n_loops/above_loop_index are unchanged. + arch = _StubArch({"DistBuf": _StubStorage(False, set())}) + compat = self._from(arch) + (t,) = compat.tensors + self.assertEqual(t.above_loop_index, len(t.loops)) + self.assertEqual(t.physical_spatial_loops, ()) + + +if __name__ == "__main__": + unittest.main() From 449c39c2004e66357e6f71d97f4fedf55e5e0966 Mon Sep 17 00:00:00 2001 From: Michael Gilbert Date: Fri, 26 Jun 2026 13:11:26 -0400 Subject: [PATCH 2/5] [FFM] Use component object in Reservation instead of passing flattened_arch everywhere --- accelforge/frontend/mapping/mapping.py | 4 ++ .../FFM/_join_pmappings/compatibility.py | 24 ++++----- .../make_pmapping_templates.py | 1 - .../_looptree/reuse/symbolic/_symbolic.py | 1 + accelforge/model/main.py | 1 - tests/test_compatibility_physical_spatial.py | 49 +++++++------------ 6 files changed, 33 insertions(+), 47 deletions(-) diff --git a/accelforge/frontend/mapping/mapping.py b/accelforge/frontend/mapping/mapping.py index 25a80acb..440529e6 100755 --- a/accelforge/frontend/mapping/mapping.py +++ b/accelforge/frontend/mapping/mapping.py @@ -1534,6 +1534,10 @@ class Reservation(MappingNode): """ Tensors for which this reservation is reserving the tensor's backing storage. """ + _component_object: "arch.Component | None" = None + """ The arch component backing the reserved resource, taken from the Storage node + this Reservation was created from. Used internally by the Mapper; do not set. """ + persistent: bool = False """ Whether this reservation is persistent. Persistent reservations can't be tiled and diff --git a/accelforge/mapper/FFM/_join_pmappings/compatibility.py b/accelforge/mapper/FFM/_join_pmappings/compatibility.py index 89ee6453..30e5d8b0 100755 --- a/accelforge/mapper/FFM/_join_pmappings/compatibility.py +++ b/accelforge/mapper/FFM/_join_pmappings/compatibility.py @@ -565,16 +565,16 @@ def from_mapping( mapping: Mapping, tensors: set[TensorName], einsum: Einsum, - flattened_arch=None, ) -> "Compatibility": """ Create Compatibility from a mapping, a set of fusable tensors, and the workload. - If ``flattened_arch`` is given, spatial loops below a physically-distributed - storage are recorded in each ``TensorReservation.physical_spatial_loops`` so that - fused Einsums sharing a tensor must agree on its physical placement. If it is - ``None`` (the default), no such loops are recorded and behavior is unchanged. + Spatial loops below a physically-distributed storage are recorded in each + ``TensorReservation.physical_spatial_loops`` so that fused Einsums sharing a + tensor must agree on its physical placement. Whether a storage is distributed is + determined from the reservation node's component object (set when the reservation + is created); reservations without one record no such loops. """ if not isinstance(einsum, Einsum): raise TypeError(f"einsum should be an Einsum, but {type(einsum)} instead") @@ -635,18 +635,12 @@ def make_loops(above_index: int, tensor_name: TensorName) -> list[MappingLoop]: return tuple(loops) def make_physical_spatial_loops( - above_index: int, tensor_name: TensorName, resource: str + above_index: int, tensor_name: TensorName, storage ) -> tuple[Loop, ...]: # Spatial loops below a distributed storage fix which physical instance holds # which slice of the tensor; fused Einsums must agree on them. Only relevant - # when the backing storage has a physical stride. - if flattened_arch is None: - return () - try: - storage = flattened_arch[resource] - except (KeyError, ValueError): - return () - if not getattr(storage, "_is_distributed", lambda: False)(): + # when the backing storage is physically distributed. + if storage is None or not storage._is_distributed(): return () out = [] for n in mapping.nodes[above_index + 1 :]: @@ -684,7 +678,7 @@ def make_physical_spatial_loops( resource_name=mapping.nodes[i].resource, persistent=mapping.nodes[i].persistent, physical_spatial_loops=make_physical_spatial_loops( - i, mapping.nodes[i].purpose, mapping.nodes[i].resource + i, mapping.nodes[i].purpose, mapping.nodes[i]._component_object ), ) for i in tensor_indices diff --git a/accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py b/accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py index afba294d..6adbfe7f 100755 --- a/accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py +++ b/accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py @@ -455,7 +455,6 @@ def make_pmapping_templates(job: Job, print_progress: bool = True) -> SameEinsum mapping_with_reservations, new_job.fusable_tensors, new_job.einsum, - flattened_arch=new_job.flattened_arch, ) jobs.append(new_job) diff --git a/accelforge/model/_looptree/reuse/symbolic/_symbolic.py b/accelforge/model/_looptree/reuse/symbolic/_symbolic.py index d86ef2a2..77fa2790 100755 --- a/accelforge/model/_looptree/reuse/symbolic/_symbolic.py +++ b/accelforge/model/_looptree/reuse/symbolic/_symbolic.py @@ -437,6 +437,7 @@ def insert_reservation_nodes( node = Reservation(purposes=[buffet.tensor], resource=buffet.level) node.persistent = tracker.node.persistent node._backing = tracker.node._backing + node._component_object = tracker.node.component_object if ( buffet.tensor not in info.tensor_to_reservation_backer_id diff --git a/accelforge/model/main.py b/accelforge/model/main.py index c51cea75..7f9cb5e5 100644 --- a/accelforge/model/main.py +++ b/accelforge/model/main.py @@ -183,7 +183,6 @@ def evaluate_mapping( job.mapping, job.fusable_tensors, einsum, - flattened_arch=job.flattened_arch, ) symbol_renames, compatibility = compatibility.make_fused_loop_symbols( einsum_name diff --git a/tests/test_compatibility_physical_spatial.py b/tests/test_compatibility_physical_spatial.py index 0dbf4ae2..41f4b536 100644 --- a/tests/test_compatibility_physical_spatial.py +++ b/tests/test_compatibility_physical_spatial.py @@ -2,6 +2,10 @@ Tests for recording spatial-for loops below a physically-distributed storage in ``Compatibility`` (see ``TensorReservation.physical_spatial_loops``). +Whether a storage is distributed is read from the reservation node's +``_component_object`` (set when the reservation is created), so these tests stub that +object directly on the reservation node rather than passing an architecture. + pytest tests/test_compatibility_physical_spatial.py -v python -m unittest tests.test_compatibility_physical_spatial """ @@ -29,14 +33,6 @@ def _has_physical_dim(self, dim_name): return dim_name in self._physical_dims -class _StubArch: - def __init__(self, name_to_storage): - self._n2s = name_to_storage - - def __getitem__(self, name): - return self._n2s[name] # raises KeyError if absent, like FlattenedArch - - def _matmul_einsum(): spec = Spec.from_yaml( af.examples.arches.simple, @@ -61,26 +57,25 @@ def setUp(self): self.sp_n0 = Spatial(rank_variable="n0", name="X", component="DistBuf", tile_shape=4) self.mapping = _mapping([self.res, self.sp_m, self.sp_n1, self.sp_n0]) - def _from(self, arch): - return Compatibility.from_mapping( - self.mapping, {"T1"}, self.einsum, flattened_arch=arch - ) + def _from(self, storage): + # The reservation carries the arch component it reserves; that object is what + # from_mapping consults to decide whether the storage is distributed. + self.res._component_object = storage + return Compatibility.from_mapping(self.mapping, {"T1"}, self.einsum) def _phys(self, compat): (t,) = compat.tensors return t.physical_spatial_loops - def test_no_arch_records_nothing(self): + def test_no_component_object_records_nothing(self): self.assertEqual(self._phys(self._from(None)), ()) def test_non_distributed_records_nothing(self): - arch = _StubArch({"DistBuf": _StubStorage(False, {"X", "Y"})}) - self.assertEqual(self._phys(self._from(arch)), ()) + self.assertEqual(self._phys(self._from(_StubStorage(False, {"X", "Y"}))), ()) def test_dim_filter(self): # Only X is a physical dim: only the 'm' loop (over X, indexing T1) is recorded. - arch = _StubArch({"DistBuf": _StubStorage(True, {"X"})}) - phys = self._phys(self._from(arch)) + phys = self._phys(self._from(_StubStorage(True, {"X"}))) self.assertEqual(len(phys), 1) loop = phys[0] self.assertTrue(loop.is_spatial) @@ -91,8 +86,7 @@ def test_dim_filter(self): def test_excludes_loop_not_indexing_tensor(self): # X and Y physical: 'm'->X and 'n1'->Y recorded; 'n0' (over X) doesn't index T1. - arch = _StubArch({"DistBuf": _StubStorage(True, {"X", "Y"})}) - phys = self._phys(self._from(arch)) + phys = self._phys(self._from(_StubStorage(True, {"X", "Y"}))) self.assertEqual(len(phys), 2) self.assertEqual({str(l.rank_name) for l in phys}, {"M", "N1"}) self.assertEqual({l.spatial_dim for l in phys}, {"X", "Y"}) @@ -100,21 +94,17 @@ def test_excludes_loop_not_indexing_tensor(self): def test_stops_at_next_storage_level(self): # A lower reservation ends the region; spatial loops below it are not recorded. lower = Reservation(purposes=["T1"], resource="DistBuf") + self.res._component_object = _StubStorage(True, {"X"}) mapping = _mapping([self.res, lower, self.sp_m]) - compat = Compatibility.from_mapping( - mapping, {"T1"}, self.einsum, - flattened_arch=_StubArch({"DistBuf": _StubStorage(True, {"X"})}), - ) + compat = Compatibility.from_mapping(mapping, {"T1"}, self.einsum) # Backing reservation (first) has the loop below the lower one excluded. backing = compat.get_reservation_of_tensor("T1") self.assertEqual(backing.physical_spatial_loops, ()) def test_distribution_difference_breaks_equality(self): - arch_x = _StubArch({"DistBuf": _StubStorage(True, {"X"})}) - arch_xy = _StubArch({"DistBuf": _StubStorage(True, {"X", "Y"})}) - a = self._from(arch_x) - b = self._from(arch_xy) - c = self._from(arch_x) + a = self._from(_StubStorage(True, {"X"})) + b = self._from(_StubStorage(True, {"X", "Y"})) + c = self._from(_StubStorage(True, {"X"})) self.assertEqual(a, c) self.assertEqual(hash(a), hash(c)) self.assertNotEqual(a, b) @@ -126,8 +116,7 @@ def test_distribution_difference_breaks_equality(self): def test_regular_loops_unaffected(self): # A reservation with above-storage loops keeps physical_spatial_loops empty # when the storage is not distributed, and n_loops/above_loop_index are unchanged. - arch = _StubArch({"DistBuf": _StubStorage(False, set())}) - compat = self._from(arch) + compat = self._from(_StubStorage(False, set())) (t,) = compat.tensors self.assertEqual(t.above_loop_index, len(t.loops)) self.assertEqual(t.physical_spatial_loops, ()) From 2d8f6bcf7516cd21ecfc80526cfb7bf39559f9b0 Mon Sep 17 00:00:00 2001 From: Michael Gilbert Date: Mon, 29 Jun 2026 11:00:06 -0400 Subject: [PATCH 3/5] [FFM] Clean up compatibility code --- accelforge/frontend/renames.py | 4 +++ .../FFM/_join_pmappings/compatibility.py | 26 +++++++++---------- 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/accelforge/frontend/renames.py b/accelforge/frontend/renames.py index 063ebb33..5f89706b 100755 --- a/accelforge/frontend/renames.py +++ b/accelforge/frontend/renames.py @@ -15,6 +15,10 @@ Rank: TypeAlias = str EinsumName: TypeAlias = str +# The "don't care" rank is useful when specifying in a binding that a loop can +# be over any rank. +RANK_DONT_CARE: Rank = "DONT_CARE" + class Rename(EvalableModel): """ diff --git a/accelforge/mapper/FFM/_join_pmappings/compatibility.py b/accelforge/mapper/FFM/_join_pmappings/compatibility.py index 30e5d8b0..15333a20 100755 --- a/accelforge/mapper/FFM/_join_pmappings/compatibility.py +++ b/accelforge/mapper/FFM/_join_pmappings/compatibility.py @@ -15,7 +15,7 @@ TilePattern, Loop as MappingLoop, ) -from accelforge.frontend.renames import Rank, RankVariable, TensorName +from accelforge.frontend.renames import Rank, RankVariable, TensorName, RANK_DONT_CARE from accelforge.frontend.workload import Einsum from accelforge.mapper.FFM._pareto_df.df_convention import ( is_fused_loop_col, @@ -60,8 +60,8 @@ class Loop(Updatable): rank_name: Rank tile_pattern: TilePattern | None is_spatial: bool - spatial_dim: str | None = None # The architecture spatial dimension (e.g. "X", "Y") this loop fans out over. + spatial_dim: str | None = None def __post_init__(self): assert isinstance(self.rank_name, Rank) @@ -169,7 +169,7 @@ class TensorReservation(Updatable): persistent: bool = False # Spatial loops *below* this storage that distribute the tensor across physical # instances - physical_spatial_loops: tuple[Loop, ...] = () + physical_spatial_loops: tuple[Loop] = () def __post_init__(self): if self.persistent: @@ -614,6 +614,11 @@ def from_mapping( ), f"Tensors {backing_remaining} not found in mapping" def get_rank(rank_variable, tensor): + """ + Return rank in tensor indexed by rank_variable or + Rank("NO RANK.RECOMPUTED") if rank not in tensor. + """ + # TODO: shouldn't this whole logic use relevancy from workload? rv = rank_variable_to_ranks[tensor].get(rank_variable, oset()) assert ( len(rv) <= 1 @@ -636,7 +641,7 @@ def make_loops(above_index: int, tensor_name: TensorName) -> list[MappingLoop]: def make_physical_spatial_loops( above_index: int, tensor_name: TensorName, storage - ) -> tuple[Loop, ...]: + ) -> tuple[Loop]: # Spatial loops below a distributed storage fix which physical instance holds # which slice of the tensor; fused Einsums must agree on them. Only relevant # when the backing storage is physically distributed. @@ -649,24 +654,17 @@ def make_physical_spatial_loops( break if not isinstance(n, Spatial): continue - # Only dimensions the storage is physically distributed along matter. - if not storage._has_physical_dim(n.name): - continue rank = get_rank(n.rank_variable, tensor_name) - # Skip loops that don't partition this tensor (don't index it). + # If the rank is irrelevant, the binding could be any rank if rank == Rank("NO RANK. RECOMPUTED."): - continue - # Clear symbolic tile patterns: below-storage spatial loops never get - # DataFrame columns, so they're matched structurally (rank + dim, plus a - # concrete fanout when known; symbolic fanouts are pinned by the fixed - # hardware fanout along the dimension). + rank = RANK_DONT_CARE out.append( Loop( rank_name=rank, tile_pattern=n.tile_pattern._symbol2str(), is_spatial=True, spatial_dim=n.name, - ).clear_symbolic_tile_patterns() + ) ) return tuple(out) From 774b4ab4a87ee37036c7dc548d13fcf2ec8f26a5 Mon Sep 17 00:00:00 2001 From: Michael Gilbert Date: Mon, 29 Jun 2026 11:47:09 -0400 Subject: [PATCH 4/5] [FFM] Clean up comments --- accelforge/mapper/FFM/_join_pmappings/compatibility.py | 10 +--------- .../test_compatibility_physical_spatial.py | 0 2 files changed, 1 insertion(+), 9 deletions(-) rename tests/{ => vibe_see_readme_in_this_dir}/test_compatibility_physical_spatial.py (100%) diff --git a/accelforge/mapper/FFM/_join_pmappings/compatibility.py b/accelforge/mapper/FFM/_join_pmappings/compatibility.py index 15333a20..ee149326 100755 --- a/accelforge/mapper/FFM/_join_pmappings/compatibility.py +++ b/accelforge/mapper/FFM/_join_pmappings/compatibility.py @@ -569,12 +569,6 @@ def from_mapping( """ Create Compatibility from a mapping, a set of fusable tensors, and the workload. - - Spatial loops below a physically-distributed storage are recorded in each - ``TensorReservation.physical_spatial_loops`` so that fused Einsums sharing a - tensor must agree on its physical placement. Whether a storage is distributed is - determined from the reservation node's component object (set when the reservation - is created); reservations without one record no such loops. """ if not isinstance(einsum, Einsum): raise TypeError(f"einsum should be an Einsum, but {type(einsum)} instead") @@ -642,9 +636,7 @@ def make_loops(above_index: int, tensor_name: TensorName) -> list[MappingLoop]: def make_physical_spatial_loops( above_index: int, tensor_name: TensorName, storage ) -> tuple[Loop]: - # Spatial loops below a distributed storage fix which physical instance holds - # which slice of the tensor; fused Einsums must agree on them. Only relevant - # when the backing storage is physically distributed. + """Make data binding of physically distributed storages.""" if storage is None or not storage._is_distributed(): return () out = [] diff --git a/tests/test_compatibility_physical_spatial.py b/tests/vibe_see_readme_in_this_dir/test_compatibility_physical_spatial.py similarity index 100% rename from tests/test_compatibility_physical_spatial.py rename to tests/vibe_see_readme_in_this_dir/test_compatibility_physical_spatial.py From b0ecf121c6219d46fdb31d5df4dcc9522db59abb Mon Sep 17 00:00:00 2001 From: Michael Gilbert Date: Mon, 29 Jun 2026 14:04:30 -0400 Subject: [PATCH 5/5] [tests/vibe] Claude added a nonsense test that it fails --- .../test_compatibility_physical_spatial.py | 126 ------------------ 1 file changed, 126 deletions(-) delete mode 100644 tests/vibe_see_readme_in_this_dir/test_compatibility_physical_spatial.py diff --git a/tests/vibe_see_readme_in_this_dir/test_compatibility_physical_spatial.py b/tests/vibe_see_readme_in_this_dir/test_compatibility_physical_spatial.py deleted file mode 100644 index 41f4b536..00000000 --- a/tests/vibe_see_readme_in_this_dir/test_compatibility_physical_spatial.py +++ /dev/null @@ -1,126 +0,0 @@ -""" -Tests for recording spatial-for loops below a physically-distributed storage in -``Compatibility`` (see ``TensorReservation.physical_spatial_loops``). - -Whether a storage is distributed is read from the reservation node's -``_component_object`` (set when the reservation is created), so these tests stub that -object directly on the reservation node rather than passing an architecture. - - pytest tests/test_compatibility_physical_spatial.py -v - python -m unittest tests.test_compatibility_physical_spatial -""" - -import types -import unittest - -import accelforge as af -from accelforge.frontend.spec import Spec -from accelforge.frontend.mapping import Reservation, Spatial -from accelforge.mapper.FFM._join_pmappings.compatibility import Compatibility, Loop - - -class _StubStorage: - """A storage that is distributed along a fixed set of physical dimensions.""" - - def __init__(self, distributed, physical_dims): - self._distributed = distributed - self._physical_dims = set(physical_dims) - - def _is_distributed(self): - return self._distributed - - def _has_physical_dim(self, dim_name): - return dim_name in self._physical_dims - - -def _matmul_einsum(): - spec = Spec.from_yaml( - af.examples.arches.simple, - af.examples.workloads.matmuls, - jinja_parse_data={"N_EINSUMS": 2, "M": 64, "KN": 64}, - ) - name = list(spec.workload.einsum_names)[0] - return spec.workload.einsums[name] - - -def _mapping(nodes): - return types.SimpleNamespace(nodes=list(nodes)) - - -class TestPhysicalSpatialLoops(unittest.TestCase): - def setUp(self): - self.einsum = _matmul_einsum() - # T1 is indexed by rank variables m -> M and n1 -> N1; n0 -> N0 does NOT index it. - self.res = Reservation(purposes=["T1"], resource="DistBuf") - self.sp_m = Spatial(rank_variable="m", name="X", component="DistBuf", tile_shape=4) - self.sp_n1 = Spatial(rank_variable="n1", name="Y", component="DistBuf", tile_shape=4) - self.sp_n0 = Spatial(rank_variable="n0", name="X", component="DistBuf", tile_shape=4) - self.mapping = _mapping([self.res, self.sp_m, self.sp_n1, self.sp_n0]) - - def _from(self, storage): - # The reservation carries the arch component it reserves; that object is what - # from_mapping consults to decide whether the storage is distributed. - self.res._component_object = storage - return Compatibility.from_mapping(self.mapping, {"T1"}, self.einsum) - - def _phys(self, compat): - (t,) = compat.tensors - return t.physical_spatial_loops - - def test_no_component_object_records_nothing(self): - self.assertEqual(self._phys(self._from(None)), ()) - - def test_non_distributed_records_nothing(self): - self.assertEqual(self._phys(self._from(_StubStorage(False, {"X", "Y"}))), ()) - - def test_dim_filter(self): - # Only X is a physical dim: only the 'm' loop (over X, indexing T1) is recorded. - phys = self._phys(self._from(_StubStorage(True, {"X"}))) - self.assertEqual(len(phys), 1) - loop = phys[0] - self.assertTrue(loop.is_spatial) - self.assertEqual(str(loop.rank_name), "M") - self.assertEqual(loop.spatial_dim, "X") - # Concrete fanout is preserved through clearing. - self.assertEqual(loop.tile_pattern.tile_shape, 4) - - def test_excludes_loop_not_indexing_tensor(self): - # X and Y physical: 'm'->X and 'n1'->Y recorded; 'n0' (over X) doesn't index T1. - phys = self._phys(self._from(_StubStorage(True, {"X", "Y"}))) - self.assertEqual(len(phys), 2) - self.assertEqual({str(l.rank_name) for l in phys}, {"M", "N1"}) - self.assertEqual({l.spatial_dim for l in phys}, {"X", "Y"}) - - def test_stops_at_next_storage_level(self): - # A lower reservation ends the region; spatial loops below it are not recorded. - lower = Reservation(purposes=["T1"], resource="DistBuf") - self.res._component_object = _StubStorage(True, {"X"}) - mapping = _mapping([self.res, lower, self.sp_m]) - compat = Compatibility.from_mapping(mapping, {"T1"}, self.einsum) - # Backing reservation (first) has the loop below the lower one excluded. - backing = compat.get_reservation_of_tensor("T1") - self.assertEqual(backing.physical_spatial_loops, ()) - - def test_distribution_difference_breaks_equality(self): - a = self._from(_StubStorage(True, {"X"})) - b = self._from(_StubStorage(True, {"X", "Y"})) - c = self._from(_StubStorage(True, {"X"})) - self.assertEqual(a, c) - self.assertEqual(hash(a), hash(c)) - self.assertNotEqual(a, b) - # Structural (cleared) compatibilities also differ. - self.assertNotEqual( - a.clear_symbolic_tile_patterns(), b.clear_symbolic_tile_patterns() - ) - - def test_regular_loops_unaffected(self): - # A reservation with above-storage loops keeps physical_spatial_loops empty - # when the storage is not distributed, and n_loops/above_loop_index are unchanged. - compat = self._from(_StubStorage(False, set())) - (t,) = compat.tensors - self.assertEqual(t.above_loop_index, len(t.loops)) - self.assertEqual(t.physical_spatial_loops, ()) - - -if __name__ == "__main__": - unittest.main()