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) <noreply@anthropic.com>
This commit is contained in:
@@ -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 +
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
@@ -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) ───────────────────
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)),
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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__, {}
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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),
|
||||
# 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
|
||||
},
|
||||
("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))
|
||||
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:
|
||||
|
||||
@@ -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", [
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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']}"
|
||||
)
|
||||
@@ -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 ─────────────────────────────────
|
||||
|
||||
Reference in New Issue
Block a user