OR-1 dataflow CPU sketch
1"""Tests for code generation.
2
3Tests verify:
4- or1-asm.AC8.1: Direct mode produces valid PEConfig with correct IRAM contents
5- or1-asm.AC8.2: Direct mode produces valid SMConfig with initial cell values
6- or1-asm.AC8.3: Direct mode produces seed MonadTokens for const nodes with no incoming edges
7- or1-asm.AC8.4: Direct mode PEConfig includes route restrictions matching edge analysis
8- or1-asm.AC8.5: Token stream mode emits SM init tokens before ROUTE_SET tokens
9- or1-asm.AC8.6: Token stream mode emits ROUTE_SET tokens before LOAD_INST tokens
10- or1-asm.AC8.7: Token stream mode emits LOAD_INST tokens before seed tokens
11- or1-asm.AC8.8: Token stream mode produces valid tokens consumable by emulator
12- or1-asm.AC8.9: Program with no data_defs produces empty SM init section
13- or1-asm.AC8.10: Single PE program produces ROUTE_SET with only self-routes
14"""
15
16from asm.codegen import generate_direct, generate_tokens, AssemblyResult
17from asm.ir import (
18 IRGraph,
19 IRNode,
20 IREdge,
21 IRDataDef,
22 SystemConfig,
23 SourceLoc,
24 ResolvedDest,
25)
26from cm_inst import ALUInst, Addr, ArithOp, CfgOp, MemOp, Port, RoutingOp, SMInst
27from tokens import CfgToken, LoadInstToken, MonadToken, RouteSetToken, SMToken
28from emu.types import PEConfig, SMConfig
29from sm_mod import Presence
30
31
32class TestDirectMode:
33 """AC8.1, AC8.2, AC8.3, AC8.4: Direct mode code generation."""
34
35 def test_ac81_simple_alu_instructions(self):
36 """AC8.1: Two ALU nodes on PE0 produce PEConfig with correct IRAM.
37
38 Tests that:
39 - ALU instructions are correctly converted to ALUInst
40 - They are placed in IRAM at assigned offsets
41 """
42 # Create two simple dyadic ALU nodes
43 add_node = IRNode(
44 name="&add",
45 opcode=ArithOp.ADD,
46 pe=0,
47 iram_offset=0,
48 ctx=0,
49 loc=SourceLoc(1, 1),
50 )
51 sub_node = IRNode(
52 name="&sub",
53 opcode=ArithOp.SUB,
54 pe=0,
55 iram_offset=1,
56 ctx=0,
57 loc=SourceLoc(2, 1),
58 )
59 system = SystemConfig(pe_count=1, sm_count=1)
60 graph = IRGraph(
61 {"&add": add_node, "&sub": sub_node},
62 system=system,
63 )
64
65 result = generate_direct(graph)
66
67 assert len(result.pe_configs) == 1
68 pe_config = result.pe_configs[0]
69 assert pe_config.pe_id == 0
70 assert len(pe_config.iram) == 2
71 assert 0 in pe_config.iram
72 assert 1 in pe_config.iram
73
74 # Check the instruction types
75 inst_0 = pe_config.iram[0]
76 inst_1 = pe_config.iram[1]
77 assert isinstance(inst_0, ALUInst) # Is an ALUInst
78 assert isinstance(inst_1, ALUInst) # Is an ALUInst
79 assert inst_0.op == ArithOp.ADD
80 assert inst_1.op == ArithOp.SUB
81
82 def test_ac82_data_defs_to_smconfig(self):
83 """AC8.2: Data definitions produce SMConfig with initial cell values.
84
85 Tests that:
86 - Data defs with SM placement are converted to SMConfig
87 - initial_cells dict contains correct (Presence.FULL, value) tuples
88 """
89 data_def = IRDataDef(
90 name="@val",
91 sm_id=0,
92 cell_addr=5,
93 value=42,
94 loc=SourceLoc(1, 1),
95 )
96 graph = IRGraph({}, data_defs=[data_def], system=SystemConfig(1, 1))
97
98 result = generate_direct(graph)
99
100 assert len(result.sm_configs) == 1
101 sm_config = result.sm_configs[0]
102 assert sm_config.sm_id == 0
103 assert sm_config.initial_cells is not None
104 assert 5 in sm_config.initial_cells
105 pres, val = sm_config.initial_cells[5]
106 assert pres == Presence.FULL
107 assert val == 42
108
109 def test_ac83_const_node_seed_token(self):
110 """AC8.3: CONST node with no incoming edges produces seed MonadToken.
111
112 Tests that:
113 - CONST nodes are detected
114 - Nodes with no incoming edges are marked as seeds
115 - MonadToken has correct target PE, offset, ctx, data
116 """
117 const_node = IRNode(
118 name="&seed",
119 opcode=RoutingOp.CONST,
120 pe=0,
121 iram_offset=2,
122 ctx=0,
123 const=99,
124 loc=SourceLoc(1, 1),
125 )
126 graph = IRGraph({"&seed": const_node}, system=SystemConfig(1, 1))
127
128 result = generate_direct(graph)
129
130 assert len(result.seed_tokens) == 1
131 token = result.seed_tokens[0]
132 assert isinstance(token, MonadToken)
133 assert token.target == 0
134 assert token.offset == 2
135 assert token.ctx == 0
136 assert token.data == 99
137 assert token.inline == False
138
139 def test_ac84_route_restrictions(self):
140 """AC8.4: Cross-PE edges produce correct allowed_pe_routes.
141
142 Tests that:
143 - Edges from PE0 to PE1 add PE1 to PE0's allowed_pe_routes
144 - Self-routes are always included
145 """
146 # PE0 node connecting to PE1 node
147 node_pe0 = IRNode(
148 name="&a",
149 opcode=ArithOp.ADD,
150 pe=0,
151 iram_offset=0,
152 ctx=0,
153 dest_l=ResolvedDest(
154 name="&b",
155 addr=Addr(a=0, port=Port.L, pe=1),
156 ),
157 loc=SourceLoc(1, 1),
158 )
159 node_pe1 = IRNode(
160 name="&b",
161 opcode=ArithOp.ADD,
162 pe=1,
163 iram_offset=0,
164 ctx=0,
165 loc=SourceLoc(2, 1),
166 )
167 edge = IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(1, 1))
168 system = SystemConfig(pe_count=2, sm_count=1)
169 graph = IRGraph(
170 {"&a": node_pe0, "&b": node_pe1},
171 edges=[edge],
172 system=system,
173 )
174
175 result = generate_direct(graph)
176
177 assert len(result.pe_configs) == 2
178 pe0_config = next(c for c in result.pe_configs if c.pe_id == 0)
179 pe1_config = next(c for c in result.pe_configs if c.pe_id == 1)
180
181 # PE0 should have routes to {0, 1}
182 assert 0 in pe0_config.allowed_pe_routes
183 assert 1 in pe0_config.allowed_pe_routes
184
185 # PE1 should have route to {1} (self only, no incoming cross-PE edges)
186 assert 1 in pe1_config.allowed_pe_routes
187
188 def test_sm_instructions_in_iram(self):
189 """Verify SMInst objects are correctly created and placed in IRAM.
190
191 Tests that MemOp instructions produce SMInst in IRAM.
192 """
193 sm_node = IRNode(
194 name="&read",
195 opcode=MemOp.READ,
196 pe=0,
197 iram_offset=0,
198 ctx=0,
199 sm_id=0,
200 const=42,
201 dest_l=ResolvedDest(
202 name="&out",
203 addr=Addr(a=1, port=Port.L, pe=0),
204 ),
205 loc=SourceLoc(1, 1),
206 )
207 graph = IRGraph({"&read": sm_node}, system=SystemConfig(1, 1))
208
209 result = generate_direct(graph)
210
211 assert len(result.pe_configs) == 1
212 pe_config = result.pe_configs[0]
213 assert 0 in pe_config.iram
214 inst = pe_config.iram[0]
215 assert isinstance(inst, SMInst) # Is an SMInst
216 assert inst.op == MemOp.READ
217 assert inst.sm_id == 0
218 assert inst.const == 42
219
220
221class TestTokenStream:
222 """AC8.5, AC8.6, AC8.7, AC8.8: Token stream generation and ordering."""
223
224 def test_ac85_ac86_ac87_token_ordering(self):
225 """AC8.5-8.7: Tokens are emitted in correct order.
226
227 Tests that:
228 - SM init tokens come first
229 - ROUTE_SET tokens come next
230 - LOAD_INST tokens come next
231 - Seed tokens come last
232 """
233 # Create a multi-PE graph with data_defs
234 data_def = IRDataDef(
235 name="@val",
236 sm_id=0,
237 cell_addr=5,
238 value=42,
239 loc=SourceLoc(1, 1),
240 )
241 node1 = IRNode(
242 name="&a",
243 opcode=ArithOp.ADD,
244 pe=0,
245 iram_offset=0,
246 ctx=0,
247 loc=SourceLoc(1, 1),
248 )
249 node2 = IRNode(
250 name="&b",
251 opcode=RoutingOp.CONST,
252 pe=0,
253 iram_offset=1,
254 ctx=0,
255 const=10,
256 loc=SourceLoc(2, 1),
257 )
258 system = SystemConfig(pe_count=1, sm_count=1)
259 graph = IRGraph(
260 {"&a": node1, "&b": node2},
261 data_defs=[data_def],
262 system=system,
263 )
264
265 tokens = generate_tokens(graph)
266
267 # Find positions of token types
268 smtoken_indices = [
269 i for i, t in enumerate(tokens)
270 if isinstance(t, SMToken)
271 ]
272 route_set_indices = [
273 i for i, t in enumerate(tokens)
274 if isinstance(t, CfgToken) and t.op == CfgOp.ROUTE_SET
275 ]
276 load_inst_indices = [
277 i for i, t in enumerate(tokens)
278 if isinstance(t, CfgToken) and t.op == CfgOp.LOAD_INST
279 ]
280 seed_indices = [
281 i for i, t in enumerate(tokens) if isinstance(t, MonadToken)
282 ]
283
284 # Verify order: SM < ROUTE_SET < LOAD_INST < seed
285 assert smtoken_indices, "Should have at least one SM token"
286 assert route_set_indices, "Should have at least one ROUTE_SET token"
287 assert load_inst_indices, "Should have at least one LOAD_INST token"
288 assert seed_indices, "Should have at least one seed token"
289 assert max(smtoken_indices) < min(route_set_indices), "SM tokens should come before ROUTE_SET"
290 assert max(route_set_indices) < min(load_inst_indices), "ROUTE_SET should come before LOAD_INST"
291 assert max(load_inst_indices) < min(seed_indices), "LOAD_INST should come before seed tokens"
292
293 def test_ac88_tokens_are_valid(self):
294 """AC8.8: Generated tokens are valid and consumable by emulator.
295
296 Tests that:
297 - All tokens have required fields set
298 - Token structure matches emulator expectations
299 - Tokens can be injected into an emulator System and execution completes
300 """
301 from emu.network import build_topology
302 import simpy
303
304 data_def = IRDataDef(
305 name="@val",
306 sm_id=0,
307 cell_addr=5,
308 value=42,
309 loc=SourceLoc(1, 1),
310 )
311 node = IRNode(
312 name="&add",
313 opcode=ArithOp.ADD,
314 pe=0,
315 iram_offset=0,
316 ctx=0,
317 loc=SourceLoc(1, 1),
318 )
319 system = SystemConfig(pe_count=1, sm_count=1)
320 graph = IRGraph(
321 {"&add": node},
322 data_defs=[data_def],
323 system=system,
324 )
325
326 result = generate_direct(graph)
327 tokens = generate_tokens(graph)
328
329 # Build emulator system from AssemblyResult configs
330 env = simpy.Environment()
331 emu_system = build_topology(
332 env,
333 result.pe_configs,
334 result.sm_configs,
335 fifo_capacity=16,
336 )
337
338 # Inject tokens into the system following the sequence:
339 # 1. SM init tokens (paired with SM ID)
340 # 2. ROUTE_SET and LOAD_INST CfgTokens
341 # 3. Seed MonadTokens
342 for token in tokens:
343 if isinstance(token, SMToken):
344 emu_system.inject(token)
345 elif isinstance(token, CfgToken):
346 emu_system.inject(token)
347 elif isinstance(token, MonadToken):
348 emu_system.inject(token)
349
350 # Run the simulation for enough steps to complete initialization
351 env.run(until=1000)
352
353 # Verify token structure
354 for token in tokens:
355 if isinstance(token, SMToken):
356 assert isinstance(token.target, int)
357 assert isinstance(token.addr, int)
358 assert isinstance(token.op, MemOp)
359 assert token.op == MemOp.WRITE
360 elif isinstance(token, RouteSetToken):
361 assert isinstance(token.target, int)
362 assert isinstance(token.op, CfgOp)
363 assert isinstance(token.pe_routes, frozenset)
364 assert isinstance(token.sm_routes, frozenset)
365 elif isinstance(token, LoadInstToken):
366 assert isinstance(token.target, int)
367 assert isinstance(token.op, CfgOp)
368 assert isinstance(token.instructions, tuple)
369 elif isinstance(token, MonadToken):
370 assert isinstance(token.target, int)
371 assert isinstance(token.offset, int)
372 assert isinstance(token.ctx, int)
373 assert isinstance(token.data, int)
374
375
376class TestEdgeCases:
377 """AC8.9, AC8.10: Edge cases for code generation."""
378
379 def test_ac89_no_data_defs(self):
380 """AC8.9: Program with no data_defs produces no SMConfig or SM tokens.
381
382 Tests that:
383 - sm_configs list is empty
384 - Token stream has no SMTokens
385 """
386 node = IRNode(
387 name="&add",
388 opcode=ArithOp.ADD,
389 pe=0,
390 iram_offset=0,
391 ctx=0,
392 loc=SourceLoc(1, 1),
393 )
394 system = SystemConfig(pe_count=1, sm_count=1)
395 graph = IRGraph({"&add": node}, system=system)
396
397 result = generate_direct(graph)
398 assert len(result.sm_configs) == 0
399
400 tokens = generate_tokens(graph)
401 sm_tokens = [t for t in tokens if isinstance(t, SMToken)]
402 assert len(sm_tokens) == 0
403
404 def test_ac810_single_pe_self_route(self):
405 """AC8.10: Single-PE program ROUTE_SET contains only self-route.
406
407 Tests that:
408 - allowed_pe_routes contains only the PE's own ID
409 - Single-PE graph produces ROUTE_SET with pe_routes=[pe_id]
410 """
411 node = IRNode(
412 name="&add",
413 opcode=ArithOp.ADD,
414 pe=0,
415 iram_offset=0,
416 ctx=0,
417 loc=SourceLoc(1, 1),
418 )
419 system = SystemConfig(pe_count=1, sm_count=1)
420 graph = IRGraph({"&add": node}, system=system)
421
422 result = generate_direct(graph)
423 assert len(result.pe_configs) == 1
424 pe_config = result.pe_configs[0]
425 assert pe_config.allowed_pe_routes == {0}
426
427 tokens = generate_tokens(graph)
428 route_set_tokens = [
429 t for t in tokens
430 if isinstance(t, RouteSetToken)
431 ]
432 assert len(route_set_tokens) == 1
433 token = route_set_tokens[0]
434 assert token.pe_routes == frozenset({0})
435
436 def test_multiple_data_defs_same_sm(self):
437 """Multiple data_defs targeting same SM produce single SMConfig.
438
439 Tests that:
440 - Multiple data_defs for SM0 are merged into single SMConfig
441 - initial_cells contains all entries
442 """
443 data_def1 = IRDataDef(
444 name="@val1",
445 sm_id=0,
446 cell_addr=5,
447 value=42,
448 loc=SourceLoc(1, 1),
449 )
450 data_def2 = IRDataDef(
451 name="@val2",
452 sm_id=0,
453 cell_addr=10,
454 value=99,
455 loc=SourceLoc(2, 1),
456 )
457 system = SystemConfig(pe_count=1, sm_count=1)
458 graph = IRGraph({}, data_defs=[data_def1, data_def2], system=system)
459
460 result = generate_direct(graph)
461
462 assert len(result.sm_configs) == 1
463 sm_config = result.sm_configs[0]
464 assert sm_config.sm_id == 0
465 assert len(sm_config.initial_cells) == 2
466 assert sm_config.initial_cells[5] == (Presence.FULL, 42)
467 assert sm_config.initial_cells[10] == (Presence.FULL, 99)
468
469 def test_const_node_with_incoming_edge_not_seed(self):
470 """CONST node with incoming edge is not a seed token.
471
472 Tests that:
473 - Only CONST nodes with NO incoming edges produce seed_tokens
474 """
475 source_node = IRNode(
476 name="&src",
477 opcode=ArithOp.ADD,
478 pe=0,
479 iram_offset=0,
480 ctx=0,
481 loc=SourceLoc(1, 1),
482 )
483 const_node = IRNode(
484 name="&const",
485 opcode=RoutingOp.CONST,
486 pe=0,
487 iram_offset=1,
488 ctx=0,
489 const=5,
490 loc=SourceLoc(2, 1),
491 )
492 edge = IREdge(source="&src", dest="&const", port=Port.L, loc=SourceLoc(1, 1))
493 system = SystemConfig(pe_count=1, sm_count=1)
494 graph = IRGraph(
495 {"&src": source_node, "&const": const_node},
496 edges=[edge],
497 system=system,
498 )
499
500 result = generate_direct(graph)
501
502 # The CONST node has an incoming edge, so it should NOT be a seed
503 assert len(result.seed_tokens) == 0
504
505
506class TestMultiPERouting:
507 """Extended tests for multi-PE routing scenarios."""
508
509 def test_multi_pe_route_computation(self):
510 """Multi-PE graph with multiple cross-PE edges.
511
512 Tests route computation across multiple PEs with various edge patterns.
513 """
514 # Create a 3-PE system with cross-PE edges
515 node_pe0 = IRNode(
516 name="&a",
517 opcode=ArithOp.ADD,
518 pe=0,
519 iram_offset=0,
520 ctx=0,
521 dest_l=ResolvedDest(
522 name="&b",
523 addr=Addr(a=0, port=Port.L, pe=1),
524 ),
525 loc=SourceLoc(1, 1),
526 )
527 node_pe1 = IRNode(
528 name="&b",
529 opcode=ArithOp.SUB,
530 pe=1,
531 iram_offset=0,
532 ctx=0,
533 dest_l=ResolvedDest(
534 name="&c",
535 addr=Addr(a=0, port=Port.L, pe=2),
536 ),
537 loc=SourceLoc(2, 1),
538 )
539 node_pe2 = IRNode(
540 name="&c",
541 opcode=ArithOp.INC,
542 pe=2,
543 iram_offset=0,
544 ctx=0,
545 loc=SourceLoc(3, 1),
546 )
547 edge1 = IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(1, 1))
548 edge2 = IREdge(source="&b", dest="&c", port=Port.L, loc=SourceLoc(2, 1))
549 system = SystemConfig(pe_count=3, sm_count=1)
550 graph = IRGraph(
551 {"&a": node_pe0, "&b": node_pe1, "&c": node_pe2},
552 edges=[edge1, edge2],
553 system=system,
554 )
555
556 result = generate_direct(graph)
557
558 pe0_config = next(c for c in result.pe_configs if c.pe_id == 0)
559 pe1_config = next(c for c in result.pe_configs if c.pe_id == 1)
560 pe2_config = next(c for c in result.pe_configs if c.pe_id == 2)
561
562 # PE0 -> PE1
563 assert 1 in pe0_config.allowed_pe_routes
564 # PE1 -> PE2
565 assert 2 in pe1_config.allowed_pe_routes
566 # PE2 has no outgoing edges
567 assert pe2_config.allowed_pe_routes == {2}