OR-1 dataflow CPU sketch
1"""Tests for auto-placement in the placement pass.
2
3Tests verify:
4- or1-asm.AC10.1: Unplaced nodes are assigned to PEs without exceeding limits
5- or1-asm.AC10.2: Explicitly placed nodes are not moved by auto-placement
6- or1-asm.AC10.3: Connected nodes prefer co-location on same PE (locality heuristic)
7- or1-asm.AC10.4: Program too large for available PEs produces error with per-PE utilization breakdown
8"""
9
10from asm.place import place
11from asm.ir import IRGraph, IRNode, IREdge, SystemConfig, SourceLoc, IRRegion, RegionKind
12import asm.ir
13from asm.errors import ErrorCategory
14from cm_inst import ArithOp, LogicOp, Port
15
16
17class TestBasicAutoPlacement:
18 """AC10.1: Unplaced nodes are assigned to PEs without exceeding limits."""
19
20 def test_four_unplaced_nodes_two_pes(self):
21 """Four unplaced nodes with SystemConfig(pe_count=2) get assigned."""
22 nodes = {
23 "&a": IRNode(name="&a", opcode=ArithOp.ADD, pe=None, loc=SourceLoc(1, 1)),
24 "&b": IRNode(name="&b", opcode=ArithOp.SUB, pe=None, loc=SourceLoc(2, 1)),
25 "&c": IRNode(name="&c", opcode=ArithOp.INC, pe=None, loc=SourceLoc(3, 1)),
26 "&d": IRNode(name="&d", opcode=ArithOp.DEC, pe=None, loc=SourceLoc(4, 1)),
27 }
28 system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
29 graph = IRGraph(nodes, system=system)
30 result = place(graph)
31
32 # Should have no errors
33 assert len(result.errors) == 0, f"Expected no errors, got: {[e.message for e in result.errors]}"
34
35 # All nodes should have PE assignments
36 for node_name in nodes.keys():
37 assert result.nodes[node_name].pe is not None, f"{node_name} still unplaced"
38 assert 0 <= result.nodes[node_name].pe < 2, f"{node_name} on invalid PE"
39
40 def test_monadic_node_placement(self):
41 """Monadic nodes take up only 1 IRAM slot."""
42 # CONST is monadic (RoutingOp.CONST)
43 from cm_inst import RoutingOp
44 nodes = {
45 "&c1": IRNode(name="&c1", opcode=RoutingOp.CONST, pe=None, const=5, loc=SourceLoc(1, 1)),
46 "&c2": IRNode(name="&c2", opcode=RoutingOp.CONST, pe=None, const=10, loc=SourceLoc(2, 1)),
47 }
48 system = SystemConfig(pe_count=1, sm_count=1, iram_capacity=64, frame_count=4)
49 graph = IRGraph(nodes, system=system)
50 result = place(graph)
51
52 assert len(result.errors) == 0
53 for node in result.nodes.values():
54 assert node.pe == 0
55
56
57class TestExplicitPreserved:
58 """AC10.2: Explicitly placed nodes are not moved by auto-placement."""
59
60 def test_mixed_placed_and_unplaced(self):
61 """Explicitly placed nodes keep their PE; unplaced ones get assigned."""
62 nodes = {
63 "&placed0": IRNode(name="&placed0", opcode=ArithOp.ADD, pe=0, loc=SourceLoc(1, 1)),
64 "&placed1": IRNode(name="&placed1", opcode=ArithOp.SUB, pe=1, loc=SourceLoc(2, 1)),
65 "&unplaced1": IRNode(name="&unplaced1", opcode=ArithOp.INC, pe=None, loc=SourceLoc(3, 1)),
66 "&unplaced2": IRNode(name="&unplaced2", opcode=ArithOp.DEC, pe=None, loc=SourceLoc(4, 1)),
67 }
68 system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
69 graph = IRGraph(nodes, system=system)
70 result = place(graph)
71
72 assert len(result.errors) == 0
73 # Explicitly placed nodes keep their PE
74 assert result.nodes["&placed0"].pe == 0
75 assert result.nodes["&placed1"].pe == 1
76 # Unplaced nodes get assigned
77 assert result.nodes["&unplaced1"].pe is not None
78 assert result.nodes["&unplaced2"].pe is not None
79
80 def test_all_explicit_placement(self):
81 """All nodes explicitly placed - no auto-placement needed."""
82 nodes = {
83 "&a": IRNode(name="&a", opcode=ArithOp.ADD, pe=0, loc=SourceLoc(1, 1)),
84 "&b": IRNode(name="&b", opcode=ArithOp.SUB, pe=1, loc=SourceLoc(2, 1)),
85 }
86 system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
87 graph = IRGraph(nodes, system=system)
88 result = place(graph)
89
90 assert len(result.errors) == 0
91 assert result.nodes["&a"].pe == 0
92 assert result.nodes["&b"].pe == 1
93
94
95class TestLocalityHeuristic:
96 """AC10.3: Connected nodes prefer co-location on same PE."""
97
98 def test_two_connected_unplaced_nodes(self):
99 """Two connected unplaced nodes end up on the same PE."""
100 nodes = {
101 "&a": IRNode(name="&a", opcode=ArithOp.ADD, pe=None, loc=SourceLoc(1, 1)),
102 "&b": IRNode(name="&b", opcode=ArithOp.SUB, pe=None, loc=SourceLoc(2, 1)),
103 }
104 edges = [
105 IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(3, 1)),
106 ]
107 system = SystemConfig(pe_count=4, sm_count=1, iram_capacity=64, frame_count=4)
108 graph = IRGraph(nodes, edges=edges, system=system)
109 result = place(graph)
110
111 assert len(result.errors) == 0
112 # Both nodes should be on the same PE
113 assert result.nodes["&a"].pe == result.nodes["&b"].pe
114
115 def test_cluster_of_three_interconnected_nodes(self):
116 """Cluster of 3 interconnected nodes end up on the same PE."""
117 nodes = {
118 "&a": IRNode(name="&a", opcode=ArithOp.ADD, pe=None, loc=SourceLoc(1, 1)),
119 "&b": IRNode(name="&b", opcode=ArithOp.SUB, pe=None, loc=SourceLoc(2, 1)),
120 "&c": IRNode(name="&c", opcode=ArithOp.INC, pe=None, loc=SourceLoc(3, 1)),
121 }
122 edges = [
123 IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(4, 1)),
124 IREdge(source="&b", dest="&c", port=Port.L, loc=SourceLoc(5, 1)),
125 IREdge(source="&a", dest="&c", port=Port.R, loc=SourceLoc(6, 1)),
126 ]
127 system = SystemConfig(pe_count=4, sm_count=1, iram_capacity=64, frame_count=4)
128 graph = IRGraph(nodes, edges=edges, system=system)
129 result = place(graph)
130
131 assert len(result.errors) == 0
132 # All three nodes should be on the same PE
133 pe_a = result.nodes["&a"].pe
134 pe_b = result.nodes["&b"].pe
135 pe_c = result.nodes["&c"].pe
136 assert pe_a == pe_b == pe_c
137
138 def test_locality_with_mixed_placed_unplaced(self):
139 """Unplaced node prefers PE of its connected placed neighbour."""
140 nodes = {
141 "&placed": IRNode(name="&placed", opcode=ArithOp.ADD, pe=1, loc=SourceLoc(1, 1)),
142 "&unplaced": IRNode(name="&unplaced", opcode=ArithOp.SUB, pe=None, loc=SourceLoc(2, 1)),
143 }
144 edges = [
145 IREdge(source="&placed", dest="&unplaced", port=Port.L, loc=SourceLoc(3, 1)),
146 ]
147 system = SystemConfig(pe_count=4, sm_count=1, iram_capacity=64, frame_count=4)
148 graph = IRGraph(nodes, edges=edges, system=system)
149 result = place(graph)
150
151 assert len(result.errors) == 0
152 # Unplaced node should prefer PE1 (where placed node is)
153 assert result.nodes["&unplaced"].pe == 1
154
155
156class TestOverflow:
157 """AC10.4: Program too large produces error with per-PE utilization breakdown."""
158
159 def test_200_nodes_overflow_error(self):
160 """200 unplaced nodes with limited resources produce overflow error."""
161 # Create 200 monadic nodes (each takes 1 IRAM slot)
162 nodes = {}
163 for i in range(200):
164 nodes[f"&n{i}"] = IRNode(
165 name=f"&n{i}",
166 opcode=ArithOp.INC,
167 pe=None,
168 loc=SourceLoc(i + 1, 1),
169 )
170
171 system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
172 graph = IRGraph(nodes, system=system)
173 result = place(graph)
174
175 # Should have error(s) about placement failure
176 assert len(result.errors) > 0
177 error = result.errors[0]
178 assert error.category == ErrorCategory.PLACEMENT
179 assert "Cannot place" in error.message or "full" in error.message.lower()
180 # Error message should include per-PE utilization info
181 assert "PE0" in error.message or "IRAM" in error.message
182
183 def test_overflow_error_includes_breakdown(self):
184 """Overflow error includes per-PE slot utilization breakdown."""
185 # 130 monadic nodes: on 2 PEs with 64 IRAM capacity each, this overflows
186 nodes = {}
187 for i in range(130):
188 nodes[f"&n{i}"] = IRNode(
189 name=f"&n{i}",
190 opcode=ArithOp.INC,
191 pe=None,
192 loc=SourceLoc(i + 1, 1),
193 )
194
195 system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
196 graph = IRGraph(nodes, system=system)
197 result = place(graph)
198
199 assert len(result.errors) > 0
200 error_msg = result.errors[0].message
201 # Should mention PE utilization
202 assert any(x in error_msg for x in ["PE", "IRAM", "full", "capacity"])
203
204
205class TestFunctionScopedNodesPlacement:
206 """Verify auto-placed nodes inside function scopes receive PE assignments."""
207
208 def test_function_scoped_node_gets_pe_assignment(self):
209 """Nodes inside function regions should get PE assignments after place()."""
210 # Create a graph with a function region containing unplaced nodes
211 func_nodes = {
212 "$main.&add": IRNode(
213 name="$main.&add",
214 opcode=ArithOp.ADD,
215 pe=None,
216 loc=SourceLoc(1, 1),
217 ),
218 "$main.&inc": IRNode(
219 name="$main.&inc",
220 opcode=ArithOp.INC,
221 pe=None,
222 loc=SourceLoc(2, 1),
223 ),
224 }
225
226 func_region = IRGraph(nodes=func_nodes, edges=[], regions=[], data_defs=[], errors=[])
227
228 main_region = asm.ir.IRRegion(
229 tag="$main",
230 kind=asm.ir.RegionKind.FUNCTION,
231 body=func_region,
232 loc=SourceLoc(0, 1),
233 )
234
235 graph = IRGraph(
236 nodes={},
237 edges=[],
238 regions=[main_region],
239 data_defs=[],
240 system=SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4),
241 errors=[],
242 )
243
244 result = place(graph)
245
246 # Verify no errors
247 assert len(result.errors) == 0, f"Expected no errors, got: {[e.message for e in result.errors]}"
248
249 # Verify nodes inside the function region received PE assignments
250 assert len(result.regions) == 1
251 result_func = result.regions[0]
252 assert result_func.tag == "$main"
253
254 # Check that nodes in the function body have PE assignments
255 for node_name, node in result_func.body.nodes.items():
256 assert node.pe is not None, f"Node {node_name} in function scope still has pe=None"
257 assert 0 <= node.pe < 2, f"Node {node_name} has invalid PE {node.pe}"