ywkang
08812eda58
Implement VA/MMU layer (ADR-0011 Phase 1) enabling Triton kernels to use contiguous virtual addresses on sharded tensors. Key changes: - PE_MMU component: hybrid inbox (MmuMapMsg) + sync translate() for PE_DMA - VirtualAllocator + PEMemAllocator: free-list with coalescing - MmuMapMsg/MmuUnmapMsg fabric path with SIP-level routing - DPPolicy-based mapping: replicate=local, sharded=broadcast - Tensor lifecycle: del + weakref cleanup, context manager - Rename: TensorHandle.pa→addr, DmaReadCmd.src_pa→src_addr, ctx→torch Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
"""Tests for VirtualAllocator: device-wide VA space management.
|
|
|
|
Validates:
|
|
T7. Basic VA allocation (contiguous, non-overlapping)
|
|
T8. VA free + reallocation (free-list reuse)
|
|
T9. VA space exhaustion raises AllocationError
|
|
"""
|
|
import pytest
|
|
|
|
from kernbench.policy.address.va_allocator import VirtualAllocator
|
|
|
|
_KB = 1024
|
|
_MB = 1024 * 1024
|
|
_GB = 1024 * 1024 * 1024
|
|
|
|
|
|
# ── T7. Basic VA allocation ──────────────────────────────────────────
|
|
|
|
|
|
def test_alloc_returns_aligned_va():
|
|
"""First allocation returns va_base."""
|
|
va = VirtualAllocator(va_base=0x1_0000_0000, va_size=1 * _GB, page_size=4096)
|
|
addr = va.alloc(4096)
|
|
assert addr == 0x1_0000_0000
|
|
|
|
|
|
def test_alloc_sequential_non_overlapping():
|
|
"""Two allocations return contiguous, non-overlapping VA ranges."""
|
|
va = VirtualAllocator(va_base=0x1_0000_0000, va_size=1 * _GB, page_size=4096)
|
|
a1 = va.alloc(4096)
|
|
a2 = va.alloc(8192)
|
|
assert a1 == 0x1_0000_0000
|
|
assert a2 == 0x1_0000_1000 # a1 + 4096
|
|
# No overlap
|
|
assert a2 >= a1 + 4096
|
|
|
|
|
|
def test_alloc_page_aligned():
|
|
"""Allocations are page-aligned even if requested size is not page-multiple."""
|
|
va = VirtualAllocator(va_base=0x1_0000_0000, va_size=1 * _GB, page_size=4096)
|
|
a1 = va.alloc(100) # < 1 page, but occupies 1 page
|
|
a2 = va.alloc(100)
|
|
assert a2 == 0x1_0000_1000 # aligned to next page
|
|
|
|
|
|
def test_alloc_large_contiguous():
|
|
"""Large allocation (multiple pages) is contiguous."""
|
|
va = VirtualAllocator(va_base=0x0, va_size=1 * _GB, page_size=2 * _MB)
|
|
addr = va.alloc(8 * _MB) # 4 pages
|
|
assert addr == 0x0
|
|
# Next alloc starts after 8 MB
|
|
addr2 = va.alloc(2 * _MB)
|
|
assert addr2 == 8 * _MB
|
|
|
|
|
|
# ── T8. VA free + reallocation ───────────────────────────────────────
|
|
|
|
|
|
def test_free_and_realloc():
|
|
"""Freed VA range can be reused by subsequent allocation."""
|
|
va = VirtualAllocator(va_base=0x1_0000_0000, va_size=1 * _GB, page_size=4096)
|
|
a1 = va.alloc(4096)
|
|
a2 = va.alloc(4096)
|
|
va.free(a1, 4096)
|
|
|
|
# New alloc should reuse a1's range
|
|
a3 = va.alloc(4096)
|
|
assert a3 == a1
|
|
|
|
|
|
def test_free_coalesce():
|
|
"""Freeing adjacent blocks allows larger reallocation."""
|
|
va = VirtualAllocator(va_base=0x0, va_size=1 * _GB, page_size=4096)
|
|
a1 = va.alloc(4096)
|
|
a2 = va.alloc(4096)
|
|
a3 = va.alloc(4096)
|
|
|
|
# Free first two (adjacent)
|
|
va.free(a1, 4096)
|
|
va.free(a2, 4096)
|
|
|
|
# Should be able to allocate 8192 contiguous from the freed region
|
|
a4 = va.alloc(8192)
|
|
assert a4 == a1 # reuses coalesced region
|
|
|
|
|
|
def test_free_out_of_order():
|
|
"""Free in non-sequential order still works."""
|
|
va = VirtualAllocator(va_base=0x0, va_size=1 * _GB, page_size=4096)
|
|
a1 = va.alloc(4096)
|
|
a2 = va.alloc(4096)
|
|
a3 = va.alloc(4096)
|
|
|
|
va.free(a2, 4096) # free middle
|
|
a4 = va.alloc(4096)
|
|
assert a4 == a2 # reuses middle slot
|
|
|
|
|
|
# ── T9. VA space exhaustion ──────────────────────────────────────────
|
|
|
|
|
|
def test_alloc_exhaustion():
|
|
"""Allocation beyond VA space raises AllocationError."""
|
|
va = VirtualAllocator(va_base=0x0, va_size=8192, page_size=4096)
|
|
va.alloc(4096)
|
|
va.alloc(4096)
|
|
with pytest.raises(Exception, match="[Aa]lloc|[Ee]xhaust|[Oo]ut of"):
|
|
va.alloc(4096)
|
|
|
|
|
|
def test_alloc_after_partial_free():
|
|
"""After freeing some, can allocate again within freed space."""
|
|
va = VirtualAllocator(va_base=0x0, va_size=8192, page_size=4096)
|
|
a1 = va.alloc(4096)
|
|
a2 = va.alloc(4096)
|
|
|
|
# Space is full
|
|
with pytest.raises(Exception):
|
|
va.alloc(4096)
|
|
|
|
# Free one, now can allocate again
|
|
va.free(a1, 4096)
|
|
a3 = va.alloc(4096)
|
|
assert a3 == a1
|
|
|
|
|
|
# ── Stats / inspection ───────────────────────────────────────────────
|
|
|
|
|
|
def test_used_and_total():
|
|
"""used and total properties reflect allocation state."""
|
|
va = VirtualAllocator(va_base=0x0, va_size=1 * _MB, page_size=4096)
|
|
assert va.used == 0
|
|
assert va.total == 1 * _MB
|
|
|
|
va.alloc(4096)
|
|
assert va.used == 4096
|
|
|
|
va.alloc(8192)
|
|
assert va.used == 4096 + 8192 # page-aligned: 4096 + 8192 = 12288
|