From 47e2c78c665abd8a4631d240394e78800c694bde Mon Sep 17 00:00:00 2001 From: Yangwook Date: Wed, 10 Jun 2026 19:18:04 -0700 Subject: [PATCH] gqa(adr-0064/0065): flat-ops CompositeCmd (P1) + structural dispatch cost (ADR-0064 Rev2); promote ADR-0064 ADR-0065 P1: CompositeCmd -> flat ordered ops list (drop legacy op/a/b/out_addr fields); OpSpec.operands dict + out handle. Meaning-preserving (op_log byte-equal); pe_scheduler + op_log read the head op. ADR-0064 Rev2: replace Rev1 per-op cost table with structural FIXED + logical_bytes*R formula. logical_bytes on every PeCommand; new common/pe_cost_model.py; cost centralized in TLContext._emit (load/recv_async charge explicitly); pe_cpu/kernel_runner wire the per-PE model + clock. D7: cap exceeded -> ValueError (no auto-segmentation). Remove Rev1 cpu_issue_cost.py + its tests. No goldens churn. Promote ADR-0064 Rev2 Proposed->Accepted (docs/adr/ + docs/adr-ko/); amend D7 (error not segmentation) + record P1-before-P0 ordering in ADR-0064/0065 Migration notes (EN+KO). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../ADR-0064-perf-cpu-issue-cost-model.md} | 0 ...rog-flat-ops-composite-softmax-merge-ko.md | 11 +- ...5-prog-flat-ops-composite-softmax-merge.md | 12 +- .../ADR-0064-perf-cpu-issue-cost-model.md | 0 src/kernbench/common/cpu_issue_cost.py | 52 --- src/kernbench/common/ipcq_types.py | 12 + src/kernbench/common/pe_commands.py | 84 +++- src/kernbench/common/pe_cost_model.py | 57 +++ src/kernbench/components/builtin/pe_cpu.py | 14 +- .../components/builtin/pe_scheduler.py | 22 +- src/kernbench/sim_engine/op_log.py | 34 +- src/kernbench/triton_emu/kernel_runner.py | 9 +- src/kernbench/triton_emu/tl_context.py | 155 +++---- tests/test_composite_epilogue.py | 6 +- tests/test_composite_flat_ops_refactor.py | 127 ++++++ tests/test_pe_cost_model.py | 268 ++++++++++++ tests/test_phase_e_cpu_issue_cost.py | 382 ------------------ tests/test_triton_emu.py | 8 +- 18 files changed, 703 insertions(+), 550 deletions(-) rename docs/{adr-proposed/ADR-0064-perf-cpu-issue-cost-model-ko.md => adr-ko/ADR-0064-perf-cpu-issue-cost-model.md} (100%) rename docs/{adr-proposed => adr}/ADR-0064-perf-cpu-issue-cost-model.md (100%) delete mode 100644 src/kernbench/common/cpu_issue_cost.py create mode 100644 src/kernbench/common/pe_cost_model.py create mode 100644 tests/test_composite_flat_ops_refactor.py create mode 100644 tests/test_pe_cost_model.py delete mode 100644 tests/test_phase_e_cpu_issue_cost.py diff --git a/docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model-ko.md b/docs/adr-ko/ADR-0064-perf-cpu-issue-cost-model.md similarity index 100% rename from docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model-ko.md rename to docs/adr-ko/ADR-0064-perf-cpu-issue-cost-model.md diff --git a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md index 3702146..ff1d7c5 100644 --- a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md +++ b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md @@ -372,7 +372,16 @@ scratch 에서 읽는 형태 — 순환 tile-loop 의존성. **기각 (incorrect ## Migration -Land 순서: +> **구현 순서 노트.** 아래의 ADR-0064-우선 순서가 원래 계획이었음. +> 실제 구현은 **1↔2 단계를 뒤집어** — ADR-0065 Phase 1 (평평화 +> `CompositeCmd`) 이 ADR-0064 Rev2 *보다 먼저* land. 이유: `logical_bytes` +> (ADR-0064 D2) 가 평평화 형태 위에서 자연스럽게 정의되므로, refactor 를 +> 먼저 하면 버려지는 transitional `logical_bytes` 를 피함. 각 단계는 +> 독립적으로 안전: P1 은 기존 cost 경로 하에서 의미 보존 refactor +> (op_log byte-equal); ADR-0064 Rev2 가 이미 평평화된 형태 위에 구조적 +> 공식을 교체. 최종 상태는 어느 쪽이든 동일. + +Land 순서 (원래 계획; as-built 1↔2 뒤집힘은 위 노트 참조): 1. **ADR-0064 Revision 2** (별도 PR): 구조적 dispatch cost + `logical_bytes` property + topology config override + 일회성 골든 재생성. 2. **ADR-0065 Phase 1** (본 ADR, PR a): `CompositeCmd` 평평화 refactor + diff --git a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md index f26780c..b9ccad7 100644 --- a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md +++ b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md @@ -413,7 +413,17 @@ input from #1's epilogue scratch — circular tile-loop dependency. ## Migration -Land order: +> **Implementation ordering note.** The ADR-0064-first ordering below was +> the original plan. The actual implementation **reversed steps 1 and 2** — +> ADR-0065 Phase 1 (flat-ops `CompositeCmd`) landed *before* ADR-0064 Rev2. +> Reason: `logical_bytes` (ADR-0064 D2) is naturally defined on the flat-ops +> shape, so doing the refactor first avoids a throwaway transitional +> `logical_bytes`. Each step is independently safe: P1 is a meaning- +> preserving refactor (op_log byte-equal) under the pre-existing cost path; +> ADR-0064 Rev2 then swaps in the structural formula on the already-flat +> shape. The end state is identical either way. + +Land order (original plan; see note above for the as-built reversal of 1↔2): 1. **ADR-0064 Revision 2** (separate PR): structural dispatch cost + `logical_bytes` property + topology config override + one-time goldens regeneration. diff --git a/docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model.md b/docs/adr/ADR-0064-perf-cpu-issue-cost-model.md similarity index 100% rename from docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model.md rename to docs/adr/ADR-0064-perf-cpu-issue-cost-model.md diff --git a/src/kernbench/common/cpu_issue_cost.py b/src/kernbench/common/cpu_issue_cost.py deleted file mode 100644 index a9371a6..0000000 --- a/src/kernbench/common/cpu_issue_cost.py +++ /dev/null @@ -1,52 +0,0 @@ -"""Per-op-type CPU issue cost table (ADR-0064 D1). - -Replaces the single uniform ``dispatch_cycles`` scalar with a cost table -keyed by command kind. Charged on PE_CPU at issue time (before the command -is dispatched to PE_SCHEDULER) so the hybrid's CPU-saturation win -(ADR-0060 §1) becomes measurable. - -The table is consulted by ``TLContext._emit_dispatch_overhead(kind)``; -live PE_CPU paths (greenlet via ``kernel_runner.py``, legacy replay via -``pe_cpu.py:_execute_legacy``) construct TLContext with -``issue_cost_table=DEFAULT_CPU_ISSUE_COST`` so all benches see the cost. - -Absolute ns values are provisional (ADR-0064 review item #1). The -defensible claim is the **ratio** — composite ≫ primitive. -""" -from __future__ import annotations - -from typing import Literal - -OpKind = Literal[ - "composite", - "load", - "store", - "dot", - "math", - "ipcq_send", - "ipcq_recv", - "copy_to", -] - - -DEFAULT_CPU_ISSUE_COST: dict[str, int] = { - "composite": 40, - "load": 5, - "store": 5, - "dot": 5, - "math": 5, - "ipcq_send": 5, - "ipcq_recv": 5, - "copy_to": 5, -} - - -def get_issue_cost(kind: str, table: dict[str, int] | None = None) -> int: - """Return per-op-type CPU issue cost in ns. - - Unknown kinds return 0 (no charge) so adding a new ``tl.*`` op kind - doesn't accidentally over-charge before the table is updated. - """ - if table is None: - table = DEFAULT_CPU_ISSUE_COST - return table.get(kind, 0) diff --git a/src/kernbench/common/ipcq_types.py b/src/kernbench/common/ipcq_types.py index 200e138..f5f0b56 100644 --- a/src/kernbench/common/ipcq_types.py +++ b/src/kernbench/common/ipcq_types.py @@ -123,6 +123,12 @@ class IpcqSendCmd: data: Any = None data_op: bool = True # ADR-0020 op_log recording flag + @property + def logical_bytes(self) -> int: + # framing + direction enum + src_addr + src_space enum + nbytes + # + shape (len marker + 4·rank) + dtype tag (ADR-0064 D2). + return 4 + 1 + 4 + 1 + 4 + (1 + 4 * len(self.shape)) + 1 + # ── D12: IpcqRecvCmd (PE_CPU → PE_IPCQ) ────────────────────────────── @@ -155,6 +161,12 @@ class IpcqRecvCmd: data_op: bool = True consume: bool = True # DIAGNOSTIC: see docstring + @property + def logical_bytes(self) -> int: + # framing + direction enum + shape (len marker + 4·rank) + dtype tag + # + dst_addr + dst_space enum (ADR-0064 D2). + return 4 + 1 + (1 + 4 * len(self.shape)) + 1 + 4 + 1 + # ── D12: IpcqDmaToken (PE_IPCQ → PE_DMA, vc_comm) ─────────────────── diff --git a/src/kernbench/common/pe_commands.py b/src/kernbench/common/pe_commands.py index 12ea7e3..05284cc 100644 --- a/src/kernbench/common/pe_commands.py +++ b/src/kernbench/common/pe_commands.py @@ -10,7 +10,7 @@ from __future__ import annotations from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: import simpy @@ -22,6 +22,24 @@ class Scope(Enum): KERNEL = "kernel" +def _extra_bytes(v: Any) -> int: + """Type-aware HW-logical byte count for an OpSpec.extra value (ADR-0064 D2). + + bool→1, int/float→4 (scalar), tuple/list→1 + 4·len (length marker + 4 + per element, e.g. shape/axes), str→1 (opcode-like tag). Default→4. + ``bool`` is checked first because it subclasses ``int``. + """ + if isinstance(v, bool): + return 1 + if isinstance(v, (int, float)): + return 4 + if isinstance(v, (tuple, list)): + return 1 + 4 * len(v) + if isinstance(v, str): + return 1 + return 4 + + @dataclass(frozen=True) class OpSpec: """One operation in a multi-op composite (head + epilogue, ADR-0014 D3.3). @@ -33,9 +51,22 @@ class OpSpec: kind: str # "gemm" | "bias" | "relu" | ... scope: "Scope" = Scope.OUTPUT_TILE - operands: tuple[Any, ...] = () # tuple[TensorHandle, ...] + operands: dict[str, Any] = field(default_factory=dict) # name → TensorHandle scalar: float | None = None extra: dict[str, Any] = field(default_factory=dict) + out: "TensorHandle | None" = None # explicit write-back handle + + @property + def logical_bytes(self) -> int: + """HW-logical byte size (ADR-0064 D2). ``scalar`` is a transitional + field not part of the flat-ops model (ADR-0065 D2) — excluded, + pending its removal.""" + return ( + 1 + 1 # opcode + scope enum + + 1 + 8 * len(self.operands) # len marker + handles + + (8 if self.out is not None else 0) # out handle + + 1 + sum(_extra_bytes(v) for v in self.extra.values()) + ) # Epilogue op contracts: kind → (required field names, default scope). @@ -103,6 +134,10 @@ class DmaReadCmd: nbytes: int data_op: bool = True + @property + def logical_bytes(self) -> int: + return 4 + 8 + 4 + 4 # framing + handle + src_addr + nbytes + @dataclass(frozen=True) class DmaWriteCmd: @@ -113,6 +148,10 @@ class DmaWriteCmd: nbytes: int data_op: bool = True + @property + def logical_bytes(self) -> int: + return 4 + 8 + 4 + 4 # framing + handle + dst_addr + nbytes + @dataclass(frozen=True) class GemmCmd: @@ -129,6 +168,10 @@ class GemmCmd: n: int data_op: bool = True + @property + def logical_bytes(self) -> int: + return 4 + 8 * 3 + 4 * 3 # framing + 3 handles + m/k/n scalars + @dataclass(frozen=True) class MathCmd: @@ -145,6 +188,13 @@ class MathCmd: axis: int | None = None # for reductions data_op: bool = True + @property + def logical_bytes(self) -> int: + return ( + 4 + 1 + 1 + 8 * len(self.inputs) + 8 # framing+opcode+len+inputs+out + + (4 if self.axis is not None else 0) + ) + @dataclass(frozen=True) class CopyCmd: @@ -161,6 +211,10 @@ class CopyCmd: nbytes: int data_op: bool = True + @property + def logical_bytes(self) -> int: + return 4 + 8 * 2 + 4 # framing + src/dst handles + nbytes + @dataclass(frozen=True) class CompositeCmd: @@ -168,20 +222,26 @@ class CompositeCmd: Non-blocking — submitted to PE_SCHEDULER which manages tile splitting and pipeline overlaps (ADR-0014 D3.2). + + Flat-ops shape (ADR-0065 D1): ``ops`` is an ordered list of OpSpecs. + The GEMM op (if any, ≤1) drives the tile loop; preceding/following + OpSpecs are placed by position + scope. ``rw_handles`` carries + cross-composite hazard metadata (ADR-0065 D6.3); unused in P1. """ completion: CompletionHandle - op: Literal["gemm", "math"] - a: TensorHandle - b: TensorHandle | None - out_addr: int - out_nbytes: int - math_op: str | None = None # for op="math": which math operation - data_op: bool = True - # Multi-op composite (ADR-0014 D3.3): when non-empty, ops[0] is the - # head and ops[1:] are epilogue stages with explicit scope. When empty, - # the legacy single-op semantics (op/a/b/math_op) apply. ops: tuple[OpSpec, ...] = () + rw_handles: tuple["TensorHandle", ...] = () + data_op: bool = True + + @property + def logical_bytes(self) -> int: + """HW-logical byte size (ADR-0064 D2). Per-op summation, no dedup.""" + return ( + 4 # framing + + 1 + sum(op.logical_bytes for op in self.ops) + + 1 + 8 * len(self.rw_handles) + ) @dataclass(frozen=True) diff --git a/src/kernbench/common/pe_cost_model.py b/src/kernbench/common/pe_cost_model.py new file mode 100644 index 0000000..433e2e5 --- /dev/null +++ b/src/kernbench/common/pe_cost_model.py @@ -0,0 +1,57 @@ +"""Structural PE_CPU dispatch cost model (ADR-0064 Revision 2). + +Replaces the Rev1 per-op-type calibration table (``cpu_issue_cost.py``, +removed) with a structural formula derived from each command's +``logical_bytes`` (ADR-0064 D2): + + dispatch_cycles(cmd) = FIXED_PER_CMD + cmd.logical_bytes * R + +The per-command ``FIXED_PER_CMD`` term models the queue-tail update / +MMIO-class RTT / completion-event registration; the byte term ``R`` models +queue-write bandwidth. The primary signal the model exposes is +**command-count reduction** (FIXED-dominated); the byte term is a secondary +refinement (ADR-0064 Context). + +Knobs are cycle-domain only; cycle→ns uses the PE node's ``clock_freq_ghz`` +(ADR-0064 D3). Defaults anchor a typical 1-OpSpec GEMM composite (≈54 bytes) +at ≈43 ns on a 16 B/cycle on-die descriptor queue. + +A composite whose ``logical_bytes`` exceeds ``max_composite_logical_bytes`` +is rejected at emit time with a ``ValueError`` (ADR-0064 D7, revised: hard +cap, no auto-segmentation). +""" +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class PeCostModel: + """Cycle-domain dispatch cost knobs (ADR-0064 D3/D4).""" + + fixed_per_cmd_cycles: float = 40.0 + byte_cycles_recip: float = 0.0625 # = 16 bytes/cycle + max_composite_logical_bytes: int = 1024 # D7 hard cap + + def dispatch_cycles(self, logical_bytes: int) -> float: + """FIXED + logical_bytes × R (ADR-0064 D1).""" + return self.fixed_per_cmd_cycles + logical_bytes * self.byte_cycles_recip + + +DEFAULT_PE_COST_MODEL = PeCostModel() + + +def from_node_attrs(attrs: dict) -> PeCostModel: + """Build a PeCostModel from a PE node's ``pe_cost_model:`` attrs block + (ADR-0064 D4). Missing keys fall back to defaults.""" + block = attrs.get("pe_cost_model", {}) or {} + d = DEFAULT_PE_COST_MODEL + return PeCostModel( + fixed_per_cmd_cycles=float( + block.get("fixed_per_cmd_cycles", d.fixed_per_cmd_cycles)), + byte_cycles_recip=float( + block.get("byte_cycles_recip", d.byte_cycles_recip)), + max_composite_logical_bytes=int( + block.get("max_composite_logical_bytes", + d.max_composite_logical_bytes)), + ) diff --git a/src/kernbench/components/builtin/pe_cpu.py b/src/kernbench/components/builtin/pe_cpu.py index 786788f..7b923fe 100644 --- a/src/kernbench/components/builtin/pe_cpu.py +++ b/src/kernbench/components/builtin/pe_cpu.py @@ -66,6 +66,14 @@ class PeCpuComponent(ComponentBase): | (self._cube_idx << 32) | (self._pe_idx << 24) ) + # ADR-0064 Rev2: structural dispatch cost model. Reads an optional + # `pe_cost_model:` block from this PE_CPU node's attrs (D4); missing + # keys fall back to defaults. Clock for cycle→ns is this node's + # `clock_freq_ghz` (D3), same attr PE_GEMM/PE_MATH use. + from kernbench.common.pe_cost_model import from_node_attrs + + self._cost_model = from_node_attrs(node.attrs) + self._clock_freq_ghz = float(node.attrs.get("clock_freq_ghz", 1.0)) def _find_shard(self, shards: tuple) -> Any: """Find shard matching this PE's (sip, cube, pe). Fallback to positional index.""" @@ -176,6 +184,8 @@ class PeCpuComponent(ComponentBase): store=store, scratch_base=self._tl_scratch_base, scratch_size=self._tl_scratch_size, + cost_model=self._cost_model, + clock_freq_ghz=self._clock_freq_ghz, ) yield from runner.run(env, kernel_fn, kernel_args, num_programs) return getattr(runner, "_composite_results", []) @@ -184,7 +194,6 @@ class PeCpuComponent(ComponentBase): self, env, kernel_fn, kernel_args, num_programs, scheduler_id, ) -> Generator: """Legacy Phase 0 + replay: generate command list, then dispatch.""" - from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST from kernbench.common.pe_commands import ( CompositeCmd, PeCpuOverheadCmd, PeInternalTxn, WaitCmd, ) @@ -194,7 +203,8 @@ class PeCpuComponent(ComponentBase): pe_id=self._pe_idx, num_programs=num_programs, cube_id=self._cube_idx, num_cubes=self._num_cubes, dispatch_cycles=0, - issue_cost_table=DEFAULT_CPU_ISSUE_COST, + cost_model=self._cost_model, + clock_freq_ghz=self._clock_freq_ghz, ) run_kernel(kernel_fn, tl, *kernel_args) commands = tl.commands diff --git a/src/kernbench/components/builtin/pe_scheduler.py b/src/kernbench/components/builtin/pe_scheduler.py index 8cd17ac..6b6e3ec 100644 --- a/src/kernbench/components/builtin/pe_scheduler.py +++ b/src/kernbench/components/builtin/pe_scheduler.py @@ -156,19 +156,21 @@ class PeSchedulerComponent(ComponentBase): pp = self._pe_prefix bpe = 2 # default bytes per element (f16) - if cmd.op == "gemm" and cmd.b is not None: - a = cmd.a - b = cmd.b + # Flat-ops (ADR-0065 D1): ops[0] is the head, ops[1:] are epilogue + # specs placed by scope. The head's kind selects the engine path. + head = cmd.ops[0] + epi_specs = tuple(cmd.ops[1:]) + + if head.kind == "gemm" and "b" in head.operands: + a = head.operands["a"] + b = head.operands["b"] M, K = a.shape[-2], a.shape[-1] N = b.shape[-1] - # When CompositeCmd.ops is populated, ops[0] is the head and - # ops[1:] is the epilogue spec list. Empty ops → legacy path. - epi_specs = tuple(cmd.ops[1:]) if cmd.ops else () return generate_gemm_plan( M=M, K=K, N=N, tile_m=self.TILE_M, tile_k=self.TILE_K, tile_n=self.TILE_N, bytes_per_element=bpe, - A_addr=a.addr, B_addr=b.addr, C_addr=cmd.out_addr, + A_addr=a.addr, B_addr=b.addr, C_addr=head.out.addr, pe_prefix=pp, a_pinned=getattr(a, "pinned", False), b_pinned=getattr(b, "pinned", False), @@ -176,14 +178,14 @@ class PeSchedulerComponent(ComponentBase): ) else: # Math composite - a = cmd.a + a = head.operands["a"] M = a.shape[-2] if len(a.shape) >= 2 else a.shape[0] N = a.shape[-1] if len(a.shape) >= 2 else 1 return generate_math_plan( M=M, N=N, tile_m=self.TILE_M, tile_n=self.TILE_N, bytes_per_element=bpe, - math_op=cmd.math_op or "identity", - src_addr=a.addr, dst_addr=cmd.out_addr, + math_op=head.extra.get("math_op") or "identity", + src_addr=a.addr, dst_addr=head.out.addr, pe_prefix=pp, ) diff --git a/src/kernbench/sim_engine/op_log.py b/src/kernbench/sim_engine/op_log.py index 922e9a2..2a71473 100644 --- a/src/kernbench/sim_engine/op_log.py +++ b/src/kernbench/sim_engine/op_log.py @@ -248,27 +248,33 @@ def _extract_op_info(msg: Any) -> tuple[str, str, dict[str, Any]]: "nbytes": msg.nbytes, } if isinstance(msg, CompositeCmd): + # Flat-ops (ADR-0065 D1): the head op carries op kind, operands, and + # the write-back handle that previously lived on the cmd directly. + head = msg.ops[0] + op = head.kind params: dict[str, Any] = { - "op": msg.op, - "out_addr": msg.out_addr, - "out_nbytes": msg.out_nbytes, + "op": op, + "out_addr": head.out.addr, + "out_nbytes": head.out.nbytes, } # ADR-0027: preserve operand info so Phase 2 DataExecutor can replay # the composite's numerical effect (treat it like a GemmCmd). - if msg.op == "gemm" and msg.a is not None and msg.b is not None: + a = head.operands.get("a") + b = head.operands.get("b") + if op == "gemm" and a is not None and b is not None: params.update({ - "src_a_addr": msg.a.addr, - "src_b_addr": msg.b.addr, - "shape_a": msg.a.shape, - "shape_b": msg.b.shape, - "dtype_in": msg.a.dtype, - "dtype_out": msg.a.dtype, - "src_a_space": getattr(msg.a, "space", "hbm"), - "src_b_space": getattr(msg.b, "space", "hbm"), + "src_a_addr": a.addr, + "src_b_addr": b.addr, + "shape_a": a.shape, + "shape_b": b.shape, + "dtype_in": a.dtype, + "dtype_out": a.dtype, + "src_a_space": getattr(a, "space", "hbm"), + "src_b_space": getattr(b, "space", "hbm"), "dst_space": "hbm", # dst_addr alias so DataExecutor._execute_gemm picks it up. - "dst_addr": msg.out_addr, + "dst_addr": head.out.addr, }) - return "gemm" if msg.op == "gemm" else "math", f"composite_{msg.op}", params + return "gemm" if op == "gemm" else "math", f"composite_{op}", params # Fallback for unknown data_op messages return "unknown", type(msg).__name__, {} diff --git a/src/kernbench/triton_emu/kernel_runner.py b/src/kernbench/triton_emu/kernel_runner.py index 6b703ad..2ffd3de 100644 --- a/src/kernbench/triton_emu/kernel_runner.py +++ b/src/kernbench/triton_emu/kernel_runner.py @@ -55,6 +55,8 @@ class KernelRunner: ipcq_id: str | None = None, scratch_base: int = 0, scratch_size: int = 1 << 20, + cost_model: Any = None, + clock_freq_ghz: float = 1.0, ) -> None: self._pe_prefix = pe_prefix self._pe_idx = pe_idx @@ -72,6 +74,9 @@ class KernelRunner: # ops produce a result that may later be used as a send/store source. self._scratch_base = scratch_base self._scratch_size = scratch_size + # ADR-0064 Rev2 structural dispatch cost model (+ clock for cycle→ns). + self._cost_model = cost_model + self._clock_freq_ghz = clock_freq_ghz def run( self, @@ -89,7 +94,6 @@ class KernelRunner: 4. Dispatches each command through SimPy components 5. Returns results to the kernel """ - from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST from kernbench.triton_emu.tl_context import TLContext self._parent = greenlet.getcurrent() @@ -103,7 +107,8 @@ class KernelRunner: runner=self, scratch_base=self._scratch_base, scratch_size=self._scratch_size, - issue_cost_table=DEFAULT_CPU_ISSUE_COST, + cost_model=self._cost_model, + clock_freq_ghz=self._clock_freq_ghz, ) self._tl = tl # exposed so switch_to_simpy can re-set on restore diff --git a/src/kernbench/triton_emu/tl_context.py b/src/kernbench/triton_emu/tl_context.py index c38ad5c..2bd7c4b 100644 --- a/src/kernbench/triton_emu/tl_context.py +++ b/src/kernbench/triton_emu/tl_context.py @@ -93,17 +93,16 @@ class TLContext: Args: pe_id: program instance index (returned by program_id). num_programs: total number of program instances. - dispatch_cycles: uniform PE_CPU overhead per tl API call. Used as - a fallback when ``issue_cost_table`` is None (ADR-0046 §D6 - back-compat). When ``issue_cost_table`` is provided, the - per-kind table value is used instead. - issue_cost_table: optional per-op-type CPU issue cost table - (ADR-0064 D1). When provided, each ``tl.*`` call charges the - table value keyed by op kind ("composite", "load", "store", - "dot", "math", "ipcq_send", "ipcq_recv", "copy_to"). Unknown - kinds fall back to ``dispatch_cycles``. Live PE_CPU paths - construct TLContext with ``DEFAULT_CPU_ISSUE_COST`` so the + 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__( @@ -116,14 +115,16 @@ class TLContext: num_cubes: int = 1, scratch_base: int = 0, scratch_size: int = 1 << 20, # 1 MiB per kernel invocation - issue_cost_table: dict[str, int] | None = None, + 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._issue_cost_table = issue_cost_table + self._cost_model = cost_model + self._clock_freq_ghz = clock_freq_ghz self._commands: list[PeCommand] = [] self._handle_counter = 0 self._completion_counter = 0 @@ -193,23 +194,30 @@ class TLContext: def _nbytes(self, shape: tuple[int, ...], dtype: str) -> int: return math.prod(shape) * self._dtype_bytes(dtype) - def _emit_dispatch_overhead(self, kind: str | None = None) -> None: - """Charge per-op-type CPU issue cost (ADR-0064 D1). + def _charge_dispatch(self, cmd: PeCommand) -> None: + """Charge structural PE_CPU dispatch cost (ADR-0064 Rev2 D1). - When ``issue_cost_table`` was provided, look up the per-kind cost - and emit ``PeCpuOverheadCmd(cycles=N)`` if N > 0. Unknown kinds - fall back to the uniform ``dispatch_cycles`` for forward-compat - when a new ``tl.*`` op is added before the table is updated. - - When ``issue_cost_table`` is None, preserve the ADR-0046 §D6 - contract: emit ``PeCpuOverheadCmd(dispatch_cycles)`` if positive. + 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._issue_cost_table is not None and kind is not None: - cycles = self._issue_cost_table.get(kind, self._dispatch_cycles) - else: - cycles = self._dispatch_cycles - if cycles > 0: - self._emit(PeCpuOverheadCmd(cycles=cycles)) + 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(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, @@ -255,7 +263,15 @@ class TLContext: # ── Data Movement (blocking, DMA engine) ────────────────────── def _emit(self, cmd: PeCommand) -> Any: - """Emit command: greenlet switch if runner available, else append to list.""" + """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) @@ -276,7 +292,6 @@ class TLContext: attaches a LoadFuture for structural compatibility (its event stays None — no engine, no SimPy event to wait on). """ - self._emit_dispatch_overhead("load") 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 @@ -292,6 +307,9 @@ class TLContext: ) 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. @@ -320,7 +338,6 @@ class TLContext: def store(self, ptr: int, handle: TensorHandle) -> None: """Store tensor from TCM to HBM.""" self._await_pending(handle) - self._emit_dispatch_overhead("store") cmd = DmaWriteCmd(handle=handle, dst_addr=ptr, nbytes=handle.nbytes) self._emit(cmd) @@ -358,7 +375,6 @@ class TLContext: "reads from HBM go through tl.load" ) self._await_pending(src) - self._emit_dispatch_overhead("copy_to") self._emit(CopyCmd(src=src, dst=dst, nbytes=src.nbytes)) # ── GEMM Engine (blocking) ──────────────────────────────────── @@ -378,7 +394,6 @@ class TLContext: out_dtype = a.dtype out = self._make_compute_out(shape=out_shape, dtype=out_dtype) self._await_pending(a, b) - self._emit_dispatch_overhead("dot") self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n)) return out @@ -387,7 +402,6 @@ class TLContext: 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_dispatch_overhead("math") self._emit(MathCmd(op=op, inputs=(x,), out=out)) return out @@ -421,7 +435,6 @@ class TLContext: out_shape[axis] = 1 out = self._make_compute_out(shape=tuple(out_shape), dtype=x.dtype) self._await_pending(x) - self._emit_dispatch_overhead("math") self._emit(MathCmd(op=op, inputs=(x,), out=out, axis=axis)) return out @@ -441,7 +454,6 @@ class TLContext: ) -> TensorHandle: out = self._make_compute_out(shape=a.shape, dtype=a.dtype) self._await_pending(a, b) - self._emit_dispatch_overhead("math") self._emit(MathCmd(op=op, inputs=(a, b), out=out)) return out @@ -450,7 +462,6 @@ class TLContext: ) -> TensorHandle: out = self._make_compute_out(shape=a.shape, dtype=a.dtype) self._await_pending(cond, a, b) - self._emit_dispatch_overhead("math") self._emit(MathCmd(op="where", inputs=(cond, a, b), out=out)) return out @@ -468,7 +479,6 @@ class TLContext: """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_dispatch_overhead("math") self._emit(MathCmd(op="fma", inputs=(a, b, c), out=out)) return out @@ -481,7 +491,6 @@ class TLContext: """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_dispatch_overhead("math") self._emit(MathCmd(op="clamp", inputs=(x, min, max), out=out)) return out @@ -494,7 +503,6 @@ class TLContext: """ out = self._make_compute_out(shape=x.shape, dtype=x.dtype) self._await_pending(x) - self._emit_dispatch_overhead("math") self._emit(MathCmd(op="softmax", inputs=(x,), out=out, axis=axis)) return out @@ -597,7 +605,6 @@ class TLContext: # 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 - self._emit_dispatch_overhead("ipcq_send") cmd = IpcqSendCmd( direction=dir, src_addr=src_addr, src_space=space, @@ -631,7 +638,6 @@ class TLContext: arrived. In greenlet/runner mode, ``handle.data`` carries the actual ndarray; in command-list mode the handle is a placeholder. """ - self._emit_dispatch_overhead("ipcq_recv") if dst_addr is not None and dst_space is not None: cmd = IpcqRecvCmd( direction=dir, @@ -682,7 +688,6 @@ class TLContext: they receive. This API is segregated from ``tl.recv`` so the diagnostic flag can never accidentally be set in real workloads. """ - self._emit_dispatch_overhead("ipcq_recv") cmd = IpcqRecvCmd( direction=dir, shape=shape, dtype=dtype, @@ -711,13 +716,15 @@ class TLContext: dtype: str = "f16", ) -> "RecvFuture": """Non-blocking recv. Returns a future to pass into ``tl.wait``.""" - self._emit_dispatch_overhead("ipcq_recv") 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)) @@ -749,39 +756,49 @@ class TLContext: # 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) - # Compute output size based on op + # Compute output geometry based on op. 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 - ops_tuple: tuple[OpSpec, ...] = () - if epilogue is not None: - head_operands = (a, b) if (op == "gemm" and b is not None) else (a,) - head_spec = OpSpec( - kind=op, scope=Scope.OUTPUT_TILE, operands=head_operands, - extra={ - k: v for k, v in ( - ("acc_dtype", acc_dtype), - ("tile_shape", tile_shape), - ("math_op", math_op), - ) if v is not None - }, - ) - epi_specs = tuple(self._build_epilogue_spec(e, i) - for i, e in enumerate(epilogue)) - ops_tuple = (head_spec, *epi_specs) + # Head op's write-back handle carries the output address (ADR-0065 D1). + out_handle = TensorHandle( + id=self._next_handle_id(), + addr=out_ptr, shape=out_shape, dtype=out_dtype, + nbytes=out_nbytes, space="tcm", + ) + + # 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 = (head_spec, *epi_specs) completion = CompletionHandle(id=self._next_completion_id()) - self._emit_dispatch_overhead("composite") - self._emit(CompositeCmd( - completion=completion, op=op, - a=a, b=b, out_addr=out_ptr, out_nbytes=out_nbytes, - math_op=math_op, ops=ops_tuple, - )) + self._emit(CompositeCmd(completion=completion, ops=ops_tuple)) return completion @staticmethod @@ -805,13 +822,13 @@ class TLContext: f"{', '.join(missing)}" ) scope = Scope(entry["scope"]) if "scope" in entry else default_scope - operands: list = [] + operands: dict = {} scalar: float | None = None extra: dict = {} for f in required: v = entry[f] if isinstance(v, TensorHandle): - operands.append(v) + operands[f] = v elif isinstance(v, (int, float)): if scalar is None: scalar = float(v) @@ -821,7 +838,7 @@ class TLContext: extra[f] = v return OpSpec( kind=kind, scope=scope, - operands=tuple(operands), scalar=scalar, extra=extra, + operands=operands, scalar=scalar, extra=extra, ) def wait(self, handle: "CompletionHandle | RecvFuture | None" = None) -> Any: diff --git a/tests/test_composite_epilogue.py b/tests/test_composite_epilogue.py index b4a4a5f..fdfaf62 100644 --- a/tests/test_composite_epilogue.py +++ b/tests/test_composite_epilogue.py @@ -56,12 +56,14 @@ def test_composite_epilogue_roundtrip(): ("scale", Scope.OUTPUT_TILE), ] - # Single-op call (no epilogue) keeps the legacy code path: ops stays empty. + # Single-op call (no epilogue): flat-ops still emits one head op + # (ADR-0065 D1 — every composite carries a head op, no legacy empty path). tl2 = TLContext() tl2.composite(op="gemm", a=a, b=b, out_ptr=0x2000) cmd2 = tl2._commands[-1] assert isinstance(cmd2, CompositeCmd) - assert cmd2.ops == () + assert len(cmd2.ops) == 1 + assert cmd2.ops[0].kind == "gemm" @pytest.mark.parametrize("bad,match", [ diff --git a/tests/test_composite_flat_ops_refactor.py b/tests/test_composite_flat_ops_refactor.py new file mode 100644 index 0000000..962db98 --- /dev/null +++ b/tests/test_composite_flat_ops_refactor.py @@ -0,0 +1,127 @@ +"""Phase 1 spec tests for ADR-0065 P1 — flat-ops ``CompositeCmd`` refactor. + +P1 restructures ``CompositeCmd`` from the legacy +``(op, a, b, out_addr, out_nbytes, math_op, ops)`` shape into a flat ordered +op list ``(completion, ops, rw_handles, data_op)`` (ADR-0065 D1/D2). The +user-facing ``tl.composite(op=..., a, b, epilogue=[...])`` API is preserved +(D6.4) and lowers internally to the flat shape. + +This is a *meaning-preserving* refactor: every existing bench's op_log must +stay byte-equal (the existing integration suite is the regression net). The +tests here pin the new lowering *shape* only. + +Phase 1 (this commit): tests only. All FAIL until the P2 production refactor: + - current ``CompositeCmd`` requires ``op/a/b/out_addr`` (no flat-only ctor), + - current ``OpSpec.operands`` is a tuple and there is no ``out`` field. +""" +from __future__ import annotations + +from kernbench.common.pe_commands import CompositeCmd, OpSpec, Scope, TensorHandle +from kernbench.triton_emu.tl_context import TLContext + +_LEGACY_FIELDS = ("op", "a", "b", "out_addr", "out_nbytes", "math_op") + + +def _tl() -> TLContext: + return TLContext(pe_id=0, num_programs=1, dispatch_cycles=0) + + +def _composites(tl: TLContext) -> list[CompositeCmd]: + return [c for c in tl.commands if isinstance(c, CompositeCmd)] + + +# ── plain GEMM composite (no epilogue) ─────────────────────────────── + + +def test_plain_gemm_composite_lowers_to_single_head_op(): + """A composite with no epilogue still emits exactly one flat head op.""" + tl = _tl() + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + tl.composite("gemm", a, b, out_ptr=0x3000) + + comps = _composites(tl) + assert len(comps) == 1, f"expected 1 composite; got {len(comps)}" + cmd = comps[0] + + assert len(cmd.ops) == 1, f"plain gemm must have 1 head op; got {len(cmd.ops)}" + head = cmd.ops[0] + assert head.kind == "gemm" + assert head.operands == {"a": a, "b": b}, head.operands + assert head.out is not None and head.out.addr == 0x3000, head.out + assert cmd.rw_handles == () + + +def test_composite_has_no_legacy_fields(): + """Legacy ``op/a/b/out_addr/out_nbytes/math_op`` are gone from the cmd.""" + tl = _tl() + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + tl.composite("gemm", a, b, out_ptr=0x3000) + cmd = _composites(tl)[0] + for f in _LEGACY_FIELDS: + assert not hasattr(cmd, f), f"legacy field {f!r} must be removed" + + +# ── GEMM + epilogue ────────────────────────────────────────────────── + + +def test_gemm_composite_with_epilogue_lowers_to_flat_ops(): + """``epilogue=[relu, bias]`` lowers to flat ops after the head, with the + epilogue operand migrated from positional tuple to a named dict.""" + tl = _tl() + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + bias = tl.ref(0x4000, shape=(8, 8), dtype="f16") + tl.composite( + "gemm", a, b, out_ptr=0x3000, + epilogue=[{"op": "relu"}, {"op": "bias", "bias": bias}], + ) + cmd = _composites(tl)[0] + + assert [o.kind for o in cmd.ops] == ["gemm", "relu", "bias"] + assert cmd.ops[0].scope == Scope.OUTPUT_TILE # head scope preserved + assert cmd.ops[1].kind == "relu" + assert cmd.ops[1].scope == Scope.OUTPUT_TILE + # bias handle migrated to a named dict keyed by the EPILOGUE_OPS field. + assert cmd.ops[2].operands == {"bias": bias}, cmd.ops[2].operands + + +# ── MATH composite ─────────────────────────────────────────────────── + + +def test_math_composite_lowers_to_flat_ops(): + """``op="math", math_op="exp"`` lowers to a single head op carrying the + math op kind in ``extra`` and its input as a named operand.""" + tl = _tl() + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + tl.composite("math", a, math_op="exp", out_ptr=0x3000) + cmd = _composites(tl)[0] + + assert len(cmd.ops) == 1 + head = cmd.ops[0] + assert head.kind == "math" + assert head.extra.get("math_op") == "exp", head.extra + assert head.operands == {"a": a}, head.operands + assert head.out is not None and head.out.addr == 0x3000 + + +# ── dataclass shape (unit) ─────────────────────────────────────────── + + +def test_opspec_accepts_dict_operands_and_out_handle(): + h = TensorHandle(id="t1", addr=0x10, shape=(4, 4), dtype="f16", nbytes=32) + op = OpSpec(kind="gemm", scope=Scope.OUTPUT_TILE, operands={"a": h}, out=h) + assert op.operands == {"a": h} + assert op.out is h + + +def test_compositecmd_flat_ctor_and_rw_handles_default(): + from kernbench.common.pe_commands import CompletionHandle + + h = TensorHandle(id="t1", addr=0x10, shape=(4, 4), dtype="f16", nbytes=32) + head = OpSpec(kind="math", scope=Scope.KERNEL, operands={"a": h}, out=h) + cmd = CompositeCmd(completion=CompletionHandle(id="c1"), ops=(head,)) + assert cmd.ops == (head,) + assert cmd.rw_handles == () + assert cmd.data_op is True diff --git a/tests/test_pe_cost_model.py b/tests/test_pe_cost_model.py new file mode 100644 index 0000000..f66bdc8 --- /dev/null +++ b/tests/test_pe_cost_model.py @@ -0,0 +1,268 @@ +"""Phase 1 spec tests for ADR-0064 Revision 2 — structural CPU dispatch cost. + +Rev2 replaces the Rev1 per-op-type cost table (``cpu_issue_cost.py``, +``DEFAULT_CPU_ISSUE_COST``) with a structural formula + + dispatch_cycles(cmd) = FIXED_PER_CMD + cmd.logical_bytes * R + +where ``logical_bytes`` is a structural property of each Pe command +(ADR-0064 D2). Tests are written against the *formula* (D1) and the byte +rule (D2), not specific absolute latencies, so they survive recalibration. + +Assumed post-Phase-2 surface: + - ``kernbench.common.pe_cost_model`` exports ``PeCostModel`` + (fields: ``fixed_per_cmd_cycles``, ``byte_cycles_recip``, + ``max_composite_logical_bytes``; method ``dispatch_cycles(lb)``) and + ``DEFAULT_PE_COST_MODEL`` (FIXED=40, R=0.0625, cap=1024). + - ``OpSpec``/``CompositeCmd``/``DmaReadCmd``/… expose ``logical_bytes``. + - ``TLContext(cost_model=...)`` charges ``FIXED + lb*R`` (÷ clock) as a + ``PeCpuOverheadCmd`` before each dispatched command; ``tl.cycles(n)`` + and ``WaitCmd`` are bypassed (D5). + - Per the user-revised D7: a composite whose ``logical_bytes`` exceeds + ``max_composite_logical_bytes`` raises a validation error at emit time + (no auto-segmentation). + +Phase 1 (this commit): tests only. All FAIL until Phase 2: + - ``pe_cost_model`` module does not exist yet, + - ``logical_bytes`` properties do not exist, + - ``TLContext`` has no ``cost_model`` kwarg. +""" +from __future__ import annotations + +import pytest + +from kernbench.common.pe_commands import ( + CompletionHandle, + CompositeCmd, + OpSpec, + PeCpuOverheadCmd, + Scope, + TensorHandle, +) +from kernbench.triton_emu.tl_context import TLContext + + +def _h(addr: int = 0x10, shape: tuple[int, ...] = (64, 64)) -> TensorHandle: + return TensorHandle(id=f"h{addr:x}", addr=addr, shape=shape, + dtype="f16", nbytes=2 * shape[0] * shape[1]) + + +def _gemm_opspec(out_h: TensorHandle) -> OpSpec: + """The ADR-0064 D3 worked example: a single-OpSpec GEMM head.""" + a, b = _h(0x1000), _h(0x2000) + return OpSpec( + kind="gemm", scope=Scope.OUTPUT_TILE, + operands={"a": a, "b": b}, out=out_h, + extra={"m": 64, "k": 64, "n": 64}, + ) + + +# ── PeCostModel defaults (D3) ──────────────────────────────────────── + + +def test_cost_model_defaults(): + from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL + + m = DEFAULT_PE_COST_MODEL + assert m.fixed_per_cmd_cycles == 40 + assert m.byte_cycles_recip == 0.0625 # = 16 bytes/cycle + assert m.max_composite_logical_bytes == 1024 + + +def test_dispatch_cycles_formula(): + """dispatch_cycles(lb) == FIXED + lb*R (ADR-0064 D1).""" + from kernbench.common.pe_cost_model import PeCostModel + + m = PeCostModel(fixed_per_cmd_cycles=40, byte_cycles_recip=0.0625) + assert m.dispatch_cycles(54) == pytest.approx(43.375) # D3 anchor + assert m.dispatch_cycles(0) == 40 + + +# ── logical_bytes byte rule (D2 / D3 worked example) ───────────────── + + +def test_opspec_logical_bytes_single_gemm_is_40(): + """ADR-0064 D3: opcode1+scope1 + len1 + 2 handles16 + out8 + len1 + + m/k/n 12 = 40.""" + op = _gemm_opspec(_h(0x3000)) + assert op.logical_bytes == 40 + + +def test_composite_logical_bytes_single_gemm_is_54(): + """ADR-0064 D3: framing4 + ops-len1 + GEMM 40 + rw-len1 + rw 8 = 54.""" + out_h = _h(0x3000) + comp = CompositeCmd( + completion=CompletionHandle(id="c1"), + ops=(_gemm_opspec(out_h),), + rw_handles=(out_h,), + ) + assert comp.logical_bytes == 54 + + +def test_per_op_summation_no_dedup(): + """ADR-0064 D2 counting rule: each operand handle reference is counted + independently — the same handle in 3 OpSpecs + rw_handles is counted + 4 times, not deduplicated.""" + O = _h(0x9000, shape=(8, 8)) + op = OpSpec(kind="mul", scope=Scope.KERNEL, operands={"s": O}, out=O) + # op.logical_bytes: 1+1 + 1+8(one operand) + 8(out) + 1+0(no extra) = 20 + assert op.logical_bytes == 20 + + with_rw = CompositeCmd( + completion=CompletionHandle(id="c1"), ops=(op, op, op), rw_handles=(O,), + ) + without_rw = CompositeCmd( + completion=CompletionHandle(id="c2"), ops=(op, op, op), rw_handles=(), + ) + # rw_handles entry adds exactly 8 bytes; O is NOT deduplicated against + # the 3 operand references in ops. + assert with_rw.logical_bytes - without_rw.logical_bytes == 8 + # Full identity: 4 + 1 + 3*20 + 1 + 8*1 = 74. + assert with_rw.logical_bytes == 74 + + +def test_primitive_commands_have_logical_bytes(): + """Every dispatched Pe command exposes a positive structural size.""" + from kernbench.common.pe_commands import ( + CopyCmd, DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd, + ) + + a, b, out = _h(0x10), _h(0x20), _h(0x30) + assert DmaReadCmd(handle=a, src_addr=0x10, nbytes=128).logical_bytes > 0 + assert DmaWriteCmd(handle=a, dst_addr=0x10, nbytes=128).logical_bytes > 0 + assert GemmCmd(a=a, b=b, out=out, m=4, k=4, n=4).logical_bytes > 0 + assert MathCmd(op="exp", inputs=(a,), out=out).logical_bytes > 0 + assert CopyCmd(src=a, dst=b, nbytes=128).logical_bytes > 0 + + +# ── formula wiring in TLContext (D1) ───────────────────────────────── + + +def test_tlcontext_charges_formula_per_command(): + """With a cost_model set, each dispatched command is preceded by a + PeCpuOverheadCmd carrying FIXED + lb*R cycles. Using R=0 isolates the + per-command FIXED term (= command-count signal) independent of lb.""" + from kernbench.common.pe_cost_model import PeCostModel + + tl = TLContext( + pe_id=0, num_programs=1, + cost_model=PeCostModel(fixed_per_cmd_cycles=100, byte_cycles_recip=0.0), + ) + tl.load(0x1000, shape=(4, 4), dtype="f16") + overheads = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)] + assert len(overheads) == 1 + assert overheads[0].cycles == 100 + + +def test_tl_cycles_bypasses_dispatch_cost(): + """ADR-0064 Test #4 (D5 bypass): manual tl.cycles(n) issues exactly n, + not n + dispatch_cycles.""" + from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL + + tl = TLContext(pe_id=0, num_programs=1, cost_model=DEFAULT_PE_COST_MODEL) + tl.cycles(7) + overheads = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)] + assert overheads == [PeCpuOverheadCmd(cycles=7)] + + +# ── D7 (user-revised): hard cap → error, no segmentation ───────────── + + +def test_composite_over_cap_raises(): + """A composite whose logical_bytes exceeds max_composite_logical_bytes + raises a validation error at emit time (revised D7 — no auto-split).""" + from kernbench.common.pe_cost_model import PeCostModel + + tl = TLContext( + pe_id=0, num_programs=1, + cost_model=PeCostModel(max_composite_logical_bytes=100), + ) + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + with pytest.raises(ValueError, match="logical_bytes"): + tl.composite( + "gemm", a, b, out_ptr=0x3000, + epilogue=[{"op": "relu"}] * 50, # pushes logical_bytes past 100 + ) + + +def test_composite_within_cap_ok(): + """A normal composite (well under the default 1024 cap) emits fine.""" + from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL + + tl = TLContext(pe_id=0, num_programs=1, cost_model=DEFAULT_PE_COST_MODEL) + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + tl.composite("gemm", a, b, out_ptr=0x3000) # must not raise + comps = [c for c in tl.commands if isinstance(c, CompositeCmd)] + assert len(comps) == 1 + assert comps[0].logical_bytes < 1024 + + +# ── live PE_CPU wiring / path parity (ADR-0064 Test #7) ────────────── + + +def test_live_pe_cpu_applies_cost_model(monkeypatch): + """The live PE_CPU path constructs TLContext with a PeCostModel (not the + removed Rev1 table). Restores the wiring assertion lost with Rev1 and + covers review item #4 / Test #7 (both live paths read the same model).""" + from pathlib import Path + + from kernbench.common.pe_cost_model import PeCostModel + from kernbench.policy.address.phyaddr import PhysAddr + from kernbench.runtime_api.kernel import KernelLaunchMsg, KernelRef + from kernbench.sim_engine.engine import GraphEngine + from kernbench.sim_engine.transaction import Transaction + from kernbench.topology.builder import load_topology + from kernbench.triton_emu import tl_context as _tlc + from kernbench.triton_emu.registry import clear_registry, register_kernel + + captured: list[dict] = [] + real_init = _tlc.TLContext.__init__ + + def spy_init(self, *args, **kwargs): + captured.append(dict(kwargs)) + real_init(self, *args, **kwargs) + + monkeypatch.setattr(_tlc.TLContext, "__init__", spy_init) + + clear_registry() + slice_bytes = 48 * (1 << 30) // 8 + hbm_pa = PhysAddr.pe_hbm_addr( + sip_id=0, die_id=0, pe_id=0, + pe_local_hbm_offset=0x1000, slice_size_bytes=slice_bytes, + ).encode() + + def single_load_kernel(tl): + tl.load(hbm_pa, shape=(4, 4), dtype="f16") + + register_kernel("test_rev2_cost_model_wiring", single_load_kernel) + + topo = Path(__file__).parent.parent / "topology.yaml" + engine = GraphEngine(load_topology(topo)) + pe_cpu_id = "sip0.cube0.pe0.pe_cpu" + done = engine._env.event() + txn = Transaction( + request=KernelLaunchMsg( + correlation_id="t", request_id="r", + kernel_ref=KernelRef(name="test_rev2_cost_model_wiring", kind="builtin"), + args=(), + ), + path=[pe_cpu_id], step=0, nbytes=0, done=done, + ) + + def inject(): + yield engine._components[pe_cpu_id]._inbox.put(txn) + yield done + + engine._env.process(inject()) + engine._env.run() + clear_registry() + + pe_ctx_calls = [c for c in captured + if c.get("pe_id") == 0 and "cost_model" in c] + assert len(pe_ctx_calls) >= 1, ( + f"live PE_CPU must construct TLContext with a cost_model kwarg; " + f"captures={captured}" + ) + assert isinstance(pe_ctx_calls[-1]["cost_model"], PeCostModel) diff --git a/tests/test_phase_e_cpu_issue_cost.py b/tests/test_phase_e_cpu_issue_cost.py deleted file mode 100644 index 023b9a9..0000000 --- a/tests/test_phase_e_cpu_issue_cost.py +++ /dev/null @@ -1,382 +0,0 @@ -"""Phase 1 spec tests for ADR-0064 (per-op-type CPU issue cost model). - -Phase E lands the cost-table machinery and turns it on by default so the -hybrid's CPU-saturation lever (ADR-0060 §1) becomes measurable instead of -modelled away. Today every ``tl.*`` call goes through -``_emit_dispatch_overhead()`` with a single uniform ``dispatch_cycles`` -scalar that is hard-coded to 0 on the live PE_CPU paths -(``pe_cpu.py:_execute_legacy`` and ``kernel_runner.py:run``). - -These tests assume the post-Phase-2 surface: - -- ``kernbench.common.cpu_issue_cost`` exports ``OpKind``, - ``DEFAULT_CPU_ISSUE_COST`` (the ADR-0064 D1 table) and - ``get_issue_cost(kind, table=None) -> int``. -- ``TLContext.__init__`` accepts ``issue_cost_table: dict[str, int] | None``. - When provided, ``_emit_dispatch_overhead(kind)`` looks up the per-kind - cost; when absent, it falls back to the uniform ``dispatch_cycles`` - (back-compat with the existing ADR-0046 §D6 contract). -- Live PE_CPU paths (greenlet + legacy replay) construct TLContext with - ``issue_cost_table=DEFAULT_CPU_ISSUE_COST``. - -Phase 1 (this commit): tests only. All tests FAIL until Phase 2. -""" -from __future__ import annotations - -from pathlib import Path - -import pytest - -from kernbench.common.pe_commands import ( - CompositeCmd, - DmaReadCmd, - DmaWriteCmd, - GemmCmd, - MathCmd, - PeCpuOverheadCmd, -) -from kernbench.policy.address.phyaddr import PhysAddr -from kernbench.runtime_api.kernel import KernelLaunchMsg, KernelRef -from kernbench.sim_engine.engine import GraphEngine -from kernbench.sim_engine.transaction import Transaction -from kernbench.topology.builder import load_topology -from kernbench.triton_emu.registry import clear_registry, register_kernel -from kernbench.triton_emu.tl_context import TLContext, run_kernel - -TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml" - - -def _engine(): - return GraphEngine(load_topology(TOPOLOGY_PATH)) - - -def _hbm_pa(sip: int = 0, cube: int = 0, pe_id: int = 0) -> int: - slice_bytes = 48 * (1 << 30) // 8 - pa = PhysAddr.pe_hbm_addr( - sip_id=sip, die_id=cube, pe_id=pe_id, - pe_local_hbm_offset=0x1000, slice_size_bytes=slice_bytes, - ) - return pa.encode() - - -# ── T1: default table shape ────────────────────────────────────── - - -def test_default_cost_table_has_expected_keys(): - """ADR-0064 D1 table — 8 keys with composite ≫ primitive ratio.""" - from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST - - expected_keys = { - "composite", "load", "store", "dot", "math", - "ipcq_send", "ipcq_recv", "copy_to", - } - assert set(DEFAULT_CPU_ISSUE_COST.keys()) == expected_keys, ( - f"DEFAULT_CPU_ISSUE_COST keys must match ADR-0064 D1; " - f"got {set(DEFAULT_CPU_ISSUE_COST.keys())}" - ) - # Composite is the lever — ratio against primitives must be ≥ 4×. - assert DEFAULT_CPU_ISSUE_COST["composite"] == 40 - for primitive in ("load", "store", "dot", "math", - "ipcq_send", "ipcq_recv", "copy_to"): - assert DEFAULT_CPU_ISSUE_COST[primitive] == 5, ( - f"primitive {primitive!r} default cost must be 5 ns" - ) - - -# ── T2: get_issue_cost lookup ──────────────────────────────────── - - -def test_get_issue_cost_lookup(): - """Helper returns table value; unknown kind returns 0 (no charge).""" - from kernbench.common.cpu_issue_cost import ( - DEFAULT_CPU_ISSUE_COST, - get_issue_cost, - ) - - assert get_issue_cost("composite") == 40 - assert get_issue_cost("load") == 5 - assert get_issue_cost("unknown_kind") == 0 - # Custom table override - custom = {"composite": 100, "load": 1} - assert get_issue_cost("composite", table=custom) == 100 - assert get_issue_cost("load", table=custom) == 1 - assert get_issue_cost("store", table=custom) == 0 - - -# ── T3: TLContext consumes a passed cost table ─────────────────── - - -def test_tlcontext_accepts_issue_cost_table(): - """TLContext(issue_cost_table=...) → tl.load emits the per-kind cycles.""" - tl = TLContext( - pe_id=0, num_programs=1, - dispatch_cycles=0, - issue_cost_table={"load": 7, "store": 3, "dot": 2, "math": 2}, - ) - tl.load(0x1000, shape=(4, 4), dtype="f16") - overheads = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)] - assert len(overheads) == 1, ( - f"expected exactly one PeCpuOverheadCmd before the DmaReadCmd; " - f"got {len(overheads)} (cmds={[type(c).__name__ for c in tl.commands]})" - ) - assert overheads[0].cycles == 7, ( - f"load issue cost from table must be 7; got {overheads[0].cycles}" - ) - - -# ── T4: composite ≫ primitive issue cost differential ──────────── - - -def test_tlcontext_composite_vs_primitive_charges_differ(): - """Composite kernel charges once (40); primitive sequence charges per-op.""" - from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST - - # Composite kernel: 1 load + 1 composite = 5 + 40 = 45 cycles of issue cost. - tl_comp = TLContext( - pe_id=0, num_programs=1, - dispatch_cycles=0, - issue_cost_table=DEFAULT_CPU_ISSUE_COST, - ) - a = tl_comp.load(0x1000, shape=(8, 16), dtype="f16") - b_ref = tl_comp.ref(0x2000, shape=(16, 8), dtype="f16") - tl_comp.composite("gemm", a, b_ref, out_ptr=0x3000) - comp_cycles = sum( - c.cycles for c in tl_comp.commands if isinstance(c, PeCpuOverheadCmd) - ) - - # Primitive kernel: 2 loads + 1 dot + 1 math (exp) = 5 + 5 + 5 + 5 = 20 cycles. - tl_prim = TLContext( - pe_id=0, num_programs=1, - dispatch_cycles=0, - issue_cost_table=DEFAULT_CPU_ISSUE_COST, - ) - a = tl_prim.load(0x1000, shape=(8, 16), dtype="f16") - b = tl_prim.load(0x2000, shape=(16, 8), dtype="f16") - c_out = tl_prim.dot(a, b) - tl_prim.exp(c_out) - prim_cycles = sum( - c.cycles for c in tl_prim.commands if isinstance(c, PeCpuOverheadCmd) - ) - - # ADR-0064 ratio: composite issue cost >> primitive issue cost per op. - # For these specific kernels: composite=45 (5+40), primitive=20 (4×5). - assert comp_cycles == 45, f"composite kernel: expected 45, got {comp_cycles}" - assert prim_cycles == 20, f"primitive kernel: expected 20, got {prim_cycles}" - # Headline assertion: a single composite charges more than all 3 - # post-load primitives combined (40 > 3×5) — the ADR-0060 §1 lever. - assert comp_cycles - 5 > 3 * 5, ( - f"composite issue charge {comp_cycles - 5} must exceed " - f"3× primitive issue charge {3 * 5} (ADR-0064 ratio)" - ) - - -# ── T5: cost table is purely additive (Q2 — no double-count) ───── - - -def test_cost_table_is_additive_only(): - """Q2 invariant: the cost table is additive on PE_CPU, NOT folded into - DMA/GEMM/MATH command shapes. Two TLContexts running the same kernel - with different cost tables must produce identical non-overhead command - sequences (same DmaReadCmd, GemmCmd, MathCmd, addrs, shapes, dtypes). - Only the PeCpuOverheadCmd ``cycles`` field is allowed to differ. - """ - from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST - - def kernel(tl): - a = tl.load(0x1000, shape=(8, 16), dtype="f16") - b = tl.load(0x2000, shape=(16, 8), dtype="f16") - c = tl.dot(a, b) - d = tl.exp(c) - tl.store(0x3000, d) - - tl_zero = TLContext( - pe_id=0, num_programs=1, - dispatch_cycles=0, - issue_cost_table={}, # empty table → every kind = 0 cost - ) - run_kernel(kernel, tl_zero) - - tl_default = TLContext( - pe_id=0, num_programs=1, - dispatch_cycles=0, - issue_cost_table=DEFAULT_CPU_ISSUE_COST, - ) - run_kernel(kernel, tl_default) - - # Strip PeCpuOverheadCmd from both streams; what remains must match. - non_overhead_zero = [ - c for c in tl_zero.commands if not isinstance(c, PeCpuOverheadCmd) - ] - non_overhead_default = [ - c for c in tl_default.commands if not isinstance(c, PeCpuOverheadCmd) - ] - assert len(non_overhead_zero) == len(non_overhead_default), ( - f"non-overhead command count differs: " - f"{len(non_overhead_zero)} vs {len(non_overhead_default)}" - ) - for a_cmd, b_cmd in zip(non_overhead_zero, non_overhead_default): - assert type(a_cmd) is type(b_cmd), ( - f"non-overhead command type changed under cost table: " - f"{type(a_cmd).__name__} vs {type(b_cmd).__name__}" - ) - # Total overhead under zero-table must be 0; under default must be > 0. - cycles_zero = sum( - c.cycles for c in tl_zero.commands if isinstance(c, PeCpuOverheadCmd) - ) - cycles_default = sum( - c.cycles for c in tl_default.commands if isinstance(c, PeCpuOverheadCmd) - ) - assert cycles_zero == 0, f"empty table must add 0 cycles; got {cycles_zero}" - assert cycles_default > 0, ( - f"default table must add > 0 cycles; got {cycles_default}" - ) - - -# ── T6: greenlet vs legacy replay use same cost table (review #6) ─ - - -def test_greenlet_and_legacy_path_parity(): - """Both PE_CPU execution paths read the same cost table. - - The greenlet path (kernel_runner.py:run) and the legacy replay path - (pe_cpu.py:_execute_legacy) must construct TLContext with the same - default cost table so the same kernel produces identical PE_CPU - overhead cycles via either route. This is ADR-0064 review item #6. - """ - from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST - - # Build the command list for a representative kernel via TLContext - # using the default table — this is what both live paths should see. - def kernel(tl): - a = tl.load(0x1000, shape=(4, 4), dtype="f16") - b = tl.load(0x2000, shape=(4, 4), dtype="f16") - c = tl.dot(a, b) - tl.store(0x3000, c) - - tl1 = TLContext( - pe_id=0, num_programs=1, - dispatch_cycles=0, - issue_cost_table=DEFAULT_CPU_ISSUE_COST, - ) - run_kernel(kernel, tl1) - cycles1 = sum( - c.cycles for c in tl1.commands if isinstance(c, PeCpuOverheadCmd) - ) - - tl2 = TLContext( - pe_id=0, num_programs=1, - dispatch_cycles=0, - issue_cost_table=DEFAULT_CPU_ISSUE_COST, - ) - run_kernel(kernel, tl2) - cycles2 = sum( - c.cycles for c in tl2.commands if isinstance(c, PeCpuOverheadCmd) - ) - - assert cycles1 == cycles2, ( - f"same kernel via same default table must produce identical " - f"overhead cycles; got {cycles1} vs {cycles2}" - ) - # Concrete expected for this kernel: 2×load(5) + 1×dot(5) + 1×store(5) = 20. - assert cycles1 == 20, ( - f"expected 20 cycles total (2 load + 1 dot + 1 store at 5 ns); " - f"got {cycles1}" - ) - - -# ── T7: back-compat — no table → uniform dispatch_cycles ───────── - - -def test_back_compat_no_table_uses_dispatch_cycles(): - """ADR-0046 §D6 contract preserved when no issue_cost_table is passed. - - Existing call sites doing ``TLContext(dispatch_cycles=0)`` must - continue to emit zero overhead. Existing call sites doing - ``TLContext(dispatch_cycles=1)`` must continue to emit uniform 1. - """ - # Zero path (most existing tests use this). - tl_zero = TLContext(pe_id=0, num_programs=1, dispatch_cycles=0) - tl_zero.load(0x1000, shape=(4, 4), dtype="f16") - overheads = [c for c in tl_zero.commands if isinstance(c, PeCpuOverheadCmd)] - assert overheads == [], ( - f"dispatch_cycles=0 with no table must emit no overhead; " - f"got {[c.cycles for c in overheads]}" - ) - - # Uniform path (test_dispatch_overhead_inserted relies on this). - tl_one = TLContext(pe_id=0, num_programs=1, dispatch_cycles=1) - tl_one.load(0x1000, shape=(4, 4), dtype="f16") - overheads = [c for c in tl_one.commands if isinstance(c, PeCpuOverheadCmd)] - assert overheads == [PeCpuOverheadCmd(cycles=1)], ( - f"dispatch_cycles=1 with no table must emit cycles=1; " - f"got {[c.cycles for c in overheads]}" - ) - - -# ── T8: live PE_CPU constructs TLContext with the default table ── - - -def test_live_pe_cpu_uses_default_table(monkeypatch): - """End-to-end: live PE_CPU greenlet path constructs TLContext with - ``issue_cost_table=DEFAULT_CPU_ISSUE_COST``. This is the wiring assertion - for ADR-0064 D4 "active by default" + review item #6 (path parity). - - We patch ``TLContext.__init__`` to record its kwargs and run a - single-load kernel through the live PE_CPU. The recorded - ``issue_cost_table`` must equal ``DEFAULT_CPU_ISSUE_COST``. - """ - from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST - from kernbench.triton_emu import tl_context as _tlc - - captured: list[dict] = [] - real_init = _tlc.TLContext.__init__ - - def spy_init(self, *args, **kwargs): - captured.append(dict(kwargs)) - real_init(self, *args, **kwargs) - - monkeypatch.setattr(_tlc.TLContext, "__init__", spy_init) - - clear_registry() - hbm_pa = _hbm_pa(sip=0, cube=0, pe_id=0) - - def single_load_kernel(tl): - tl.load(hbm_pa, shape=(4, 4), dtype="f16") - - register_kernel("test_active_default_table", single_load_kernel) - - engine = _engine() - pe_cpu_id = "sip0.cube0.pe0.pe_cpu" - done = engine._env.event() - txn = Transaction( - request=KernelLaunchMsg( - correlation_id="t", request_id="r", - kernel_ref=KernelRef(name="test_active_default_table", kind="builtin"), - args=(), - ), - path=[pe_cpu_id], step=0, nbytes=0, done=done, - ) - - def inject(): - yield engine._components[pe_cpu_id]._inbox.put(txn) - yield done - - engine._env.process(inject()) - engine._env.run() - clear_registry() - - # Live PE_CPU constructs TLContext on either the greenlet path - # (kernel_runner.py:run) or the legacy replay path - # (pe_cpu.py:_execute_legacy) — whichever is chosen depends on whether - # MemoryStore is wired. Both must pass DEFAULT_CPU_ISSUE_COST. - pe_ctx_calls = [c for c in captured if c.get("pe_id") == 0 - and "issue_cost_table" in c] - assert len(pe_ctx_calls) >= 1, ( - f"expected at least one TLContext constructed by a live PE_CPU " - f"path with an issue_cost_table kwarg; got {len(pe_ctx_calls)} " - f"(all captures: {captured})" - ) - last = pe_ctx_calls[-1] - assert last["issue_cost_table"] == DEFAULT_CPU_ISSUE_COST, ( - f"live PE_CPU must use DEFAULT_CPU_ISSUE_COST; got {last['issue_cost_table']}" - ) diff --git a/tests/test_triton_emu.py b/tests/test_triton_emu.py index 160ec4b..85ba537 100644 --- a/tests/test_triton_emu.py +++ b/tests/test_triton_emu.py @@ -178,9 +178,11 @@ def test_tl_composite_nonblocking(): assert isinstance(h, CompletionHandle) comp_cmds = [c for c in tl.commands if isinstance(c, CompositeCmd)] assert len(comp_cmds) == 1 - assert comp_cmds[0].op == "gemm" - assert comp_cmds[0].out_addr == 0x3000 - assert comp_cmds[0].out_nbytes == 32 * 32 * 2 # M×N×dtype_bytes + # Flat-ops shape (ADR-0065 D1): head op carries kind + write-back handle. + head = comp_cmds[0].ops[0] + assert head.kind == "gemm" + assert head.out.addr == 0x3000 + assert head.out.nbytes == 32 * 32 * 2 # M×N×dtype_bytes # ── 8. tl.wait(handle) → WaitCmd ─────────────────────────────────