kernbench2/src/kernbench/components/builtin/pe_scheduler.py

"""PE_SCHEDULER: plan generation + tile dispatch (ADR-0014 D6).

Receives PeInternalTxn from PE_CPU, routes to engines:
  - Simple commands (DmaReadCmd, GemmCmd, etc.) → direct dispatch to engine
  - CompositeCmd → generate TilePlan, feed tiles via _feed_loop

Composite pipeline uses token self-routing (ADR-0014 D6):
  Scheduler only does initial dispatch + completion tracking.
  Tiles chain through components based on their plan's stage sequence.
"""
from __future__ import annotations

from collections.abc import Generator
from typing import TYPE_CHECKING, Any

import simpy

from kernbench.components.base import ComponentBase

if TYPE_CHECKING:
    from kernbench.common.pe_commands import PeInternalTxn
    from kernbench.components.context import ComponentContext
    from kernbench.topology.types import Node


class PeSchedulerComponent(ComponentBase):
    """PE_SCHEDULER: sole dispatcher inside a PE (ADR-0014 D1, D6).

    Simple commands are forwarded to the appropriate engine.
    CompositeCmd creates a TilePlan and feeds tiles into the pipeline.

    Single _feed_loop process per scheduler ensures FIFO command ordering.
    """

    TILE_M = 32
    TILE_K = 64
    TILE_N = 32

    _CMD_DISPATCH: dict[type, str] = {}

    @classmethod
    def _ensure_dispatch_table(cls) -> None:
        if cls._CMD_DISPATCH:
            return
        from kernbench.common.pe_commands import (
            CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd,
        )
        cls._CMD_DISPATCH = {
            DmaReadCmd: "pe_dma",
            DmaWriteCmd: "pe_dma",
            GemmCmd: "pe_gemm",
            MathCmd: "pe_math",
            # ADR-0063 §D3.1: tl.copy_to → vector engine.
            CopyCmd: "pe_math",
        }

    def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
        super().__init__(node, ctx)
        self._pe_prefix = node.id.rsplit(".", 1)[0]
        self._ensure_dispatch_table()
        self._pending_feeds: simpy.Store | None = None
        self._pipeline_counter = 0

    def start(self, env: simpy.Environment) -> None:
        self._pending_feeds = simpy.Store(env)
        super().start(env)
        env.process(self._feed_loop(env))

    def run(self, env: simpy.Environment, nbytes: int) -> Generator:
        overhead_ns = float(self.node.attrs.get("overhead_ns", 0.0))
        yield env.timeout(overhead_ns)

    def _worker(self, env: simpy.Environment) -> Generator:
        from kernbench.common.pe_commands import PeInternalTxn

        while True:
            msg: Any = yield self._inbox.get()
            if isinstance(msg, PeInternalTxn):
                env.process(self._dispatch(env, msg))
            else:
                yield from self._forward_txn(env, msg)

    def _dispatch(self, env: simpy.Environment, pe_txn: PeInternalTxn) -> Generator:
        from kernbench.common.pe_commands import CompositeCmd, PeCpuOverheadCmd

        yield from self.run(env, 0)  # scheduler overhead

        cmd = pe_txn.command

        # Simple command dispatch
        engine_suffix = self._CMD_DISPATCH.get(type(cmd))
        if engine_suffix is not None:
            yield self.out_ports[f"{self._pe_prefix}.{engine_suffix}"].put(pe_txn)
            return

        # CompositeCmd: generate plan and feed
        if isinstance(cmd, CompositeCmd):
            yield from self._dispatch_composite(env, pe_txn, cmd)
            return

        if isinstance(cmd, PeCpuOverheadCmd):
            yield env.timeout(cmd.cycles)
            pe_txn.done.succeed()
            return

        pe_txn.done.succeed()

    def _dispatch_composite(
        self, env: simpy.Environment, pe_txn: Any, cmd: Any,
    ) -> Generator:
        """Generate plan and enqueue to feeder. Non-blocking (ADR-0014 D6)."""
        from kernbench.components.builtin.pe_types import PipelineContext

        plan = self._generate_plan(cmd)

        self._pipeline_counter += 1
        ctx = PipelineContext(
            id=f"p{self._pipeline_counter}",
            total_tiles=len(plan.tiles),
            done_event=pe_txn.done,
        )

        # Enqueue to feeder — scheduler worker returns immediately
        assert self._pending_feeds is not None
        yield self._pending_feeds.put((plan, ctx))

    def _feed_loop(self, env: simpy.Environment) -> Generator:
        """Single feeder process: FIFO command ordering (ADR-0014 D6).

        No tile feed interleaving between commands.
        Queue full → only this process blocks.
        """
        from kernbench.components.builtin.pe_types import TileToken

        assert self._pending_feeds is not None
        while True:
            plan, ctx = yield self._pending_feeds.get()
            for tile in plan.tiles:
                first_stage = tile.stages[0]
                token = TileToken(
                    tile_id=tile.tile_id,
                    pipeline_ctx=ctx,
                    plan=tile,
                    stage_idx=0,
                    params=first_stage.params,
                )
                yield self.out_ports[first_stage.component].put(token)

    def _generate_plan(self, cmd: Any) -> Any:
        """Generate a PipelinePlan from CompositeCmd."""
        from kernbench.components.builtin.tiling import (
            generate_gemm_plan,
            generate_math_plan,
        )

        pp = self._pe_prefix
        bpe = 2  # default bytes per element (f16)

        if cmd.op == "gemm" and cmd.b is not None:
            a = cmd.a
            b = cmd.b
            M, K = a.shape[-2], a.shape[-1]
            N = b.shape[-1]
            # When CompositeCmd.ops is populated, ops[0] is the head and
            # ops[1:] is the epilogue spec list. Empty ops → legacy path.
            epi_specs = tuple(cmd.ops[1:]) if cmd.ops else ()
            return generate_gemm_plan(
                M=M, K=K, N=N,
                tile_m=self.TILE_M, tile_k=self.TILE_K, tile_n=self.TILE_N,
                bytes_per_element=bpe,
                A_addr=a.addr, B_addr=b.addr, C_addr=cmd.out_addr,
                pe_prefix=pp,
                a_pinned=getattr(a, "pinned", False),
                b_pinned=getattr(b, "pinned", False),
                epilogue_specs=epi_specs,
            )
        else:
            # Math composite
            a = cmd.a
            M = a.shape[-2] if len(a.shape) >= 2 else a.shape[0]
            N = a.shape[-1] if len(a.shape) >= 2 else 1
            return generate_math_plan(
                M=M, N=N,
                tile_m=self.TILE_M, tile_n=self.TILE_N,
                bytes_per_element=bpe,
                math_op=cmd.math_op or "identity",
                src_addr=a.addr, dst_addr=cmd.out_addr,
                pe_prefix=pp,
            )