Model latency and distributed occupancy

Author: gilbertmike

accelforge/model/_looptree/latency/memory.py

@@ -79,6 +79,17 @@ def component_latency(
                 f"Component {component} is not a TensorHolder or Compute"
             )
 
+    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
+
+        actions["hops_actions"] += network_stats.max_hops
+
     longest_compute_latency = Max(
         0, *[s.max_latency for s in looptree_results.compute_stats.values()]
     )

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

@@ -818,6 +818,10 @@ def inherit_add(attr: str, default_value: Any = fills) -> Any:
                             dim_fanout[d] / component_object._get_physical_stride_along(d)
                         )
 
+        # ==========================
+        # Recalculate usage of distributed buffers
+        stats.max_occupancy /= n_active_physical_units
+
         # ==========================
         # Data exchanges with parent
         if count_downward_movement[tensor]:  # Parent -> Me