from __future__ import annotations from abc import ABC, abstractmethod from collections.abc import Generator from typing import TYPE_CHECKING, Any import simpy if TYPE_CHECKING: from kernbench.components.context import ComponentContext from kernbench.topology.types import Node class ComponentBase(ABC): """Base class for all SimPy component implementations (ADR-0007 D3, ADR-0015). Each component corresponds to one node in the compiled topology graph. It models the processing overhead at that node as a SimPy generator, allowing future implementations to add queueing and contention. Port model (ADR-0015 D1): in_ports[src_node_id] — SimPy Store for incoming messages from src out_ports[dst_node_id] — SimPy Store for outgoing messages to dst Ports are wired by GraphEngine at initialization; wire processes model propagation delay between connected ports (ADR-0015 D2). Context (ADR-0015 D4): ctx — ComponentContext with router and resolver. """ def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None: self.node = node self.ctx = ctx self.in_ports: dict[str, simpy.Store] = {} self.out_ports: dict[str, simpy.Store] = {} self._op_logger: Any | None = None # OpLogger, set by GraphEngine if enabled def start(self, env: simpy.Environment) -> None: """Called once after all ports are wired. Default: starts a fan-in collector and a generic forwarding worker. The worker calls self.run() for per-component latency, then routes the Transaction to the next hop or signals done (duck-typed; no direct Transaction import to avoid circular dependencies). Override in components that need custom fan-out / aggregation logic (e.g. MCpuComponent, IoCpuComponent for kernel launch). """ if not self.in_ports: return self._inbox: simpy.Store = simpy.Store(env) for port in self.in_ports.values(): env.process(self._fan_in(port)) env.process(self._worker(env)) # ADR-0033 Phase 2c: flit-aware components consume Flits directly; # non-flit-aware components reassemble Flits into the parent # Transaction before delivery to _inbox. Default False preserves # legacy single-msg semantics during incremental rollout. _FLIT_AWARE: bool = False def _fan_in(self, port: simpy.Store) -> Generator: """Relay messages from in_port to _inbox. For non-flit-aware components (default), Flits are accumulated by parent Transaction and only the reassembled Transaction is placed on _inbox once ``is_last`` arrives. Step is updated to this component's path position for legacy step-based routing.""" from kernbench.sim_engine.transaction import Flit if self._FLIT_AWARE: while True: msg = yield port.get() yield self._inbox.put(msg) return flit_buffers: dict[int, list[Any]] = {} while True: msg = yield port.get() if isinstance(msg, Flit): tid = id(msg.txn) flit_buffers.setdefault(tid, []).append(msg) if msg.is_last: flit_buffers.pop(tid, None) self._update_step(msg.txn) yield self._inbox.put(msg.txn) else: yield self._inbox.put(msg) def _update_step(self, txn: Any) -> None: """Set txn.step to this component's index in txn.path (if found). Allows legacy step-based routing to work even when flit-aware upstream components don't call txn.advance().""" my_id = self.node.id path = getattr(txn, "path", None) if not path: return for i, n in enumerate(path): if n == my_id: txn.step = i return def _worker(self, env: simpy.Environment) -> Generator: """Generic forwarding worker: spawns _forward_txn per message (pipeline).""" while True: txn: Any = yield self._inbox.get() env.process(self._forward_txn(env, txn)) def _on_process_start(self, env: simpy.Environment, msg: Any) -> None: """Op log hook: record service start for data_op messages (ADR-0020 D2).""" if self._op_logger and getattr(msg, "data_op", False): self._op_logger.record_start(env.now, self.node.id, msg) def _on_process_end(self, env: simpy.Environment, msg: Any) -> None: """Op log hook: record service end for data_op messages (ADR-0020 D2).""" if self._op_logger and getattr(msg, "data_op", False): self._op_logger.record_end(env.now, self.node.id, msg) def _forward_txn(self, env: simpy.Environment, txn: Any) -> Generator: """Apply run() latency, then forward to next hop or drain at terminal.""" self._on_process_start(env, txn) yield from self.run(env, txn.nbytes) self._on_process_end(env, txn) next_hop = txn.next_hop # duck-typed: Transaction.next_hop if next_hop: yield self.out_ports[next_hop].put(txn.advance()) else: drain = getattr(txn, "drain_ns", 0.0) if drain > 0: yield env.timeout(drain) txn.done.succeed() @abstractmethod def run(self, env: simpy.Environment, nbytes: int) -> Generator: """SimPy process: yield one or more events for this node's processing. Subclasses yield env.timeout(overhead_ns) or compute latency dynamically. Called by _forward_txn and subclass-specific handlers. """ ... class PeEngineBase(ComponentBase): """Base class for PE-internal engines (PE_DMA, PE_GEMM, PE_MATH). Provides: - ``_pe_prefix``: extracted from node.id (e.g. "sip0.cube0.pe0") - Dual-message ``_worker``: dispatches PeInternalTxn to ``handle_command()`` and Transaction to inherited ``_forward_txn()``. - ``init_resources(env)``: hook for subclass resource initialization, called by ``start()`` before the worker is spawned. Subclass contract: 1. Override ``handle_command(env, pe_txn)`` — process a PeInternalTxn. 2. Override ``run(env, nbytes)`` — yield component latency. 3. Optionally override ``init_resources(env)`` for DMA channels, etc. """ def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None: super().__init__(node, ctx) self._pe_prefix: str = node.id.rsplit(".", 1)[0] def start(self, env: simpy.Environment) -> None: self.init_resources(env) super().start(env) def init_resources(self, env: simpy.Environment) -> None: """Hook for subclass resource initialization. Called before worker spawn.""" def _worker(self, env: simpy.Environment) -> Generator: """Dual-message dispatch: PeInternalTxn → handle_command, Transaction → _forward_txn.""" from kernbench.common.pe_commands import PeInternalTxn while True: msg: Any = yield self._inbox.get() if isinstance(msg, PeInternalTxn): env.process(self._handle_with_hooks(env, msg)) else: env.process(self._forward_txn(env, msg)) def _handle_with_hooks(self, env: simpy.Environment, pe_txn: Any) -> Generator: """Wrap handle_command with op log hooks on the inner command. Subclasses that need to defer record_start until after a resource wait (e.g. pe_dma's DMA-channel acquire) set ``_DEFER_RECORD_START = True`` and call ``self._on_process_start(env, pe_txn.command)`` themselves at the post-wait moment. record_end still fires here. """ if not getattr(self, "_DEFER_RECORD_START", False): self._on_process_start(env, pe_txn.command) yield from self.handle_command(env, pe_txn) self._on_process_end(env, pe_txn.command) @abstractmethod def handle_command(self, env: simpy.Environment, pe_txn: Any) -> Generator: """Process a PE-internal command (PeInternalTxn). Subclass must: - Perform engine-specific work (acquire resources, compute, etc.) - Call ``pe_txn.done.succeed()`` on completion. """ ... class ComponentRegistry: """DI registry: maps node.impl strings to ComponentBase subclasses. Resolution order for ComponentRegistry.create(node, overrides, ctx): 1. overrides[node.impl] — caller-injected override 2. _registry[node.impl] — globally registered impl (lazy import) 3. Error — no fallback; every node must have an impl Registry is populated from components.yaml via load_components_yaml(). Manual register() is still supported for tests and overrides. """ _registry: dict[str, type[ComponentBase]] = {} _lazy: dict[str, str] = {} # impl → "module.path:ClassName" _loaded: bool = False @classmethod def register(cls, impl: str, component_cls: type[ComponentBase]) -> None: cls._registry[impl] = component_cls @classmethod def load_components_yaml(cls, path: str | None = None) -> None: """Load impl→class mappings from components.yaml. Lazy imports on first use.""" if cls._loaded: return import yaml from pathlib import Path if path is None: # Search: project root (cwd), then relative to this file candidates = [ Path.cwd() / "components.yaml", Path(__file__).parent.parent.parent.parent / "components.yaml", ] for p in candidates: if p.exists(): path = str(p) break if path is None: return with open(path) as f: spec = yaml.safe_load(f) for impl, class_path in (spec.get("components") or {}).items(): cls._lazy[impl] = class_path cls._loaded = True @classmethod def _resolve(cls, impl: str) -> type[ComponentBase] | None: """Resolve impl name: check _registry first, then lazy import from _lazy.""" if impl in cls._registry: return cls._registry[impl] if not cls._loaded: cls.load_components_yaml() class_path = cls._lazy.get(impl) if class_path is None: return None import importlib module_path, class_name = class_path.rsplit(":", 1) mod = importlib.import_module(module_path) component_cls = getattr(mod, class_name) cls._registry[impl] = component_cls # cache for next lookup return component_cls @classmethod def create( cls, node: Node, overrides: dict[str, type[ComponentBase]] | None = None, ctx: ComponentContext | None = None, ) -> ComponentBase: if overrides and node.impl in overrides: return overrides[node.impl](node, ctx) component_cls = cls._resolve(node.impl) if component_cls is not None: return component_cls(node, ctx) raise ValueError( f"No component registered for impl '{node.impl}' (node: {node.id}). " f"Add it to components.yaml or call ComponentRegistry.register()." )