[frontend,examples,notebooks,tests] Finish Network spec implementation

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,7 @@
 Arch.model_rebuild()
 Fork.model_rebuild()
 Hierarchical.model_rebuild()
+Flat.model_rebuild()
 Network.model_rebuild()
 
 __all__ = [
@@ -28,6 +29,7 @@
     "Compute",
     "Container",
     "Fork",
+    "Flat",
     "Hierarchical",
     "Leaf",
     "Memory",

accelforge/frontend/arch/components.py

@@ -1000,5 +1000,11 @@ def _render_node_color(self) -> str:
         return "#E0EEFF"
 
 
-class Network(Component, Branch):
-    shape: EvalableList[Spatial]
+class Network(Component, Leaf):
+    shape: EvalableList[Spatial] = EvalableList()
+
+    def _render_node_shape(self) -> str:
+        return "Msquare"
+
+    def _render_node_color(self) -> str:
+        return "#FAF8C8"

accelforge/frontend/arch/structure.py

@@ -68,6 +68,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, Flat)):
+            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)}"
@@ -164,6 +187,7 @@ class Branch(ArchNode):
                 Annotated["Container", Tag("Container")],
                 Annotated["Network", Tag("Network")],
                 Annotated["Hierarchical", Tag("Hierarchical")],
+                Annotated["Flat", Tag("Flat")],
                 Annotated["Fork", Tag("Fork")],
             ],
             Discriminator(_get_tag),
@@ -236,6 +260,68 @@ def _power_gating(
         return result, non_power_gated_porp
 
 
+class Flat(Branch):
+    def _parent2child_names(
+        self, parent_name: str = None
+    ) -> tuple[list[tuple[str, str]], str]:
+        from accelforge.frontend.arch.components import Compute
+
+        edges = []
+        current_parent_name = parent_name
+
+        for node in self.nodes:
+            if isinstance(node, (Hierarchical, Fork)):
+                child_edges, _last_child_name = node._parent2child_names(
+                    current_parent_name
+                )
+                edges.extend(child_edges)
+            elif isinstance(node, Flat):
+                child_edges, _last_child_name = node._parent2child_names(
+                    current_parent_name
+                )
+                edges.extend(child_edges)
+            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()))
+            else:
+                if current_parent_name is not None:
+                    edges.append((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 = []
+        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()
+        print(edges)
+        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,
@@ -300,6 +386,13 @@ def _parent2child_names(
                 edges.extend(child_edges)
                 if last_child_name is not None:
                     current_parent_name = last_child_name
+            elif isinstance(node, Flat):
+                child_edges, last_child_name = node._parent2child_names(
+                    current_parent_name
+                )
+                edges.extend(child_edges)
+                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:

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
 

examples/arches/networked/flat.yaml

@@ -12,10 +12,14 @@ arch:
 
   - !Network
     name: NoC
+    area: 0
+    leak_power: 0
     shape:
     - {name: X, fanout: 4}
     - {name: Y, fanout: 4}
     actions: []
+
+  - !Flat
     nodes:
     - !Memory
       name: GlobalBuffer
@@ -80,5 +84,7 @@ arch:
 
       - !Compute
         name: MAC
+        area: 0
+        leak_power: 0
         actions:
         - {name: compute, energy: 0, latency: 0}

…/arches/networked/hierarchical_only.yaml …mples/arches/networked/hierarchical.yamlexamples/arches/networked/hierarchical_only.yaml renamed to examples/arches/networked/hierarchical.yaml examples/arches/networked/hierarchical_only.yaml renamed to examples/arches/networked/hierarchical.yaml

@@ -0,0 +1,46 @@
+arch:
+  nodes:
+  - !Memory
+    name: MainMemory
+    size: inf
+    area: 0
+    leak_power: 0
+    tensors: {keep: All}
+    actions:
+    - {name: read, energy: 0, latency: 0}
+    - {name: write, energy: 0, latency: 0}
+
+  - !Network
+    name: NoC
+    area: 0
+    leak_power: 0
+    actions: []
+
+  - !Memory
+    name: HBM
+    size: inf
+    area: 0
+    leak_power: 0
+    tensors: {keep: ~MainMemory, may_keep: All}
+    actions:
+    - {name: read, energy: 0, latency: 0}
+    - {name: write, energy: 0, latency: 0}
+    spatial:
+    - {name: X, fanout: 4}
+    - {name: Y, fanout: 4}
+
+  - !Memory
+    name: Buffer
+    size: inf
+    area: 0
+    leak_power: 0
+    actions:
+    - {name: read, energy: 0, latency: 0}
+    - {name: write, energy: 0, latency: 0}
+
+  - !Compute
+    name: MAC
+    area: 0
+    leak_power: 0
+    actions:
+    - {name: compute, energy: 0, latency: 0}