"""TLContext: fake Triton Language module for kernel performance simulation. Passed as the `tl` parameter to kernel functions. Each API call records a PeCommand in the internal trace. After the kernel returns, PE_CPU replays the command list through SimPy. Kernel code looks like standard Python — no yield, no async: def my_kernel(a_ptr, b_ptr, out_ptr, tl): pid = tl.program_id(0) a = tl.load(a_ptr, shape=(32, 64), dtype="f16") b = tl.load(b_ptr + pid * stride, shape=(64, 32), dtype="f16") tl.composite(op="gemm", a=a, b=b, out_ptr=out_ptr) """ from __future__ import annotations import math from typing import Literal from kernbench.common.ipcq_types import IpcqRecvCmd, IpcqSendCmd, RecvFuture from kernbench.common.pe_commands import ( EPILOGUE_OPS, CompletionHandle, CompositeCmd, CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd, OpSpec, PeCommand, PeCpuOverheadCmd, Scope, TensorHandle, WaitCmd, ) _DTYPE_BYTES: dict[str, int] = { "f16": 2, "f32": 4, "f64": 8, "bf16": 2, "i8": 1, "i16": 2, "i32": 4, "i64": 8, "u8": 1, "u16": 2, "u32": 4, "u64": 8, } class LoadFuture: """Lazy ``tl.load`` future (ADR-0062 §D2). Mirrors ``RecvFuture`` for IPCQ, generalised to HBM loads. Carries the originating ``DmaReadCmd``, the SimPy completion event (set by the runner once the DMA has been issued), and a resolved flag. Consumer ops auto-wait via ``TLContext._await_pending(handle)``, which yields ``event`` if ``resolved`` is False, then reads the DMA's bytes into ``handle.data`` and marks the future resolved. """ __slots__ = ("cmd", "event", "resolved", "data") def __init__(self, cmd: DmaReadCmd) -> None: self.cmd = cmd self.event: object | None = None # simpy.Event set by runner self.resolved: bool = False self.data: object = None class CompositeFuture: """Pending marker on a composite's output handle (ADR-0065 D8). Carries the composite's ``CompletionHandle`` so a downstream op (via ``_await_pending``) or ``tl.wait`` can await it. Distinct from ``LoadFuture``: composite completion is tracked by completion id in the runner's pending dict, so awaiting just re-issues a ``WaitCmd`` (a second wait on an already-drained completion is a harmless no-op). """ __slots__ = ("completion", "resolved") def __init__(self, completion: "CompletionHandle") -> None: self.completion = completion self.resolved: bool = False class _ScratchScope: """Context manager that recycles per-tile scratch (ADR-0063 D1). ``__enter__`` snapshots ``_scratch_cursor``; ``__exit__`` restores it, so every handle allocated inside the ``with``-block has its address freed for the next iteration. Persistent state (running ``(m, ℓ, O)``, prefetch buffers) lives outside the scope per ADR-0063 D3. """ def __init__(self, ctx: "TLContext") -> None: self._ctx = ctx self._save: int | None = None def __enter__(self) -> "_ScratchScope": self._save = self._ctx._scratch_cursor return self def __exit__(self, *exc_info: object) -> bool: if self._save is not None: self._ctx._scratch_cursor = self._save return False class TLContext: """Fake Triton Language context. Args: pe_id: program instance index (returned by program_id). num_programs: total number of program instances. dispatch_cycles: uniform PE_CPU overhead per dispatched command. Back-compat fallback used only when ``cost_model`` is None (ADR-0046 §D6). cost_model: structural dispatch cost model (ADR-0064 Rev2). When provided, every dispatched PeCommand is charged ``FIXED + cmd.logical_bytes × R`` cycles (÷ ``clock_freq_ghz`` for ns) as a ``PeCpuOverheadCmd`` emitted just before it. Live PE_CPU paths construct TLContext with the per-PE model so the hybrid's CPU-saturation lever is measurable. clock_freq_ghz: cycle→ns conversion for the cost model (ADR-0064 D3). """ def __init__( self, pe_id: int = 0, num_programs: int = 1, dispatch_cycles: int = 1, runner: Any = None, cube_id: int = 0, num_cubes: int = 1, scratch_base: int = 0, scratch_size: int = 1 << 20, # 1 MiB per kernel invocation cost_model: "PeCostModel | None" = None, clock_freq_ghz: float = 1.0, ) -> None: self._pe_id = pe_id self._num_programs = num_programs self._cube_id = cube_id self._num_cubes = num_cubes self._dispatch_cycles = dispatch_cycles self._cost_model = cost_model self._clock_freq_ghz = clock_freq_ghz self._commands: list[PeCommand] = [] self._handle_counter = 0 self._completion_counter = 0 self._runner = runner # KernelRunner for greenlet mode (ADR-0020 D3) # PE-local scratch allocator for math/compute output handles. # Each binary/unary/reduction op auto-allocates a unique addr from # this pool so the resulting TensorHandle can be the source of a # later tl.send / tl.store. Cursor resets on every kernel invocation. self._scratch_base = scratch_base self._scratch_size = scratch_size self._scratch_cursor = 0 def scratch_scope(self) -> _ScratchScope: """Per-tile scratch recycling context manager (ADR-0063). Usage: with tl.scratch_scope(): s = tl.dot(q, k_t) # per-tile temporaries — p = tl.softmax(s) # their scratch is rewound o_j = tl.dot(p, v) # on __exit__ Persistent state (running ``(m, ℓ, O)``, lazy-load prefetch buffers) must be allocated **outside** the scope; only handles allocated inside are recycled. """ return _ScratchScope(self) def _scratch_alloc(self, nbytes: int) -> int: """Allocate a unique scratch address for an output TensorHandle. Returns 0 if no scratch base was configured (e.g. command-list mode); in that case the resulting handle has addr=0 and cannot be used as a send/store source. Greenlet/runner mode always supplies a base. """ if self._scratch_base == 0: return 0 # 16-byte alignment aligned = (nbytes + 15) & ~15 addr = self._scratch_base + self._scratch_cursor self._scratch_cursor += aligned if self._scratch_cursor > self._scratch_size: raise RuntimeError( f"TLContext scratch overflow: requested {nbytes}B, " f"used {self._scratch_cursor}/{self._scratch_size}B" ) return addr @property def commands(self) -> list[PeCommand]: """Return the recorded command trace.""" return self._commands # ── helpers ──────────────────────────────────────────────────── def _next_handle_id(self) -> str: self._handle_counter += 1 return f"t{self._handle_counter}" def _next_completion_id(self) -> str: self._completion_counter += 1 return f"c{self._completion_counter}" def _dtype_bytes(self, dtype: str) -> int: return _DTYPE_BYTES.get(dtype, 2) def _nbytes(self, shape: tuple[int, ...], dtype: str) -> int: return math.prod(shape) * self._dtype_bytes(dtype) def _charge_dispatch(self, cmd: PeCommand) -> None: """Charge structural PE_CPU dispatch cost (ADR-0064 Rev2 D1). With a ``cost_model``: emit ``PeCpuOverheadCmd`` carrying ``round((FIXED + cmd.logical_bytes × R) / clock_freq_ghz)`` ns just before ``cmd``. A CompositeCmd over the descriptor-size cap is rejected (D7). Without a cost_model, fall back to the uniform ``dispatch_cycles`` contract (ADR-0046 §D6). """ if self._cost_model is not None: if isinstance(cmd, CompositeCmd): lb = cmd.logical_bytes cap = self._cost_model.max_composite_logical_bytes if lb > cap: raise ValueError( f"CompositeCmd logical_bytes {lb} exceeds " f"max_composite_logical_bytes {cap} (ADR-0064 D7)" ) cycles = self._cost_model.dispatch_cycles( type(cmd), cmd.logical_bytes) ns = round(cycles / self._clock_freq_ghz) if ns > 0: self._emit(PeCpuOverheadCmd(cycles=ns)) elif self._dispatch_cycles > 0: self._emit(PeCpuOverheadCmd(cycles=self._dispatch_cycles)) def _make_handle( self, addr: int, shape: tuple[int, ...], dtype: str, space: str = "tcm", pinned: bool = False, ) -> TensorHandle: return TensorHandle( id=self._next_handle_id(), addr=addr, shape=shape, dtype=dtype, nbytes=self._nbytes(shape, dtype), space=space, pinned=pinned, ) def _make_compute_out( self, shape: tuple[int, ...], dtype: str, ) -> TensorHandle: """Allocate an output TensorHandle in PE-local scratch (TCM space). Used by math/compute ops so the result has a real address that can be the source of a later send/store. The data field stays None in Phase 1 — Phase 2 DataExecutor fills the actual ndarray. """ nbytes = self._nbytes(shape, dtype) addr = self._scratch_alloc(nbytes) return TensorHandle( id=self._next_handle_id(), addr=addr, shape=shape, dtype=dtype, nbytes=nbytes, space="tcm", # TCM-resident: a downstream composite operand reads it on-chip, # not via a DMA_READ of its (bit-61) scratch address — same as # the composite auto-output and recipe scratch handles below. pinned=True, ) # ── Reference (no DMA, metadata only) ──────────────────────── def ref( self, ptr: int, shape: tuple[int, ...], dtype: str = "f16", ) -> TensorHandle: """Create a TensorHandle referencing HBM data without issuing DMA. Used when the scheduler will stream data per-tile (e.g., tensor b in a composite GEMM) or as a composite's HBM output handle (ADR-0065 D8). The handle's ``space="hbm"`` — it references HBM data. No command is generated. (Operand-input DMA stays ``pinned``-based per ADR-0065 D4, so this label does not change input streaming.) """ return self._make_handle(addr=ptr, shape=shape, dtype=dtype, space="hbm") # ── Data Movement (blocking, DMA engine) ────────────────────── def _emit(self, cmd: PeCommand) -> Any: """Emit command: greenlet switch if runner available, else append to list. Each dispatched PeCommand is preceded by a ``PeCpuOverheadCmd`` carrying its structural dispatch cost (ADR-0064 Rev2 D1). ``PeCpuOverheadCmd`` (incl. manual ``tl.cycles``) and ``WaitCmd`` bypass the charge (D5). """ if not isinstance(cmd, (PeCpuOverheadCmd, WaitCmd)): self._charge_dispatch(cmd) if self._runner is not None: return self._runner.switch_to_simpy(cmd) self._commands.append(cmd) return None def load( self, ptr: int, shape: tuple[int, ...], dtype: str = "f16", ) -> TensorHandle: """Load tensor from HBM — **lazy** (ADR-0062 §D1/§D2). Posts the ``DmaReadCmd`` to PE_DMA and returns immediately with a ``TensorHandle`` whose ``pending`` field references a fresh ``LoadFuture``. The actual DMA completion is awaited at the first consuming op (``tl.dot``, MATH, ``tl.store``, ``tl.send``, ``tl.copy_to``, ``tl.composite``) via ``_await_pending``. Command-list mode: emits the DmaReadCmd to ``self._commands``, attaches a LoadFuture for structural compatibility (its event stays None — no engine, no SimPy event to wait on). """ nbytes = self._nbytes(shape, dtype) # LoadFuture is mutable; create it first, attach to the handle, # then point its ``cmd`` at the DmaReadCmd that references the # *final* handle. This guarantees ``handle.pending.cmd.handle is handle``. future = LoadFuture.__new__(LoadFuture) future.event = None future.resolved = False future.data = None handle = TensorHandle( id=self._next_handle_id(), addr=ptr, shape=shape, dtype=dtype, nbytes=nbytes, data=None, space="hbm", pinned=True, pending=future, ) cmd = DmaReadCmd(handle=handle, src_addr=ptr, nbytes=nbytes) future.cmd = cmd # Lazy load bypasses _emit (it posts via a load_issue future), so it # charges its own dispatch cost here (ADR-0064 D1). self._charge_dispatch(cmd) if self._runner is not None: # Lazy: runner posts the DmaReadCmd, sets future.event, then # switches back immediately. No yield on completion here. self._runner.switch_to_simpy(("load_issue", future)) else: self._commands.append(cmd) return handle def _await_pending(self, *handles: TensorHandle | None) -> None: """Auto-wait at first use (ADR-0062 §D2). For each handle carrying an unresolved ``LoadFuture``, yield on the DMA completion event (greenlet → runner). Command-list mode is a no-op. """ if self._runner is None: return for h in handles: if h is None: continue pending = getattr(h, "pending", None) if pending is None or pending.resolved: continue if isinstance(pending, CompositeFuture): # Composite output: wait on its completion (ADR-0065 D8). self._emit(WaitCmd(handle=pending.completion)) pending.resolved = True else: self._runner.switch_to_simpy(("load_await", h)) def store(self, ptr: int, handle: TensorHandle) -> None: """Store tensor from TCM to HBM.""" self._await_pending(handle) cmd = DmaWriteCmd(handle=handle, dst_addr=ptr, nbytes=handle.nbytes) self._emit(cmd) def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None: """TCM-to-TCM byte copy (ADR-0063 §D3.1). Writes ``src``'s bytes into ``dst``'s address. Both handles must live in TCM and have matching shape and dtype. Symmetric to ``tl.store`` (which targets HBM) but stays on-chip so it doesn't emit a DMA entry into op_log. Used inside ``tl.scratch_scope()`` to persist a scoped result — typically an updated running ``(m, ℓ, O)`` — to an outside-scope (persistent) handle so its bytes survive the scope's ``__exit__`` cursor rewind (ADR-0063 §D3 two-arena pattern). """ if src.shape != dst.shape: raise ValueError( f"tl.copy_to: shape mismatch — src.shape={src.shape} " f"vs dst.shape={dst.shape}" ) if src.dtype != dst.dtype: raise ValueError( f"tl.copy_to: dtype mismatch — src.dtype={src.dtype!r} " f"vs dst.dtype={dst.dtype!r}" ) if dst.space != "tcm": raise ValueError( f"tl.copy_to: dst must be in TCM (got space={dst.space!r}); " "writes to HBM go through tl.store" ) if src.space != "tcm": raise ValueError( f"tl.copy_to: src must be in TCM (got space={src.space!r}); " "reads from HBM go through tl.load" ) self._await_pending(src) self._emit(CopyCmd(src=src, dst=dst, nbytes=src.nbytes)) # ── GEMM Engine (blocking) ──────────────────────────────────── def dot(self, a: TensorHandle, b: TensorHandle) -> TensorHandle: """Matrix multiply: out = a @ b. Both operands must be in TCM. a: (M, K), b: (K, N) → out: (M, N) """ if len(a.shape) < 2 or len(b.shape) < 2: raise ValueError("dot requires 2D tensors") m, k = a.shape[-2], a.shape[-1] k2, n = b.shape[-2], b.shape[-1] if k != k2: raise ValueError(f"dot shape mismatch: a.K={k} != b.K={k2}") out_shape = (*a.shape[:-2], m, n) out_dtype = a.dtype out = self._make_compute_out(shape=out_shape, dtype=out_dtype) self._await_pending(a, b) self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n)) return out # ── MATH Engine: unary (blocking) ───────────────────────────── def _unary_math(self, op: str, x: TensorHandle) -> TensorHandle: out = self._make_compute_out(shape=x.shape, dtype=x.dtype) self._await_pending(x) self._emit(MathCmd(op=op, inputs=(x,), out=out)) return out def exp(self, x: TensorHandle) -> TensorHandle: return self._unary_math("exp", x) def log(self, x: TensorHandle) -> TensorHandle: return self._unary_math("log", x) def sqrt(self, x: TensorHandle) -> TensorHandle: return self._unary_math("sqrt", x) def abs(self, x: TensorHandle) -> TensorHandle: return self._unary_math("abs", x) def sigmoid(self, x: TensorHandle) -> TensorHandle: return self._unary_math("sigmoid", x) def cos(self, x: TensorHandle) -> TensorHandle: return self._unary_math("cos", x) def sin(self, x: TensorHandle) -> TensorHandle: return self._unary_math("sin", x) # ── MATH Engine: reduction (blocking) ───────────────────────── def _reduction( self, op: str, x: TensorHandle, axis: int, ) -> TensorHandle: out_shape = list(x.shape) out_shape[axis] = 1 out = self._make_compute_out(shape=tuple(out_shape), dtype=x.dtype) self._await_pending(x) self._emit(MathCmd(op=op, inputs=(x,), out=out, axis=axis)) return out def sum(self, x: TensorHandle, axis: int) -> TensorHandle: return self._reduction("sum", x, axis) def max(self, x: TensorHandle, axis: int) -> TensorHandle: return self._reduction("max", x, axis) def min(self, x: TensorHandle, axis: int) -> TensorHandle: return self._reduction("min", x, axis) # ── MATH Engine: binary (blocking) ──────────────────────────── def _binary_math( self, op: str, a: TensorHandle, b: TensorHandle, ) -> TensorHandle: out = self._make_compute_out(shape=a.shape, dtype=a.dtype) self._await_pending(a, b) self._emit(MathCmd(op=op, inputs=(a, b), out=out)) return out def where( self, cond: TensorHandle, a: TensorHandle, b: TensorHandle, ) -> TensorHandle: out = self._make_compute_out(shape=a.shape, dtype=a.dtype) self._await_pending(cond, a, b) self._emit(MathCmd(op="where", inputs=(cond, a, b), out=out)) return out def maximum(self, a: TensorHandle, b: TensorHandle) -> TensorHandle: """Element-wise max of two tensors (real Triton: tl.maximum).""" return self._binary_math("maximum", a, b) def minimum(self, a: TensorHandle, b: TensorHandle) -> TensorHandle: """Element-wise min of two tensors (real Triton: tl.minimum).""" return self._binary_math("minimum", a, b) def fma( self, a: TensorHandle, b: TensorHandle, c: TensorHandle, ) -> TensorHandle: """Fused multiply-add: a * b + c (real Triton: tl.fma).""" out = self._make_compute_out(shape=a.shape, dtype=a.dtype) self._await_pending(a, b, c) self._emit(MathCmd(op="fma", inputs=(a, b, c), out=out)) return out def clamp( self, x: TensorHandle, min: TensorHandle, max: TensorHandle, ) -> TensorHandle: """Clamp x to [min, max] (real Triton: tl.clamp).""" out = self._make_compute_out(shape=x.shape, dtype=x.dtype) self._await_pending(x, min, max) self._emit(MathCmd(op="clamp", inputs=(x, min, max), out=out)) return out def softmax(self, x: TensorHandle, axis: int = -1) -> TensorHandle: """Numerically-stable softmax along ``axis`` (real Triton: tl.softmax). Implemented as a single MathCmd (op="softmax") so timing accounts for one MATH dispatch; Phase 2 DataExecutor expands it to the canonical (x - max) → exp → sum → div sequence. """ out = self._make_compute_out(shape=x.shape, dtype=x.dtype) self._await_pending(x) self._emit(MathCmd(op="softmax", inputs=(x,), out=out, axis=axis)) return out # ── Scalar helpers (real Triton: tl.cdiv etc.) ──────────────── @staticmethod def cdiv(a: int, b: int) -> int: """Ceiling division: (a + b - 1) // b (real Triton: tl.cdiv). Used by host/kernel grid math; not a tensor op, so no MathCmd is emitted. Mirrors triton.cdiv. """ return -(-int(a) // int(b)) # ── Index / Scalar (PE_CPU, no engine) ──────────────────────── def program_id(self, axis: int = 0) -> int: """Return program instance index (ADR-0022). axis=0: local PE id within cube. axis=1: cube id. """ if axis == 1: return self._cube_id return self._pe_id def num_programs(self, axis: int = 0) -> int: """Return total number of program instances (ADR-0022). axis=0: num PEs per cube. axis=1: num cubes. """ if axis == 1: return self._num_cubes return self._num_programs def arange(self, start: int, end: int, dtype: str = "i32") -> TensorHandle: """Create index range tensor in TCM.""" n = end - start return self._make_handle(addr=0, shape=(n,), dtype=dtype) def zeros(self, shape: tuple[int, ...], dtype: str = "f16") -> TensorHandle: """Create zero-filled tensor in TCM.""" return self._make_handle(addr=0, shape=shape, dtype=dtype) def full( self, shape: tuple[int, ...], value: float | int, dtype: str = "f16", ) -> TensorHandle: """Create constant-filled tensor in TCM.""" return self._make_handle(addr=0, shape=shape, dtype=dtype) # ── Metadata (no compute, no DMA) ───────────────────────────── def trans(self, x: TensorHandle) -> TensorHandle: """Transpose — shape change only, no command generated.""" if len(x.shape) < 2: raise ValueError("trans requires at least 2D tensor") new_shape = (*x.shape[:-2], x.shape[-1], x.shape[-2]) return TensorHandle( id=x.id, addr=x.addr, shape=new_shape, dtype=x.dtype, nbytes=x.nbytes, data=x.data, ) # ── IPCQ (CCL) collective primitives (ADR-0023 D4) ──────────── def send( self, dir: str, src: TensorHandle | None = None, *, src_addr: int | None = None, nbytes: int | None = None, shape: tuple[int, ...] | None = None, dtype: str = "f16", space: str = "tcm", ) -> None: """Send tensor data to the peer in the given direction. Two calling forms: tl.send(dir, handle) # use handle's metadata tl.send(dir, src_addr=..., nbytes=..., shape=..., dtype=..., space=...) Blocking: returns when PE_IPCQ has accepted the request and forwarded the IpcqDmaToken to PE_DMA. Backpressure may apply. """ if src is not None: src_addr = src.addr nbytes = src.nbytes shape = src.shape dtype = src.dtype space = getattr(src, "space", space) if src_addr is None or nbytes is None or shape is None: raise ValueError("tl.send: provide either a TensorHandle or src_addr/nbytes/shape") # ADR-0062: if the source is a lazy-loaded handle, await first so # the data snapshot below sees the real bytes. self._await_pending(src) # Carry the handle's .data snapshot (if available). When the source # is a recv slot, .data holds the numpy array that was read from # MemoryStore at recv-time. This prevents a Phase 1 race where a # later IPCQ inbound overwrites the slot before the outbound # PE_DMA reads it. handle_data = getattr(src, "data", None) if src is not None else None cmd = IpcqSendCmd( direction=dir, src_addr=src_addr, src_space=space, nbytes=nbytes, shape=shape, dtype=dtype, handle_id=self._next_handle_id(), data=handle_data, ) self._emit(cmd) def recv( self, dir: str | None = None, shape: tuple[int, ...] = (), dtype: str = "f16", space: str = "tcm", dst_addr: int | None = None, dst_space: str | None = None, ) -> TensorHandle: """Receive tensor data from a peer. Args: dir: specific direction (e.g. "W"), or None for round-robin. shape, dtype: expected tensor metadata. dst_addr / dst_space: if both are provided, the slot data is copied to (dst_space, dst_addr) before the handle is returned ("copy_to_dst" mode). Otherwise the slot address is returned directly ("return_slot" mode). Returns: TensorHandle pointing to the slot (or dst) where the data has arrived. In greenlet/runner mode, ``handle.data`` carries the actual ndarray; in command-list mode the handle is a placeholder. """ if dst_addr is not None and dst_space is not None: cmd = IpcqRecvCmd( direction=dir, shape=shape, dtype=dtype, handle_id=self._next_handle_id(), recv_mode="copy_to_dst", dst_addr=dst_addr, dst_space=dst_space, ) else: cmd = IpcqRecvCmd( direction=dir, shape=shape, dtype=dtype, handle_id=self._next_handle_id(), ) result = self._emit(cmd) if isinstance(result, dict): slot_addr = int(result.get("src_addr", 0)) slot_space = str(result.get("src_space", "tcm")) data = result.get("data") return TensorHandle( id=self._next_handle_id(), addr=slot_addr, shape=shape, dtype=dtype, 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) def recv_no_consume( self, dir: str | None = None, shape: tuple[int, ...] = (), dtype: str = "f16", ) -> TensorHandle: """DIAGNOSTIC ONLY — recv that blocks for arrival but skips slot read. Same blocking semantics as ``tl.recv``: the kernel waits until the payload has landed in the IPCQ slot. Differs from ``tl.recv`` by skipping the slot-read latency charge (slot-IO + PE↔bank fabric drain) on DST. This entry point exists solely so the pe2pe overview plot can draw an apples-to-apples comparison against ``tl.store`` (a one-sided fabric write that pays no read on DST). Production kernels MUST use ``tl.recv`` — they need to consume the data they receive. This API is segregated from ``tl.recv`` so the diagnostic flag can never accidentally be set in real workloads. """ cmd = IpcqRecvCmd( direction=dir, shape=shape, dtype=dtype, handle_id=self._next_handle_id(), consume=False, ) result = self._emit(cmd) # type: ignore[arg-type] if isinstance(result, dict): slot_addr = int(result.get("src_addr", 0)) slot_space = str(result.get("src_space", "tcm")) return TensorHandle( id=self._next_handle_id(), addr=slot_addr, shape=shape, dtype=dtype, nbytes=self._nbytes(shape, dtype), data=None, space=slot_space, pinned=slot_space != "hbm", ) return self._make_handle(addr=0, shape=shape, dtype=dtype) def recv_async( self, dir: str, shape: tuple[int, ...] = (), dtype: str = "f16", ) -> "RecvFuture": """Non-blocking recv. Returns a future to pass into ``tl.wait``.""" cmd = IpcqRecvCmd( direction=dir, shape=shape, dtype=dtype, handle_id=self._next_handle_id(), blocking=False, ) # recv_async bypasses _emit (posts via a recv_async future), so it # charges its own dispatch cost here (ADR-0064 D1). self._charge_dispatch(cmd) future = RecvFuture(cmd=cmd) if self._runner is not None: self._runner.switch_to_simpy(("recv_async", future)) return future # ── Composite + Control ─────────────────────────────────────── def composite( self, op: Literal["gemm", "math"], a: TensorHandle | None = None, b: TensorHandle | None = None, math_op: str | None = None, *, out: TensorHandle | None = None, out_ptr: int | None = None, prologue: list[dict] | None = None, epilogue: list[dict] | None = None, acc_dtype: str | None = None, tile_shape: tuple[int, int, int] | None = None, ) -> TensorHandle: """Submit a composite command (non-blocking, tiled pipeline). ``prologue`` (ADR-0065 D5) is an ordered list of recipe dicts — each with an ``"op"`` naming a ``RECIPE_DESCRIPTORS`` entry plus its operands. TLContext expands each into flat MATH OpSpecs that run before the head op and auto-binds the recipe's primary output into the head GEMM's ``a`` operand (D6.6). ``epilogue`` is an ordered list of ``EPILOGUE_OPS`` dicts. ``out`` (ADR-0065 D8) is a ``TensorHandle``: an HBM ref (``tl.ref(addr, shape)`` → ``space="hbm"`` → DMA_WRITE) or an in-place TCM handle (e.g. a recipe accumulator → STORE only). ``out_ptr`` (int) is shorthand for an HBM output — equivalent to ``out=tl.ref(out_ptr, )``. If neither is given, a TCM scratch is auto-allocated (chainable, like ``tl.dot``). Returns the **output TensorHandle** — auto-awaited by downstream ops (and by ``tl.wait(handle)``) via its ``pending`` field. """ # ADR-0062: composite operand DMA paths still need their inputs # to be resolved before the composite reads them via PE_SCHEDULER. self._await_pending(a, b) # Prologue recipes → flat MATH OpSpecs + primary-output handle (D5). prologue_ops: tuple[OpSpec, ...] = () primary_out: TensorHandle | None = None rw_handles: tuple[TensorHandle, ...] = () if prologue: prologue_ops, primary_out, rw_handles = self._expand_prologue(prologue) # Await any prologue operand that is a pending composite output # (e.g. the score tile produced by a prior #1 composite) so it is # computed before this composite reads it (RAW across composites). for _item in prologue: self._await_pending( *[v for v in _item.values() if isinstance(v, TensorHandle)] ) # Auto-bind (D6.6): the recipe's primary output fills the head GEMM's # `a` unless the kernel supplied `a` explicitly (then it's ambiguous). if op == "gemm" and primary_out is not None: if a is not None: raise ValueError( "auto-bind conflict: head GEMM operand 'a' was provided " "explicitly but the prologue recipe declares a primary_out " "(ADR-0065 D6.6) — omit 'a' to let the recipe bind it" ) a = primary_out if a is None: raise ValueError( "composite requires operand 'a' (or a prologue recipe with a " "primary_out)" ) # Output geometry. if op == "gemm" and b is not None: m, k = a.shape[-2], a.shape[-1] n = b.shape[-1] out_dtype = a.dtype out_shape: tuple[int, ...] = (m, n) out_nbytes = m * n * self._dtype_bytes(out_dtype) else: out_dtype = a.dtype out_shape = a.shape out_nbytes = a.nbytes # Write-back handle (ADR-0065 D8). Explicit `out` (HBM ref or an # in-place TCM accumulator) wins; `out_ptr` is HBM shorthand; else # auto-allocate a TCM scratch — the chainable default, like tl.dot. # pinned=True on the auto output: it is TCM-resident, so a downstream # op reading it needs no DMA_READ. if out is not None: out_handle = out elif out_ptr is not None: out_handle = TensorHandle( id=self._next_handle_id(), addr=out_ptr, shape=out_shape, dtype=out_dtype, nbytes=out_nbytes, space="hbm", ) else: out_handle = TensorHandle( id=self._next_handle_id(), addr=self._scratch_alloc(out_nbytes), shape=out_shape, dtype=out_dtype, nbytes=out_nbytes, space="tcm", pinned=True, ) # Head op (flat-ops, ADR-0065 D2): GEMM takes named a/b and carries # m/k/n in extra; MATH takes named a and carries math_op in extra. if op == "gemm" and b is not None: head_operands: dict[str, Any] = {"a": a, "b": b} head_extra: dict[str, Any] = {"m": m, "k": k, "n": n} else: head_operands = {"a": a} head_extra = {} for _k, _v in (("acc_dtype", acc_dtype), ("tile_shape", tile_shape), ("math_op", math_op)): if _v is not None: head_extra[_k] = _v head_spec = OpSpec( kind=op, scope=Scope.OUTPUT_TILE, operands=head_operands, extra=head_extra, out=out_handle, ) epi_specs = tuple(self._build_epilogue_spec(e, i) for i, e in enumerate(epilogue or [])) ops_tuple = (*prologue_ops, head_spec, *epi_specs) self._validate_composite_ops(ops_tuple, prologue_ops) completion = CompletionHandle(id=self._next_completion_id()) self._emit(CompositeCmd( completion=completion, ops=ops_tuple, rw_handles=rw_handles, )) # Attach the completion to the output handle so downstream ops (and # tl.wait) auto-await it (ADR-0065 D8 / ADR-0062 lazy pattern), and # return it so the result chains like tl.dot's (frozen dataclass → # set pending via object.__setattr__). object.__setattr__(out_handle, "pending", CompositeFuture(completion)) return out_handle def _expand_prologue( self, prologue: list[dict], ) -> tuple[tuple[OpSpec, ...], TensorHandle | None, tuple[TensorHandle, ...]]: """Expand prologue recipe items into flat MATH OpSpecs (ADR-0065 D5). Returns ``(math_ops, primary_out_handle, rw_handles)``. PE_SCHEDULER never sees the recipe — only the flat ops. """ from kernbench.triton_emu.tl_recipes import ( PRIMARY_OUT_SLOT, RECIPE_DESCRIPTORS, ) math_ops: list[OpSpec] = [] primary_out: TensorHandle | None = None rw_handles: list[TensorHandle] = [] for item in prologue: if not isinstance(item, dict) or "op" not in item: raise ValueError("prologue entry must be a dict with an 'op' key") name = item["op"] recipe = RECIPE_DESCRIPTORS.get(name) if recipe is None: known = ", ".join(sorted(RECIPE_DESCRIPTORS)) raise ValueError( f"unknown prologue recipe {name!r} (known: {known})" ) # Bind recipe operands from the item. slots: dict[str, TensorHandle] = {} for opd_name in recipe.operands: if opd_name not in item: raise ValueError( f"prologue recipe {name!r} missing operand {opd_name!r}" ) slots[opd_name] = item[opd_name] # Expand engine_seq sequentially — later ops read earlier dsts. for eop in recipe.engine_seq: operands: dict[str, Any] = {} for key in ("src", "src_a", "src_b", "src_c"): slot_name = getattr(eop, key) if slot_name is None: continue if slot_name not in slots: raise ValueError( f"recipe {name!r} op {eop.op_kind!r} reads unbound " f"slot {slot_name!r}" ) operands[key] = slots[slot_name] dst_handle = self._resolve_recipe_dst(eop, recipe, slots, operands) extra: dict[str, Any] = {} if eop.reduce_axis is not None: extra["reduce_axis"] = eop.reduce_axis if eop.bcast_axis is not None: extra["bcast_axis"] = eop.bcast_axis math_ops.append(OpSpec( kind=eop.op_kind, scope=Scope.KERNEL, operands=operands, extra=extra, out=dst_handle, )) slots[eop.dst] = dst_handle if recipe.primary_out is not None: primary_out = slots.get(PRIMARY_OUT_SLOT) for opd_name, rw in recipe.operands.items(): if rw == "RW": rw_handles.append(slots[opd_name]) return tuple(math_ops), primary_out, tuple(rw_handles) def _resolve_recipe_dst( self, eop: Any, recipe: Any, slots: dict, operands: dict, ) -> TensorHandle: """Resolve an EngineOp's dst to a handle: reuse an existing slot (RW operand / already-allocated), else allocate fresh TCM scratch with an inferred shape.""" from kernbench.triton_emu.tl_recipes import PRIMARY_OUT_SLOT dst_name = eop.dst if dst_name in slots: return slots[dst_name] # write in place if recipe.primary_out is not None and dst_name == PRIMARY_OUT_SLOT: ref = slots[recipe.primary_out.from_shape] shape, dtype = ref.shape, ref.dtype elif eop.reduce_axis is not None: # Reductions keep the reduced axis as size 1 (keepdims), so the # result broadcasts against the running (m, l) — e.g. rmax of # s=(G,TILE) over axis -1 → (G,1), not (G,) (ADR-0065 N2). ref = next(iter(operands.values())) dims = list(ref.shape) if len(dims) > 1: ax = eop.reduce_axis % len(dims) dims[ax] = 1 shape = tuple(dims) dtype = ref.dtype else: ref = next(iter(operands.values())) shape, dtype = ref.shape, ref.dtype nbytes = self._nbytes(shape, dtype) addr = self._scratch_alloc(nbytes) # pinned=True: recipe scratch / primary-out live in TCM already, so # the head GEMM's auto-bound `a` (= primary-out P) is consumed in # place — no DMA_READ (ADR-0065 P3, pinned-based DMA decision). return TensorHandle( id=self._next_handle_id(), addr=addr, shape=shape, dtype=dtype, nbytes=nbytes, space="tcm", pinned=True, ) @staticmethod def _validate_composite_ops( ops: tuple[OpSpec, ...], prologue_ops: tuple[OpSpec, ...], ) -> None: """Emit-time invariants (ADR-0065 D6.1 / D6.7). D6.1 (GEMM count ≤ 1) applies to the whole composite. D6.7 (MATH operands TCM-only) applies to **prologue recipe ops only** — the head op (gemm *or* math) and epilogue ops keep the existing DMA-staged-from-HBM behavior (a math head's `a` is streamed by PE_SCHEDULER, like a GEMM operand). """ gemm_count = sum(1 for o in ops if o.kind == "gemm") if gemm_count > 1: raise ValueError( f"composite must carry at most 1 GEMM op (ADR-0065 D6.1); " f"got {gemm_count}" ) for o in prologue_ops: for h in o.operands.values(): if getattr(h, "space", "tcm") != "tcm": raise ValueError( f"MATH operand must be TCM-resident (ADR-0065 D6.7); " f"prologue op {o.kind!r} operand has space=" f"{getattr(h, 'space', None)!r}" ) @staticmethod def _build_epilogue_spec(entry: dict, idx: int) -> OpSpec: if not isinstance(entry, dict) or "op" not in entry: raise ValueError( f"epilogue[{idx}]: each entry must be a dict with an 'op' key" ) kind = entry["op"] if kind not in EPILOGUE_OPS: known = ", ".join(sorted(EPILOGUE_OPS)) raise ValueError( f"epilogue[{idx}]: unknown op {kind!r} " f"(known ops: {known})" ) required, default_scope = EPILOGUE_OPS[kind] missing = [f for f in required if f not in entry] if missing: raise ValueError( f"epilogue[{idx}] op {kind!r} missing required field(s): " f"{', '.join(missing)}" ) scope = Scope(entry["scope"]) if "scope" in entry else default_scope operands: dict = {} scalar: float | None = None extra: dict = {} for f in required: v = entry[f] if isinstance(v, TensorHandle): operands[f] = v elif isinstance(v, (int, float)): if scalar is None: scalar = float(v) else: extra[f] = v else: extra[f] = v return OpSpec( kind=kind, scope=scope, operands=operands, scalar=scalar, extra=extra, ) def wait(self, handle: "CompletionHandle | RecvFuture | None" = None) -> Any: """Wait for a composite, a recv future, or all pending composites. - ``CompletionHandle`` (or None): wait for composite completion. - ``RecvFuture``: wait for a non-blocking ``recv_async`` to finish. Returns the resolved ``TensorHandle``. """ if isinstance(handle, RecvFuture): if handle.resolved: return handle.result if self._runner is None: raise RuntimeError( "tl.wait(RecvFuture) requires runner mode (greenlet)" ) result_dict = self._runner.switch_to_simpy(("recv_wait", handle)) slot_addr = int(result_dict.get("src_addr", 0)) slot_space = str(result_dict.get("src_space", "tcm")) data = result_dict.get("data") th = TensorHandle( id=self._next_handle_id(), addr=slot_addr, shape=handle.cmd.shape, dtype=handle.cmd.dtype, nbytes=self._nbytes(handle.cmd.shape, handle.cmd.dtype), data=data, space=slot_space, pinned=slot_space != "hbm", ) handle.resolved = True handle.result = th return th # Composite output handle (ADR-0065 D8): wait on its completion. if isinstance(handle, TensorHandle): pending = getattr(handle, "pending", None) if isinstance(pending, CompositeFuture): if not pending.resolved: self._emit(WaitCmd(handle=pending.completion)) pending.resolved = True return None # Composite path (CompletionHandle or None — existing behaviour). self._emit(WaitCmd(handle=handle)) return None def cycles(self, n: int) -> None: """Declare PE_CPU scalar execution overhead (cycles).""" self._emit(PeCpuOverheadCmd(cycles=n)) # ── TensorHandle arithmetic operators ───────────────────────────── # Enables: a + b, a * b, a - b, a / b in kernel code. # Each creates a MathCmd via a module-level helper that requires a # TLContext. We attach the context to handles via a closure approach. def _enable_tensor_ops() -> None: """Patch TensorHandle with arithmetic operators. Called once at module load. Operators create MathCmd entries via a thread-local TLContext reference set during kernel execution. """ import threading _local = threading.local() def set_active_context(ctx: TLContext | None) -> None: _local.ctx = ctx def get_active_context() -> TLContext: ctx = getattr(_local, "ctx", None) if ctx is None: raise RuntimeError("TensorHandle ops require an active TLContext") return ctx def _binop(op: str): def method(self: TensorHandle, other: TensorHandle) -> TensorHandle: ctx = get_active_context() return ctx._binary_math(op, self, other) return method # Patch TensorHandle class with operators TensorHandle.__add__ = _binop("add") # type: ignore[attr-defined] TensorHandle.__sub__ = _binop("sub") # type: ignore[attr-defined] TensorHandle.__mul__ = _binop("mul") # type: ignore[attr-defined] TensorHandle.__truediv__ = _binop("div") # type: ignore[attr-defined] # Expose context management TLContext._set_active = staticmethod(set_active_context) # type: ignore[attr-defined] TLContext._get_active = staticmethod(get_active_context) # type: ignore[attr-defined] _enable_tensor_ops() def run_kernel( kernel_fn, tl_ctx: TLContext, *args, **kwargs, ) -> list[PeCommand]: """Execute a kernel function with the given TLContext and return commands. Sets tl_ctx as the active context for TensorHandle operators, calls the kernel, then clears the context. """ TLContext._set_active(tl_ctx) # type: ignore[attr-defined] try: kernel_fn(*args, tl=tl_ctx, **kwargs) finally: TLContext._set_active(None) # type: ignore[attr-defined] return tl_ctx.commands