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 → IRAM writes → seeds)
7
8Reference: Phase 6 design doc, Tasks 1-2.
9"""
10
11from dataclasses import dataclass
12from collections import defaultdict
13
14from asm.errors import AssemblyError, ErrorCategory
15from asm.ir import (
16 IRGraph, IRNode, IREdge, ResolvedDest, collect_all_nodes_and_edges, collect_all_data_defs,
17 DEFAULT_IRAM_CAPACITY, DEFAULT_CTX_SLOTS
18)
19from asm.opcodes import is_dyadic
20from cm_inst import ALUInst, MemOp, Port, RoutingOp, SMInst
21from emu.types import PEConfig, SMConfig
22from tokens import DyadToken, IRAMWriteToken, MonadToken, SMToken
23from sm_mod import Presence
24
25
26@dataclass(frozen=True)
27class AssemblyResult:
28 """Result of code generation in direct mode.
29
30 Attributes:
31 pe_configs: List of PEConfig objects, one per PE
32 sm_configs: List of SMConfig objects, one per SM with data_defs
33 seed_tokens: List of MonadTokens for const nodes with no incoming edges
34 """
35 pe_configs: list[PEConfig]
36 sm_configs: list[SMConfig]
37 seed_tokens: list[MonadToken]
38
39
40
41
42def _build_iram_for_pe(
43 nodes_on_pe: list[IRNode],
44 all_nodes: dict[str, IRNode],
45 all_edges: list[IREdge],
46) -> dict[int, ALUInst | SMInst]:
47 """Build IRAM instruction dict for a single PE.
48
49 Args:
50 nodes_on_pe: List of IRNodes on this PE
51 all_nodes: All nodes in graph (for lookups)
52 all_edges: All edges in graph (for ctx_override detection)
53
54 Returns:
55 Dict mapping IRAM offset to ALUInst or SMInst
56 """
57 iram = {}
58
59 # Build edge map for quick lookup: node name -> list of outgoing edges
60 edges_by_source = defaultdict(list)
61 for edge in all_edges:
62 edges_by_source[edge.source].append(edge)
63
64 for node in nodes_on_pe:
65 if node.iram_offset is None:
66 # Node not allocated, skip
67 continue
68
69 if isinstance(node.opcode, MemOp):
70 # Memory operation -> SMInst
71 ret_addr = node.dest_l.addr if isinstance(node.dest_l, ResolvedDest) else None
72 ret_dyadic = False
73 if isinstance(node.dest_l, ResolvedDest):
74 dest_node = all_nodes.get(node.dest_l.name)
75 if dest_node is not None:
76 ret_dyadic = is_dyadic(dest_node.opcode, dest_node.const)
77 inst = SMInst(
78 op=node.opcode,
79 sm_id=node.sm_id,
80 const=node.const,
81 ret=ret_addr,
82 ret_dyadic=ret_dyadic,
83 )
84 else:
85 # ALU operation -> ALUInst
86 # Extract Addr from ResolvedDest or keep None
87 dest_l_addr = None
88 dest_r_addr = None
89
90 if node.dest_l is not None and isinstance(node.dest_l, ResolvedDest):
91 dest_l_addr = node.dest_l.addr
92
93 if node.dest_r is not None and isinstance(node.dest_r, ResolvedDest):
94 dest_r_addr = node.dest_r.addr
95
96 # Check if this node has ctx_override edges (AC5.2, AC5.3)
97 ctx_mode = 0
98 packed_const = node.const
99 node_edges = edges_by_source.get(node.name, [])
100 has_ctx_override = any(edge.ctx_override for edge in node_edges)
101
102 if has_ctx_override:
103 # AC5.3: Conflict detection - node with both const and ctx_override
104 if node.const is not None:
105 error = AssemblyError(
106 category=ErrorCategory.VALUE,
107 message=f"Node '{node.name}' requires both const operand and CTX_OVRD — expected expand pass to insert trampoline",
108 loc=node.loc,
109 )
110 raise ValueError(f"Codegen error: {error.message}")
111
112 ctx_mode = 1
113 # Pack const field: target_ctx and target_gen
114 # Find the first ctx_override edge to get the target context
115 # The allocator should have set ctx on destination nodes
116 for edge in node_edges:
117 if edge.ctx_override:
118 dest_node = all_nodes.get(edge.dest)
119 if dest_node is not None and dest_node.ctx is not None:
120 target_ctx = dest_node.ctx
121 target_gen = 0 # Initial generation
122 # Pack: ((target_ctx & 0xF) << 4) | ((target_gen & 0x3) << 2)
123 # Upper 8 bits must be zero (reserved)
124 packed_const = ((target_ctx & 0xF) << 4) | ((target_gen & 0x3) << 2)
125 break
126
127 # Defensive guard: ensure packed_const is set for CTX_OVRD
128 if packed_const is None:
129 raise ValueError(
130 f"Codegen error: Node '{node.name}' has ctx_override edge but destination context is not resolved. "
131 f"Allocator should have assigned ctx to destination nodes."
132 )
133
134 inst = ALUInst(
135 op=node.opcode,
136 dest_l=dest_l_addr,
137 dest_r=dest_r_addr,
138 const=packed_const,
139 ctx_mode=ctx_mode,
140 )
141
142 iram[node.iram_offset] = inst
143
144 return iram
145
146
147def _compute_route_restrictions(
148 nodes_by_pe: dict[int, list[IRNode]],
149 all_edges: list[IREdge],
150 all_nodes: dict[str, IRNode],
151 pe_id: int,
152) -> tuple[set[int], set[int]]:
153 """Compute allowed PE and SM routes for a given PE.
154
155 Analyzes all edges involving nodes on this PE to determine which other
156 PEs and SMs it can route to. Includes self-routes.
157
158 Args:
159 nodes_by_pe: Dict mapping PE ID to list of nodes on that PE
160 all_edges: List of all edges in graph
161 all_nodes: Dict of all nodes
162 pe_id: The PE we're computing routes for
163
164 Returns:
165 Tuple of (allowed_pe_routes set, allowed_sm_routes set)
166 """
167 nodes_on_pe_set = {node.name for node in nodes_by_pe.get(pe_id, [])}
168
169 pe_routes = {pe_id} # Always include self-route
170 sm_routes = set()
171
172 # Scan all edges for those sourced from this PE
173 for edge in all_edges:
174 if edge.source in nodes_on_pe_set:
175 # This edge originates from our PE
176 dest_node = all_nodes.get(edge.dest)
177 if dest_node is not None:
178 if dest_node.pe is not None:
179 pe_routes.add(dest_node.pe)
180
181 # Scan all nodes on this PE for SM instructions
182 for node in nodes_by_pe.get(pe_id, []):
183 if isinstance(node.opcode, MemOp) and node.sm_id is not None:
184 sm_routes.add(node.sm_id)
185
186 return pe_routes, sm_routes
187
188
189def generate_direct(graph: IRGraph) -> AssemblyResult:
190 """Generate PEConfig, SMConfig, and seed tokens from an allocated IRGraph.
191
192 Args:
193 graph: A fully allocated IRGraph (after allocate pass)
194
195 Returns:
196 AssemblyResult with pe_configs, sm_configs, and seed_tokens
197 """
198 all_nodes, all_edges = collect_all_nodes_and_edges(graph)
199 all_data_defs = collect_all_data_defs(graph)
200
201 # Group nodes by PE
202 nodes_by_pe: dict[int, list[IRNode]] = defaultdict(list)
203 for node in all_nodes.values():
204 if node.pe is not None:
205 nodes_by_pe[node.pe].append(node)
206
207 # Build PEConfigs
208 pe_configs = []
209 for pe_id in sorted(nodes_by_pe.keys()):
210 nodes_on_pe = nodes_by_pe[pe_id]
211
212 # Build IRAM for this PE
213 iram = _build_iram_for_pe(nodes_on_pe, all_nodes, all_edges)
214
215 # Compute route restrictions
216 allowed_pe_routes, allowed_sm_routes = _compute_route_restrictions(
217 nodes_by_pe, all_edges, all_nodes, pe_id
218 )
219
220 # Create PEConfig
221 config = PEConfig(
222 pe_id=pe_id,
223 iram=iram,
224 ctx_slots=graph.system.ctx_slots if graph.system else DEFAULT_CTX_SLOTS,
225 offsets=graph.system.iram_capacity if graph.system else DEFAULT_IRAM_CAPACITY,
226 allowed_pe_routes=allowed_pe_routes,
227 allowed_sm_routes=allowed_sm_routes,
228 )
229 pe_configs.append(config)
230
231 # Build SMConfigs from data_defs
232 sm_configs_by_id: dict[int, dict[int, tuple[Presence, int]]] = defaultdict(dict)
233 for data_def in all_data_defs:
234 if data_def.sm_id is not None and data_def.cell_addr is not None:
235 sm_configs_by_id[data_def.sm_id][data_def.cell_addr] = (
236 Presence.FULL, data_def.value
237 )
238
239 sm_count = max(1, graph.system.sm_count if graph.system else 1)
240 for sm_id in range(sm_count):
241 if sm_id not in sm_configs_by_id:
242 sm_configs_by_id[sm_id] = {}
243
244 sm_configs = []
245 for sm_id in sorted(sm_configs_by_id.keys()):
246 initial_cells = sm_configs_by_id[sm_id]
247 config = SMConfig(
248 sm_id=sm_id,
249 initial_cells=initial_cells if initial_cells else None,
250 )
251 sm_configs.append(config)
252
253 # Detect seed tokens
254 seed_tokens = []
255
256 # Build edge indices
257 edges_by_dest = defaultdict(list)
258 edges_by_source = defaultdict(list)
259 for edge in all_edges:
260 edges_by_dest[edge.dest].append(edge)
261 edges_by_source[edge.source].append(edge)
262
263 for node in all_nodes.values():
264 if node.seed:
265 # Seed node: generate token(s) targeted at destination(s)
266 out_edges = edges_by_source.get(node.name, [])
267 for edge in out_edges:
268 dest_node = all_nodes.get(edge.dest)
269 if dest_node is None or dest_node.pe is None:
270 continue
271 dest_is_dyadic = is_dyadic(dest_node.opcode, dest_node.const)
272 if dest_is_dyadic:
273 token = DyadToken(
274 target=dest_node.pe,
275 offset=dest_node.iram_offset if dest_node.iram_offset is not None else 0,
276 ctx=dest_node.ctx if dest_node.ctx is not None else 0,
277 data=node.const if node.const is not None else 0,
278 port=edge.port,
279 gen=0,
280 wide=False,
281 )
282 else:
283 token = MonadToken(
284 target=dest_node.pe,
285 offset=dest_node.iram_offset if dest_node.iram_offset is not None else 0,
286 ctx=dest_node.ctx if dest_node.ctx is not None else 0,
287 data=node.const if node.const is not None else 0,
288 inline=False,
289 )
290 seed_tokens.append(token)
291 elif node.opcode == RoutingOp.CONST:
292 # Triggerable constant: CONST node in IRAM with no incoming edges
293 if node.name not in edges_by_dest:
294 token = MonadToken(
295 target=node.pe if node.pe is not None else 0,
296 offset=node.iram_offset if node.iram_offset is not None else 0,
297 ctx=node.ctx if node.ctx is not None else 0,
298 data=node.const if node.const is not None else 0,
299 inline=False,
300 )
301 seed_tokens.append(token)
302
303 return AssemblyResult(
304 pe_configs=pe_configs,
305 sm_configs=sm_configs,
306 seed_tokens=seed_tokens,
307 )
308
309
310def generate_tokens(graph: IRGraph) -> list:
311 """Generate bootstrap token sequence from an allocated IRGraph.
312
313 Produces tokens in order: SM init → IRAM writes → seeds
314
315 Args:
316 graph: A fully allocated IRGraph (after allocate pass)
317
318 Returns:
319 List of tokens (SMToken, IRAMWriteToken, MonadToken) in bootstrap order
320 """
321 # Use direct mode to get configs and seeds
322 result = generate_direct(graph)
323
324 tokens = []
325
326 # 1. SM init tokens
327 all_data_defs = collect_all_data_defs(graph)
328 for data_def in all_data_defs:
329 if data_def.sm_id is not None and data_def.cell_addr is not None:
330 token = SMToken(
331 target=data_def.sm_id,
332 addr=data_def.cell_addr,
333 op=MemOp.WRITE,
334 flags=None,
335 data=data_def.value,
336 ret=None,
337 )
338 tokens.append(token)
339
340 # 2. IRAM write tokens
341 for pe_config in result.pe_configs:
342 offsets = sorted(pe_config.iram.keys())
343 iram_instructions = [pe_config.iram[offset] for offset in offsets]
344 token = IRAMWriteToken(
345 target=pe_config.pe_id,
346 offset=0,
347 ctx=0,
348 data=0,
349 instructions=tuple(iram_instructions),
350 )
351 tokens.append(token)
352
353 # 3. Seed tokens
354 tokens.extend(result.seed_tokens)
355
356 return tokens