nonbinary.computer
OR-1 dataflow CPU sketch
1"""Code generation for OR1 assembly.
2
3Converts fully allocated IRGraphs to emulator-ready configuration objects and
4token streams. Two output modes:
51. Direct mode: Produces PEConfig/SMConfig lists + seed tokens (for direct setup)
62. Token stream mode: Produces bootstrap sequence (SM init → ROUTE_SET → LOAD_INST → seeds)
7
8Reference: Phase 6 design doc, Tasks 1-2.
9"""
10
11from dataclasses import dataclass
12from collections import defaultdict
13
14from asm.ir import (
15 IRGraph, IRNode, IREdge, ResolvedDest, collect_all_nodes_and_edges, collect_all_data_defs,
16 DEFAULT_IRAM_CAPACITY, DEFAULT_CTX_SLOTS
17)
18from cm_inst import ALUInst, CfgOp, MemOp, RoutingOp, SMInst
19from emu.types import PEConfig, SMConfig
20from tokens import LoadInstToken, MonadToken, RouteSetToken, SMToken
21from sm_mod import Presence
22
23
24@dataclass(frozen=True)
25class AssemblyResult:
26 """Result of code generation in direct mode.
27
28 Attributes:
29 pe_configs: List of PEConfig objects, one per PE
30 sm_configs: List of SMConfig objects, one per SM with data_defs
31 seed_tokens: List of MonadTokens for const nodes with no incoming edges
32 """
33 pe_configs: list[PEConfig]
34 sm_configs: list[SMConfig]
35 seed_tokens: list[MonadToken]
36
37
38
39
40def _build_iram_for_pe(
41 nodes_on_pe: list[IRNode],
42 all_nodes: dict[str, IRNode],
43) -> dict[int, ALUInst | SMInst]:
44 """Build IRAM instruction dict for a single PE.
45
46 Args:
47 nodes_on_pe: List of IRNodes on this PE
48 all_nodes: All nodes in graph (for lookups)
49
50 Returns:
51 Dict mapping IRAM offset to ALUInst or SMInst
52 """
53 iram = {}
54
55 for node in nodes_on_pe:
56 if node.iram_offset is None:
57 # Node not allocated, skip
58 continue
59
60 if isinstance(node.opcode, MemOp):
61 # Memory operation -> SMInst
62 inst = SMInst(
63 op=node.opcode,
64 sm_id=node.sm_id,
65 const=node.const,
66 ret=node.dest_l.addr if isinstance(node.dest_l, ResolvedDest) else None,
67 )
68 else:
69 # ALU operation -> ALUInst
70 # Extract Addr from ResolvedDest or keep None
71 dest_l_addr = None
72 dest_r_addr = None
73
74 if node.dest_l is not None and isinstance(node.dest_l, ResolvedDest):
75 dest_l_addr = node.dest_l.addr
76
77 if node.dest_r is not None and isinstance(node.dest_r, ResolvedDest):
78 dest_r_addr = node.dest_r.addr
79
80 inst = ALUInst(
81 op=node.opcode,
82 dest_l=dest_l_addr,
83 dest_r=dest_r_addr,
84 const=node.const,
85 )
86
87 iram[node.iram_offset] = inst
88
89 return iram
90
91
92def _compute_route_restrictions(
93 nodes_by_pe: dict[int, list[IRNode]],
94 all_edges: list[IREdge],
95 all_nodes: dict[str, IRNode],
96 pe_id: int,
97) -> tuple[set[int], set[int]]:
98 """Compute allowed PE and SM routes for a given PE.
99
100 Analyzes all edges involving nodes on this PE to determine which other
101 PEs and SMs it can route to. Includes self-routes.
102
103 Args:
104 nodes_by_pe: Dict mapping PE ID to list of nodes on that PE
105 all_edges: List of all edges in graph
106 all_nodes: Dict of all nodes
107 pe_id: The PE we're computing routes for
108
109 Returns:
110 Tuple of (allowed_pe_routes set, allowed_sm_routes set)
111 """
112 nodes_on_pe_set = {node.name for node in nodes_by_pe.get(pe_id, [])}
113
114 pe_routes = {pe_id} # Always include self-route
115 sm_routes = set()
116
117 # Scan all edges for those sourced from this PE
118 for edge in all_edges:
119 if edge.source in nodes_on_pe_set:
120 # This edge originates from our PE
121 dest_node = all_nodes.get(edge.dest)
122 if dest_node is not None:
123 if dest_node.pe is not None:
124 pe_routes.add(dest_node.pe)
125
126 # Scan all nodes on this PE for SM instructions
127 for node in nodes_by_pe.get(pe_id, []):
128 if isinstance(node.opcode, MemOp) and node.sm_id is not None:
129 sm_routes.add(node.sm_id)
130
131 return pe_routes, sm_routes
132
133
134def generate_direct(graph: IRGraph) -> AssemblyResult:
135 """Generate PEConfig, SMConfig, and seed tokens from an allocated IRGraph.
136
137 Args:
138 graph: A fully allocated IRGraph (after allocate pass)
139
140 Returns:
141 AssemblyResult with pe_configs, sm_configs, and seed_tokens
142 """
143 all_nodes, all_edges = collect_all_nodes_and_edges(graph)
144 all_data_defs = collect_all_data_defs(graph)
145
146 # Group nodes by PE
147 nodes_by_pe: dict[int, list[IRNode]] = defaultdict(list)
148 for node in all_nodes.values():
149 if node.pe is not None:
150 nodes_by_pe[node.pe].append(node)
151
152 # Build PEConfigs
153 pe_configs = []
154 for pe_id in sorted(nodes_by_pe.keys()):
155 nodes_on_pe = nodes_by_pe[pe_id]
156
157 # Build IRAM for this PE
158 iram = _build_iram_for_pe(nodes_on_pe, all_nodes)
159
160 # Compute route restrictions
161 allowed_pe_routes, allowed_sm_routes = _compute_route_restrictions(
162 nodes_by_pe, all_edges, all_nodes, pe_id
163 )
164
165 # Create PEConfig
166 config = PEConfig(
167 pe_id=pe_id,
168 iram=iram,
169 ctx_slots=graph.system.ctx_slots if graph.system else DEFAULT_CTX_SLOTS,
170 offsets=graph.system.iram_capacity if graph.system else DEFAULT_IRAM_CAPACITY,
171 allowed_pe_routes=allowed_pe_routes,
172 allowed_sm_routes=allowed_sm_routes,
173 )
174 pe_configs.append(config)
175
176 # Build SMConfigs from data_defs
177 sm_configs_by_id: dict[int, dict[int, tuple[Presence, int]]] = defaultdict(dict)
178 for data_def in all_data_defs:
179 if data_def.sm_id is not None and data_def.cell_addr is not None:
180 sm_configs_by_id[data_def.sm_id][data_def.cell_addr] = (
181 Presence.FULL, data_def.value
182 )
183
184 sm_configs = []
185 for sm_id in sorted(sm_configs_by_id.keys()):
186 initial_cells = sm_configs_by_id[sm_id]
187 config = SMConfig(
188 sm_id=sm_id,
189 initial_cells=initial_cells if initial_cells else None,
190 )
191 sm_configs.append(config)
192
193 # Detect seed tokens: CONST nodes with no incoming edges
194 seed_tokens = []
195
196 # Build index of edges by destination
197 edges_by_dest = defaultdict(list)
198 for edge in all_edges:
199 edges_by_dest[edge.dest].append(edge)
200
201 for node in all_nodes.values():
202 # Check if this is a CONST node
203 if node.opcode == RoutingOp.CONST:
204 # Check if it has no incoming edges
205 if node.name not in edges_by_dest:
206 # This is a seed token
207 token = MonadToken(
208 target=node.pe if node.pe is not None else 0,
209 offset=node.iram_offset if node.iram_offset is not None else 0,
210 ctx=node.ctx if node.ctx is not None else 0,
211 data=node.const if node.const is not None else 0,
212 inline=False,
213 )
214 seed_tokens.append(token)
215
216 return AssemblyResult(
217 pe_configs=pe_configs,
218 sm_configs=sm_configs,
219 seed_tokens=seed_tokens,
220 )
221
222
223def generate_tokens(graph: IRGraph) -> list:
224 """Generate bootstrap token sequence from an allocated IRGraph.
225
226 Produces tokens in order: SM init → ROUTE_SET → LOAD_INST → seeds
227
228 Args:
229 graph: A fully allocated IRGraph (after allocate pass)
230
231 Returns:
232 List of tokens (SMToken, CfgToken, MonadToken) in bootstrap order
233 """
234 # Use direct mode to get configs and seeds
235 result = generate_direct(graph)
236
237 tokens = []
238
239 # 1. SM init tokens
240 all_data_defs = collect_all_data_defs(graph)
241 for data_def in all_data_defs:
242 if data_def.sm_id is not None and data_def.cell_addr is not None:
243 token = SMToken(
244 target=data_def.sm_id,
245 addr=data_def.cell_addr,
246 op=MemOp.WRITE,
247 flags=None,
248 data=data_def.value,
249 ret=None,
250 )
251 tokens.append(token)
252
253 # 2. ROUTE_SET tokens
254 for pe_config in result.pe_configs:
255 token = RouteSetToken(
256 target=pe_config.pe_id,
257 addr=None,
258 op=CfgOp.ROUTE_SET,
259 pe_routes=frozenset(pe_config.allowed_pe_routes or ()),
260 sm_routes=frozenset(pe_config.allowed_sm_routes or ()),
261 )
262 tokens.append(token)
263
264 # 3. LOAD_INST tokens
265 for pe_config in result.pe_configs:
266 # Get instructions in offset order
267 offsets = sorted(pe_config.iram.keys())
268 iram_instructions = [pe_config.iram[offset] for offset in offsets]
269 token = LoadInstToken(
270 target=pe_config.pe_id,
271 addr=0,
272 op=CfgOp.LOAD_INST,
273 instructions=tuple(iram_instructions),
274 )
275 tokens.append(token)
276
277 # 4. Seed tokens
278 tokens.extend(result.seed_tokens)
279
280 return tokens