ywkang
e8d6c283d8
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>
97 lines
3.8 KiB
Python
97 lines
3.8 KiB
Python
"""Phase 1 spec tests for ADR-0065 D8 — composite output handle + output-space DMA.
|
|
|
|
`tl.composite(...)` returns the **output TensorHandle** (not a bare
|
|
CompletionHandle) so its result chains into downstream ops like `tl.dot`'s.
|
|
`out` is always a handle: `out=tl.ref(addr, shape)` for HBM output (DMA_WRITE
|
|
inside the composite), `out=<TCM handle>` for in-place, or omitted → TLContext
|
|
auto-allocates a TCM scratch. The output handle's `space` decides the tile
|
|
loop's write-back (tcm → STORE only; hbm → STORE + DMA_WRITE). `tl.ref`
|
|
returns `space="hbm"`.
|
|
|
|
Phase 1 (this commit): tests only. FAIL until D8:
|
|
- `tl.ref` returns space="tcm"; `composite` returns CompletionHandle;
|
|
output always DMA_WRITEs; `tl.wait` rejects a TensorHandle.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from kernbench.common.pe_commands import OpSpec, Scope, TensorHandle
|
|
from kernbench.triton_emu.tl_context import TLContext
|
|
|
|
|
|
def _tl() -> TLContext:
|
|
return TLContext(pe_id=0, num_programs=1, scratch_base=0x200000,
|
|
scratch_size=1 << 20)
|
|
|
|
|
|
def test_ref_returns_hbm_space():
|
|
h = _tl().ref(0x1000, shape=(4, 4), dtype="f16")
|
|
assert h.space == "hbm"
|
|
|
|
|
|
def test_composite_returns_handle_auto_tcm():
|
|
tl = _tl()
|
|
a = tl.ref(0x1000, shape=(8, 16), dtype="f16")
|
|
b = tl.ref(0x2000, shape=(16, 8), dtype="f16")
|
|
out = tl.composite(op="gemm", a=a, b=b) # out omitted → auto TCM
|
|
assert isinstance(out, TensorHandle)
|
|
assert out.space == "tcm"
|
|
assert out.addr >= 0x200000 # auto-allocated scratch
|
|
assert out.shape == (8, 8)
|
|
|
|
|
|
def test_composite_returns_the_explicit_out_handle():
|
|
tl = _tl()
|
|
a = tl.ref(0x1000, shape=(8, 16), dtype="f16")
|
|
b = tl.ref(0x2000, shape=(16, 8), dtype="f16")
|
|
O = TensorHandle(id="O", addr=0x5000, shape=(8, 8), dtype="f16",
|
|
nbytes=128, space="tcm", pinned=True)
|
|
out = tl.composite(op="gemm", a=a, b=b, out=O)
|
|
assert out is O
|
|
|
|
|
|
def test_composite_output_chains_into_next_op():
|
|
tl = _tl()
|
|
a = tl.ref(0x1000, shape=(8, 16), dtype="f16")
|
|
b = tl.ref(0x2000, shape=(16, 8), dtype="f16")
|
|
scores = tl.composite(op="gemm", a=a, b=b) # handle
|
|
out2 = tl.composite(op="math", a=scores, math_op="exp")
|
|
assert isinstance(out2, TensorHandle)
|
|
|
|
|
|
def test_wait_accepts_tensor_handle():
|
|
tl = _tl()
|
|
a = tl.ref(0x1000, shape=(8, 16), dtype="f16")
|
|
b = tl.ref(0x2000, shape=(16, 8), dtype="f16")
|
|
out = tl.composite(op="gemm", a=a, b=b)
|
|
tl.wait(out) # must not raise
|
|
|
|
|
|
def test_output_space_drives_dma_write():
|
|
from kernbench.components.builtin.pe_types import StageType
|
|
from kernbench.components.builtin.tiling import generate_plan_from_ops
|
|
|
|
a = TensorHandle(id="a", addr=0x1000, shape=(64, 64), dtype="f16",
|
|
nbytes=8192, space="tcm", pinned=True)
|
|
b = TensorHandle(id="b", addr=0x2000, shape=(64, 64), dtype="f16",
|
|
nbytes=8192, space="hbm")
|
|
|
|
def _gemm(out):
|
|
return OpSpec(kind="gemm", scope=Scope.OUTPUT_TILE,
|
|
operands={"a": a, "b": b}, out=out,
|
|
extra={"m": 64, "k": 64, "n": 64})
|
|
|
|
o_tcm = TensorHandle(id="ot", addr=0x300000, shape=(64, 64), dtype="f16",
|
|
nbytes=8192, space="tcm", pinned=True)
|
|
o_hbm = TensorHandle(id="oh", addr=0x9000, shape=(64, 64), dtype="f16",
|
|
nbytes=8192, space="hbm")
|
|
|
|
def _has_write(ops):
|
|
plan = generate_plan_from_ops(ops, tile_m=64, tile_k=64, tile_n=64,
|
|
bytes_per_element=2,
|
|
pe_prefix="sip0.cube0.pe0")
|
|
return any(s.stage_type == StageType.DMA_WRITE
|
|
for s in plan.tiles[0].stages)
|
|
|
|
assert not _has_write((_gemm(o_tcm),)), "tcm out → no DMA_WRITE"
|
|
assert _has_write((_gemm(o_hbm),)), "hbm out → DMA_WRITE"
|