[network] Refactor network cost to handle different topologies

Author: gilbertmike

accelforge/frontend/arch/components.py

@@ -1294,7 +1294,7 @@ def _render_node_color(self) -> str:
         return "#E0EEFF"
 
 
-class TopologySpec(str, enum.Enum):
+class TopologySpec(enum.StrEnum):
     MESH = "mesh"
 
 

accelforge/model/_looptree/latency/memory.py

@@ -109,6 +109,7 @@ def component_latency(
 
     network_to_max_link_traffic = defaultdict(lambda: defaultdict(lambda: 0))
     network_to_max_hops = defaultdict(lambda: [])
+    # Aggregates across tensors
     for network, network_stats in looptree_results.network_stats.items():
         component = network.component
         if component not in name2component:

accelforge/model/_looptree/reuse/symbolic/_network.py

@@ -1,6 +1,11 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
 from accelforge.frontend.mapping import (
     Spatial
 )
+from accelforge.frontend.arch.components import TopologySpec
 from accelforge.frontend._workload_isl._symbolic import (
     compute_dense_tile_occupancy,
     Irrelevant,
@@ -14,22 +19,178 @@
 from ._stats import NetworkStats, SymbolicAnalysisOutput
 
 
-class NetworkAnalyzer:
-    def __init__(self, network_stats):
+@dataclass
+class PerLoopTransferCost:
+    """The per-spatial-loop cost contributed by a single network, as computed
+    by a :class:`TopologyModel`."""
+
+    total_cost: Any
+    """Total hops contributed by data movement over this spatial loop."""
+    max_hops: Any
+    """Hops added to the longest route by this spatial loop."""
+    max_traffic: Any
+    """Maximum traffic (in actions) on any single link along this dimension."""
+
+
+class TopologyModel(ABC):
+    """Computes the cost of moving data across a network of a given topology.
+
+    Subclasses encapsulate everything topology-specific about how a tensor's
+    data is delivered across a spatial fanout. :class:`NetworkAnalyzer` selects
+    the model for each network from its component's
+    :class:`~accelforge.frontend.arch.components.TopologySpec` and remains
+    agnostic to the topology itself.
+
+    Instances are stateful: they accumulate per-network max hops across the
+    repeated spatial-loop iterations of a single :class:`NetworkAnalyzer`, so a
+    fresh model is constructed for each analyzer (see :func:`get_topology_model`).
+    """
+
+    def __init__(self):
+        # Running total of max hops per network, accumulated across the
+        # repeated spatial-loop iterations handled by one NetworkAnalyzer.
         self.overall_max_hops: dict = {}
+
+    def accumulate_max_hops(self, network, max_hops):
+        """Add this loop's ``max_hops`` to ``network``'s running total and
+        return the updated total.
+
+        Each call to :meth:`NetworkAnalyzer.accumulate_child_result` (i.e., over
+        a different iteration of a spatial loop) adds more to the max hops.
+        """
+        self.overall_max_hops[network] = (
+            self.overall_max_hops.get(network, 0) + max_hops
+        )
+        return self.overall_max_hops[network]
+
+    @abstractmethod
+    def per_loop_transfer_cost(
+        self,
+        relevancy,
+        *,
+        shape_repeats,
+        last_fanout,
+        volume,
+        src_component,
+        dim_name: str,
+    ) -> PerLoopTransferCost:
+        """Return the :class:`PerLoopTransferCost` for moving ``volume`` of data across one
+        spatial loop.
+
+        Args:
+            relevancy: The relevancy of the spatial loop's rank variable to the
+                tensor (``Irrelevant``, ``Relevant``, or ``PartiallyRelevant``).
+            shape_repeats: The number of iterations of this spatial loop.
+            last_fanout: The fanout in this dimension among mapping nodes below
+                (i.e., the stride).
+            volume: The data volume (in actions) moved per destination.
+            src_component: The flattened-arch component sourcing the data, used
+                to query physical fanout/stride.
+            dim_name: The name of the spatial dimension (e.g., ``X`` or ``Y``).
+        """
+        raise NotImplementedError
+
+
+class MeshTopologyModel(TopologyModel):
+    """Cost model for a mesh network.
+
+    Data travels link-by-link along one axis of the mesh. Multicast delivers a
+    value to every point along the dimension; unicast delivers a distinct value
+    to each point. When the source is physically distributed, data is bound as
+    locally as possible across the physical buffers.
+    """
+
+    def per_loop_transfer_cost(
+        self,
+        relevancy,
+        *,
+        shape_repeats,
+        last_fanout,
+        volume,
+        src_component,
+        dim_name,
+    ) -> PerLoopTransferCost:
+        if isinstance(relevancy, Irrelevant):
+            # The volume travels through link by link in one axis of the mesh
+            # Distributed or not, the amount of total cost is the same.
+            # However, the accesses now come from different physical memories
+            total_cost = multicast_cost(shape_repeats, last_fanout) * volume
+            max_hops = shape_repeats * last_fanout
+            max_traffic = volume
+        elif isinstance(relevancy, Relevant):
+            # If distributed, then we bind data as locally as possible in the
+            # physical buffers
+            if src_component._get_physical_fanout_along(dim_name) > 1:
+                physical_stride = src_component._get_physical_stride_along(dim_name)
+
+                n_dsts_per_physical = MinGeqZero(
+                    # if last_fanout > physical_stride, set n_dst to 1, which results in 0 hops
+                    # later (which is correct because the set of destinations always overlap
+                    # the set of sources).
+                    MaxGeqZero(physical_stride / last_fanout, 1),
+                    shape_repeats
+                )
+                n_activated_physical = MaxGeqZero(shape_repeats * last_fanout / physical_stride, 1)
+                total_cost = (
+                    n_activated_physical
+                    *
+                    unicast_cost(n_dsts_per_physical, last_fanout)
+                    *
+                    volume
+                )
+                max_hops = MinGeqZero((n_dsts_per_physical - 1) * last_fanout, physical_stride)
+                max_traffic = (n_dsts_per_physical - 1) * volume
+            else:
+                total_cost = unicast_cost(shape_repeats, last_fanout) * volume
+                max_hops = shape_repeats * last_fanout
+                max_traffic = (shape_repeats - 1) * volume
+        elif isinstance(relevancy, PartiallyRelevant):
+            raise NotImplementedError()
+        else:
+            raise RuntimeError(f"unhandled relevancy type {relevancy}")
+
+        return PerLoopTransferCost(total_cost=total_cost, max_hops=max_hops, max_traffic=max_traffic)
+
+
+# Registry mapping each topology to the model class that costs its data
+# movement. Classes (not instances) are stored because models are stateful and
+# each NetworkAnalyzer needs its own.
+TOPOLOGY_MODELS: dict[TopologySpec, type[TopologyModel]] = {
+    TopologySpec.MESH: MeshTopologyModel,
+}
+
+
+def get_topology_model(topology) -> TopologyModel:
+    """Construct a fresh :class:`TopologyModel` for the given topology."""
+    return TOPOLOGY_MODELS[topology]()
+
+
+class NetworkAnalyzer:
+    def __init__(self, network_stats, info: AnalysisInfo, einsum_name, node: Spatial):
         self.network_stats = network_stats
+        # These don't change across calls to accumulate_child_result.
+        self.info = info
+        self.einsum_name = einsum_name
+        self.node = node
+        # Each network gets its own topology model, since different networks may
+        # have different topologies. Models are constructed lazily, the first
+        # time a network needs costing, and reused for the analyzer's lifetime so
+        # their accumulated max hops persist.
+        self.topology_models: dict = {}
+
+    def _get_topology_model(self, network, topology) -> TopologyModel:
+        if network not in self.topology_models:
+            self.topology_models[network] = get_topology_model(topology)
+        return self.topology_models[network]
 
     def accumulate_child_result(
         self,
         child_result: SymbolicAnalysisOutput,
-        info: AnalysisInfo,
         shape_repeats,
-        einsum_name,
         child_shape,
-        node: Spatial,
     ):
         """This function is called for every repeated shape."""
-        flattened_arch = info.job.flattened_arch
+        flattened_arch = self.info.job.flattened_arch
 
         for network, child_network_stats in child_result.network_stats.items():
             src_component = flattened_arch[network.source.level]
@@ -39,7 +200,7 @@ def accumulate_child_result(
 
             # We only need to update the summary if the spatial loop is for
             # a component higher than the network of interest
-            if flattened_arch.is_above(node.component, network.component):
+            if flattened_arch.is_above(self.node.component, network.component):
                 accumulated_network_stats.total_hops += (
                     child_network_stats.total_hops * shape_repeats
                 )
@@ -54,71 +215,51 @@ def accumulate_child_result(
                     )
                 continue
 
-            volume = self._get_data_volume(network, einsum_name, info, child_shape)
+            volume = self._get_data_volume(network, child_shape)
 
-            relevancy = info.tensor_to_relevancy[network.tensor][node.rank_variable]
+            relevancy = self.info.tensor_to_relevancy[network.tensor][self.node.rank_variable]
 
             # The fanout in this dimension in mapping nodes below, i.e., the stride
-            last_fanout = child_result.fanout.get((node.component, einsum_name), {})
-            last_fanout = last_fanout.get(node.name, 1)
-            if isinstance(relevancy, Irrelevant):
-                # The volume travels through link by link in one axis of the mesh
-                # Distributed or not, the amount of total cost is the same.
-                # However, the accesses now come from different physical memories
-                total_cost = multicast_cost(shape_repeats, last_fanout)*volume
-                max_hops = shape_repeats*last_fanout
-                max_traffic = volume
-            elif isinstance(relevancy, Relevant):
-                # If distributed, then we bind data as locally as possible in the
-                # physical buffers
-                if src_component._get_physical_fanout_along(node.name) > 1:
-                    physical_stride = src_component._get_physical_stride_along(node.name)
-
-                    n_dsts_per_physical = MinGeqZero(
-                        # if last_fanout > physical_stride, set n_dst to 1, which results in 0 hops
-                        # later (which is correct because the set of destinations always overlap
-                        # the set of sources).
-                        MaxGeqZero(physical_stride / last_fanout, 1),
-                        shape_repeats
-                    )
-                    n_activated_physical = MaxGeqZero(shape_repeats*last_fanout/physical_stride, 1)
-                    total_cost = (
-                        n_activated_physical
-                        *
-                        unicast_cost(n_dsts_per_physical, last_fanout)
-                        *
-                        volume
-                    )
-                    max_hops = MinGeqZero((n_dsts_per_physical-1)*last_fanout, physical_stride)
-                    max_traffic = (n_dsts_per_physical-1)*volume
-                else:
-                    total_cost = unicast_cost(shape_repeats, last_fanout)*volume
-                    max_hops = shape_repeats * last_fanout
-                    max_traffic = (shape_repeats-1)*volume
-            elif isinstance(relevancy, PartiallyRelevant):
-                raise NotImplementedError()
-            else:
-                raise RuntimeError(f"unhandled relevancy type {relevancy}")
+            last_fanout = child_result.fanout.get((self.node.component, self.einsum_name), {})
+            last_fanout = last_fanout.get(self.node.name, 1)
 
-            # Each subsequent call to this function (i.e., over different iterations of a spatial loop)
-            # adds more to the max hops
-            self.overall_max_hops[network] = self.overall_max_hops.get(network, 0) + max_hops
+            topology_model = self._get_topology_model(
+                network, flattened_arch[network.component].topology
+            )
+            per_loop_transfer_cost = topology_model.per_loop_transfer_cost(
+                relevancy,
+                shape_repeats=shape_repeats,
+                last_fanout=last_fanout,
+                volume=volume,
+                src_component=src_component,
+                dim_name=self.node.name,
+            )
+
+            overall_max_hops = topology_model.accumulate_max_hops(
+                network, per_loop_transfer_cost.max_hops
+            )
 
             accumulated_network_stats.total_hops += (
-                total_cost + child_network_stats.total_hops*shape_repeats
+                per_loop_transfer_cost.total_cost
+                + child_network_stats.total_hops * shape_repeats
             )
             accumulated_network_stats.max_hops = MaxGeqZero(
                 accumulated_network_stats.max_hops,
-                self.overall_max_hops[network] + child_network_stats.max_hops,
+                overall_max_hops + child_network_stats.max_hops,
             )
-            accumulated_network_stats.max_traffic[node.name] = MaxGeqZero(
-                accumulated_network_stats.max_traffic.get(node.name, 0),
-                max_traffic + child_network_stats.max_traffic.get(node.name, 0)
+            accumulated_network_stats.max_traffic[self.node.name] = MaxGeqZero(
+                accumulated_network_stats.max_traffic.get(self.node.name, 0),
+                per_loop_transfer_cost.max_traffic + child_network_stats.max_traffic.get(self.node.name, 0)
             )
 
-        return self.overall_max_hops
+        overall_max_hops = {}
+        for model in self.topology_models.values():
+            overall_max_hops.update(model.overall_max_hops)
+        return overall_max_hops
 
-    def _get_data_volume(self, network, einsum_name, info, child_shape):
+    def _get_data_volume(self, network, child_shape):
+        info = self.info
+        einsum_name = self.einsum_name
         flattened_arch = info.job.flattened_arch
         projection = info.einsum_tensor_to_projection[(einsum_name, network.tensor)]
         component_object = flattened_arch[network.component]
@@ -153,4 +294,4 @@ def unicast_cost(n_dsts, stride):
 
 
 def arithmetic_sum(n):
-    return 0.5 * (n+1) * n
+    return 0.5 * (n+1) * n

accelforge/model/_looptree/reuse/symbolic/_symbolic.py

@@ -592,7 +592,9 @@ def analyze_spatial(node_idx, current_shape, info: AnalysisInfo):
 
     result_accumulator = SymbolicAnalysisOutput()
 
-    network_analyzer = NetworkAnalyzer(result_accumulator.network_stats)
+    network_analyzer = NetworkAnalyzer(
+        result_accumulator.network_stats, info, einsum_name, node
+    )
 
     def handle_repeated_value(repeated_shape):
         shape_value = repeated_shape.value
@@ -633,7 +635,7 @@ def handle_repeated_value(repeated_shape):
             )
 
         network_analyzer.accumulate_child_result(
-            child_result, info, shape_repeats, einsum_name, child_shape, node
+            child_result, shape_repeats, child_shape
         )
 
         for einsum, child_steps in child_result.temporal_steps.items():