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>

src/kernbench/components/builtin/tiling.py

@@ -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