or1-design / emu / network.py

import simpy

from emu.pe import ProcessingElement
from emu.sm import StructureMemory
from emu.types import PEConfig, SMConfig
from sm_mod import SMCell
from tokens import PEToken, SMToken, Token


class System:
    def __init__(
        self,
        env: simpy.Environment,
        pes: dict[int, ProcessingElement],
        sms: dict[int, StructureMemory],
    ):
        self.env = env
        self.pes = pes
        self.sms = sms

    def inject(self, token: Token) -> None:
        """Inject a token into the appropriate element's input store (direct append).

        Routes by type: SMToken → target SM, PEToken → target PE.
        Uses direct list append (bypasses SimPy put), suitable for pre-simulation setup.
        """
        store = self._target_store(token)
        store.items.append(token)

    def send(self, token: Token):
        """Inject a token via SimPy store.put() with 1-cycle delivery delay (generator, yields).

        Same routing as inject() but adds network latency and respects FIFO backpressure.
        Must be called from within a SimPy process or env.process().
        """
        store = self._target_store(token)
        yield self.env.timeout(1)  # 1-cycle network delivery latency
        yield store.put(token)

    def load(self, tokens: list[Token]) -> None:
        """Spawn a SimPy process that sends each token via store.put().

        Tokens are injected in order, respecting FIFO capacity.
        """
        def _loader():
            for token in tokens:
                yield from self.send(token)
        self.env.process(_loader())

    def _target_store(self, token: Token) -> simpy.Store:
        """Resolve the destination store for a token."""
        if isinstance(token, SMToken):
            return self.sms[token.target].input_store
        if isinstance(token, PEToken):
            return self.pes[token.target].input_store
        raise TypeError(f"Unknown token type: {type(token).__name__}")


def build_topology(
    env: simpy.Environment,
    pe_configs: list[PEConfig],
    sm_configs: list[SMConfig],
    fifo_capacity: int = 8,
) -> System:
    pes: dict[int, ProcessingElement] = {}
    sms: dict[int, StructureMemory] = {}

    for cfg in pe_configs:
        pe = ProcessingElement(
            env=env,
            pe_id=cfg.pe_id,
            config=cfg,
        )
        pes[cfg.pe_id] = pe

    t0_store: list[int] = []

    for cfg in sm_configs:
        sm = StructureMemory(
            env=env,
            sm_id=cfg.sm_id,
            cell_count=cfg.cell_count,
            fifo_capacity=fifo_capacity,
            tier_boundary=cfg.tier_boundary,
            on_event=cfg.on_event,
        )
        sm.t0_store = t0_store
        if cfg.initial_cells is not None:
            for addr, (pres, data) in cfg.initial_cells.items():
                sm.cells[addr] = SMCell(pres, data, None)
        sms[cfg.sm_id] = sm

    pe_stores: dict[int, simpy.Store] = {pe_id: pe.input_store for pe_id, pe in pes.items()}
    sm_stores: dict[int, simpy.Store] = {sm_id: sm.input_store for sm_id, sm in sms.items()}

    for pe in pes.values():
        pe.route_table.update(pe_stores)
        pe.sm_routes.update(sm_stores)

    for sm in sms.values():
        sm.route_table.update(pe_stores)

    # Apply route restrictions based on PEConfig allowed_pe_routes and allowed_sm_routes
    for cfg in pe_configs:
        pe = pes[cfg.pe_id]
        if cfg.allowed_pe_routes is not None:
            pe.route_table = {
                pid: store for pid, store in pe.route_table.items()
                if pid in cfg.allowed_pe_routes
            }
        if cfg.allowed_sm_routes is not None:
            pe.sm_routes = {
                sid: store for sid, store in pe.sm_routes.items()
                if sid in cfg.allowed_sm_routes
            }

    system = System(env, pes, sms)
    for sm in sms.values():
        sm.system = system
    return system