ywkang
3d4a3d43c4
data_executor._compute_math gains the 5 recipe ops (rmax/rsum as keepdims reductions, max_elem/exp_diff/mul_bcast as binary; numpy broadcasting covers bcast_axis). tiling._math_stage now carries operand+output addrs/shapes/spaces + axis on the recipe prologue MATH stages. op_log.record_end promotes a MATH stage to op_kind='math' ONLY when it carries input_addrs -> the DataExecutor runs it; legacy epilogue MATH stages (bias/relu, no addrs) stay op_kind-opaque -> byte-equal. Fixed a latent P2 lowering bug exposed by data mode: _resolve_recipe_dst gave reductions shape (G,) (axis removed) but max_elem broadcasts them against the running m=(G,1); reductions now keep the reduced axis as size 1 (keepdims) -> (G,1). Verified end-to-end: running the recipe's 8 MATH ops through the DataExecutor produces m, l (fully updated online-softmax) and O (rescaled by corr) matching a numpy reference. Suite 815 pass / 3 pre-existing fail. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
96 lines
4.2 KiB
Python
96 lines
4.2 KiB
Python
"""Data-mode numeric tests for the softmax_merge recipe MATH ops (ADR-0065 N2).
|
|
|
|
N2 makes the recipe's 8 MATH ops compute in data mode: `_compute_math` gains
|
|
rmax/rsum/max_elem/exp_diff/mul_bcast, the recipe MATH stages carry operand
|
|
addresses (`_math_stage`), and op_log `record_end` promotes an address-
|
|
carrying MATH stage to op_kind="math". The MATH ops fully compute the online-
|
|
softmax update of `m` and `l` (and rescale `O`); the `+P·V` of `O` is the
|
|
GEMM-accumulate step (N3).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from kernbench.common.pe_commands import CompositeCmd, TensorHandle
|
|
from kernbench.components.builtin.tiling import generate_plan_from_ops
|
|
from kernbench.sim_engine.data_executor import DataExecutor, _compute_math
|
|
from kernbench.sim_engine.memory_store import MemoryStore
|
|
from kernbench.sim_engine.op_log import OpRecord
|
|
from kernbench.triton_emu.tl_context import TLContext
|
|
|
|
G, TILE, D = 8, 64, 128
|
|
|
|
|
|
def test_compute_math_recipe_ops():
|
|
x = np.random.randn(4, 6).astype(np.float32)
|
|
y = np.random.randn(4, 1).astype(np.float32)
|
|
assert np.allclose(_compute_math("rmax", [x], -1), x.max(-1, keepdims=True))
|
|
assert np.allclose(_compute_math("rsum", [x], -1), x.sum(-1, keepdims=True))
|
|
assert np.allclose(_compute_math("max_elem", [y, x.max(-1, keepdims=True)], None),
|
|
np.maximum(y, x.max(-1, keepdims=True)))
|
|
assert np.allclose(_compute_math("exp_diff", [x, y], None), np.exp(x - y))
|
|
assert np.allclose(_compute_math("mul_bcast", [x, y], None), x * y)
|
|
|
|
|
|
def _tcm(addr, shape):
|
|
return TensorHandle(id=f"h{addr:x}", addr=addr, shape=shape, dtype="f16",
|
|
nbytes=2 * int(np.prod(shape)), space="tcm")
|
|
|
|
|
|
def test_softmax_merge_computes_m_l_O_rescale():
|
|
"""Run the recipe's prologue MATH ops through the DataExecutor and check
|
|
m, l (fully updated) and O (rescaled by corr) match a numpy reference."""
|
|
rng = np.random.default_rng(0)
|
|
s = rng.standard_normal((G, TILE)).astype(np.float16)
|
|
m0 = rng.standard_normal((G, 1)).astype(np.float16)
|
|
l0 = np.abs(rng.standard_normal((G, 1))).astype(np.float16)
|
|
O0 = rng.standard_normal((G, D)).astype(np.float16)
|
|
|
|
sh = _tcm(0x1000, (G, TILE))
|
|
mh = _tcm(0x2000, (G, 1))
|
|
lh = _tcm(0x3000, (G, 1))
|
|
Oh = _tcm(0x4000, (G, D))
|
|
V = TensorHandle(id="V", addr=0x5000, shape=(TILE, D), dtype="f16",
|
|
nbytes=2 * TILE * D, space="hbm")
|
|
|
|
tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000,
|
|
scratch_size=1 << 20)
|
|
tl.composite(
|
|
prologue=[{"op": "softmax_merge", "s": sh, "m": mh, "l": lh, "O": Oh}],
|
|
op="gemm", b=V, out=Oh,
|
|
)
|
|
cmd = [c for c in tl.commands if isinstance(c, CompositeCmd)][-1]
|
|
plan = generate_plan_from_ops(cmd.ops, tile_m=32, tile_k=32, tile_n=32,
|
|
bytes_per_element=2,
|
|
pe_prefix="sip0.cube0.pe0")
|
|
|
|
store = MemoryStore()
|
|
store.write("tcm", sh.addr, s)
|
|
store.write("tcm", mh.addr, m0)
|
|
store.write("tcm", lh.addr, l0)
|
|
store.write("tcm", Oh.addr, O0)
|
|
|
|
records = [
|
|
OpRecord(t_start=float(i), t_end=float(i), component_id="pe_math",
|
|
op_kind="math", op_name=st.params["op"], params=dict(st.params))
|
|
for i, st in enumerate(plan.prologue_stages)
|
|
]
|
|
DataExecutor(records, store).run()
|
|
|
|
# numpy reference (online-softmax merge of a single new tile into m/l/O).
|
|
sf = s.astype(np.float32)
|
|
m_loc = sf.max(-1, keepdims=True)
|
|
m_new = np.maximum(m0.astype(np.float32), m_loc)
|
|
corr = np.exp(m0.astype(np.float32) - m_new)
|
|
P = np.exp(sf - m_new)
|
|
l_new = l0.astype(np.float32) * corr + P.sum(-1, keepdims=True)
|
|
O_resc = O0.astype(np.float32) * corr
|
|
|
|
m_got = store.read("tcm", mh.addr, shape=(G, 1), dtype="f16").astype(np.float32)
|
|
l_got = store.read("tcm", lh.addr, shape=(G, 1), dtype="f16").astype(np.float32)
|
|
O_got = store.read("tcm", Oh.addr, shape=(G, D), dtype="f16").astype(np.float32)
|
|
|
|
assert np.allclose(m_got, m_new, atol=2e-2), f"m: {m_got.ravel()[:3]} vs {m_new.ravel()[:3]}"
|
|
assert np.allclose(l_got, l_new, rtol=5e-2, atol=5e-2), "l mismatch"
|
|
assert np.allclose(O_got, O_resc, rtol=5e-2, atol=5e-2), "O rescale mismatch"
|