ywkang
5bbe293bea
Two fixes make the recipe's P.V GEMM fold into the rescaled accumulator: (1) op_log marks the composite gemm record accumulate=True when an 'add' epilogue targets the GEMM's own output; data_executor._execute_gemm then does O = O_existing + a@b instead of overwriting. (2) PE_SCHEDULER records the composite's numeric op at COMPLETION (in _retire_on_done) instead of at dispatch, so the GEMM's t_start sorts AFTER its prologue MATH ops -> it reads the computed P and the rescaled O. Verified end-to-end (DataExecutor): the full recipe produces O = O_old*corr + P@V matching a numpy flash-attention reference. Suite 816 pass / 3 pre-existing fail. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
263 lines
10 KiB
Python
263 lines
10 KiB
Python
"""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)
|
|
# Retire from the hazard tracker AND record the composite's numeric
|
|
# op — both at completion (see _retire_on_done).
|
|
env.process(self._retire_on_done(env, pe_txn.done, cmd))
|
|
|
|
plan = self._generate_plan(cmd)
|
|
|
|
self._pipeline_counter += 1
|
|
# Prologue / post-loop single-shot stages each count as a tile for
|
|
# completion (ADR-0065 D3). Empty for legacy composites → unchanged.
|
|
n_pre = len(getattr(plan, "prologue_stages", ()))
|
|
n_post = len(getattr(plan, "epilogue_stages", ()))
|
|
ctx = PipelineContext(
|
|
id=f"p{self._pipeline_counter}",
|
|
total_tiles=len(plan.tiles) + n_pre + n_post,
|
|
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, cmd: Any,
|
|
) -> Generator:
|
|
"""On composite completion: retire from the RW hazard tracker AND
|
|
record its numeric op (ADR-0027 / ADR-0065 N1+N3).
|
|
|
|
Recording at completion (not dispatch) gives the composite's GEMM a
|
|
t_start AFTER its prologue MATH ops, so in data mode it reads their
|
|
results (P, the rescaled O) instead of stale values. No-op without an
|
|
op_logger → latency-only mode unchanged.
|
|
"""
|
|
yield done_event
|
|
self._hazard.retire(cmd.completion.id)
|
|
self._on_process_start(env, cmd)
|
|
self._on_process_end(env, cmd)
|
|
|
|
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()
|
|
tid = 0
|
|
# Prologue single-shot stages first (ADR-0065 D3) — each a
|
|
# standalone 1-stage tile that runs on its component and
|
|
# completes (no inter-component routing).
|
|
for st in getattr(plan, "prologue_stages", ()):
|
|
yield from self._feed_single_stage(st, ctx, tid)
|
|
tid += 1
|
|
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)
|
|
for st in getattr(plan, "epilogue_stages", ()):
|
|
yield from self._feed_single_stage(st, ctx, tid)
|
|
tid += 1
|
|
|
|
def _feed_single_stage(self, stage: Any, ctx: Any, tile_id: int) -> Generator:
|
|
"""Feed one standalone stage as a 1-stage tile (prologue/post-loop)."""
|
|
from kernbench.components.builtin.pe_types import TilePlan, TileToken
|
|
|
|
token = TileToken(
|
|
tile_id=tile_id, pipeline_ctx=ctx,
|
|
plan=TilePlan(tile_id=tile_id, stages=(stage,)),
|
|
stage_idx=0, params=stage.params,
|
|
)
|
|
yield self.out_ports[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,
|
|
)
|