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

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>
2026-06-10 20:24:33 -07:00
parent 184f654295
commit 55f025c4b1
7 changed files with 304 additions and 90 deletions
@@ -0,0 +1,129 @@
+"""Phase 1 spec tests for ADR-0065 P3 — flat-ops PE_SCHEDULER plan.
+
+P3 rewrites tile-plan generation to scan the flat ops list for the GEMM
+(≤1) and place pre-GEMM KERNEL ops as single-shot prologue MATH stages
+(ADR-0065 D3). A new tiling entry ``generate_plan_from_ops(ops, ...)``
+produces a ``PipelinePlan`` with ``prologue_stages`` / ``epilogue_stages``.
+
+DMA insertion keeps the existing **``pinned``** signal (already-in-TCM →
+skip DMA; not-pinned → stream from HBM), NOT a ``space`` flip — so existing
+bench Stage sequences stay byte-equal. Recipe TCM scratch / primary-out
+handles are marked ``pinned=True`` so the head GEMM's auto-bound ``a`` (= the
+recipe's ``P``) is consumed in place, while ``b`` (a ``tl.ref``) is streamed.
+
+These tests target the tiling entry directly (no engine), building OpSpecs
+by hand so each operand's ``pinned`` flag is controlled precisely.
+
+Phase 1 (this commit): tests only. FAIL until P3:
+  - ``generate_plan_from_ops`` does not exist,
+  - ``PipelinePlan`` has no ``prologue_stages`` field.
+"""
+from __future__ import annotations
+
+import math
+
+from kernbench.common.pe_commands import OpSpec, Scope, TensorHandle
+
+PE = "sip0.cube0.pe0"
+TM = TK = TN = 64
+
+
+def _h(addr: int, shape: tuple[int, ...], pinned: bool = False) -> TensorHandle:
+    return TensorHandle(
+        id=f"h{addr:x}", addr=addr, shape=shape, dtype="f16",
+        nbytes=2 * math.prod(shape), space="tcm", pinned=pinned,
+    )
+
+
+def _gemm_op(a, b, out) -> OpSpec:
+    return OpSpec(
+        kind="gemm", scope=Scope.OUTPUT_TILE,
+        operands={"a": a, "b": b}, out=out,
+        extra={"m": a.shape[-2], "k": a.shape[-1], "n": b.shape[-1]},
+    )
+
+
+def _plan(ops):
+    from kernbench.components.builtin.tiling import generate_plan_from_ops
+
+    return generate_plan_from_ops(
+        ops=tuple(ops), tile_m=TM, tile_k=TK, tile_n=TN,
+        bytes_per_element=2, pe_prefix=PE,
+    )
+
+
+def _stage_types(stages):
+    return [s.stage_type.name for s in stages]
+
+
+# ── DMA from the pinned flag (byte-equal with generate_gemm_plan) ─────
+
+
+def test_unpinned_operand_gets_dma_read():
+    from kernbench.components.builtin.pe_types import StageType
+
+    a = _h(0x1000, (TM, TK), pinned=True)            # already in TCM
+    b = _h(0x2000, (TK, TN), pinned=False)           # streamed from HBM
+    out = _h(0x3000, (TM, TN))
+    tile = _plan([_gemm_op(a, b, out)]).tiles[0]
+    reads = [s for s in tile.stages if s.stage_type == StageType.DMA_READ]
+    assert len(reads) == 1, _stage_types(tile.stages)   # only B
+
+
+def test_pinned_operands_skip_dma():
+    from kernbench.components.builtin.pe_types import StageType
+
+    a = _h(0x1000, (TM, TK), pinned=True)
+    b = _h(0x2000, (TK, TN), pinned=True)
+    out = _h(0x3000, (TM, TN))
+    tile = _plan([_gemm_op(a, b, out)]).tiles[0]
+    assert not [s for s in tile.stages if s.stage_type == StageType.DMA_READ]
+
+
+# ── prologue placement (ADR-0065 D3) ─────────────────────────────────
+
+
+def test_recipe_prologue_math_ops_become_prologue_stages():
+    from kernbench.components.builtin.pe_types import StageType
+
+    # 8 KERNEL-scope MATH ops, then the head GEMM whose `a` is the recipe's
+    # primary-out P (pinned → in place) and `b` is a streamed ref.
+    sc = _h(0x100000, (8, 64), pinned=True)
+    pre = [
+        OpSpec(kind=k, scope=Scope.KERNEL, operands={"src": sc}, out=sc)
+        for k in ("rmax", "max_elem", "exp_diff", "exp_diff",
+                  "rsum", "fma", "mul_bcast", "copy")
+    ]
+    P = _h(0x100400, (8, 64), pinned=True)           # primary-out, in TCM
+    V = _h(0x2000, (64, 128), pinned=False)          # streamed
+    out = _h(0x3000, (8, 128))
+    plan = _plan([*pre, _gemm_op(P, V, out)])
+
+    # 8 single-shot MATH prologue stages, none of them DMA.
+    assert len(plan.prologue_stages) == 8, _stage_types(plan.prologue_stages)
+    assert all(s.stage_type == StageType.MATH for s in plan.prologue_stages)
+    assert not any(s.stage_type in (StageType.DMA_READ, StageType.DMA_WRITE)
+                   for s in plan.prologue_stages)
+    # Head GEMM drives the tile loop; only V (not pinned) gets a DMA_READ.
+    reads = [s for s in plan.tiles[0].stages
+             if s.stage_type == StageType.DMA_READ]
+    assert len(reads) == 1, _stage_types(plan.tiles[0].stages)
+
+
+def test_math_only_composite_has_no_gemm_tiles():
+    """A legacy MATH composite (no GEMM op) → math tile plan, no GEMM stage.
+
+    Shaped like ``tl.composite(op="math", a=..., math_op="exp")`` — head
+    operand keyed ``"a"`` with the op kind in ``extra["math_op"]``."""
+    from kernbench.components.builtin.pe_types import StageType
+
+    src = _h(0x1000, (TM, TN), pinned=False)
+    out = _h(0x3000, (TM, TN))
+    op = OpSpec(kind="math", scope=Scope.OUTPUT_TILE,
+                operands={"a": src}, extra={"math_op": "exp"}, out=out)
+    plan = _plan([op])
+    all_stages = []
+    for t in plan.tiles:
+        all_stages.extend(t.stages)
+    assert plan.tiles, "math-only composite must still produce tiles"
+    assert not any(s.stage_type == StageType.GEMM for s in all_stages)