OR-1 dataflow CPU sketch
at main 152 lines 4.5 kB view raw
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 )