commit - release 1

2026-03-18 11:47:48 -07:00
commit 6f43807900
109 changed files with 14909 additions and 0 deletions
@@ -0,0 +1,268 @@
+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)