diff --git a/src/kernbench/triton_emu/tl_context.py b/src/kernbench/triton_emu/tl_context.py index 773f75f..a6d7ecf 100644 --- a/src/kernbench/triton_emu/tl_context.py +++ b/src/kernbench/triton_emu/tl_context.py @@ -697,6 +697,10 @@ class TLContext: nbytes=self._nbytes(shape, dtype), data=data, space=slot_space, + # On-chip (tcm/sram) IPCQ slot: read in place as a composite + # operand, not DMA_READ of its non-HBM slot address. An HBM + # slot is a valid PA, so it stays unpinned (DMA streams it). + pinned=slot_space != "hbm", ) return self._make_handle(addr=0, shape=shape, dtype=dtype) @@ -738,6 +742,7 @@ class TLContext: nbytes=self._nbytes(shape, dtype), data=None, space=slot_space, + pinned=slot_space != "hbm", ) return self._make_handle(addr=0, shape=shape, dtype=dtype) @@ -1105,6 +1110,7 @@ class TLContext: nbytes=self._nbytes(handle.cmd.shape, handle.cmd.dtype), data=data, space=slot_space, + pinned=slot_space != "hbm", ) handle.resolved = True handle.result = th diff --git a/tests/test_composite_recv_operand.py b/tests/test_composite_recv_operand.py new file mode 100644 index 0000000..0d4b297 --- /dev/null +++ b/tests/test_composite_recv_operand.py @@ -0,0 +1,75 @@ +"""Regression: a recv handle used as a composite ``b`` operand must be read +in place (on-chip), NOT DMA_READ from its non-HBM slot address. + +Before the recv-pinning fix, recv handles were unpinned → the composite +scheduler emitted DMA_READ(B) from the IPCQ slot address (non-HBM) → +PhysAddr.decode raised PhysAddrError. On-chip (TCM/SRAM) recv slots are +now pinned (ADR-0065 D4), so a composite operand reads them in place. + +Isolates operand ``b``: ``a`` is a pinned ``tl.load``; only ``b`` is the +recv handle. Mirrors the decode row-chain IPCQ directions (intra_E recv / +intra_W send) so the recv actually lands without deadlock. +""" +from __future__ import annotations + +from pathlib import Path + +from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config +from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip +from kernbench.policy.placement.dp import DPPolicy +from kernbench.runtime_api.bench_runner import run_bench +from kernbench.runtime_api.types import resolve_device +from kernbench.sim_engine.engine import GraphEngine +from kernbench.topology.builder import resolve_topology + +TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[1] / "topology.yaml" +P = 8 + + +def _ccl_cfg(): + return resolve_algorithm_config( + load_ccl_config(), name="lrab_hierarchical_allreduce", + ) + + +def _engine_factory(t, d): + return GraphEngine(getattr(t, "topology_obj", t), enable_data=True) + + +def _recv_operand_kernel(a_ptr, b_ptr, o_ptr, M, K, N, P, *, tl): + pe_id = tl.program_id(axis=0) + pe_in_group = pe_id % P + group_cols = min(4, P) + pe_col = pe_in_group % group_cols + + A = tl.load(a_ptr, shape=(M, K), dtype="f16") # pinned (no DMA) + if pe_col < group_cols - 1: + B = tl.recv(dir="intra_E", shape=(K, N), dtype="f16") # recv handle + out = tl.composite("gemm", a=A, b=B) # ← b = recv slot + tl.store(o_ptr, out) + if pe_col > 0: + B_send = tl.load(b_ptr, shape=(K, N), dtype="f16") + tl.send(dir="intra_W", src=B_send) + + +def test_recv_handle_as_composite_operand_reads_in_place(): + M = K = N = 64 + topo = resolve_topology(str(TOPOLOGY_DEFAULT)) + + def _bench_fn(ctx): + configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg()) + dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=P) + a = ctx.zeros((M, K), dtype="f16", dp=dp, name="recv_operand_a") + b = ctx.zeros((K, N), dtype="f16", dp=dp, name="recv_operand_b") + o = ctx.empty((M, N), dtype="f16", dp=dp, name="recv_operand_o") + ctx.launch( + "recv_operand_composite", _recv_operand_kernel, + a, b, o, M, K, N, P, + _auto_dim_remap=False, + ) + + r = run_bench( + topology=topo, bench_fn=_bench_fn, + device=resolve_device(None), engine_factory=_engine_factory, + ) + assert r.completion.ok, f"recv-operand composite completion: {r.completion}"