OR-1 dataflow CPU sketch
1"""State snapshot capture from a live simulation System.
2
3Provides a frozen dataclass representation of PE and SM state at a moment in time,
4suitable for serialization and network transport.
5"""
6
7from __future__ import annotations
8
9from dataclasses import dataclass
10from typing import Optional
11
12from cm_inst import Instruction, FrameSlotValue, Port
13from emu.network import System
14from sm_mod import Presence
15from tokens import Token
16
17
18@dataclass(frozen=True)
19class PESnapshot:
20 """Frozen snapshot of a Processing Element's state at a moment in time."""
21
22 pe_id: int
23 iram: dict[int, Instruction]
24 frames: tuple[tuple[FrameSlotValue, ...], ...]
25 tag_store: dict[int, tuple[int, int]]
26 presence: tuple[tuple[tuple[bool, ...], ...], ...]
27 port_store: tuple[tuple[tuple[Port | None, ...], ...], ...]
28 match_data: tuple[tuple[tuple[int | None, ...], ...], ...]
29 free_frames: tuple[int, ...]
30 lane_count: int
31 input_queue: tuple[Token, ...]
32 output_log: tuple[Token, ...]
33
34
35@dataclass(frozen=True)
36class SMCellSnapshot:
37 """Frozen snapshot of a Structure Memory cell's state."""
38
39 pres: Presence
40 data_l: Optional[int]
41 data_r: Optional[int]
42
43
44@dataclass(frozen=True)
45class SMSnapshot:
46 """Frozen snapshot of a Structure Memory's state at a moment in time."""
47
48 sm_id: int
49 cells: dict[int, SMCellSnapshot]
50 deferred_read: Optional[dict]
51 t0_store: tuple[int, ...]
52 input_queue: tuple[Token, ...]
53
54
55@dataclass(frozen=True)
56class StateSnapshot:
57 """Complete state snapshot of all PEs and SMs at a moment in time."""
58
59 sim_time: float
60 next_time: float
61 pes: dict[int, PESnapshot]
62 sms: dict[int, SMSnapshot]
63
64
65def capture(system: System) -> StateSnapshot:
66 """Capture a frozen snapshot of the entire system state at the current moment.
67
68 Reads from live PE and SM attributes (frames, tag_store, presence, port_store,
69 free_frames, iram, input_store.items, output_log, cells, deferred_read, t0_store).
70
71 Args:
72 system: A System instance from emu.network.build_topology()
73
74 Returns:
75 StateSnapshot containing frozen copies of all PE and SM state.
76 """
77 pes = {}
78 for pe_id, pe in system.pes.items():
79 frames = tuple(
80 tuple(slot for slot in frame)
81 for frame in pe.frames
82 )
83 tag_store = dict(pe.tag_store)
84 presence = tuple(
85 tuple(
86 tuple(lane_val for lane_val in offset_lanes)
87 for offset_lanes in frame_presence
88 )
89 for frame_presence in pe.presence
90 )
91 port_store = tuple(
92 tuple(
93 tuple(lane_val for lane_val in offset_lanes)
94 for offset_lanes in frame_ports
95 )
96 for frame_ports in pe.port_store
97 )
98 match_data = tuple(
99 tuple(
100 tuple(lane_val for lane_val in offset_lanes)
101 for offset_lanes in frame_match
102 )
103 for frame_match in pe.match_data
104 )
105 free_frames = tuple(pe.free_frames)
106
107 pes[pe_id] = PESnapshot(
108 pe_id=pe_id,
109 iram=dict(pe.iram),
110 frames=frames,
111 tag_store=tag_store,
112 presence=presence,
113 port_store=port_store,
114 match_data=match_data,
115 free_frames=free_frames,
116 lane_count=pe.lane_count,
117 input_queue=tuple(pe.input_store.items),
118 output_log=tuple(pe.output_log),
119 )
120
121 sms = {}
122 for sm_id, sm in system.sms.items():
123 cells = {}
124 for i, cell in enumerate(sm.cells):
125 if cell.pres != Presence.EMPTY or cell.data_l is not None:
126 cells[i] = SMCellSnapshot(
127 pres=cell.pres,
128 data_l=cell.data_l,
129 data_r=cell.data_r,
130 )
131
132 dr = None
133 if sm.deferred_read is not None:
134 dr = {
135 "cell_addr": sm.deferred_read.cell_addr,
136 "return_route": sm.deferred_read.return_route,
137 }
138
139 sms[sm_id] = SMSnapshot(
140 sm_id=sm_id,
141 cells=cells,
142 deferred_read=dr,
143 t0_store=tuple(sm.t0_store),
144 input_queue=tuple(sm.input_store.items),
145 )
146
147 return StateSnapshot(
148 sim_time=system.env.now,
149 next_time=system.env.peek(),
150 pes=pes,
151 sms=sms,
152 )