From 55f025c4b178a46718f2b961089fddf3d04a8a46 Mon Sep 17 00:00:00 2001 From: Yangwook Date: Wed, 10 Jun 2026 20:24:33 -0700 Subject: [PATCH] =?UTF-8?q?gqa(adr-0065):=20P3=20=E2=80=94=20flat-ops=20PE?= =?UTF-8?q?=5FSCHEDULER=20plan=20(position-scan=20+=20prologue=20stages)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- ...rog-flat-ops-composite-softmax-merge-ko.md | 52 ++++--- ...5-prog-flat-ops-composite-softmax-merge.md | 61 +++++---- .../components/builtin/pe_scheduler.py | 54 ++------ src/kernbench/components/builtin/pe_types.py | 6 + src/kernbench/components/builtin/tiling.py | 86 +++++++++++- src/kernbench/triton_emu/tl_context.py | 6 +- tests/test_pe_scheduler_flat_ops_plan.py | 129 ++++++++++++++++++ 7 files changed, 304 insertions(+), 90 deletions(-) create mode 100644 tests/test_pe_scheduler_flat_ops_plan.py diff --git a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md index ff1d7c5..172ce6f 100644 --- a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md +++ b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md @@ -115,24 +115,32 @@ PE_SCHEDULER 가 `cmd.ops` 에서 GEMM op 검색 (composite 당 ≤ 1, D6 참조 GEMM op 가 없는 composite (예: MATH-only) 는 모든 op 를 KERNEL-scope 직렬로 처리 — phase 가 "순서대로 single-shot" 으로 collapse. -### D4. PE_SCHEDULER 가 operand `space` 보고 DMA 자동 삽입 +### D4. PE_SCHEDULER 가 operand `pinned` 플래그 보고 DMA 자동 삽입 -PE_SCHEDULER 가 GEMM 의 `operands` 와 `out` 검사. 각각: -- `space == "hbm"` → FETCH/GEMM Stage 앞에 DMA_READ Stage 삽입 (`out` 은 DMA_WRITE) -- `space == "tcm"` → DMA Stage 없음, in-place 소비 +PE_SCHEDULER 가 GEMM 의 `operands` 검사. 각각: +- **not `pinned`** (데이터가 HBM 에 있어 스트리밍 필요) → FETCH/GEMM Stage + 앞에 DMA_READ Stage 삽입 +- **`pinned`** (이전 `tl.load` 로 이미 TCM 에 staged, 또는 recipe 의 TCM + scratch / primary-out `P`) → DMA Stage 없음, in-place 소비 -오늘 이미 `a_pinned`/`b_pinned` 로 부분 모델링; D4 가 룰을 균일화 + -가시화: **커널은 composite operand 의 명시적 DMA cmd 를 절대 emit 안 함; -PE_SCHEDULER 가 핸들에서 추론**. `tl.ref(addr, shape)` → `space="hbm"`; -`tl.zeros` / `tl.full` / `tl.load` 결과 → `space="tcm"`. +커널은 composite operand 의 명시적 DMA cmd 를 절대 emit 안 함; +PE_SCHEDULER 가 핸들에서 추론. `tl.ref` operand 는 not pinned (DMA_READ); +`tl.load` 결과와 recipe scratch / primary-out 은 pinned (in-place). -**비 GEMM op 은 TCM-only (Phase 1 한계, TLContext emit 시 강제).** -모든 non-GEMM (MATH) OpSpec 의 모든 operand 는 **반드시** -`space == "tcm"`. 커널이 `space == "hbm"` 핸들을 MATH operand 로 전달하면 -TLContext 가 emit 시 validation error. Phase 1 에서 DMA 자동 삽입을 -GEMM operand + head 출력 핸들로 제한; MATH-with-HBM 경로는 미래 확장 -(명시적 `tl.load` + MATH composite 분리, 또는 새 "DMA-as-prologue" recipe -variant). 이 invariant 는 D6 #7 에 재기술. +> **as-built 노트 (P3).** D4 의 초안은 DMA 결정을 `space` 태그 +> (`hbm`→DMA, `tcm`→in-place) 에서 도출하며 현재 `space` 값 +> (오늘 `tl.load`=`hbm`, `tl.ref`=`tcm`) 을 뒤집어야 했음. 구현은 기존 +> **`pinned` 플래그**를 DMA 신호로 그대로 유지 — `pinned` 이 이미 +> "이미 TCM 에 있음, DMA skip" 을 인코딩하고, `space` 뒤집기는 기존 bench +> Stage 시퀀스를 바꾸면서 기능적 이득이 없음. `space` 통일 + `pinned` +> 폐기는 후순위 cleanup 으로 연기. + +**Prologue recipe op 은 TCM-only (Phase 1 한계, TLContext emit 시 강제).** +*prologue recipe* (비 GEMM) OpSpec 의 모든 operand 는 **반드시** TCM 상주; +HBM 핸들을 recipe operand 로 전달하면 emit 시 validation error. Phase 1 +에서 DMA staging 을 head op 의 operand 로 제한. **head op 자체** (gemm +또는 math) 는 기존 DMA-staged-from-HBM 동작 유지 — D4 의 TCM-only 룰은 +head 에 적용 **안 됨**. 이 invariant 는 D6 #7 에 재기술. ### D5. RECIPE_DESCRIPTORS — TLContext 내부, `pe_commands.py` 가 아님 @@ -210,8 +218,9 @@ op="gemm", ...)` 에서: error. 커널은 recipe 가 바인딩하게 두거나 (operand 생략) **혹은** 명시적으로 지정 (prologue 생략, 또는 `primary_out` 없는 recipe 사용) 해야 함. "어느 값이 이기느냐" 의 모호함 방지. -7. **MATH operand TCM-only.** D4 의 재기술: 모든 비 GEMM OpSpec 의 모든 - operand 는 `space == "tcm"`. Phase 1 에서 TLContext emit 시 강제. +7. **Prologue MATH operand TCM-only.** D4 의 재기술: *prologue recipe* + (비 GEMM) OpSpec 의 모든 operand 는 TCM 상주. head op (gemm 또는 math) + 은 예외 — HBM 스트리밍 허용. Phase 1 에서 TLContext emit 시 강제. ### D7. Boundary 요약 (compiler vs scheduler vs engine) @@ -328,10 +337,11 @@ scratch 에서 읽는 형태 — 순환 tile-loop 의존성. **기각 (incorrect 시작 시에만). 4. **Strict-FIFO RW 직렬화.** `rw_handle` 공유하는 두 연속 composite 가 dispatch 순서대로 완료; stage 가 interleave 안 함. -5. **`space` 에서 DMA 자동 삽입.** `b=tl.ref(...)` (space=hbm) 의 GEMM - composite 가 DMA_READ Stage emit; `b` 가 `tl.zeros(...)` (space=tcm) 면 - emit 안 함. 출력 핸들 `space=hbm` 이면 DMA_WRITE emit; `space=tcm` 이면 - emit 안 함. +5. **`pinned` 에서 DMA 자동 삽입.** not-`pinned` operand (예: `b=tl.ref(...)`) + 의 GEMM composite 가 DMA_READ Stage emit; `pinned` operand (예: recipe + 의 TCM scratch / primary-out, 또는 `tl.load` 결과) 는 emit 안 함. + (as-built D4 노트: DMA 결정은 `space` 태그가 아니라 기존 `pinned` + 플래그 사용.) 6. **opt2 가 opt3 와 수치 동등.** data mode 에서 opt2 의 최종 `(m, l, O)` 가 opt3 와 fp tolerance 안. 7. **opt2 dispatch ratio (ADR-0064 Rev2 이후) — robust.** opt3 vs opt2 diff --git a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md index b9ccad7..8514fdc 100644 --- a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md +++ b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md @@ -125,28 +125,35 @@ D6). Calling its index `g`: A composite with no GEMM op (e.g., MATH-only composite) treats all ops as KERNEL-scope sequential — phase collapses to "single-shot in order". -### D4. PE_SCHEDULER auto-inserts DMAs from operand `space` +### D4. PE_SCHEDULER auto-inserts DMAs from the operand `pinned` flag -PE_SCHEDULER scans the GEMM's `operands` and `out`. For each: -- `space == "hbm"` → emit DMA_READ Stage (or DMA_WRITE for `out`) before - the FETCH/GEMM Stages -- `space == "tcm"` → no DMA Stage, operand consumed in place +PE_SCHEDULER scans the GEMM's `operands`. For each: +- **not `pinned`** (data lives in HBM, must be streamed) → emit DMA_READ + Stage before the FETCH/GEMM Stages +- **`pinned`** (already staged in TCM via a prior `tl.load`, or a recipe's + TCM scratch / primary-out `P`) → no DMA Stage, consumed in place -This is already partially modelled today via `a_pinned`/`b_pinned`; D4 -makes the rule uniform and visible: **the kernel never emits explicit -DMA cmds for composite operands; PE_SCHEDULER infers them from -handles**. `tl.ref(addr, shape)` returns a handle with `space="hbm"`; -`tl.zeros` / `tl.full` / `tl.load` outputs return `space="tcm"`. +The kernel never emits explicit DMA cmds for composite operands; +PE_SCHEDULER infers them from handles. `tl.ref` operands are not pinned +(DMA_READ); `tl.load` results and recipe scratch / primary-out are pinned +(in place). -**Non-GEMM ops are TCM-only (Phase 1 limit, enforced at TLContext emit).** -Every operand of every non-GEMM (MATH) OpSpec **must** have -`space == "tcm"`. If a kernel passes a handle with `space == "hbm"` as a -MATH operand, TLContext raises a validation error at emit time. This -keeps DMA auto-insertion bounded to GEMM operands and head output -handles in Phase 1; a MATH-with-HBM path would require a future -extension (split into explicit `tl.load` + MATH composite, or a new -"DMA-as-prologue" recipe variant). This invariant is also restated in -D6 #7. +> **As-built note (P3).** An earlier draft of D4 derived the DMA decision +> from a `space` tag (`hbm`→DMA, `tcm`→in place) and would have flipped the +> current `space` values (today `tl.load`=`hbm`, `tl.ref`=`tcm`). The +> implementation **keeps the existing `pinned` flag** as the DMA signal +> instead — `pinned` already encodes "already in TCM, skip DMA", flipping +> `space` would change existing bench Stage sequences for no functional +> gain. Unifying `space` and retiring `pinned` is a deferred low-priority +> cleanup. + +**Prologue recipe ops are TCM-only (Phase 1 limit, enforced at TLContext +emit).** Every operand of a *prologue recipe* (non-GEMM) OpSpec **must** be +TCM-resident; passing an HBM handle as a recipe operand raises a validation +error at emit time. This bounds DMA staging to the head op's operands in +Phase 1. The **head op itself** (gemm *or* math) keeps the existing +DMA-staged-from-HBM behavior — D4's TCM-only rule does **not** apply to it. +This invariant is also restated in D6 #7. ### D5. RECIPE_DESCRIPTORS — TLContext-internal, NOT in `pe_commands.py` @@ -236,9 +243,10 @@ sees only the flat ops list. kernel must either let the recipe bind (omit the operand) **or** specify it explicitly (omit the prologue, or use a recipe without `primary_out`). This prevents the ambiguity of "which value wins". -7. **MATH operands TCM-only.** Restatement of D4: every operand of a - non-GEMM OpSpec must have `space == "tcm"`. Phase 1 enforced at - TLContext emit. +7. **Prologue MATH operands TCM-only.** Restatement of D4: every operand + of a *prologue recipe* (non-GEMM) OpSpec must be TCM-resident. The head + op (gemm or math) is exempt — it may stream from HBM. Phase 1 enforced + at TLContext emit. ### D7. Boundary summary (compiler vs scheduler vs engine) @@ -368,10 +376,11 @@ input from #1's epilogue scratch — circular tile-loop dependency. 4. **Strict-FIFO RW serialization.** Two consecutive composites sharing any `rw_handle` complete in dispatch order; their stages do not interleave. -5. **DMA auto-insertion from `space`.** A GEMM composite with - `b=tl.ref(...)` (space=hbm) emits a DMA_READ Stage; with `b` from - `tl.zeros(...)` (space=tcm) it does not. Output handle with - `space=hbm` emits DMA_WRITE; `space=tcm` does not. +5. **DMA auto-insertion from `pinned`.** A GEMM composite with a + not-`pinned` operand (e.g. `b=tl.ref(...)`) emits a DMA_READ Stage; a + `pinned` operand (e.g. a recipe's TCM scratch / primary-out, or a + `tl.load` result) does not. (As-built D4 note: the DMA decision uses + the existing `pinned` flag, not a `space` tag.) 6. **opt2 numeric parity with opt3.** In data mode, opt2's final `(m, l, O)` matches opt3's within fp tolerance. 7. **opt2 dispatch ratio (after ADR-0064 Rev2) — robust.** opt3 vs diff --git a/src/kernbench/components/builtin/pe_scheduler.py b/src/kernbench/components/builtin/pe_scheduler.py index 6b6e3ec..b72f4e8 100644 --- a/src/kernbench/components/builtin/pe_scheduler.py +++ b/src/kernbench/components/builtin/pe_scheduler.py @@ -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, - ) diff --git a/src/kernbench/components/builtin/pe_types.py b/src/kernbench/components/builtin/pe_types.py index 64cfb26..8cd79b9 100644 --- a/src/kernbench/components/builtin/pe_types.py +++ b/src/kernbench/components/builtin/pe_types.py @@ -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 ───────────────────────────────────────────────── diff --git a/src/kernbench/components/builtin/tiling.py b/src/kernbench/components/builtin/tiling.py index 0704ea3..519d1f8 100644 --- a/src/kernbench/components/builtin/tiling.py +++ b/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 diff --git a/src/kernbench/triton_emu/tl_context.py b/src/kernbench/triton_emu/tl_context.py index ca6d36b..c25cf2b 100644 --- a/src/kernbench/triton_emu/tl_context.py +++ b/src/kernbench/triton_emu/tl_context.py @@ -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 diff --git a/tests/test_pe_scheduler_flat_ops_plan.py b/tests/test_pe_scheduler_flat_ops_plan.py new file mode 100644 index 0000000..06b2535 --- /dev/null +++ b/tests/test_pe_scheduler_flat_ops_plan.py @@ -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)