[frontend,mapper,model] Refactor flattened arch into a class

Author: gilbertmike

accelforge/frontend/arch/_flattened_arch.py

@@ -0,0 +1,53 @@
+class FlattenedArch:
+    """
+    A flattened arch is an architecture spec that has been
+    flattened into a hierarchy for the purpose of mapping
+    a particular Einsum.
+
+    Several steps (may not be exhaustive) are applied when
+    an arch is flattened:
+    - A compute unit is selected, and a flattened arch
+      is a path from the root to that compute unit
+    - Expressions have been evaluated in the context of an
+      Einsum.
+    - Non-hierarchical arch nodes (e.g., ones inside the
+      `nodes` key of an `Array` has been inserted into
+      a hierarchy.)
+
+    This class should only be relevant to the model and
+    mapper. That is, the user should generally not define
+    a flattened arch directly. So, unlike other classes in
+    `frontend`, this one is intentionally *not* an
+    `EvalableModel`.
+    """
+    def __init__(self, nodes: list["Leaf"]):
+        self.nodes = nodes
+
+    def __getitem__(self, idx: int | str | slice):
+        if isinstance(idx, (int, slice)):
+            return self.nodes[idx]
+        elif isinstance(idx, str):
+            for node in self.nodes:
+                if node.name == idx:
+                    return node
+            raise KeyError(f"arch node with name {idx} not found")
+        raise ValueError(f"idx should be int or str, but instead {type(idx)}")
+
+    def __iter__(self):
+        for node in self.nodes:
+            yield node
+
+    def index(self, name: str):
+        for i, node in enumerate(self.nodes):
+            if node.name == name:
+                return i
+        raise ValueError(f"no node found with name {name}")
+
+    def is_above(self, name_a: str, name_b: str):
+        """
+        Returns True if node with name_a is above node with name_b.
+        Raises ValueError if either is not found.
+        """
+        idx_a = self.index(name_a)
+        idx_b = self.index(name_b)
+        return idx_a < idx_b

accelforge/frontend/arch/structure.py

@@ -21,6 +21,7 @@
 from accelforge.util.exceptions import EvaluationError
 
 from accelforge.frontend.arch.spatialable import Spatialable
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 
 from pydantic import Discriminator
 from accelforge.util._basetypes import _uninstantiable
@@ -328,7 +329,7 @@ def _flatten(
         compute_node: str,
         fanout: int = 1,
         return_fanout: bool = False,
-    ):
+    ) -> FlattenedArch:
         from accelforge.frontend.arch.components import Compute
 
         nodes = []
@@ -350,6 +351,7 @@ def _flatten(
                 e.add_field(node)
                 raise e
 
+        nodes = FlattenedArch(nodes)
         if return_fanout:
             return nodes, fanout
         return nodes
@@ -418,7 +420,7 @@ def _flatten(
         compute_node: str,
         fanout: int = 1,
         return_fanout: bool = False,
-    ):
+    ) -> FlattenedArch:
         from accelforge.frontend.arch.components import Compute
 
         nodes = []
@@ -466,6 +468,7 @@ def _flatten(
                 e.add_field(node)
                 raise e
 
+        nodes = FlattenedArch(nodes)
         if return_fanout:
             return nodes, fanout
         return nodes

accelforge/frontend/mapping/mapping.py

@@ -33,6 +33,7 @@
 from pydantic import ConfigDict, Discriminator, Tag, computed_field
 import sympy
 
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 from accelforge.frontend.renames import EinsumName, TensorName
 from accelforge.util._basetypes import (
     # Parsing helpers for the input files.
@@ -1285,7 +1286,7 @@ def compact_str(self) -> str:
     def _get_single_tensor_mapping(
         self,
         tensor_name: TensorName,
-        flattened_arch: list[arch.Leaf],
+        flattened_arch: FlattenedArch,
         tensor_rank_variables: set[str],
     ) -> Self:
         """

accelforge/frontend/spec.py

@@ -12,6 +12,7 @@
     Container,
     Spatialable,
 )
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 
 from accelforge.frontend.workload import Workload
 from accelforge.frontend.variables import Variables
@@ -276,7 +277,7 @@ def _get_flattened_architecture(
         self,
         compute_node: str | Compute | None = None,
         einsum_name: EinsumName | None = None,
-    ) -> list[list[Leaf]] | list[Leaf]:
+    ) -> list[FlattenedArch] | FlattenedArch:
         """
         Return the architecture as paths of ``Leaf`` instances from the highest-level
         node to each ``Compute`` node. Parses arithmetic expressions in the

accelforge/mapper/FFM/_make_pmappings/contraints/constraints.py

@@ -2,6 +2,7 @@
 import logging
 from typing import List
 from accelforge._accelerated_imports import np
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 from accelforge.frontend._workload_isl._symbolic import PartiallyRelevant, Relevant
 import accelforge.frontend.arch as arch
 from accelforge.frontend.arch import (
@@ -196,7 +197,7 @@ def constrained_loops(
 
 
 def get_constraints(
-    flattened_arch: list[arch.Leaf],
+    flattened_arch: FlattenedArch,
     mapping: List[MappingNode],
     symbol_table: dict[str, InvertibleSet],
     einsum_name: EinsumName,

accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_loops.py

@@ -5,6 +5,7 @@
 from enum import Enum
 
 import accelforge.frontend.arch as arch
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 from accelforge.util._frozenset import oset
 from accelforge.frontend.mapping import (
     MappingNode,
@@ -39,15 +40,14 @@ def insert_temporal_loops(
     ranks_with_tile_pattern: set,
     workload: Workload,
     _can_lower_outermost_memory: bool,
-    flattened_arch: list[arch.Leaf],
+    flattened_arch: FlattenedArch,
     max_fused_loops: int,
     fanouts: dict[str, int],
     fusable_tensors: set[TensorName],
     intermediate_tensors: set[TensorName],
     let_non_intermediate_tensors_respawn_in_backing_storage: bool,
     explore_loop_orders: bool,
 ):
-    arch_node_names = [n.name for n in flattened_arch]
     # First establish insertion points. Insertion points are:
     # - Below the last instance of the first memory
     # - Between any two TensorHolder nodes
@@ -88,7 +88,7 @@ def insert_temporal_loops(
     for s in split_mapping:
         # Within each split mapping group, sort by arch levels.
         # This can help create places to put spatial loops
-        s.sort(key=lambda tensor_holder: arch_node_names.index(tensor_holder.component))
+        s.sort(key=lambda tensor_holder: flattened_arch.index(tensor_holder.component))
 
     if sum(map(len, split_mapping)) != len(mapping):
         raise RuntimeError("BUG: number of storage nodes post-split != original")
@@ -333,11 +333,10 @@ def _get_next_storages(i: int, toll_allowed: bool = False) -> list[TensorHolder]
 def insert_spatial_loops(
     mapping: list[MappingNode],
     einsum: Einsum,
-    flattened_arch: list[arch.Memory],
+    flattened_arch: FlattenedArch,
     intermediate_tensors: set[TensorName],
 ):
     nodes_with_fanout = [n for n in flattened_arch if n.get_fanout() > 1]
-    arch_node_names = [n.name for n in flattened_arch]
     tensor2fully_relevant_rank_vars = einsum.tensor2directly_indexing_rank_variables
     simple_rank_variables = einsum._simple_rank_variables
 
@@ -346,7 +345,7 @@ def insert_spatial_loops(
         # above the fanout in the arch, and below any temporal loops in the
         # same block.
         insertion_point = _idx_below_lowest_tensor_holder_with_component_above_fanout(
-            node, mapping, arch_node_names
+            node, mapping, flattened_arch
         )
         while insertion_point < len(mapping) and isinstance(
             mapping[insertion_point], Temporal
@@ -386,16 +385,16 @@ def _tensors_seen_above_point(idx, mapping):
 
 
 def _idx_below_lowest_tensor_holder_with_component_above_fanout(
-    fanout_node, mapping, arch_node_names
+    fanout_node, mapping, flattened_arch: FlattenedArch
 ):
     """Return the index right after the lowest TensorHolder whose component
     is above the fanout in the arch. If none found, returns len(mapping)."""
-    fanout_arch_idx = arch_node_names.index(fanout_node.name)
+    fanout_arch_idx = flattened_arch.index(fanout_node.name)
     result = 0
     for i in range(len(mapping)):
         if not isinstance(mapping[i], TensorHolder):
             continue
-        if arch_node_names.index(mapping[i].component) < fanout_arch_idx:
+        if flattened_arch.index(mapping[i].component) < fanout_arch_idx:
             result = i + 1
     return result
 

accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_pmapping_templates.py

@@ -11,6 +11,7 @@
 from tqdm import tqdm
 
 import accelforge.frontend.arch as arch
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 from accelforge.util._frozenset import oset
 from accelforge.frontend.mapping import (
     Compute,
@@ -86,7 +87,7 @@ def unpack_loops_to_rank_variables(mapping: List[MappingNode]):
 # Iterate over mappings
 # =================================================================================================
 def place_missing_temporal_loops(
-    mapping: List[MappingNode], einsum: Einsum, flattened_arch: list[arch.Leaf]
+    mapping: List[MappingNode], einsum: Einsum, flattened_arch: FlattenedArch
 ):
     """
     Adds temporal loops to the mapping to fill in any rank variables that are missing.
@@ -132,7 +133,7 @@ def place_missing_temporal_loops(
 
 def remove_unordered_spatial_temporal_loops(
     mapping: list[MappingNode],
-    flattened_arch: list[arch.Leaf],
+    flattened_arch: FlattenedArch,
     einsum: Einsum,
     explore_unordered_spatial_loops: bool = True,
 ):
@@ -258,7 +259,7 @@ def assert_proper_fusion_labeling(
 def iterate_mappings_no_constraints(
     spec: Spec,
     einsum_name: str,
-    flattened_arch: list[arch.Leaf],
+    flattened_arch: FlattenedArch,
     rank_variable_bounds: dict[RankVariable, int],
     job: Job,
 ) -> Iterator[tuple[Mapping, SymbolTable, arch.Compute, int]]:
@@ -331,7 +332,7 @@ def iterate_mappings_no_constraints(
 def iterate_mappings_constraints(
     spec: Spec,
     einsum_names: list[str] | str,
-    flattened_arch: list[arch.Leaf],
+    flattened_arch: FlattenedArch,
     rank_variable_bounds: dict[RankVariable, int],
     tensor_to_relevancy: dict[
         TensorName, dict[RankVariable, Relevant | PartiallyRelevant]

accelforge/mapper/FFM/_make_pmappings/make_pmapping_templates/make_reservations.py

@@ -2,6 +2,7 @@
 from typing import Any
 
 import accelforge.frontend.arch as arch
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 from accelforge.frontend.mapping import MappingNode, Reservation, Storage, TensorHolder
 
 
@@ -22,7 +23,7 @@ def _recursive_iter_fence_positions(
 
 def get_reservation_choices(
     mapping: list[TensorHolder],
-    flattened_arch: list[arch.Leaf],
+    flattened_arch: FlattenedArch,
 ) -> Generator[tuple[list[TensorHolder], Any], None, None]:
     # Rules:
     # - In general, reservations go right under their storage node

accelforge/mapper/FFM/_make_pmappings/pmapper_job.py

@@ -6,6 +6,7 @@
 from uuid import UUID, uuid4
 
 import accelforge.frontend.arch as arch
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 from accelforge.util._frozenset import oset
 from accelforge.frontend.mapping import (
     Mapping,
@@ -57,7 +58,7 @@ class Job:
     mapping: Mapping | None = None
     constraints: MappingConstraints | None = None
     fusable_tensors: set[TensorName] | None = None
-    flattened_arch: list[arch.Leaf] | None = None
+    flattened_arch: FlattenedArch | None = None
 
     einsum_name: EinsumName | None = None
     """If the Job is for a single einsum, this is the einsum name."""

accelforge/model/_looptree/latency/memory.py

@@ -2,6 +2,7 @@
 
 from accelforge.frontend import arch
 from accelforge.frontend.arch import Leaf, Memory, TensorHolder, Component
+from accelforge.frontend.arch._flattened_arch import FlattenedArch
 from accelforge.frontend.mapping import Compute, Mapping
 from accelforge.frontend.spec import Spec
 
@@ -39,7 +40,7 @@ def isl_to_summarized(
 
 def component_latency(
     looptree_results: SymbolicAnalysisOutput,
-    flattened_arch: list[Leaf],
+    flattened_arch: FlattenedArch,
     mapping: Mapping,
     spec: Spec,
 ):