Add TieredAllocator for capacity-bounded on-chip/spill allocation

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

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)
+        }

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)