[network] Tested distributed buffers

gilbertmike · gilbertmike · commit 307196ddaf4a · 2026-05-22T15:13:12.000-04:00
diff --git a/accelforge/frontend/arch/_flattened_arch.py b/accelforge/frontend/arch/_flattened_arch.py
@@ -1,3 +1,12 @@
+from typing import TypeVar
+
+
+_FIND_SENTINEL = object()
+
+D = TypeVar("D")
+T = TypeVar("T")
+
+
 class FlattenedArch:
     """
     A flattened arch is an architecture spec that has been
@@ -51,3 +60,27 @@ def is_above(self, name_a: str, name_b: str):
         idx_a = self.index(name_a)
         idx_b = self.index(name_b)
         return idx_a < idx_b
+
+    def find_first_of_type_between(
+        self, node_type: T, name_lower: str, name_upper: str, default: D = _FIND_SENTINEL
+    ) -> T | D:
+        """
+        Returns the first node with type `node_type` above `name_lower` and under `name_upper`.
+
+        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.
+        """
+        upper_idx = self.index(name_upper)
+        lower_idx = self.index(name_lower)
+
+        for i, node in enumerate(self.nodes):
+            if not isinstance(node, node_type) or i <= upper_idx or i >= lower_idx:
+                continue
+            else:
+                return node
+        if default is not _FIND_SENTINEL:
+            return default
+        else:
+            raise ValueError(f"node with type {node_type} between {name_upper} and {name_lower} not found")
diff --git a/accelforge/model/_looptree/reuse/symbolic/_symbolic.py b/accelforge/model/_looptree/reuse/symbolic/_symbolic.py
@@ -903,19 +903,21 @@ def analyze_reservation(node_idx, current_shape, info: AnalysisInfo):
     child_result.buffet_stats[buffet] = stats
 
     # Reservation nodes are the first to produce stats for a network
-    network_node = info.job.spec_one_einsum.arch.find_first_of_type_above(
-        NetworkSpec, buffet.level, default=None
-    )
-    if network_node is not None:
-        network = Network(
-            tensor,
-            einsum_name,
-            info.data_movement_connections.get_src(buffet),
-            buffet,
-            component=network_node.name if network_node else network_node,
+    src = info.data_movement_connections.get_src(buffet)
+    if src is not None:
+        network_node = info.job.flattened_arch.find_first_of_type_between(
+            NetworkSpec, buffet.level, src.level, default=None
         )
-        assert network not in child_result.network_stats
-        child_result.network_stats[network] = NetworkStats()
+        if network_node is not None:
+            network = Network(
+                tensor,
+                einsum_name,
+                src,
+                buffet,
+                component=network_node.name if network_node else network_node,
+            )
+            assert network not in child_result.network_stats
+            child_result.network_stats[network] = NetworkStats()
 
     fanout_key = (node.resource, einsum_name)
     if fanout_key not in child_result.fanout:
@@ -965,18 +967,20 @@ def analyze_compute(
         stats.max_occupancy = 1
         result_accumulator.buffet_stats[buffet] = stats
 
-        network_node = info.job.spec_one_einsum.arch.find_first_of_type_above(
-            NetworkSpec, node.component, default=None
-        )
-        if network_node is not None:
-            network = Network(
-                tensor,
-                info.job.einsum_name,
-                info.data_movement_connections.get_src(buffet),
-                buffet,
-                component=network_node.name if network_node else network_node,
+        src = info.data_movement_connections.get_src(buffet)
+        if src is not None:
+            network_node = info.job.flattened_arch.find_first_of_type_between(
+                NetworkSpec, node.component, src.level, default=None
             )
-            result_accumulator.network_stats[network] = NetworkStats()
+            if network_node is not None:
+                network = Network(
+                    tensor,
+                    info.job.einsum_name,
+                    src,
+                    buffet,
+                    component=network_node.name if network_node else network_node,
+                )
+                result_accumulator.network_stats[network] = NetworkStats()
 
     return result_accumulator
 
diff --git a/tests/input_files/networked/flat.yaml b/tests/input_files/networked/flat.yaml
@@ -1,5 +1,3 @@
-{% set N_ROW_BUFFER = N_ROW_BUFFER | default(4) %}
-{% set N_COL_BUFFER = N_COL_BUFFER | default(4) %}
 arch:
   nodes:
   - !Memory
@@ -12,29 +10,12 @@ arch:
     - {name: read, energy: 100, latency: 0}
     - {name: write, energy: 100, latency: 0}
 
-  - !Network
-    name: NoC
-    area: 0
-    leak_power: 0
-    actions:
-    - {name: hops, energy: 1, latency: 0}
-
   - !Array
     name: Array
     spatial:
     - {name: X, fanout: 4}
     - {name: Y, fanout: 4}
     nodes:
-    - !Memory
-      name: GlobalBuffer
-      size: inf
-      area: 0
-      leak_power: 0
-      tensors: {keep: ~MainMemory, may_keep: All}
-      actions:
-      - {name: read, energy: 10, latency: 0}
-      - {name: write, energy: 10, latency: 0}
-
     - !Memory
       name: RowBuffer
       size: inf
@@ -45,7 +26,7 @@ arch:
       - {name: read, energy: 5, latency: 0}
       - {name: write, energy: 5, latency: 0}
       spatial:
-      - {name: X, fanout: {{N_ROW_BUFFER}}}
+      - {name: X, fanout: 4}
 
     - !Memory
       name: ColumnBuffer
@@ -57,7 +38,7 @@ arch:
       - {name: read, energy: 5, latency: 0}
       - {name: write, energy: 5, latency: 0}
       spatial:
-      - {name: Y, fanout: {{N_COL_BUFFER}}}
+      - {name: Y, fanout: 4}
 
     - !Memory
       name: DistributedBuffer
@@ -72,6 +53,14 @@ arch:
       - {name: X, fanout: 2}
       - {name: Y, fanout: 2}
 
+    - !Network
+      name: NoC
+      area: 0
+      leak_power: 0
+      actions:
+      - {name: hops, energy: 1, latency: 0}
+
+
   - !Memory
     name: Scratchpad
     size: inf
diff --git a/tests/input_files/networked/one_matmul_to_flat.yaml b/tests/input_files/networked/one_matmul_to_flat.yaml
@@ -0,0 +1,42 @@
+mapping:
+  nodes:
+  - !Storage
+    component: MainMemory
+    tensors: [T0, T1, W0]
+  - !Storage
+    component: DistributedBuffer
+    tensors: [W0]
+  - !Temporal
+    rank_variable: m
+    tile_shape: {{ M_TILE }}
+  - !Storage
+    component: RowBuffer
+    tensors: [T0]
+  - !Storage
+    component: ColumnBuffer
+    tensors: [T1]
+  - !Spatial
+    rank_variable: n0
+    tile_shape: {{ MAC_TILE }}
+    component: Array
+    name: X
+  - !Spatial
+    rank_variable: n1
+    tile_shape: {{ MAC_TILE }}
+    component: Array
+    name: Y
+  - !Storage
+    component: Scratchpad
+    tensors: [T0, T1, W0]
+  - !Temporal
+    rank_variable: m
+    tile_shape: 1
+  - !Temporal
+    rank_variable: n0
+    tile_shape: 1
+  - !Temporal
+    rank_variable: n1
+    tile_shape: 1
+  - !Compute
+    einsum: Matmul0
+    component: MAC
diff --git a/tests/test_network.py b/tests/test_network.py
@@ -33,7 +33,6 @@ def test_flat(self):
             {n.name for n in spec.arch.get_nodes_of_type(af.spec.Leaf)},
             {
                 "MainMemory",
-                "GlobalBuffer",
                 "NoC",
                 "RowBuffer",
                 "ColumnBuffer",
@@ -226,6 +225,71 @@ def test_hierarchical(self):
             * BITS_PER_VALUE,
         )
 
+    def test_flat(self):
+        M = 8
+        KN = 8
+        MAC_TILE = 2
+        M_TILE = 4
+        BITS_PER_VALUE = 8
+
+        spec = af.Spec.from_yaml(
+            af.examples.workloads.matmuls,
+            INPUT_FILES_DIR / "flat.yaml",
+            INPUT_FILES_DIR / "one_matmul_to_flat.yaml",
+            jinja_parse_data={
+                "N_EINSUMS": 1,
+                "M": 8,
+                "KN": 8,
+                "MAC_TILE": MAC_TILE,
+                "M_TILE": M_TILE,
+            },
+        )
+        result = spec.evaluate_mapping()
+        self.assertEqual(
+            result.data['Matmul0<SEP>action<SEP>NoC<SEP>T0<SEP>hops'].iloc[0],
+            (
+                M / M_TILE
+                *
+                (KN / MAC_TILE) * (KN / MAC_TILE - 1)   # num rows * multicast_hops
+                *
+                M_TILE * MAC_TILE  # tile shape
+                *
+                BITS_PER_VALUE
+            )
+        )
+        self.assertEqual(
+            result.data['Matmul0<SEP>action<SEP>NoC<SEP>T1<SEP>hops'].iloc[0],
+            (
+                M / M_TILE
+                *
+                (KN / MAC_TILE) * (KN / MAC_TILE - 1)   # num rows * multicast_hops
+                *
+                M_TILE * MAC_TILE  # tile shape
+                *
+                BITS_PER_VALUE
+            )
+        )
+        self.assertEqual(
+            result.data['Matmul0<SEP>action<SEP>NoC<SEP>W0<SEP>hops'].iloc[0],
+            (
+                M / M_TILE
+                *
+                (
+                    4   # a 2x2 grid of physical buffers
+                    *
+                    (
+                        sum(i for i in range(2)) * MAC_TILE  # unicast along row * tile shape
+                        +
+                        2 * sum(i for i in range(2))  # num cols * unicast down col
+                    )
+                )
+                *
+                MAC_TILE * MAC_TILE  # tile shape
+                *
+                BITS_PER_VALUE
+            )
+        )
+
 
 class TestMapper(TestCase):
     def test_hierarchical(self):
