Something went wrong. Try again.
OR-1 dataflow CPU sketch
Something went wrong. Try again.
23 kB · 698 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699"""Tests for frame-based code generation (Phase 6).
Tests verify:- pe-frame-redesign.AC6.1: AssemblyResult includes setup_tokens field- pe-frame-redesign.AC6.2: Token stream ordering: SM init → IRAM writes → ALLOC → frame slot writes → seeds- pe-frame-redesign.AC6.3: IRAM write data uses pack_instruction()- pe-frame-redesign.AC6.4: Destination frame slot writes use pack_flit1() with is_dest=True- pe-frame-redesign.AC6.5: T0 bootstrap data uses pack_token() for packed flits- pe-frame-redesign.AC6.6: Seed tokens use act_id, no gen field"""
import pytest
from asm.codegen import generate_direct, generate_tokens, AssemblyResult, _build_iram_for_pefrom asm.ir import ( IRGraph, IRNode, IREdge, IRDataDef, SystemConfig, SourceLoc, ResolvedDest, FrameLayout, FrameSlotMap,)from cm_inst import ( Instruction, OutputStyle, TokenKind, FrameDest, FrameOp, ArithOp, MemOp, Port, RoutingOp)from tokens import ( DyadToken, MonadToken, SMToken, PELocalWriteToken, FrameControlToken)from encoding import pack_instruction, pack_flit1, unpack_instruction, pack_tokenfrom emu.types import PEConfig, SMConfigfrom sm_mod import Presence
class TestTask1BuildIramForPE: """Task 1: Rewrite _build_iram_for_pe() for Instruction objects."""
def test_builds_instruction_objects(self): """Produces Instruction objects from IRNodes with mode set.""" node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, act_id=0, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=5, loc=SourceLoc(1, 1), )
iram = _build_iram_for_pe([node], {"&add": node}, [])
assert len(iram) == 1 assert 0 in iram inst = iram[0] assert isinstance(inst, Instruction) assert inst.opcode == ArithOp.ADD assert inst.output == OutputStyle.INHERIT assert inst.has_const == False assert inst.dest_count == 2 assert inst.wide == False assert inst.fref == 5
def test_excludes_seed_nodes(self): """Skips seed nodes from IRAM.""" seed_node = IRNode( name="&const", opcode=RoutingOp.CONST, pe=0, iram_offset=10, act_id=0, mode=(OutputStyle.INHERIT, True, 1), wide=False, fref=0, seed=True, # Mark as seed loc=SourceLoc(1, 1), )
iram = _build_iram_for_pe([seed_node], {"&const": seed_node}, [])
assert len(iram) == 0
def test_excludes_unallocated_nodes(self): """Skips nodes without iram_offset.""" node = IRNode( name="&unallocated", opcode=ArithOp.ADD, pe=0, iram_offset=None, # Not allocated act_id=0, mode=(OutputStyle.SINK, False, 0), wide=False, fref=0, loc=SourceLoc(1, 1), )
iram = _build_iram_for_pe([node], {"&unallocated": node}, [])
assert len(iram) == 0
def test_excludes_nodes_without_mode(self): """Skips nodes without mode (output style) allocation.""" node = IRNode( name="&no_mode", opcode=ArithOp.ADD, pe=0, iram_offset=5, act_id=0, mode=None, # No mode set wide=False, fref=0, loc=SourceLoc(1, 1), )
iram = _build_iram_for_pe([node], {"&no_mode": node}, [])
assert len(iram) == 0
def test_handles_mem_op_instructions(self): """Produces Instruction objects for MemOp opcodes.""" node = IRNode( name="&read", opcode=MemOp.READ, pe=0, iram_offset=1, act_id=0, sm_id=0, mode=(OutputStyle.CHANGE_TAG, False, 1), wide=True, fref=3, loc=SourceLoc(2, 1), )
iram = _build_iram_for_pe([node], {"&read": node}, [])
assert 1 in iram inst = iram[1] assert inst.opcode == MemOp.READ assert inst.output == OutputStyle.CHANGE_TAG assert inst.wide == True assert inst.fref == 3
class TestAC63PackInstructionValues: """AC6.3: IRAM write data uses pack_instruction() — verify known values."""
def test_known_instruction_pack_value(self): """Verify a known Instruction packs to expected 16-bit value.
Example from phase design: inst = Instruction(opcode=ArithOp.ADD, output=OutputStyle.INHERIT, has_const=False, dest_count=1, wide=False, fref=0) Should pack to 0x0000 (all bits zero). """ inst = Instruction( opcode=ArithOp.ADD, output=OutputStyle.INHERIT, has_const=False, dest_count=1, wide=False, fref=0, ) packed = pack_instruction(inst) # ADD opcode = 0, INHERIT = 0b000, has_const=False → mode=0b000, wide=0, fref=0 # [type:1=0][opcode:5=0][mode:3=0][wide:1=0][fref:6=0] = 0x0000 assert packed == 0x0000
def test_roundtrip_pack_unpack(self): """Verify pack_instruction and unpack_instruction are inverses.""" orig = Instruction( opcode=ArithOp.ADD, output=OutputStyle.INHERIT, has_const=True, dest_count=2, wide=True, fref=15, ) packed = pack_instruction(orig) unpacked = unpack_instruction(packed) assert unpacked == orig
class TestTask2FrameSetupTokens: """Task 2: Frame setup token generation."""
def test_ac61_assembly_result_has_setup_tokens(self): """AC6.1: AssemblyResult includes setup_tokens field.""" node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, act_id=0, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(1, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph({"&add": node}, system=system)
result = generate_direct(graph)
assert hasattr(result, 'setup_tokens') assert isinstance(result.setup_tokens, list)
def test_ac62_token_ordering(self): """AC6.2: Token stream ordering verified.
Order: SM init → IRAM writes → ALLOC → frame slot writes → seeds """ # Create a simple graph with data def and instruction data_def = IRDataDef( name="@data", sm_id=0, cell_addr=10, value=42, loc=SourceLoc(1, 1), ) node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, act_id=0, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(2, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph( {"&add": node}, data_defs=[data_def], system=system )
result = generate_direct(graph) tokens = result.setup_tokens
# Verify ordering: SM init should come before IRAM writes sm_tokens = [t for t in tokens if isinstance(t, SMToken)] iram_tokens = [t for t in tokens if isinstance(t, PELocalWriteToken) and t.region == 0] alloc_tokens = [t for t in tokens if isinstance(t, FrameControlToken)]
if sm_tokens and iram_tokens: assert tokens.index(sm_tokens[0]) < tokens.index(iram_tokens[0]) if iram_tokens and alloc_tokens: assert tokens.index(iram_tokens[0]) < tokens.index(alloc_tokens[0])
def test_ac63_iram_write_uses_pack_instruction(self): """AC6.3: IRAM write PELocalWriteTokens carry pack_instruction() data.""" node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=5, act_id=0, mode=(OutputStyle.INHERIT, True, 2), wide=False, fref=12, loc=SourceLoc(1, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph({"&add": node}, system=system)
result = generate_direct(graph) iram_tokens = [t for t in result.setup_tokens if isinstance(t, PELocalWriteToken) and t.region == 0]
assert len(iram_tokens) > 0 token = iram_tokens[0]
# Verify that the data matches packed instruction expected_inst = Instruction( opcode=ArithOp.ADD, output=OutputStyle.INHERIT, has_const=True, dest_count=2, wide=False, fref=12, ) expected_data = pack_instruction(expected_inst) assert token.data == expected_data
def test_alloc_tokens_per_activation(self): """ALLOC tokens generated for each activation on each PE.""" node1 = IRNode( name="&a", opcode=ArithOp.ADD, pe=0, iram_offset=0, act_id=1, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(1, 1), ) node2 = IRNode( name="&b", opcode=ArithOp.SUB, pe=0, iram_offset=1, act_id=2, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(2, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph( {"&a": node1, "&b": node2}, system=system )
result = generate_direct(graph) alloc_tokens = [t for t in result.setup_tokens if isinstance(t, FrameControlToken)]
# Should have 2 ALLOC tokens (one for act_id=1, one for act_id=2) assert len(alloc_tokens) == 2 assert all(t.op == FrameOp.ALLOC for t in alloc_tokens)
# Verify that PE configs have initial_tag_store with tuple values assert len(result.pe_configs) == 1 pe_cfg = result.pe_configs[0] assert pe_cfg.initial_tag_store, "initial_tag_store should not be empty for PE with activations" for act_id, val in pe_cfg.initial_tag_store.items(): assert isinstance(val, tuple) and len(val) == 2, \ f"initial_tag_store[{act_id}] should be (frame_id, lane) tuple, got {val}" frame_id, lane = val assert isinstance(frame_id, int), f"frame_id should be int, got {type(frame_id)}" assert isinstance(lane, int), f"lane should be int, got {type(lane)}"
class TestTask3SeedTokens: """Task 3: Seed token generation with act_id."""
def test_ac66_seed_tokens_use_act_id(self): """AC6.6: Seed tokens use act_id field (not ctx).""" seed_node = IRNode( name="&const", opcode=RoutingOp.CONST, pe=0, iram_offset=10, act_id=5, const=42, mode=(OutputStyle.INHERIT, True, 1), wide=False, fref=0, seed=True, loc=SourceLoc(1, 1), ) edge = IREdge(source="&const", dest="&consumer", port=Port.L) consumer = IRNode( name="&consumer", opcode=ArithOp.ADD, pe=1, iram_offset=0, act_id=3, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(2, 1), ) system = SystemConfig(pe_count=2, sm_count=1) graph = IRGraph( {"&const": seed_node, "&consumer": consumer}, edges=[edge], system=system )
result = generate_direct(graph) seed_tokens = result.seed_tokens
assert len(seed_tokens) > 0 token = seed_tokens[0] # Verify token uses act_id field assert hasattr(token, 'act_id') assert token.act_id == 3
def test_dyadic_seed_token_no_gen(self): """DyadToken seed tokens have no gen field.""" seed_node = IRNode( name="&const", opcode=RoutingOp.CONST, pe=0, iram_offset=10, act_id=0, const=100, mode=(OutputStyle.INHERIT, True, 1), wide=False, fref=0, seed=True, loc=SourceLoc(1, 1), ) edge = IREdge(source="&const", dest="&dyadic", port=Port.L) dyadic_node = IRNode( name="&dyadic", opcode=ArithOp.ADD, pe=1, iram_offset=0, act_id=1, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(2, 1), ) system = SystemConfig(pe_count=2, sm_count=1) graph = IRGraph( {"&const": seed_node, "&dyadic": dyadic_node}, edges=[edge], system=system )
result = generate_direct(graph) dyadic_tokens = [t for t in result.seed_tokens if isinstance(t, DyadToken)]
assert len(dyadic_tokens) > 0 token = dyadic_tokens[0] # Verify no gen field exists assert not hasattr(token, 'gen')
def test_generate_direct_produces_configs(self): """generate_direct() produces PEConfigs with Instruction IRAM.""" node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=3, act_id=0, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=7, loc=SourceLoc(1, 1), ) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph({"&add": node}, system=system)
result = generate_direct(graph)
assert len(result.pe_configs) == 1 pe_cfg = result.pe_configs[0] assert pe_cfg.pe_id == 0 assert 3 in pe_cfg.iram
inst = pe_cfg.iram[3] assert isinstance(inst, Instruction) assert inst.opcode == ArithOp.ADD
def test_generate_tokens_ordering(self): """generate_tokens() ordering: SM init → IRAM → ALLOC → frame slots → seeds.""" data_def = IRDataDef( name="@data", sm_id=0, cell_addr=5, value=99, loc=SourceLoc(1, 1), ) node = IRNode( name="&add", opcode=ArithOp.ADD, pe=0, iram_offset=0, act_id=0, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(2, 1), ) seed_node = IRNode( name="&const", opcode=RoutingOp.CONST, pe=0, iram_offset=10, act_id=0, const=42, mode=(OutputStyle.INHERIT, True, 1), wide=False, fref=0, seed=True, loc=SourceLoc(3, 1), ) edge = IREdge(source="&const", dest="&add", port=Port.L) system = SystemConfig(pe_count=1, sm_count=1) graph = IRGraph( {"&add": node, "&const": seed_node}, edges=[edge], data_defs=[data_def], system=system )
tokens = generate_tokens(graph)
# Verify that setup_tokens come before seeds setup_count = sum(1 for t in tokens if isinstance(t, (SMToken, PELocalWriteToken, FrameControlToken))) seed_count = sum(1 for t in tokens if isinstance(t, (DyadToken, MonadToken)))
# Find indices if setup_count > 0 and seed_count > 0: first_setup = next(i for i, t in enumerate(tokens) if isinstance(t, (SMToken, PELocalWriteToken, FrameControlToken))) first_seed = next(i for i, t in enumerate(tokens) if isinstance(t, (DyadToken, MonadToken))) assert first_setup < first_seed
class TestAC65T0BootstrapPacking: """AC6.5: T0 bootstrap data uses pack_token() for packed flits."""
def test_pack_token_for_t0_bootstrap(self): """Verify pack_token() can pack tokens for T0 bootstrap storage.
In T0 bootstrap, tokens are written to T0 storage as packed flit sequences. Each flit is written via SMToken(WRITE) to T0 addresses. """ # Create a sample dyadic token that might be bootstrapped to T0 token = DyadToken( target=1, offset=10, act_id=2, data=0x1234, port=Port.L, )
# Pack the token flits = pack_token(token)
# Verify we get a sequence of flits assert isinstance(flits, list) assert len(flits) >= 1 # At least flit 1 (header) assert all(isinstance(f, int) for f in flits)
# For T0 bootstrap, each flit would be written as: # SMToken(target=sm_id, addr=t0_offset + i, op=MemOp.WRITE, data=flit) for i, flit in enumerate(flits): # Verify flit is 16-bit assert 0 <= flit <= 0xFFFF, f"Flit {i} out of range: {flit}"
def test_pack_token_monad_for_t0(self): """MonadToken can also be packed for T0 bootstrap.""" token = MonadToken( target=2, offset=5, act_id=1, data=0x5678, inline=False, )
flits = pack_token(token) assert len(flits) >= 1 assert all(0 <= f <= 0xFFFF for f in flits)
class TestIntegration: """Integration tests for complete codegen pipeline."""
def test_multinode_multiactivation(self): """Complex graph with multiple nodes and activations.""" nodes = [ IRNode( name="&a", opcode=ArithOp.ADD, pe=0, iram_offset=0, act_id=0, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(1, 1), ), IRNode( name="&b", opcode=ArithOp.SUB, pe=0, iram_offset=1, act_id=0, mode=(OutputStyle.INHERIT, False, 2), wide=False, fref=0, loc=SourceLoc(2, 1), ), IRNode( name="&c", opcode=ArithOp.INC, pe=1, iram_offset=0, act_id=1, mode=(OutputStyle.INHERIT, False, 1), wide=False, fref=0, loc=SourceLoc(3, 1), ), ] edges = [ IREdge(source="&a", dest="&b", port=Port.L), IREdge(source="&b", dest="&c", port=Port.L), ] system = SystemConfig(pe_count=2, sm_count=1) graph = IRGraph( {node.name: node for node in nodes}, edges=edges, system=system )
result = generate_direct(graph)
# Check pe_configs assert len(result.pe_configs) == 2 assert result.pe_configs[0].pe_id == 0 assert result.pe_configs[1].pe_id == 1
# Check setup_tokens assert len(result.setup_tokens) > 0
# Check that ALLOC tokens are generated for act_id=1 alloc_tokens = [t for t in result.setup_tokens if isinstance(t, FrameControlToken)] assert any(t.act_id == 1 for t in alloc_tokens)
class TestAC6_4SetupTokensWithFrameDest: """AC6.4: Verify destination frame slot writes use pack_flit1() with is_dest=True."""
def test_setup_tokens_frame_writes_use_pack_flit1(self): """Verify setup_tokens use pack_flit1() for frame destination writes.
This test verifies that the codegen correctly uses pack_flit1() to encode FrameDest objects when writing to frame slots (region=1, is_dest=True). """ from dataclasses import replace from asm.allocate import allocate from asm.ir import ResolvedDest
# Create a simple graph with a node that has dest_l/dest_r source_node = IRNode( name="&source", opcode=ArithOp.ADD, pe=0, iram_offset=0, act_id=0, mode=(OutputStyle.INHERIT, False, 1), wide=False, fref=0, loc=SourceLoc(1, 1), )
dest_node = IRNode( name="&dest", opcode=ArithOp.SUB, pe=0, iram_offset=1, act_id=0, mode=(OutputStyle.SINK, False, 0), wide=False, fref=0, loc=SourceLoc(2, 1), )
edge = IREdge(source="&source", dest="&dest", port=Port.L)
system = SystemConfig(pe_count=1, sm_count=0) graph = IRGraph( nodes={"&source": source_node, "&dest": dest_node}, edges=[edge], regions=[], system=system, call_sites=[], )
# Run allocate to set up frame layouts and dest_l/dest_r allocated_graph = allocate(graph) if allocated_graph.errors: raise ValueError(f"Allocation errors: {allocated_graph.errors}")
# Now run generate_direct on the allocated graph codegen_result = generate_direct(allocated_graph)
# Find PELocalWriteToken entries with is_dest=True (frame dest writes) dest_write_tokens = [ t for t in codegen_result.setup_tokens if isinstance(t, PELocalWriteToken) and t.is_dest ]
# Should have at least one frame dest write assert len(dest_write_tokens) > 0, \ "Should have PELocalWriteToken with is_dest=True for frame destinations"
# Verify each dest write token has valid packed flit1 data for token in dest_write_tokens: assert token.region == 1, \ f"Frame dest writes should use region=1 (frame), got {token.region}" # Data should be a valid packed flit1 (16-bit) assert 0 <= token.data <= 0xFFFF, \ f"Frame dest data should be 16-bit, got {token.data}"