Network specification and start hops modeling (#31)

* Introduces a new node called `Array`, which can be used to create an array of components within a hierarchical environment, but also allows components to be in a flat organization.

* Cleanup of some modeling code.

* Create classes to track network statistics.
---------

Co-authored-by: tanner-andrulis <andrulis@mit.edu>

Author: gilbertmike

Makefile

@@ -66,6 +66,9 @@ push-arm64:
 run-docker:
 	docker-compose up
 
+clean-notebooks:
+	nb-clean clean notebooks/*.ipynb
+
 .PHONY: generate-docs
 generate-docs:
     # pip install sphinx-autobuild sphinx_autodoc_typehints sphinx-copybutton pydata-sphinx-theme

accelforge/frontend/arch/__init__.py

@@ -15,6 +15,8 @@
 Arch.model_rebuild()
 Fork.model_rebuild()
 Hierarchical.model_rebuild()
+Array.model_rebuild()
+Network.model_rebuild()
 
 __all__ = [
     "Action",
@@ -27,6 +29,7 @@
     "Compute",
     "Container",
     "Fork",
+    "Array",
     "Hierarchical",
     "Leaf",
     "Memory",

accelforge/frontend/arch/components.py

@@ -29,8 +29,8 @@
 from accelforge.util._setexpressions import InvertibleSet, eval_set_expression
 from accelforge.frontend.renames import TensorName
 from accelforge.frontend.arch.constraints import Comparison
-from accelforge.frontend.arch.structure import ArchNode, Leaf
-from accelforge.frontend.arch.spatialable import Spatialable
+from accelforge.frontend.arch.structure import ArchNode, Branch, Leaf
+from accelforge.frontend.arch.spatialable import Spatial, Spatialable
 
 from accelforge.util._basetypes import _uninstantiable
 
@@ -998,3 +998,11 @@ def _render_node_shape(self) -> str:
 
     def _render_node_color(self) -> str:
         return "#E0EEFF"
+
+
+class Network(Component, Leaf):
+    def _render_node_shape(self) -> str:
+        return "Msquare"
+
+    def _render_node_color(self) -> str:
+        return "#FAF8C8"

accelforge/frontend/arch/spatialable.py

@@ -81,3 +81,9 @@ class Spatialable(EvalableModel):
     def get_fanout(self) -> int:
         """The spatial fanout of this node."""
         return int(math.prod(x.fanout for x in self.spatial))
+
+    def _spatial_str(self, include_newline=True) -> str:
+        if not self.spatial:
+            return ""
+        result = ", ".join(f"{s.fanout}× {s.name}" for s in self.spatial)
+        return f"\n[{result}]" if include_newline else result

accelforge/frontend/arch/structure.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 import math
 from typing import (
     Any,
@@ -19,6 +20,8 @@
 
 from accelforge.util.exceptions import EvaluationError
 
+from accelforge.frontend.arch.spatialable import Spatialable
+
 from pydantic import Discriminator
 from accelforge.util._basetypes import _uninstantiable
 from accelforge.util.parallel import _SVGJupyterRender
@@ -68,6 +71,29 @@ def find(self, name: str, default: Any = _FIND_SENTINEL) -> Union["Leaf", Any]:
             return default
         raise ValueError(f"Leaf {name} not found in {self}")
 
+    def iterate_hierarchically(self, _parents=None):
+        """
+        Iterates over all Leaf nodes while also yielding the list of all Leaf
+        nodes that are hierarchical parents over the current node.
+        """
+        if _parents is None:
+            _parents = []
+
+        if isinstance(self, Leaf):
+            yield self, _parents
+            _parents.append(self)
+            return
+
+        assert isinstance(self, Branch)
+        if isinstance(self, (Fork, Array)):
+            for node in self.nodes:
+                yield from node.iterate_hierarchically(list(_parents))
+        elif isinstance(self, Hierarchical):
+            for node in self.nodes:
+                yield from node.iterate_hierarchically(_parents)
+        else:
+            raise RuntimeError("unhandled structure type")
+
     def _render_node_name(self) -> str:
         """The name for a Pydot node."""
         return f"{self.__class__.__name__}_{id(self)}"
@@ -137,12 +163,6 @@ def __str__(self) -> str:
             result = f"{result} [{spatial_str}]"
         return result
 
-    def _spatial_str(self, include_newline=True) -> str:
-        if not self.spatial:
-            return ""
-        result = ", ".join(f"{s.fanout}× {s.name}" for s in self.spatial)
-        return f"\n[{result}]" if include_newline else result
-
     def _render_node_label(self) -> str:
         return f"{self.name}" + self._spatial_str()
 
@@ -162,7 +182,9 @@ class Branch(ArchNode):
                 Annotated["Memory", Tag("Memory")],
                 Annotated["Toll", Tag("Toll")],
                 Annotated["Container", Tag("Container")],
+                Annotated["Network", Tag("Network")],
                 Annotated["Hierarchical", Tag("Hierarchical")],
+                Annotated["Array", Tag("Array")],
                 Annotated["Fork", Tag("Fork")],
             ],
             Discriminator(_get_tag),
@@ -235,6 +257,94 @@ def _power_gating(
         return result, non_power_gated_porp
 
 
+class Array(Branch, Spatialable):
+    def model_post_init(self, __context__=None) -> None:
+        for node in self.nodes:
+            if isinstance(node, Fork):
+                raise EvaluationError("cannot have fork inside array")
+
+    def _flatten(
+        self,
+        compute_node: str,
+        fanout: int = 1,
+        return_fanout: bool = False,
+    ):
+        from accelforge.frontend.arch.components import Compute
+        nodes = []
+
+        for node in self.nodes:
+            try:
+                if isinstance(node, Branch):
+                    raise RuntimeError("do not put branches inside array")
+                elif isinstance(node, Leaf):
+                    fanout *= node.get_fanout()
+                    node = deepcopy(node)
+                    node.spatial = EvalableList()
+                    nodes.append(node)
+                else:
+                    raise RuntimeError(f"unhandled structure type {node}")
+            except EvaluationError as e:
+                e.add_field(node)
+                raise e
+
+        if return_fanout:
+            return nodes, fanout
+        return nodes
+
+    def _render_node_label(self) -> str:
+        return f"Array {self._spatial_str()}"
+
+    def _render_node_color(self) -> str:
+        return "#FCC2FC"
+
+    def _parent2child_names(
+        self, parent_name: str = None
+    ) -> tuple[list[tuple[str, str, str]], str]:
+        from accelforge.frontend.arch.components import Compute
+
+        edges = []
+        current_parent_name = parent_name
+
+        for node in self.nodes:
+            if isinstance(node, Branch):
+                raise EvaluationError("do not put branch in array")
+            elif isinstance(node, Compute):
+                # Compute nodes branch off to the side like a Fork
+                if current_parent_name is not None:
+                    edges.append((current_parent_name, node._render_node_name(), "dashed"))
+            else:
+                if current_parent_name is not None:
+                    edges.append((current_parent_name, node._render_node_name(), "dashed"))
+        return edges, self._render_node_name()
+
+    def _render_make_children(self) -> list[pydot.Node]:
+        """Renders only children, not the Hierarchical node itself."""
+        result = [self._render_node()]
+        for node in self.nodes:
+            children = node._render_make_children()
+            result.extend(child for child in children if child is not None)
+        return result
+
+    def render(self) -> str:
+        """Renders the architecture as a Pydot graph."""
+        graph = _pydot_graph()
+
+        # Render all nodes (Hierarchical nodes return None and are filtered out)
+        for node in self._render_make_children():
+            if node is not None:
+                graph.add_node(node)
+
+        # Add edges
+        edges, _ = self._parent2child_names()
+        for parent_name, child_name in edges:
+            graph.add_edge(pydot.Edge(parent_name, child_name))
+
+        return _SVGJupyterRender(graph.create_svg(prog="dot").decode("utf-8"))
+
+    def _repr_svg_(self) -> str:
+        return self.render()
+
+
 class Hierarchical(Branch):
     def _flatten(
         self,
@@ -265,6 +375,18 @@ def _flatten(
                     ):
                         break
                     assert not isinstance(node, Fork)
+                elif isinstance(node, Array):
+                    new_nodes = node._flatten(
+                        compute_node, fanout, return_fanout=False
+                    )
+                    nodes.extend(new_nodes)
+                    nodes.append(node)
+                    fanout *= node.get_fanout()
+                    if any(
+                        isinstance(n, Compute) and n.name == compute_node
+                        for n in new_nodes
+                    ):
+                        break
                 elif isinstance(node, Compute):
                     if node.name == compute_node:
                         fanout *= node.get_fanout()
@@ -299,23 +421,27 @@ def _parent2child_names(
                 edges.extend(child_edges)
                 if last_child_name is not None:
                     current_parent_name = last_child_name
+            elif isinstance(node, Array):
+                child_edges, last_child_name = node._parent2child_names(
+                    node._render_node_name()
+                )
+                edges.extend(child_edges)
+                edges.append((current_parent_name, last_child_name, "solid"))
+                if last_child_name is not None:
+                    current_parent_name = last_child_name
             elif isinstance(node, Compute):
                 # Compute nodes branch off to the side like a Fork
                 if current_parent_name is not None:
-                    edges.append((current_parent_name, node._render_node_name()))
+                    edges.append((current_parent_name, node._render_node_name(), "solid"))
             else:
                 if current_parent_name is not None:
-                    edges.append((current_parent_name, node._render_node_name()))
+                    edges.append((current_parent_name, node._render_node_name(), "solid"))
 
                 # Update parent for next iteration
                 current_parent_name = node._render_node_name()
 
         return edges, current_parent_name
 
-    def _render_node(self) -> pydot.Node:
-        """Hierarchical nodes should not be rendered."""
-        return None
-
     def _render_make_children(self) -> list[pydot.Node]:
         """Renders only children, not the Hierarchical node itself."""
         result = []
@@ -335,8 +461,8 @@ def render(self) -> str:
 
         # Add edges
         edges, _ = self._parent2child_names()
-        for parent_name, child_name in edges:
-            graph.add_edge(pydot.Edge(parent_name, child_name))
+        for parent_name, child_name, style in edges:
+            graph.add_edge(pydot.Edge(parent_name, child_name, style=style))
 
         return _SVGJupyterRender(graph.create_svg(prog="dot").decode("utf-8"))
 

accelforge/frontend/spec.py

@@ -3,7 +3,7 @@
 from accelforge.frontend.mapper import FFM
 from accelforge.frontend.renames import EinsumName, Renames
 from accelforge.util._eval_expressions import EvaluationError, ParseExpressionsContext
-from accelforge.frontend.arch import Compute, Leaf, Component, Arch, Container
+from accelforge.frontend.arch import Compute, Leaf, Component, Arch, Container, Spatialable
 
 from accelforge.frontend.workload import Workload
 from accelforge.frontend.variables import Variables
@@ -194,7 +194,6 @@ def calculate_component_area_energy_latency_leak(
         if einsum_name is None and len(self.workload.einsums) > 0:
             einsum_name = self.workload.einsums[0].name
 
-        components = set()
         try:
             if not getattr(self, "_evaluated", False):
                 self = self._spec_eval_expressions(einsum_name=einsum_name)
@@ -211,36 +210,37 @@ def calculate_component_area_energy_latency_leak(
                 )
             raise
 
-        for arch in self._get_flattened_architecture():
-            fanout = 1
-            for component in arch:
-                fanout *= component.get_fanout()
-                if component.name in components or isinstance(component, Container):
-                    continue
-                assert isinstance(component, Component)
-                components.add(component.name)
-                orig: Component = self.arch.find(component.name)
-                c = component
-                if area:
-                    c = c.calculate_area(models)
-                    orig.area = c.area
-                    orig.total_area = c.area * fanout
-                if energy:
-                    c = c.calculate_action_energy(models)
-                    for a in c.actions:
-                        orig_action = orig.actions[a.name]
-                        orig_action.energy = a.energy
-                if latency:
-                    c = c.calculate_action_latency(models)
-                    for a in c.actions:
-                        orig_action = orig.actions[a.name]
-                        orig_action.latency = a.latency
-                if leak:
-                    c = c.calculate_leak_power(models)
-                    orig.leak_power = c.leak_power
-                    orig.total_leak_power = c.leak_power * fanout
-                orig.component_modeling_log.extend(c.component_modeling_log)
-                orig.component_model = c.component_model
+        for leaf, parents in self.arch.iterate_hierarchically():
+            if not isinstance(leaf, Component):
+                continue
+
+            global_fanout = 1
+            for p in parents:
+                if isinstance(p, Spatialable):
+                    global_fanout *= p.get_fanout()
+
+            orig: Component = self.arch.find(leaf.name)
+            c = leaf
+            if area:
+                c = c.calculate_area(models)
+                orig.area = c.area
+                orig.total_area = c.area * global_fanout
+            if energy:
+                c = c.calculate_action_energy(models)
+                for a in c.actions:
+                    orig_action = orig.actions[a.name]
+                    orig_action.energy = a.energy
+            if latency:
+                c = c.calculate_action_latency(models)
+                for a in c.actions:
+                    orig_action = orig.actions[a.name]
+                    orig_action.latency = a.latency
+            if leak:
+                c = c.calculate_leak_power(models)
+                orig.leak_power = c.leak_power
+                orig.total_leak_power = c.leak_power * global_fanout
+            orig.component_modeling_log.extend(c.component_modeling_log)
+            orig.component_model = c.component_model
 
         return self
 

accelforge/mapper/FFM/_join_pmappings/join_pmappings.py

@@ -618,6 +618,8 @@ def no_match_lookahead_error(
         # ======================================================================
         # If we delayed the mapping merging, do it now.
         # ======================================================================
+        import copy
+        prev_combined = copy.deepcopy(combined)
         if DELAY:
             mappings = parallel(
                 [c.mappings for c in combined],