kernbench2/src/kernbench/components/base.py

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()."
        )