commit - release 1

src/kernbench/components/base.py

@@ -0,0 +1,167 @@
+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] = {}
+
+    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))
+
+    def _fan_in(self, port: simpy.Store) -> Generator:
+        """Relay messages from one in_port into the shared inbox."""
+        while True:
+            msg = yield port.get()
+            yield self._inbox.put(msg)
+
+    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 _forward_txn(self, env: simpy.Environment, txn: Any) -> Generator:
+        """Apply run() latency, then forward to next hop or drain at terminal."""
+        yield from self.run(env, txn.nbytes)
+        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_command(env, msg))
+            else:
+                env.process(self._forward_txn(env, msg))
+
+    @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
+      3. Error                   — no fallback; every node must have an impl
+    """
+
+    _registry: dict[str, type[ComponentBase]] = {}
+
+    @classmethod
+    def register(cls, impl: str, component_cls: type[ComponentBase]) -> None:
+        cls._registry[impl] = 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)
+        if node.impl in cls._registry:
+            return cls._registry[node.impl](node, ctx)
+        raise ValueError(
+            f"No component registered for impl '{node.impl}' (node: {node.id}). "
+            f"Register it in kernbench.components.impls.__init__."
+        )