"""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 _RwHazardTracker: """Strict-FIFO cross-composite RW hazard tracker (ADR-0065 D6.3 / DDD §8). A composite is registered as in-flight by its *write* set (``rw_handles``). A new composite whose read **or** write handles intersect any in-flight write set waits on those composites' done events before being admitted. Composites with no ``rw_handles`` (legacy) are never registered and never block — existing benches are unaffected. Strict FIFO is conservative: even partial overlap waits for the prior overlapping composite to fully drain (RW-aware reorder is ADR-0065 A4, deferred). """ def __init__(self) -> None: # (completion_id, write-handle id set, done_event) self._inflight: list[tuple[str, set, Any]] = [] @staticmethod def _conflict_ids(cmd: Any) -> set: """Handle ids this composite reads (op operands) or writes (rw).""" ids = {h.id for h in cmd.rw_handles} for op in cmd.ops: for h in op.operands.values(): hid = getattr(h, "id", None) if hid is not None: ids.add(hid) return ids def admit(self, env: Any, cmd: Any, done_event: Any) -> Generator: """Block until no in-flight composite's write set overlaps this composite's read/write set, then register this composite's writes.""" cids = self._conflict_ids(cmd) while True: blocking = [ev for _cid, rw, ev in self._inflight if rw & cids] if not blocking: break yield blocking[0] if cmd.rw_handles: self._inflight.append( (cmd.completion.id, {h.id for h in cmd.rw_handles}, done_event) ) def retire(self, completion_id: str) -> None: self._inflight = [t for t in self._inflight if t[0] != completion_id] 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 self._hazard = _RwHazardTracker() # ADR-0065 D6.3 strict-FIFO RW 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). Strict-FIFO RW gate (ADR-0065 D6.3): wait for overlapping in-flight composites, register this one's write set, retire on completion. Legacy composites (no rw_handles) pass through immediately. """ from kernbench.components.builtin.pe_types import PipelineContext yield from self._hazard.admit(env, cmd, pe_txn.done) env.process(self._retire_on_done(env, pe_txn.done, cmd.completion.id)) 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 _retire_on_done( self, env: simpy.Environment, done_event: Any, completion_id: str, ) -> Generator: """Retire a composite from the RW hazard tracker on completion.""" yield done_event self._hazard.retire(completion_id) 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, )