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>
This commit is contained in:
@@ -24,6 +24,7 @@ def generate_gemm_plan(
|
||||
a_pinned: bool = False,
|
||||
b_pinned: bool = False,
|
||||
epilogue_specs: tuple = (),
|
||||
out_space: str = "hbm",
|
||||
) -> PipelinePlan:
|
||||
"""Generate GEMM tile plan: M→N→K order.
|
||||
|
||||
@@ -152,13 +153,18 @@ def generate_gemm_plan(
|
||||
"nbytes": out_bytes,
|
||||
},
|
||||
))
|
||||
stages.append(Stage(
|
||||
stage_type=StageType.DMA_WRITE,
|
||||
component=dma_id,
|
||||
params={
|
||||
"dst_addr": c_addr, "nbytes": out_bytes,
|
||||
},
|
||||
))
|
||||
# DMA_WRITE only when the output is HBM-resident
|
||||
# (ADR-0065 D8). A TCM output stays on-chip (chainable);
|
||||
# writing its high-bit scratch address as an HBM PA would
|
||||
# also fail the DMA decoder.
|
||||
if out_space == "hbm":
|
||||
stages.append(Stage(
|
||||
stage_type=StageType.DMA_WRITE,
|
||||
component=dma_id,
|
||||
params={
|
||||
"dst_addr": c_addr, "nbytes": out_bytes,
|
||||
},
|
||||
))
|
||||
|
||||
tiles.append(TilePlan(tile_id=tile_id, stages=tuple(stages)))
|
||||
tile_id += 1
|
||||
@@ -281,6 +287,7 @@ def generate_plan_from_ops(
|
||||
a_pinned=getattr(a, "pinned", False),
|
||||
b_pinned=getattr(b, "pinned", False),
|
||||
epilogue_specs=tuple(post_ops),
|
||||
out_space=getattr(head.out, "space", "hbm"),
|
||||
)
|
||||
|
||||
# Prologue (pre-GEMM KERNEL) + post-loop KERNEL ops become single-shot
|
||||
|
||||
Reference in New Issue
Block a user