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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user