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>
This commit is contained in:
2026-06-10 20:24:33 -07:00
parent 184f654295
commit 55f025c4b1
7 changed files with 304 additions and 90 deletions
@@ -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
@@ -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