ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa(adr-0065): P3 — flat-ops PE_SCHEDULER plan (position-scan + prologue stages)

Browse Source 
tiling.generate_plan_from_ops: scan flat ops for the GEMM (<=1); pre-GEMM KERNEL ops become single-shot prologue MATH stages, post-GEMM ops split by scope (K_TILE/OUTPUT_TILE epilogue + KERNEL post-loop). MATH-only composite reproduces the legacy math-head plan. Prologue/post-loop stages fold into the first/last tile so the feeder + completion counting are untouched (existing benches have neither -> byte-equal op_log).

PipelinePlan gains prologue_stages/epilogue_stages. pe_scheduler._generate_plan delegates to generate_plan_from_ops. DMA keeps the existing pinned signal (NOT a space flip); recipe scratch/primary-out handles are pinned=True so the head GEMM's auto-bound a (=P) is consumed in place.

ADR-0065 D4/D6.7/Test#5 amended (EN+KO): DMA decision from pinned, not space; prologue recipe ops TCM-only (head op exempt -- it may stream from HBM).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook
2026-06-10 20:24:33 -07:00
parent 184f654295
commit 55f025c4b1
7 changed files with 304 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/components/builtin/pe_scheduler.py
+13 -41
View File
@@ -147,45 +147,17 @@ class PeSchedulerComponent(ComponentBase):
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,
"""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,
)
pp = self._pe_prefix
bpe = 2 # default bytes per element (f16)
# Flat-ops (ADR-0065 D1): ops[0] is the head, ops[1:] are epilogue
# specs placed by scope. The head's kind selects the engine path.
head = cmd.ops[0]
epi_specs = tuple(cmd.ops[1:])
if head.kind == "gemm" and "b" in head.operands:
a = head.operands["a"]
b = head.operands["b"]
M, K = a.shape[-2], a.shape[-1]
N = b.shape[-1]
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=head.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 = head.operands["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=head.extra.get("math_op") or "identity",
src_addr=a.addr, dst_addr=head.out.addr,
pe_prefix=pp,
)
src/kernbench/components/builtin/pe_types.py
+6
View File
@@ -54,6 +54,12 @@ class PipelinePlan:
m_tiles: int = 0
k_tiles: int = 0
n_tiles: int = 0
# Flat-ops (ADR-0065 D3): single-shot MATH stages placed before / after
# the GEMM tile loop (recipe prologue / KERNEL-scope post-loop ops).
# Informational — execution folds these into the first / last tile so
# the feeder + completion counting stay unchanged.
prologue_stages: tuple[Stage, ...] = ()
epilogue_stages: tuple[Stage, ...] = ()
# ── Pipeline Context ─────────────────────────────────────────────────
src/kernbench/components/builtin/tiling.py
+85 -1
View File
@@ -5,7 +5,7 @@ Ported from pe_accel tiling.py with stage-based plan structure.
"""
from __future__ import annotations
from math import ceil
from math import ceil, prod
from kernbench.components.builtin.pe_types import (
PipelinePlan,
@@ -218,3 +218,87 @@ def generate_math_plan(
tile_id += 1
return PipelinePlan(tiles=tiles, m_tiles=M_tiles, n_tiles=N_tiles)
def _math_stage(op: object, pe_prefix: str) -> Stage:
"""Single-shot MATH stage for a KERNEL-scope flat op (ADR-0065 D3)."""
out = getattr(op, "out", None)
num_elements = prod(out.shape) if out is not None else 1
return Stage(
StageType.MATH, f"{pe_prefix}.pe_math",
{"op_kind": op.kind, "num_elements": num_elements, "scope": "kernel"},
)
def generate_plan_from_ops(
ops: tuple,
tile_m: int, tile_k: int, tile_n: int,
bytes_per_element: int,
pe_prefix: str,
) -> PipelinePlan:
"""Generate a PipelinePlan from a flat ops list (ADR-0065 D3).
Scans for the single GEMM op (≤1). Pre-GEMM KERNEL ops become single-shot
prologue MATH stages; post-GEMM ops are split by scope — K_TILE/
OUTPUT_TILE feed the GEMM tile loop's epilogue (as today), KERNEL ops
become post-loop stages. A composite with no GEMM op is MATH-only.
DMA insertion keeps the existing ``pinned`` signal (operand already in
TCM → no DMA_READ). Prologue / post-loop stages are folded into the
first / last tile so the feeder and completion counting stay unchanged.
"""
from kernbench.common.pe_commands import Scope as _Scope
gemm_idx = next((i for i, o in enumerate(ops) if o.kind == "gemm"), None)
# MATH-only composite (no GEMM): reproduce the legacy math-head plan.
if gemm_idx is None:
head = ops[0]
a = head.operands["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=tile_m, tile_n=tile_n,
bytes_per_element=bytes_per_element,
math_op=head.extra.get("math_op") or "identity",
src_addr=a.addr, dst_addr=head.out.addr, pe_prefix=pe_prefix,
)
head = ops[gemm_idx]
pre_ops = ops[:gemm_idx]
post_ops = ops[gemm_idx + 1:]
a = head.operands["a"]
b = head.operands["b"]
M, K = a.shape[-2], a.shape[-1]
N = b.shape[-1]
plan = generate_gemm_plan(
M=M, K=K, N=N,
tile_m=tile_m, tile_k=tile_k, tile_n=tile_n,
bytes_per_element=bytes_per_element,
A_addr=a.addr, B_addr=b.addr, C_addr=head.out.addr,
pe_prefix=pe_prefix,
a_pinned=getattr(a, "pinned", False),
b_pinned=getattr(b, "pinned", False),
epilogue_specs=tuple(post_ops),
)
pre_stages = tuple(_math_stage(o, pe_prefix) for o in pre_ops)
post_stages = tuple(_math_stage(o, pe_prefix) for o in post_ops
if o.scope == _Scope.KERNEL)
plan.prologue_stages = pre_stages
plan.epilogue_stages = post_stages
# Fold prologue/post-loop stages into the first/last tile so the feeder
# and completion counting are untouched (existing benches have neither,
# so their tiles are unchanged → byte-equal op_log).
if pre_stages and plan.tiles:
t0 = plan.tiles[0]
plan.tiles[0] = TilePlan(tile_id=t0.tile_id,
stages=(*pre_stages, *t0.stages))
if post_stages and plan.tiles:
tl = plan.tiles[-1]
plan.tiles[-1] = TilePlan(tile_id=tl.tile_id,
stages=(*tl.stages, *post_stages))
return plan
src/kernbench/triton_emu/tl_context.py
+5 -1
View File
@@ -934,9 +934,13 @@ class TLContext:
nbytes = self._nbytes(shape, dtype)
addr = self._scratch_alloc(nbytes)
# pinned=True: recipe scratch / primary-out live in TCM already, so
# the head GEMM's auto-bound `a` (= primary-out P) is consumed in
# place — no DMA_READ (ADR-0065 P3, pinned-based DMA decision).
return TensorHandle(
id=self._next_handle_id(),
addr=addr, shape=shape, dtype=dtype, nbytes=nbytes, space="tcm",
addr=addr, shape=shape, dtype=dtype, nbytes=nbytes,
space="tcm", pinned=True,
)
@staticmethod