Implement ADR-0020: 2-pass data execution with greenlet kernel runner
Step 1 — Foundation: - OpRecord/OpLogger: op log infrastructure with t_start stable ordering - MemoryStore: numpy ndarray tensor-granular storage (reference semantics) - data_op=True flag on DmaReadCmd, DmaWriteCmd, GemmCmd, MathCmd, CompositeCmd - numpy/greenlet dependencies added to pyproject.toml Step 2 — ComponentBase hooks: - _on_process_start/end hooks in _forward_txn (fabric messages) - _handle_with_hooks in PeEngineBase (PE-internal commands) - op_logger optional — zero overhead when disabled Step 3 — KernelRunner + greenlet: - KernelRunner: greenlet ↔ SimPy bridge in triton_emu/kernel_runner.py - TLContext: _emit() method routes to greenlet switch or command list - tl.load() returns real numpy data in greenlet mode - Dynamic control flow supported (memory-read based branching) Step 4 — PE_CPU integration: - Greenlet mode when ctx.memory_store is set, legacy fallback otherwise - Refactored into _execute_greenlet/_execute_legacy/_send_response - ComponentContext gains memory_store and op_logger fields Step 5 — DataExecutor: - Phase 2 numpy execution for GEMM/Math ops from op_log - _compute_math: all unary/binary/reduction ops - verify(): compare MemoryStore against expected with dtype tolerance 28 new tests, 366 total passing. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,188 @@
|
||||
"""Tests for DataExecutor Phase 2 execution (ADR-0020 D6)."""
|
||||
import numpy as np
|
||||
|
||||
from kernbench.sim_engine.data_executor import DataExecutor
|
||||
from kernbench.sim_engine.memory_store import MemoryStore
|
||||
from kernbench.sim_engine.op_log import OpRecord
|
||||
|
||||
|
||||
def test_gemm_execution():
|
||||
"""Phase 2 GEMM: out = a @ b with f32 accumulation."""
|
||||
store = MemoryStore()
|
||||
a = np.ones((4, 8), dtype=np.float16)
|
||||
b = np.ones((8, 4), dtype=np.float16) * 2.0
|
||||
store.write("tcm", 0x0, a)
|
||||
store.write("tcm", 0x100, b)
|
||||
|
||||
op = OpRecord(
|
||||
t_start=0.0, t_end=100.0,
|
||||
component_id="pe_gemm",
|
||||
op_kind="gemm", op_name="gemm_f16",
|
||||
params={
|
||||
"src_a_addr": 0x0, "src_b_addr": 0x100, "dst_addr": 0x200,
|
||||
"shape_a": (4, 8), "shape_b": (8, 4), "shape_out": (4, 4),
|
||||
"dtype_in": "f16", "dtype_acc": "f32", "dtype_out": "f16",
|
||||
"addr_space": "tcm",
|
||||
},
|
||||
)
|
||||
|
||||
executor = DataExecutor([op], store)
|
||||
executor.run()
|
||||
|
||||
result = store.read("tcm", 0x200)
|
||||
expected = (a.astype(np.float32) @ b.astype(np.float32)).astype(np.float16)
|
||||
assert np.allclose(result, expected)
|
||||
|
||||
|
||||
def test_math_exp():
|
||||
store = MemoryStore()
|
||||
x = np.array([0.0, 1.0, 2.0], dtype=np.float32)
|
||||
store.write("tcm", 0x0, x)
|
||||
|
||||
op = OpRecord(
|
||||
t_start=0.0, t_end=10.0,
|
||||
component_id="pe_math",
|
||||
op_kind="math", op_name="exp",
|
||||
params={
|
||||
"op": "exp",
|
||||
"input_addrs": [0x0], "input_shapes": [(3,)],
|
||||
"dst_addr": 0x100, "shape_out": (3,),
|
||||
"dtype": "f32", "axis": None, "addr_space": "tcm",
|
||||
},
|
||||
)
|
||||
|
||||
executor = DataExecutor([op], store)
|
||||
executor.run()
|
||||
|
||||
result = store.read("tcm", 0x100)
|
||||
assert np.allclose(result, np.exp(x))
|
||||
|
||||
|
||||
def test_math_add():
|
||||
store = MemoryStore()
|
||||
a = np.array([1.0, 2.0], dtype=np.float32)
|
||||
b = np.array([3.0, 4.0], dtype=np.float32)
|
||||
store.write("tcm", 0x0, a)
|
||||
store.write("tcm", 0x100, b)
|
||||
|
||||
op = OpRecord(
|
||||
t_start=0.0, t_end=5.0,
|
||||
component_id="pe_math",
|
||||
op_kind="math", op_name="add",
|
||||
params={
|
||||
"op": "add",
|
||||
"input_addrs": [0x0, 0x100], "input_shapes": [(2,), (2,)],
|
||||
"dst_addr": 0x200, "shape_out": (2,),
|
||||
"dtype": "f32", "axis": None, "addr_space": "tcm",
|
||||
},
|
||||
)
|
||||
|
||||
executor = DataExecutor([op], store)
|
||||
executor.run()
|
||||
|
||||
result = store.read("tcm", 0x200)
|
||||
assert np.array_equal(result, np.array([4.0, 6.0], dtype=np.float32))
|
||||
|
||||
|
||||
def test_math_sum_reduction():
|
||||
store = MemoryStore()
|
||||
x = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
|
||||
store.write("tcm", 0x0, x)
|
||||
|
||||
op = OpRecord(
|
||||
t_start=0.0, t_end=5.0,
|
||||
component_id="pe_math",
|
||||
op_kind="math", op_name="sum",
|
||||
params={
|
||||
"op": "sum",
|
||||
"input_addrs": [0x0], "input_shapes": [(2, 2)],
|
||||
"dst_addr": 0x100, "shape_out": (1, 2),
|
||||
"dtype": "f32", "axis": 0, "addr_space": "tcm",
|
||||
},
|
||||
)
|
||||
|
||||
executor = DataExecutor([op], store)
|
||||
executor.run()
|
||||
|
||||
result = store.read("tcm", 0x100)
|
||||
assert np.array_equal(result, np.array([[4.0, 6.0]], dtype=np.float32))
|
||||
|
||||
|
||||
def test_verify_pass():
|
||||
store = MemoryStore()
|
||||
store.write("hbm", 0x0, np.array([1.0, 2.0], dtype=np.float32))
|
||||
|
||||
executor = DataExecutor([], store)
|
||||
results = executor.verify({
|
||||
("hbm", 0x0): np.array([1.0, 2.0], dtype=np.float32),
|
||||
})
|
||||
assert results["hbm:0x0"] is True
|
||||
|
||||
|
||||
def test_verify_fail():
|
||||
store = MemoryStore()
|
||||
store.write("hbm", 0x0, np.array([1.0, 2.0], dtype=np.float32))
|
||||
|
||||
executor = DataExecutor([], store)
|
||||
results = executor.verify({
|
||||
("hbm", 0x0): np.array([9.0, 9.0], dtype=np.float32),
|
||||
})
|
||||
assert results["hbm:0x0"] is False
|
||||
|
||||
|
||||
def test_memory_ops_skipped():
|
||||
"""Memory ops in op_log should be skipped (handled in Phase 1)."""
|
||||
store = MemoryStore()
|
||||
op = OpRecord(
|
||||
t_start=0.0, t_end=5.0,
|
||||
component_id="pe_dma",
|
||||
op_kind="memory", op_name="dma_read",
|
||||
params={"src_addr": 0x0, "nbytes": 64, "handle_id": "t0"},
|
||||
)
|
||||
# Should not raise
|
||||
executor = DataExecutor([op], store)
|
||||
executor.run()
|
||||
|
||||
|
||||
def test_sequential_gemm_then_math():
|
||||
"""GEMM output feeds into math op."""
|
||||
store = MemoryStore()
|
||||
a = np.eye(2, dtype=np.float16)
|
||||
b = np.ones((2, 2), dtype=np.float16)
|
||||
store.write("tcm", 0x0, a)
|
||||
store.write("tcm", 0x100, b)
|
||||
|
||||
ops = [
|
||||
OpRecord(
|
||||
t_start=0.0, t_end=50.0,
|
||||
component_id="pe_gemm",
|
||||
op_kind="gemm", op_name="gemm_f16",
|
||||
params={
|
||||
"src_a_addr": 0x0, "src_b_addr": 0x100, "dst_addr": 0x200,
|
||||
"shape_a": (2, 2), "shape_b": (2, 2), "shape_out": (2, 2),
|
||||
"dtype_in": "f16", "dtype_acc": "f32", "dtype_out": "f32",
|
||||
"addr_space": "tcm",
|
||||
},
|
||||
),
|
||||
OpRecord(
|
||||
t_start=50.0, t_end=55.0,
|
||||
component_id="pe_math",
|
||||
op_kind="math", op_name="exp",
|
||||
params={
|
||||
"op": "exp",
|
||||
"input_addrs": [0x200], "input_shapes": [(2, 2)],
|
||||
"dst_addr": 0x300, "shape_out": (2, 2),
|
||||
"dtype": "f32", "axis": None, "addr_space": "tcm",
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
executor = DataExecutor(ops, store)
|
||||
executor.run()
|
||||
|
||||
gemm_result = store.read("tcm", 0x200)
|
||||
expected_gemm = (a.astype(np.float32) @ b.astype(np.float32)).astype(np.float32)
|
||||
assert np.allclose(gemm_result, expected_gemm)
|
||||
|
||||
exp_result = store.read("tcm", 0x300)
|
||||
assert np.allclose(exp_result, np.exp(expected_gemm))
|
||||
Reference in New Issue
Block a user