269 lines
8.3 KiB
Python
269 lines
8.3 KiB
Python
import pytest
|
|
|
|
from kernbench.policy.address.allocator import AddressConfig, AllocationError, PEMemAllocator
|
|
from kernbench.policy.address.phyaddr import PhysAddr, PhysAddrError, UnitType
|
|
|
|
_MB = 1 << 20
|
|
_GB = 1 << 30
|
|
|
|
# Topology-matching config: 48GB HBM / 8 slices / 16MB TCM / 4MB reserved / 32MB SRAM
|
|
_CFG = AddressConfig(
|
|
sip_count=2,
|
|
cubes_per_sip=16,
|
|
pes_per_cube=8,
|
|
hbm_bytes_per_cube=48 * _GB,
|
|
hbm_slices_per_cube=8,
|
|
tcm_bytes_per_pe=16 * _MB,
|
|
tcm_scheduler_reserved_bytes=4 * _MB,
|
|
sram_bytes_per_cube=32 * _MB,
|
|
)
|
|
|
|
|
|
# ── Immutability & value semantics ──────────────────────────────────
|
|
|
|
|
|
def test_physaddr_immutable():
|
|
pa = PhysAddr.hbm_addr(rack_id=0, sip_id=0, cube_id=0, hbm_offset=0)
|
|
with pytest.raises(AttributeError):
|
|
pa.rack_id = 1 # type: ignore[misc]
|
|
# hashable
|
|
{pa}
|
|
# comparable
|
|
pa2 = PhysAddr.hbm_addr(rack_id=0, sip_id=0, cube_id=0, hbm_offset=0)
|
|
assert pa == pa2
|
|
|
|
|
|
# ── HBM encode/decode roundtrip ────────────────────────────────────
|
|
|
|
|
|
def test_hbm_encode_decode_roundtrip():
|
|
pa = PhysAddr.hbm_addr(rack_id=2, sip_id=3, cube_id=5, hbm_offset=0x1000)
|
|
raw = pa.encode()
|
|
dec = PhysAddr.decode(raw)
|
|
assert dec.rack_id == 2
|
|
assert dec.sip_id == 3
|
|
assert dec.cube_id == 5
|
|
assert dec.kind == "hbm"
|
|
assert dec.hbm_offset == 0x1000
|
|
|
|
|
|
# ── PE resource encode/decode roundtrip ─────────────────────────────
|
|
|
|
|
|
def test_pe_resource_encode_decode_roundtrip():
|
|
pa = PhysAddr(
|
|
rack_id=1, sip_id=2, sip_seg=7, local_offset=0,
|
|
kind="pe_resource", cube_id=7,
|
|
unit_type=UnitType.PE, pe_id=3, ext=1, sub_offset=0xFF,
|
|
)
|
|
# manually build local_offset matching bit layout
|
|
local_offset = (UnitType.PE << 34) | (3 << 30) | (1 << 29) | 0xFF
|
|
pa2 = PhysAddr(
|
|
rack_id=1, sip_id=2, sip_seg=7, local_offset=local_offset,
|
|
kind="pe_resource", cube_id=7,
|
|
unit_type=UnitType.PE, pe_id=3, ext=1, sub_offset=0xFF,
|
|
)
|
|
raw = pa2.encode()
|
|
dec = PhysAddr.decode(raw)
|
|
assert dec.kind == "pe_resource"
|
|
assert dec.unit_type == UnitType.PE
|
|
assert dec.pe_id == 3
|
|
assert dec.ext == 1
|
|
assert dec.sub_offset == 0xFF
|
|
|
|
|
|
# ── pe_hbm_addr factory ────────────────────────────────────────────
|
|
|
|
|
|
def test_pe_hbm_addr_factory():
|
|
SLICE = 6 * (1 << 30) # 6 GB per PE slice
|
|
pa = PhysAddr.pe_hbm_addr(
|
|
rack_id=0, sip_id=0, cube_id=0,
|
|
pe_id=2, pe_local_hbm_offset=1024, slice_size_bytes=SLICE,
|
|
)
|
|
assert pa.kind == "hbm"
|
|
assert pa.cube_id == 0
|
|
assert pa.hbm_offset == 2 * SLICE + 1024
|
|
|
|
|
|
def test_pe_hbm_addr_overflow():
|
|
SLICE = 6 * (1 << 30)
|
|
with pytest.raises(PhysAddrError, match="pe_local_hbm_offset"):
|
|
PhysAddr.pe_hbm_addr(
|
|
rack_id=0, sip_id=0, cube_id=0,
|
|
pe_id=0, pe_local_hbm_offset=SLICE, slice_size_bytes=SLICE,
|
|
)
|
|
|
|
|
|
# ── Invalid unit_type decode (fix #1) ──────────────────────────────
|
|
|
|
|
|
def test_invalid_unit_type_raises():
|
|
# Craft a PE-resource address with unit_type=7 (invalid)
|
|
local_offset = (7 << 34) | (0 << 30) | 0
|
|
pa_raw = PhysAddr(
|
|
rack_id=0, sip_id=0, sip_seg=0, local_offset=local_offset,
|
|
)
|
|
raw = pa_raw.encode()
|
|
with pytest.raises(PhysAddrError, match="unit_type"):
|
|
PhysAddr.decode(raw)
|
|
|
|
|
|
# ── hbm_pe_id utility (fix #3) ─────────────────────────────────────
|
|
|
|
|
|
def test_hbm_pe_id_utility():
|
|
SLICE = 6 * (1 << 30) # 6 GB
|
|
pa = PhysAddr.pe_hbm_addr(
|
|
rack_id=0, sip_id=0, cube_id=0,
|
|
pe_id=5, pe_local_hbm_offset=256, slice_size_bytes=SLICE,
|
|
)
|
|
assert PhysAddr.hbm_pe_id(pa.hbm_offset, SLICE) == 5
|
|
|
|
|
|
# ── UnitType.SRAM exists (fix #5) ──────────────────────────────────
|
|
|
|
|
|
def test_sram_unit_type_exists():
|
|
assert UnitType.SRAM == 2
|
|
|
|
|
|
# ── cube_sram_addr factory + roundtrip ──────────────────────────────
|
|
|
|
|
|
def test_cube_sram_addr_roundtrip():
|
|
pa = PhysAddr.cube_sram_addr(
|
|
rack_id=0, sip_id=1, cube_id=3, sram_offset=0x800,
|
|
)
|
|
assert pa.kind == "pe_resource"
|
|
assert pa.unit_type == UnitType.SRAM
|
|
assert pa.cube_id == 3
|
|
assert pa.sub_offset == 0x800
|
|
# encode → decode roundtrip
|
|
dec = PhysAddr.decode(pa.encode())
|
|
assert dec.unit_type == UnitType.SRAM
|
|
assert dec.cube_id == 3
|
|
assert dec.sub_offset == 0x800
|
|
|
|
|
|
def test_cube_sram_addr_range_check():
|
|
with pytest.raises(PhysAddrError):
|
|
PhysAddr.cube_sram_addr(
|
|
rack_id=0, sip_id=0, cube_id=0,
|
|
sram_offset=(1 << 29), # exceeds 29-bit sub_offset
|
|
)
|
|
|
|
|
|
# ── pe_tcm_addr factory + roundtrip ────────────────────────────────
|
|
|
|
|
|
def test_pe_tcm_addr_roundtrip():
|
|
pa = PhysAddr.pe_tcm_addr(
|
|
rack_id=0, sip_id=0, cube_id=2, pe_id=7, tcm_offset=0x400,
|
|
)
|
|
assert pa.kind == "pe_resource"
|
|
assert pa.unit_type == UnitType.PE
|
|
assert pa.pe_id == 7
|
|
assert pa.cube_id == 2
|
|
assert pa.sub_offset == 0x400
|
|
# encode → decode roundtrip
|
|
dec = PhysAddr.decode(pa.encode())
|
|
assert dec.unit_type == UnitType.PE
|
|
assert dec.pe_id == 7
|
|
assert dec.sub_offset == 0x400
|
|
|
|
|
|
def test_pe_tcm_addr_range_check():
|
|
with pytest.raises(PhysAddrError):
|
|
PhysAddr.pe_tcm_addr(
|
|
rack_id=0, sip_id=0, cube_id=0, pe_id=0,
|
|
tcm_offset=(1 << 29), # exceeds 29-bit sub_offset
|
|
)
|
|
|
|
|
|
# ── AddressConfig ───────────────────────────────────────────────────
|
|
|
|
|
|
def test_address_config_derived_sizes():
|
|
assert _CFG.hbm_slice_bytes == 6 * _GB
|
|
assert _CFG.tcm_allocatable_bytes == 12 * _MB
|
|
|
|
|
|
# ── PEMemAllocator: HBM ────────────────────────────────────────────
|
|
|
|
|
|
def _make_alloc(pe_id: int = 0) -> PEMemAllocator:
|
|
return PEMemAllocator(rack_id=0, sip_id=0, cube_id=0, pe_id=pe_id, cfg=_CFG)
|
|
|
|
|
|
def test_allocator_hbm_basic():
|
|
a = _make_alloc(pe_id=3)
|
|
pa = a.alloc_hbm(4096)
|
|
assert pa.kind == "hbm"
|
|
assert pa.sip_id == 0
|
|
assert pa.cube_id == 0
|
|
# hbm_offset should be pe3's slice start
|
|
assert pa.hbm_offset == 3 * 6 * _GB
|
|
|
|
|
|
def test_allocator_hbm_sequential():
|
|
a = _make_alloc()
|
|
pa1 = a.alloc_hbm(1024)
|
|
pa2 = a.alloc_hbm(2048)
|
|
assert pa1.hbm_offset == 0 # pe0 slice start + 0
|
|
assert pa2.hbm_offset == 1024 # pe0 slice start + 1024
|
|
|
|
|
|
def test_allocator_hbm_overflow():
|
|
a = _make_alloc()
|
|
a.alloc_hbm(6 * _GB - 256)
|
|
with pytest.raises(AllocationError, match="HBM"):
|
|
a.alloc_hbm(512)
|
|
|
|
|
|
# ── PEMemAllocator: TCM ────────────────────────────────────────────
|
|
|
|
|
|
def test_allocator_tcm_basic():
|
|
a = _make_alloc(pe_id=5)
|
|
pa = a.alloc_tcm(256)
|
|
assert pa.kind == "pe_resource"
|
|
assert pa.unit_type == UnitType.PE
|
|
assert pa.pe_id == 5
|
|
assert pa.sub_offset == 0
|
|
|
|
|
|
def test_allocator_tcm_respects_reserved():
|
|
a = _make_alloc()
|
|
# allocatable = 12 MB, should succeed
|
|
a.alloc_tcm(12 * _MB)
|
|
assert a.tcm_used == 12 * _MB
|
|
assert a.tcm_total == 12 * _MB
|
|
|
|
|
|
def test_allocator_tcm_overflow():
|
|
a = _make_alloc()
|
|
a.alloc_tcm(12 * _MB)
|
|
with pytest.raises(AllocationError, match="TCM"):
|
|
a.alloc_tcm(1)
|
|
|
|
|
|
# ── PEMemAllocator: stats & determinism ─────────────────────────────
|
|
|
|
|
|
def test_allocator_stats():
|
|
a = _make_alloc()
|
|
a.alloc_hbm(1000)
|
|
a.alloc_tcm(500)
|
|
assert a.hbm_used == 1000
|
|
assert a.hbm_total == 6 * _GB
|
|
assert a.tcm_used == 500
|
|
assert a.tcm_total == 12 * _MB
|
|
|
|
|
|
def test_allocator_deterministic():
|
|
a1 = _make_alloc(pe_id=2)
|
|
a2 = _make_alloc(pe_id=2)
|
|
assert a1.alloc_hbm(4096) == a2.alloc_hbm(4096)
|
|
assert a1.alloc_tcm(256) == a2.alloc_tcm(256)
|