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 a flat-ops CompositeCmd (ADR-0065 D3).

        Position-scan for the GEMM (≤1) + prologue/epilogue placement lives
        in ``generate_plan_from_ops``; a legacy composite (GEMM at index 0,
        no prologue) produces the identical plan it did before.
        """
        from kernbench.components.builtin.tiling import generate_plan_from_ops

        return generate_plan_from_ops(
            ops=cmd.ops,
            tile_m=self.TILE_M, tile_k=self.TILE_K, tile_n=self.TILE_N,
            bytes_per_element=2,  # f16
            pe_prefix=self._pe_prefix,
        )