[model] Classes to hold network stats + cleanup

Author: gilbertmike

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

@@ -0,0 +1,63 @@
+from accelforge.frontend.mapping import (
+    Mapping,
+)
+
+from accelforge.frontend.mapping import (
+    Mapping,
+    TensorHolder,
+)
+from accelforge.frontend.workload import (
+    TensorName,
+)
+from accelforge.model._looptree.types import Buffet, Compute
+
+
+class DataMovementConnections:
+    """
+    For querying connection between Buffet or Compute
+    (e.g., which Buffet supplies this Compute).
+    """
+    def __init__(self):
+        self.src_to_dst = {}
+        self.dst_to_src = {}
+
+    def get_src(self, dst: Buffet | Compute) -> Buffet | None:
+        return self.dst_to_src[dst]
+
+    def get_dst(self, src: TensorHolder) -> Buffet | Compute:
+        return self.src_to_dst[src]
+
+    @classmethod
+    def from_pmapping(cls, pmapping):
+        einsum_name = pmapping[-1].einsum
+        src_to_dst = {}
+        dst_to_src = {}
+        tensor_to_last_src = {}
+        for node in pmapping:
+            if isinstance(node, TensorHolder):
+                for tensor in node.tensors:
+                    buffet = Buffet(tensor, einsum_name, node.component)
+                    if tensor in tensor_to_last_src:
+                        src = tensor_to_last_src[tensor]
+                    else:
+                        src = None
+                    tensor_to_last_src[tensor] = buffet
+                    dst_to_src[buffet] = src
+                    src_to_dst[src] = buffet
+                    src_to_dst[buffet] = None
+        result = DataMovementConnections()
+        result.src_to_dst = src_to_dst
+        result.dst_to_src = dst_to_src
+        return result
+
+
+def get_tensor_to_backer_id(mapping: Mapping):
+    tensor_to_ids: dict[TensorName, set[int]] = {}
+    for node in mapping:
+        if isinstance(node, TensorHolder):
+            for tensor in node.tensors:
+                if tensor in tensor_to_ids:
+                    continue
+                tensor_to_ids[tensor] = id(node)
+    return tensor_to_ids
+

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

@@ -1,16 +1,11 @@
 import copy
 from dataclasses import dataclass, field
-import itertools
 from accelforge.frontend.mapping import (
     Compute,
     Mapping,
-    Nested,
-    Pipeline,
     Toll,
     Reservation,
-    Sequential,
     Spatial,
-    Split,
     Storage,
     Temporal,
 )
@@ -33,9 +28,7 @@
 from accelforge.frontend.workload import (
     Workload,
     TensorName,
-    isl_expression_has_variable,
 )
-from accelforge.frontend._workload_isl._isl import get_rank_variable_bounds
 from accelforge.frontend._workload_isl._symbolic import (
     get_projection_expr,
     get_rank_variable_relevancy,
@@ -45,7 +38,11 @@
     PartiallyRelevant,
 )
 
-from accelforge.model._looptree.types import Buffet
+from accelforge.model._looptree.types import Buffet, Compute, Network
+from accelforge.model._looptree.reuse.symbolic.mapping_utils import (
+    DataMovementConnections,
+    get_tensor_to_backer_id
+)
 
 from accelforge.mapper.FFM._make_pmappings.pmapper_job import Job
 from accelforge.util._sympy.broadcast_max import Min, Max
@@ -59,12 +56,6 @@
 PRINT_FORMULAS = False
 
 
-@dataclass(eq=True, frozen=True)
-class Compute:
-    einsum: str
-    level: str
-
-
 class Uninitialized:
     def __init__(self):
         pass
@@ -106,6 +97,12 @@ def max_dict(a: dict[Any, Any], b: dict[Any, Any]) -> dict[Any, Any]:
     return new
 
 
+@dataclass
+class NetworkStats:
+    buffet_to_total_hops: Any
+    buffet_to_max_hops: Any
+
+
 @dataclass
 class BuffetStats:
     total_reads_to_parent: Any = field(default=0)
@@ -240,11 +237,11 @@ def net_max_per_unit_write_actions(self) -> Any:
             - self.min_per_unit_skipped_first_write_actions
         )
 
-
-def blank_buffet_stats() -> BuffetStats:
-    stats = BuffetStats()
-    stats.n_loops_above = None  # Inherit from whoever is added to this
-    return stats
+    @classmethod
+    def blank(cls):
+        stats = cls()
+        stats.n_loops_above = None # Inherit from whoever is added to this
+        return stats
 
 
 @dataclass
@@ -310,6 +307,8 @@ class SymbolicAnalysisOutput:
 
     buffet_stats: dict[Buffet, BuffetStats] = field(default_factory=dict)
 
+    network_stats: dict[Network, NetworkStats] = field(default_factory=dict)
+
     # Mapping [level, einsum] to the fanout
     fanout: dict[(Buffet, str), int] = field(default_factory=dict)
 
@@ -357,7 +356,7 @@ def get_child_buffet_stats(self, buffet: Buffet) -> BuffetStats:
     def sum_buffet_stats_per_level(self) -> dict[str, BuffetStats]:
         result: dict[str, BuffetStats] = {}
         for buffet, stats in self.buffet_stats.items():
-            result.setdefault(buffet.level, blank_buffet_stats())
+            result.setdefault(buffet.level, BuffetStats.blank())
             result[buffet.level] += stats
         return result
 
@@ -396,6 +395,8 @@ class AnalysisInfo:
 
     tensor_to_reservation_backer_id: dict[TensorName, int] = field(default_factory=dict)
 
+    data_movement_connections: DataMovementConnections = None
+
     # We track first latency for these nodes (should be Temporal)
     last_temporal_node_idx: int = None
     """
@@ -517,7 +518,7 @@ def analyze_reuse_and_add_reservations_to_mapping(
     - *
     - Compute
 
-    A loop spatial-for (Reg) k in [0..K) would affect the register at the point of the
+    A loop 'spatial-for (Reg) k in [0..K)' would affect the register at the point of the
     first asterisk, but the global buffer at the point of the second asterisk.
 
     To handle this, we make a separate mapping for each tensor, analyze each, and
@@ -555,6 +556,7 @@ def analyze_reuse_and_add_reservations_to_mapping(
             tensor_to_backer_id=tensor_to_backer_id,
             is_copy_operation=is_copy_operation,
             job=job,
+            src_to_dst_movement=DataMovementConnections.from_pmapping(cur_mapping.nodes),
         )
         cur_result = analyze_node(0, job.rank_variable_bounds, info)
         if result is None:
@@ -589,17 +591,6 @@ def analyze_reuse_and_add_reservations_to_mapping(
     return result
 
 
-def get_tensor_to_backer_id(mapping: Mapping):
-    tensor_to_ids: dict[TensorName, set[int]] = {}
-    for node in mapping:
-        if isinstance(node, TensorHolder):
-            for tensor in node.tensors:
-                if tensor in tensor_to_ids:
-                    continue
-                tensor_to_ids[tensor] = id(node)
-    return tensor_to_ids
-
-
 class ReservationAnalysisTracker:
     def __init__(self, buffet, node):
         self.buffet: Buffet = buffet
@@ -811,7 +802,7 @@ def handle_repeated_value(repeated_shape):
             relevancy = info.tensor_to_relevancy[buffet.tensor][node.rank_variable]
             is_fully_relevant = isinstance(relevancy, Relevant)
             accumulated_stats = accumulated_buffet_stats.setdefault(
-                buffet, blank_buffet_stats()
+                buffet, BuffetStats.blank()
             )
             accumulated_stats += stats.repeat_temporal(
                 shape_repeats, is_fully_relevant=is_fully_relevant
@@ -882,7 +873,7 @@ def handle_repeated_value(repeated_shape):
         for i, (buffet, buffet_stats) in enumerate(child_stats):
             stats = buffet_stats
             accumulated_stats = accumulated_buffet_stats.setdefault(
-                buffet, blank_buffet_stats()
+                buffet, BuffetStats.blank()
             )
             relevancy = info.tensor_to_relevancy[buffet.tensor][rank_var]
 

accelforge/model/_looptree/types.py

@@ -24,3 +24,24 @@ class Buffet:
     "An einsum operating on the tensor."
     level: ComponentName
     "The abstract hardware level the buffet resides in."
+
+
+@dataclass(eq=True, frozen=True)
+class Compute:
+    einsum: str
+    level: str
+
+
+@dataclass(eq=True, frozen=True)
+class Network:
+    """
+    A logical network that delivers a tensor.
+    """
+    tensor: TensorName
+    "The tensor held by the buffet."
+    einsum: EinsumName
+    "An einsum operating on the tensor."
+    source: ComponentName
+    "The source of data."
+    destination: ComponentName
+    "The destination of data."