ywkang
51004c311c
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>
86 lines
2.3 KiB
Python
86 lines
2.3 KiB
Python
"""Tests for MemoryStore (ADR-0020 D7)."""
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from kernbench.sim_engine.memory_store import MemoryStore
|
|
|
|
|
|
def test_write_read_reference():
|
|
"""Write and read return the same numpy array (no copy)."""
|
|
store = MemoryStore()
|
|
data = np.ones((4, 4), dtype=np.float16)
|
|
store.write("tcm", 0x0, data)
|
|
result = store.read("tcm", 0x0)
|
|
assert result is data
|
|
|
|
|
|
def test_overwrite_replaces():
|
|
"""Same addr write replaces the previous tensor."""
|
|
store = MemoryStore()
|
|
data1 = np.zeros((4,), dtype=np.float32)
|
|
data2 = np.ones((4,), dtype=np.float32)
|
|
store.write("hbm", 0x100, data1)
|
|
store.write("hbm", 0x100, data2)
|
|
result = store.read("hbm", 0x100)
|
|
assert result is data2
|
|
|
|
|
|
def test_read_missing_raises():
|
|
store = MemoryStore()
|
|
with pytest.raises(KeyError):
|
|
store.read("hbm", 0x999)
|
|
|
|
|
|
def test_read_different_space():
|
|
store = MemoryStore()
|
|
data = np.array([1, 2, 3], dtype=np.int32)
|
|
store.write("tcm", 0x0, data)
|
|
with pytest.raises(KeyError):
|
|
store.read("hbm", 0x0) # different space
|
|
|
|
|
|
def test_dtype_reinterpret():
|
|
"""Read with different dtype does view cast."""
|
|
store = MemoryStore()
|
|
data = np.array([1.0, 2.0], dtype=np.float32) # 8 bytes
|
|
store.write("tcm", 0x0, data)
|
|
result = store.read("tcm", 0x0, dtype="u8")
|
|
assert result.dtype == np.uint8
|
|
assert result.nbytes == data.nbytes
|
|
|
|
|
|
def test_reshape():
|
|
store = MemoryStore()
|
|
data = np.arange(12, dtype=np.float32)
|
|
store.write("tcm", 0x0, data)
|
|
result = store.read("tcm", 0x0, shape=(3, 4))
|
|
assert result.shape == (3, 4)
|
|
|
|
|
|
def test_shape_mismatch_raises():
|
|
store = MemoryStore()
|
|
data = np.arange(12, dtype=np.float32)
|
|
store.write("tcm", 0x0, data)
|
|
with pytest.raises(ValueError, match="Shape mismatch"):
|
|
store.read("tcm", 0x0, shape=(5, 5))
|
|
|
|
|
|
def test_has():
|
|
store = MemoryStore()
|
|
assert not store.has("tcm", 0x0)
|
|
store.write("tcm", 0x0, np.array([1]))
|
|
assert store.has("tcm", 0x0)
|
|
|
|
|
|
def test_snapshot():
|
|
store = MemoryStore()
|
|
data = np.ones((4,), dtype=np.float16)
|
|
store.write("hbm", 0x0, data)
|
|
|
|
snap = store.snapshot()
|
|
assert snap.read("hbm", 0x0) is data # same reference
|
|
|
|
# Modifying snap doesn't affect original
|
|
snap.write("hbm", 0x0, np.zeros((4,), dtype=np.float16))
|
|
assert store.read("hbm", 0x0) is data
|