Merge branch 'main' of github.com:Accelergy-Project/fastfusion

Author: tanner-andrulis

accelforge/frontend/arch/components.py

@@ -1001,8 +1001,46 @@ def _render_node_color(self) -> str:
 
 
 class Network(Component, Leaf):
+    """
+    Defines a network component.
+
+    The routing is currently defined using the mapping, the routing follows the order
+    of the spatial nodes from top to bottom.
+    """
+    bits_per_value_scale: EvalsTo[dict] = {"All": 1}
+    """
+    A scaling factor for the bits per value of the tensors in this `TensorHolder`. If
+    this is a dictionary, keys in the dictionary are evaluated as expressions and may
+    reference one or more tensors.
+    """
+
+    bits_per_action: EvalsTo[int | float | None] = None
+    """
+    The number of bits accessed in each of this component's actions. Overridden by
+    bits_per_action in any action of this component. If set here, acts as a default
+    value for the bits_per_action of all actions of this component.
+    """
+
     def _render_node_shape(self) -> str:
         return "Msquare"
 
     def _render_node_color(self) -> str:
         return "#FAF8C8"
+
+    def _eval_expressions(self, *args, **kwargs):
+        # Sometimes the same component object may appear in the mapping and the
+        # architecture, in which case we don't want parsing to happen twice.
+        if getattr(self, "_evaluated", False):
+            return super()._eval_expressions(*args, **kwargs)
+
+        class MyPostCall(_PostCall):
+            def __call__(self, field, value, evaluated, symbol_table):
+                if field == "bits_per_value_scale":
+                    evaluated = _eval_tensor2bits(
+                        evaluated,
+                        location="bits_per_value_scale",
+                        symbol_table=symbol_table,
+                    )
+                return evaluated
+
+        return super()._eval_expressions(*args, **kwargs, post_calls=(MyPostCall(),))

accelforge/frontend/arch/structure.py

@@ -29,6 +29,9 @@
 
 _FIND_SENTINEL = object()
 
+D = TypeVar("D")
+T = TypeVar("T")
+
 
 class ArchNode(EvalableModel):
     """A node in the architecture."""
@@ -71,6 +74,46 @@ def find(self, name: str, default: Any = _FIND_SENTINEL) -> Union["Leaf", Any]:
             return default
         raise ValueError(f"Leaf {name} not found in {self}")
 
+    def is_above(self, node_a: str, node_b: str) -> bool:
+        """Returns whether node_a is above node_b in a hierarchy."""
+        self.find(node_a)
+        self.find(node_b)
+        for node, parents in self.iterate_hierarchically():
+            if node.name != node_b:
+                continue
+            return any(p.name == node_a for p in parents)
+
+    def find_first_of_type_above(
+        self,
+        node_type: T,
+        name: str,
+        default: D = _FIND_SENTINEL
+    ) -> T | D:
+        """
+        Returns the first node with type `node_type` above `name`.
+
+        If `name` does not exist, raises an error.
+
+        If no node of `node_type` is found, either `default` is
+        returned (if provided) or raises an error.
+        """
+        # Check if name exists
+        # This *should* raise even if default != _FIND_SENTINEL
+        self.find(name)
+
+        for node, parents in self.iterate_hierarchically():
+            if node.name != name:
+                continue
+            for p in reversed(parents):
+                if isinstance(p, node_type):
+                    return p
+            if default is not _FIND_SENTINEL:
+                return default
+            raise ValueError(f"Parent of {name} with type {node_type} not found")
+        raise RuntimeError(
+            "BUG: find() finds node but iterate_hierarchically() does not"
+        )
+
     def iterate_hierarchically(self, _parents=None):
         """
         Iterates over all Leaf nodes while also yielding the list of all Leaf
@@ -170,9 +213,6 @@ def _render_make_children(self) -> list[pydot.Node]:
         return [self._render_node()]
 
 
-T = TypeVar("T")
-
-
 @_uninstantiable
 class Branch(ArchNode):
     nodes: ArchNodes[

accelforge/model/_looptree/energy.py

@@ -1,21 +1,17 @@
-from collections import defaultdict
 from collections.abc import Mapping as MappingABC
-from dataclasses import dataclass
 import logging
 from numbers import Number
 from numbers import Real
 
 from accelforge.frontend import arch
-from accelforge.frontend.mapping.mapping import MappingNode
 from accelforge.frontend.spec import Spec
 from accelforge.model._looptree.reuse.symbolic import SymbolicAnalysisOutput
 from accelforge.util._base_analysis_types import (
     ActionCount,
     ActionKey,
     VerboseActionKey,
 )
-from accelforge.frontend.workload import Workload
-from accelforge.frontend.mapping import Mapping
+from accelforge.model._looptree.types import Network
 
 
 def gather_actions(
@@ -29,6 +25,7 @@ def gather_actions(
 
     buffet_keyer = _get_buffet_keyer(verbose, use_name, bindings)
     compute_keyer = _get_compute_keyer(verbose, use_name, bindings)
+    network_keyer = _get_network_keyer(verbose, use_name, bindings)
 
     for buffet, accesses in looptree_results.buffet_stats.items():
         if buffet.level in compute_levels:
@@ -60,6 +57,13 @@ def gather_actions(
         actions[key].total += ops.total_ops
         actions[key].max_per_unit += ops.max_per_unit_ops
 
+    for network, stats in looptree_results.network_stats.items():
+        key = network_keyer(network, "hops")
+        if key not in actions:
+            actions[key] = ActionCount.default()
+        actions[key].total += stats.total_hops
+        actions[key].max_per_unit += stats.max_hops
+
     return actions
 
 
@@ -111,6 +115,22 @@ def compute_keyer(compute, action_name):
     return compute_keyer
 
 
+def _get_network_keyer(verbose, use_name, bindings):
+    if not verbose:
+        def network_keyer(network: Network, action_name: str):
+            component = network.component
+            if not use_name:
+                component = bindings[component]
+            return ActionKey(component, action_name)
+    else:
+        def network_keyer(network: Network, action_name: str):
+            component = network.component
+            if not use_name:
+                component = bindings[component]
+            return VerboseActionKey(component, action_name, network.tensor, network.einsum)
+    return network_keyer
+
+
 def compute_energy_from_actions(
     spec: Spec,
     action_counts: MappingABC[ActionKey, Real],
@@ -135,7 +155,7 @@ def compute_energy_from_actions(
             energy_per_ac = component_obj.actions[key.action].energy
         except KeyError as e:
             raise KeyError(
-                f"Action {key.action} not found in component {key.component}. Action occurred "
+                f"Action {key.action} not found in component {key.level}. Action occurred "
                 f"{counts.total} times."
             ) from None
         energy_result[key] = counts.total * energy_per_ac

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

@@ -5,6 +5,7 @@
 from accelforge.frontend.mapping import (
     Mapping,
     TensorHolder,
+    Compute as ComputeSpec
 )
 from accelforge.frontend.workload import (
     TensorName,
@@ -54,6 +55,13 @@ def from_pmapping(cls, pmapping):
                     if src is not None:
                         src_to_dst[src] = buffet
                     src_to_dst[buffet] = None
+            elif isinstance(node, ComputeSpec):
+                for tensor in tensor_to_last_src:
+                    dst_buffet = Buffet(tensor, einsum_name, node.component)
+                    src = tensor_to_last_src[tensor]
+                    src_to_dst[src] = dst_buffet
+                    dst_to_src[dst_buffet] = src
+
         result = DataMovementConnections()
         result.src_to_dst = src_to_dst
         result.dst_to_src = dst_to_src