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