ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa(adr-0065): D8 — composite returns output handle + output-space DMA

Browse Source 
tl.composite now returns the output TensorHandle (not a CompletionHandle) so its result chains like tl.dot's; the handle carries the completion in a CompositeFuture pending so downstream ops and tl.wait auto-await it. out is a handle: out=tl.ref(addr,shape) (HBM, DMA_WRITE inside the composite) or an in-place TCM handle (STORE only); omitted -> TLContext auto-allocates a TCM scratch. out_ptr kept as HBM shorthand (= out=tl.ref(out_ptr, shape)) to avoid churning ~30 existing call sites.

tl.ref now returns space=hbm (it references HBM data; operand-input DMA stays pinned-based per D4 so input streaming is unchanged). tiling: the tile loop's DMA_WRITE is gated on out.space==hbm (out analog of the operand pinned rule) — a TCM output stays on-chip (chainable) and its high-bit scratch address no longer hits the DMA PA decoder.

Fixes the opt2 data-mode crash: the recipe accumulator O is TCM -> no DMA_WRITE -> opt2 now RUNS end-to-end in data mode (enable_data=True). Numeric parity of the recipe MATH ops is the next step. Suite 812 pass / 3 pre-existing fail; existing composite benches use out_ptr->hbm->DMA_WRITE unchanged (byte-equal).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook
2026-06-10 22:29:10 -07:00
parent dd525bfcb7
commit e8d6c283d8
7 changed files with 267 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md
+40
View File
@@ -267,6 +267,46 @@ TLContext (compiler analog) PE_SCHEDULER (scheduler/dispatcher)
imports: nothing recipe-related
```
### D8. Composite output handle + output-space DMA
`tl.composite(...)` returns the **output `TensorHandle`** (not a bare
`CompletionHandle`), so a composite's result can feed a downstream op the
same way `tl.dot`'s result does. The handle's `pending` field references
the composite's completion; downstream consumers auto-await it
(`_await_pending`, ADR-0062 lazy pattern) and `tl.wait(handle)` works.
**`out` is always a `TensorHandle`** (never a raw address) — its location is
the kernel's choice, never guessed:
- **HBM output** → `out=tl.ref(addr, shape)`. `tl.ref` returns a handle with
`space="hbm"`, so the composite writes it back via DMA_WRITE. The
STORE + DMA_WRITE stay **inside** the composite pipeline — they are part of
the fused op, *not* a separate `tl.store`.
- **In-place TCM** → `out=<existing TCM handle>` (e.g. the recipe accumulator
`O`). `space="tcm"` → STORE only, result stays in TCM (chainable).
- **Omitted** → TLContext **auto-allocates a TCM scratch** (like `tl.dot`)
and returns it.
For ergonomics, `out_ptr: int` is accepted as **shorthand for an HBM
output** — equivalent to `out=tl.ref(out_ptr, <output shape>)` (the compiler
wraps it into an `space="hbm"` handle). When no output is specified the
compiler owns the scratch address, never the kernel. (`tl.ref` returns
`space="hbm"` — it references HBM data. The operand-**input** DMA decision
stays `pinned`-based per D4, so this label does not affect input streaming.)
The output handle's **`space` drives the tile loop's write-back** (the
output analog of the operand `pinned` rule in D4):
| `out.space` | tile-loop write-back |
| --- | --- |
| `"tcm"` | STORE only — result stays in TCM (chainable; **no DMA_WRITE**) |
| `"hbm"` | STORE + DMA_WRITE — final HBM output |
This makes a TCM-resident composite output (e.g. a recipe accumulator) legal
**and** keeps a TCM-scratch address out of the DMA engine's PA decoder
(which would otherwise reject the high-bit scratch address). Existing benches
wrap their HBM output via `tl.ref` → write-back unchanged (byte-equal).
## Alternatives
### A1. Three composites per tile (no prologue concept)