ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa(adr-0064/0065): flat-ops CompositeCmd (P1) + structural dispatch cost (ADR-0064 Rev2); promote ADR-0064

Browse Source 
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:
Yangwook
2026-06-10 19:18:04 -07:00
parent 79ddb12b42
commit 47e2c78c66
18 changed files with 703 additions and 550 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/triton_emu/tl_context.py
+86 -69
View File
@@ -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: