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