Implement (almost) proper latency model; waiting for hwcomponents latency/bandwidth update

gilbertmike · gilbertmike · commit db5eae4579f6 · 2026-05-27T17:05:40.000-04:00
diff --git a/accelforge/frontend/arch/components.py b/accelforge/frontend/arch/components.py
@@ -1174,6 +1174,20 @@ class Network(Component, Leaf):
     of the spatial nodes from top to bottom.
     """
 
+    total_latency: str | int | float = "max(max_hops*actions['hops'].latency, max_link_traffic/actions['hops'].latency)"
+    """
+    Models latency as either:
+    - *Latency-bound*, which means that the latency of the route with the most number of
+      hops dominate the overall communication latency.
+    - *Bandwidth-bound*, which means that the traffic over the most congested link
+      dominates the overall communication latency.
+
+    Keywords:
+    - `max_hops` returns the number of hops in the longest route.
+    - `max_link_traffic` returns the amount of traffic (in bits) over the most congested
+      link.
+    """
+
     bits_per_value: EvalsTo[dict] = {}
     """
     Sets the bits per value for tensors in this `TensorHolder`. Keys are evaluated as
diff --git a/accelforge/model/_looptree/latency/memory.py b/accelforge/model/_looptree/latency/memory.py
@@ -14,7 +14,7 @@
 
 from accelforge.model._looptree.reuse.symbolic import BuffetStats
 from accelforge.util._eval_expressions import MATH_FUNCS, eval_expression
-from accelforge.util._sympy.broadcast_max import Max, Min
+from accelforge.util._sympy.broadcast_max import Max, Min, MaxGeqZero
 import symengine as se
 
 
@@ -47,6 +47,10 @@ def component_latency(
     component_to_actions: dict[str, dict[str, float]] = defaultdict(
         lambda: defaultdict(lambda: 0)
     )
+    # Holds ``keywords" that do not map neatly to actions, e.g., max_hops for network
+    component_to_keywords: dict[str, dict[str, float]] = defaultdict(
+        lambda: defaultdict(lambda: 0)
+    )
     name2component: dict[str, Component] = {node.name: node for node in flattened_arch}
 
     compute_obj = flattened_arch[-1]
@@ -79,16 +83,28 @@ def component_latency(
                 f"Component {component} is not a TensorHolder or Compute"
             )
 
+    network_to_max_link_traffic = defaultdict(lambda: defaultdict(lambda: 0))
+    network_to_max_hops = defaultdict(lambda: [])
     for network, network_stats in looptree_results.network_stats.items():
         component = network.component
-        actions = component_to_actions[component]
         if component not in name2component:
             raise ValueError(f"Component {component} found in mapping but not arch")
 
-        for action in name2component[component].actions:
-            actions[f"{action.name}_actions"] += 0
+        dim_traffic = network_to_max_link_traffic[component]
+        for dim, max_traffic_in_dim in network_stats.max_traffic.items():
+            dim_traffic[dim] += max_traffic_in_dim
 
-        actions["hops_actions"] += network_stats.max_hops
+        network_to_max_hops[component].append(network_stats.max_hops)
+
+    for network, network_stats in looptree_results.network_stats.items():
+        component = network.component
+        keywords = component_to_keywords[component]
+        keywords["max_link_traffic"] = MaxGeqZero(
+            *network_to_max_link_traffic[component].values()
+        )
+        keywords["max_hops"] = MaxGeqZero(
+            *network_to_max_hops[component]
+        )
 
     longest_compute_latency = Max(
         0, *[s.max_latency for s in looptree_results.compute_stats.values()]
@@ -126,14 +142,15 @@ def component_latency(
         "sum": se.Add,
     }
 
-    for component, actions in component_to_actions.items():
-        component_obj = name2component[component]
+    for component, component_obj in name2component.items():
+        actions = component_to_actions[component]
         symbol_table = {
             "action2latency": component_to_action_latency[component],
             **symbol_table_base,
             **name2component[component].shallow_model_dump(include_None=True),
             **actions,
             **component_to_action_latency[component],
+            **component_to_keywords[component],
         }
         if name2component[component].total_latency is not None:
             component_latency[component] = eval_expression(
diff --git a/accelforge/model/_looptree/reuse/symbolic/_network.py b/accelforge/model/_looptree/reuse/symbolic/_network.py
@@ -1,9 +1,5 @@
-import copy
-from accelforge.frontend import arch
-from accelforge.frontend.arch import Network as NetworkSpec
 from accelforge.frontend.mapping import (
-    TensorHolder,
-    TensorName
+    Spatial
 )
 from accelforge.frontend._workload_isl._symbolic import (
     compute_dense_tile_occupancy,
@@ -20,7 +16,7 @@
 
 class NetworkAnalyzer:
     def __init__(self, network_stats):
-        self.overall_max_hops = 0
+        self.overall_max_hops: dict = {}
         self.network_stats = network_stats
 
     def accumulate_child_result(
@@ -30,8 +26,9 @@ def accumulate_child_result(
         shape_repeats,
         einsum_name,
         child_shape,
-        node,
+        node: Spatial,
     ):
+        """This function is called for every repeated shape."""
         flattened_arch = info.job.flattened_arch
 
         for network, child_network_stats in child_result.network_stats.items():
@@ -40,44 +37,37 @@ def accumulate_child_result(
                 self.network_stats[network] = NetworkStats()
             accumulated_network_stats = self.network_stats[network]
 
-            accumulated_network_stats.total_hops += (
-                child_network_stats.total_hops * shape_repeats
-            )
-            accumulated_network_stats.max_hops = MaxGeqZero(
-                accumulated_network_stats.max_hops,
-                child_network_stats.max_hops,
-            )
-            projection = info.einsum_tensor_to_projection[(einsum_name, network.tensor)]
-            component_object = flattened_arch[network.component]
-            workload_bpv = info.job.einsum.tensor_accesses[
-                network.tensor
-            ].bits_per_value
-            bits_per_value = component_object.bits_per_value.get(
-                network.tensor, workload_bpv
-            )
-            bits_per_action = component_object.bits_per_action
-            if bits_per_action is not None:
-                actions_per_value = bits_per_value / bits_per_action
-            else:
-                actions_per_value = bits_per_value
-            volume = (
-                compute_dense_tile_occupancy(projection, child_shape)
-                * actions_per_value
-            )
-
+            # 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):
+                accumulated_network_stats.total_hops += (
+                    child_network_stats.total_hops * shape_repeats
+                )
+                accumulated_network_stats.max_hops = MaxGeqZero(
+                    accumulated_network_stats.max_hops,
+                    child_network_stats.max_hops,
+                )
+                for k, v in child_network_stats.max_traffic.items():
+                    accumulated_network_stats.max_traffic[k] = MaxGeqZero(
+                        accumulated_network_stats.max_traffic.get(k, 0),
+                        v
+                    )
                 continue
 
+            volume = self._get_data_volume(network, einsum_name, info, child_shape)
+
             relevancy = info.tensor_to_relevancy[network.tensor][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
@@ -99,26 +89,56 @@ def accumulate_child_result(
                         *
                         volume
                     )
-                    max_hops = MinGeqZero(shape_repeats*last_fanout, physical_stride)
+                    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}")
 
-            # TODO: this is sketchy
-            self.overall_max_hops += max_hops
+            # 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
 
-            accumulated_network_stats.total_hops += total_cost
+            accumulated_network_stats.total_hops += (
+                total_cost + child_network_stats.total_hops*shape_repeats
+            )
             accumulated_network_stats.max_hops = MaxGeqZero(
                 accumulated_network_stats.max_hops,
-                self.overall_max_hops + child_network_stats.max_hops,
+                self.overall_max_hops[network] + 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)
             )
 
         return self.overall_max_hops
 
+    def _get_data_volume(self, network, einsum_name, info, child_shape):
+        flattened_arch = info.job.flattened_arch
+        projection = info.einsum_tensor_to_projection[(einsum_name, network.tensor)]
+        component_object = flattened_arch[network.component]
+        workload_bpv = info.job.einsum.tensor_accesses[
+            network.tensor
+        ].bits_per_value
+        bits_per_value = component_object.bits_per_value.get(
+            network.tensor, workload_bpv
+        )
+        bits_per_action = component_object.bits_per_action
+        if bits_per_action is not None:
+            actions_per_value = bits_per_value / bits_per_action
+        else:
+            actions_per_value = bits_per_value
+        volume = (
+            compute_dense_tile_occupancy(projection, child_shape)
+            * actions_per_value
+        )
+        return volume
+
 
 def multicast_cost(n_dsts, stride):
     """Returns total hops of multicast along a dimension."""
diff --git a/accelforge/model/_looptree/reuse/symbolic/_stats.py b/accelforge/model/_looptree/reuse/symbolic/_stats.py
@@ -21,7 +21,11 @@
 @dataclass
 class NetworkStats:
     total_hops: Any = field(default=0)
+    """Total number of hops overall. Useful to calculate energy."""
     max_hops: Any = field(default=0)
+    """Longest hops among all routes."""
+    max_traffic: dict[int | str, Any] = field(default_factory=dict)
+    """Maximum traffic occuring on any single link along a dimension."""
 
     def repeat(self, n_repeats):
         new = copy.copy(self)
@@ -32,10 +36,6 @@ def repeat(self, n_repeats):
         new.total_hops = new.total_hops * n_repeats
         return new
 
-    def combine(self, other: "NetworkStats"):
-        self.total_hops += other.total_hops
-        self.max_hops = max(self.max_hops, other.max_hops)
-
 
 @dataclass
 class BuffetStats:
diff --git a/accelforge/model/run_model.py b/accelforge/model/run_model.py
@@ -43,6 +43,7 @@ def run_model(
     )
 
     latency = component_latency(reuse, job.flattened_arch, pmapping, spec)
+    print(latency)
     try:
         overall_latency = MaxGeqZero(*latency.values())
     except Exception as e:
diff --git a/tests/input_files/networked/hierarchical_1d.yaml b/tests/input_files/networked/hierarchical_1d.yaml
@@ -24,8 +24,9 @@ arch:
     name: PeArray
     area: 0
     leak_power: 0
+    total_latency: "max_hops"
     actions:
-    - {name: hops, energy: 1, latency: 0}
+    - {name: hops, energy: 1, latency: 1}
 
   - !Memory
     name: Scratchpad
diff --git a/tests/test_network.py b/tests/test_network.py
@@ -125,6 +125,10 @@ def test_hierarchical_1d(self):
             * KN
             * BITS_PER_VALUE,
         )
+        self.assertEqual(
+            result.data["Total<SEP>latency"].iloc[0],
+            4
+        )
 
     def test_hierarchical(self):
         M = 8

Original file line number	Diff line number	Diff line change
`@@ -43,6 +43,7 @@ def run_model(`
`43`	`43`	`)`
`44`	`44`
`45`	`45`	`latency = component_latency(reuse, job.flattened_arch, pmapping, spec)`
	`46`	`+ print(latency)`
`46`	`47`	`try:`
`47`	`48`	`overall_latency = MaxGeqZero(*latency.values())`
`48`	`49`	`except Exception as e:`