diff --git a/src/python/omnimalloc/allocators/__init__.py b/src/python/omnimalloc/allocators/__init__.py index 336f32e..809bd9c 100644 --- a/src/python/omnimalloc/allocators/__init__.py +++ b/src/python/omnimalloc/allocators/__init__.py @@ -38,6 +38,8 @@ from .tabu_search import TabuSearchConfig as TabuSearchConfig from .telamalloc import TelamallocAllocator as TelamallocAllocator from .telamalloc import TelamallocConfig as TelamallocConfig +from .tiered import TieredAllocator as TieredAllocator +from .tiered import TieredConfig as TieredConfig from .utils import AVAILABLE_ALLOCATORS as AVAILABLE_ALLOCATORS from .utils import DEFAULT_ALLOCATOR as DEFAULT_ALLOCATOR from .utils import get_allocator_by_name as get_allocator_by_name diff --git a/src/python/omnimalloc/allocators/tiered.py b/src/python/omnimalloc/allocators/tiered.py new file mode 100644 index 0000000..2395d31 --- /dev/null +++ b/src/python/omnimalloc/allocators/tiered.py @@ -0,0 +1,113 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +from collections.abc import Mapping +from dataclasses import dataclass + +from omnimalloc._cpp import compute_temporal_overlaps, first_fit_place +from omnimalloc.primitives import Allocation +from omnimalloc.primitives.allocation import IdType + +from .base import BaseAllocator +from .greedy_base import ( + order_by_area, + order_by_conflict, + order_by_conflict_size, + order_by_duration, + order_by_size, + order_by_start, +) + +# Buffer priority for the fast tier: the bytes most expensive to spill are +# considered first, so largest size is the sensible default. +_ORDERS = { + "size": order_by_size, + "duration": order_by_duration, + "area": order_by_area, + "conflict": order_by_conflict, + "conflict_size": order_by_conflict_size, + "start": order_by_start, +} + + +@dataclass(frozen=True) +class TieredConfig: + """Configuration for TieredAllocator.""" + + capacity: int | None = None + order: str = "size" + + def __post_init__(self) -> None: + if self.capacity is not None and self.capacity <= 0: + raise ValueError(f"capacity must be positive, got {self.capacity}") + if self.order not in _ORDERS: + raise ValueError( + f"unknown order {self.order!r}; choose from {sorted(_ORDERS)}" + ) + + +class TieredAllocator(BaseAllocator): + """Pack buffers into a fixed-capacity fast memory and spill the rest above it. + + Models the on-chip/off-chip decision at the heart of ML accelerators. The + fast memory occupies offsets ``[0, capacity)``; every buffer first-fit can + place under that ceiling stays on-chip, and the rest are packed contiguously + above ``capacity`` (the spill region). The result is a single valid + allocation: a buffer is on-chip iff its ``height`` (offset + size) is at most + ``capacity``. With ``capacity=None`` the fast memory is unbounded, so nothing + spills and the allocator reduces to first-fit in ``order``. + + Buffers are considered for the fast tier in ``order`` (default largest size + first, so the bytes most expensive to spill are kept on-chip). + """ + + def __init__(self, config: TieredConfig | None = None) -> None: + self._config = config or TieredConfig() + + def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + if not allocations: + return allocations + overlaps = compute_temporal_overlaps(allocations) + ordered = _ORDERS[self._config.order](allocations) + placed, spilled = self._pack_fast(ordered, overlaps) + if spilled: + placed.update(self._pack_spill(spilled)) + return tuple(a.with_offset(placed[a.id]) for a in allocations) + + def _pack_fast( + self, + ordered: tuple[Allocation, ...], + overlaps: Mapping[IdType, set[IdType]], + ) -> tuple[dict[IdType, int], list[Allocation]]: + """Place each buffer at its lowest under-ceiling gap; spill when none fits.""" + capacity = self._config.capacity + sizes = {a.id: a.size for a in ordered} + placed: dict[IdType, int] = {} + spilled: list[Allocation] = [] + for alloc in ordered: + neighbors = overlaps.get(alloc.id) or () + occupied = sorted((placed[n], sizes[n]) for n in neighbors if n in placed) + cursor = 0 + fits = False + for offset, size in occupied: + if offset - cursor >= alloc.size: + fits = True + break + cursor = max(cursor, offset + size) + if not fits and (capacity is None or cursor + alloc.size <= capacity): + fits = True + if fits: + placed[alloc.id] = cursor + else: + spilled.append(alloc) + return placed, spilled + + def _pack_spill(self, spilled: list[Allocation]) -> dict[IdType, int]: + """Pack the spilled buffers with first-fit, offset above the fast region.""" + base = self._config.capacity or 0 + ordered = order_by_size(tuple(spilled)) + overlaps = compute_temporal_overlaps(ordered) + return { + a.id: base + (a.offset or 0) for a in first_fit_place(ordered, overlaps) + } diff --git a/tests/unit/allocators/test_tiered.py b/tests/unit/allocators/test_tiered.py new file mode 100644 index 0000000..d1184ec --- /dev/null +++ b/tests/unit/allocators/test_tiered.py @@ -0,0 +1,133 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +import pytest +from omnimalloc.allocators.greedy import GreedyBySizeAllocator +from omnimalloc.allocators.greedy_base import peak_memory +from omnimalloc.allocators.tiered import TieredAllocator, TieredConfig +from omnimalloc.primitives import Allocation, Pool +from omnimalloc.validate import validate_allocation + + +def _is_valid(result: tuple[Allocation, ...]) -> bool: + return validate_allocation(Pool(id="test_pool", allocations=result)) + + +def _spilled_bytes(result: tuple[Allocation, ...], capacity: int) -> int: + return sum(a.size for a in result if (a.height or 0) > capacity) + + +def test_tiered_empty() -> None: + assert TieredAllocator(TieredConfig(capacity=100)).allocate(()) == () + + +def test_tiered_unbounded_single_at_zero() -> None: + result = TieredAllocator().allocate((Allocation(id=1, size=100, start=0, end=10),)) + assert result[0].offset == 0 + + +def test_tiered_buffer_larger_than_capacity_spills() -> None: + result = TieredAllocator(TieredConfig(capacity=50)).allocate( + (Allocation(id=1, size=100, start=0, end=10),) + ) + assert result[0].offset == 50 + + +def test_tiered_non_overlapping_share_offset_on_chip() -> None: + allocs = ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=100, start=10, end=20), + ) + result = TieredAllocator(TieredConfig(capacity=100)).allocate(allocs) + by_id = {a.id: a for a in result} + assert by_id[1].offset == 0 + assert by_id[2].offset == 0 + + +def test_tiered_large_capacity_spills_nothing() -> None: + allocs = tuple(Allocation(id=i, size=100, start=0, end=10) for i in range(5)) + result = TieredAllocator(TieredConfig(capacity=10_000)).allocate(allocs) + assert _spilled_bytes(result, 10_000) == 0 + assert peak_memory(result) == 500 + + +def test_tiered_overlap_beyond_capacity_spills_above_ceiling() -> None: + allocs = ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=100, start=0, end=10), + ) + result = TieredAllocator(TieredConfig(capacity=100)).allocate(allocs) + offsets = sorted(a.offset for a in result) + assert offsets == [0, 100] + assert _is_valid(result) + + +def test_tiered_unbounded_matches_greedy_by_size() -> None: + allocs = tuple( + Allocation(id=i, size=(i % 4 + 1) * 32, start=i, end=i + 5) for i in range(20) + ) + tiered = TieredAllocator().allocate(allocs) + greedy = GreedyBySizeAllocator().allocate(allocs) + assert peak_memory(tiered) == peak_memory(greedy) + + +def test_tiered_invalid_capacity_raises() -> None: + with pytest.raises(ValueError, match="capacity must be positive"): + TieredConfig(capacity=0) + + +def test_tiered_invalid_order_raises() -> None: + with pytest.raises(ValueError, match="unknown order"): + TieredConfig(order="nope") + + +def test_tiered_deterministic() -> None: + allocs = tuple( + Allocation(id=i, size=(i * 37 % 90) + 10, start=i % 7, end=i % 7 + 6) + for i in range(30) + ) + config = TieredConfig(capacity=300) + first = TieredAllocator(config).allocate(allocs) + second = TieredAllocator(config).allocate(allocs) + assert [a.offset for a in first] == [a.offset for a in second] + + +@pytest.mark.parametrize( + "order", ["size", "duration", "area", "conflict", "conflict_size", "start"] +) +def test_tiered_every_order_is_valid(order: str) -> None: + allocs = tuple( + Allocation(id=i, size=(i * 53 % 120) + 16, start=i % 9, end=i % 9 + 7) + for i in range(40) + ) + result = TieredAllocator(TieredConfig(capacity=400, order=order)).allocate(allocs) + assert _is_valid(result) + assert {a.id for a in result} == {a.id for a in allocs} + + +def test_tiered_spill_shrinks_monotonically_with_capacity() -> None: + allocs = tuple( + Allocation(id=i, size=(i * 41 % 100) + 20, start=i % 11, end=i % 11 + 8) + for i in range(60) + ) + spilled = [] + for capacity in (200, 400, 600, 800, 1000): + result = TieredAllocator(TieredConfig(capacity=capacity)).allocate(allocs) + assert _is_valid(result) + spilled.append(_spilled_bytes(result, capacity)) + assert spilled == sorted(spilled, reverse=True) + + +def test_tiered_partition_is_consistent_and_complete() -> None: + allocs = tuple( + Allocation(id=i, size=(i * 29 % 80) + 24, start=i % 6, end=i % 6 + 5) + for i in range(50) + ) + capacity = 250 + result = TieredAllocator(TieredConfig(capacity=capacity)).allocate(allocs) + on_chip = [a for a in result if (a.height or 0) <= capacity] + spilled = [a for a in result if (a.height or 0) > capacity] + assert len(on_chip) + len(spilled) == len(allocs) + assert all(a.offset is not None and a.offset >= capacity for a in spilled) + assert _is_valid(result)