sim: pin compute outputs and on-chip recv slots for composite operands

2026-06-24 15:33:31 -07:00
parent 1971ac731c
commit a455ae61b4
2 changed files with 81 additions and 0 deletions
@@ -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}"