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:
@@ -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
|
||||
Reference in New Issue
Block a user