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

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>

tests/test_composite_output_handle.py

@@ -0,0 +1,96 @@
+"""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"

tests/test_triton_emu.py

@@ -175,7 +175,10 @@ def test_tl_composite_nonblocking():
     a = tl.load(0x1000, shape=(32, 64), dtype="f16")
     b = tl.load(0x2000, shape=(64, 32), dtype="f16")
     h = tl.composite(op="gemm", a=a, b=b, out_ptr=0x3000)
-    assert isinstance(h, CompletionHandle)
+    # ADR-0065 D8: composite returns the output TensorHandle (chainable),
+    # carrying the completion in its `pending` field.
+    assert isinstance(h, TensorHandle)
+    assert h.addr == 0x3000 and h.space == "hbm"
     comp_cmds = [c for c in tl.commands if isinstance(c, CompositeCmd)]
     assert len(comp_cmds) == 1
     # Flat-ops shape (ADR-0065 D1): head op carries kind + write-back handle.
@@ -195,7 +198,8 @@ def test_tl_wait_specific():
     tl.wait(h)
     wait_cmds = [c for c in tl.commands if isinstance(c, WaitCmd)]
     assert len(wait_cmds) == 1
-    assert wait_cmds[0].handle == h
+    # ADR-0065 D8: h is the output handle; tl.wait targets its completion.
+    assert wait_cmds[0].handle == h.pending.completion
 
 
 # ── 9. tl.wait() → WaitCmd(handle=None) ──────────────────────────