From e5bca1dda9c46aa63bc9238d7bab6f7f3448547f Mon Sep 17 00:00:00 2001 From: Fabian Peddinghaus Date: Tue, 21 Jul 2026 17:33:48 +0200 Subject: [PATCH] Accept raw allocation sequences and streamline the public API Let every public entry point (allocate, validate_allocation, save_allocation, plot_allocation) accept a raw Sequence[Allocation] in addition to System/Memory/Pool, wrapped via the new Pool.from_allocations; allocate returns the placed allocations as a tuple in input order. Pool.allocate now restores the pool's input order after allocator-internal reordering so positions keep corresponding to the request, and checks the returned set by length as well as ids. Rename Memory.capacity to Memory.size (validation message follows), and rename the analysis functions pressure/pressure_per_allocation to antichain_pressure/antichain_pressure_per_allocation, keeping the old names exported as aliases for the canonical default metric. Drop the keyword-only markers from allocator constructors, Timer, and plot_allocation so the common single-argument calls read naturally. Add any_allocated across Pool/Memory/System and use it in io. Add nanobind to the dev dependency group. --- pyproject.toml | 1 + src/python/omnimalloc/_allocate.py | 48 ++++++++--- src/python/omnimalloc/allocators/genetic.py | 1 - src/python/omnimalloc/allocators/greedy.py | 2 +- .../omnimalloc/allocators/greedy_base.py | 1 - src/python/omnimalloc/allocators/hillclimb.py | 1 - .../omnimalloc/allocators/minimalloc.py | 2 +- src/python/omnimalloc/allocators/omni.py | 2 +- src/python/omnimalloc/allocators/random.py | 2 +- .../allocators/simulated_annealing.py | 1 - .../omnimalloc/allocators/supermalloc.py | 1 - .../omnimalloc/allocators/tabu_search.py | 1 - .../omnimalloc/allocators/telamalloc.py | 1 - src/python/omnimalloc/analysis/__init__.py | 4 + src/python/omnimalloc/analysis/_pressure.py | 21 +++-- src/python/omnimalloc/benchmark/timer.py | 2 +- src/python/omnimalloc/io.py | 12 +-- src/python/omnimalloc/primitives/memory.py | 17 ++-- src/python/omnimalloc/primitives/pool.py | 37 ++++++--- src/python/omnimalloc/primitives/system.py | 5 ++ src/python/omnimalloc/primitives/utils.py | 12 +++ src/python/omnimalloc/validate.py | 33 +++++--- src/python/omnimalloc/visualize.py | 80 ++++++++++--------- tests/integration/test_omni_torture.py | 4 +- tests/unit/analysis/test_pressure.py | 7 ++ tests/unit/primitives/test_memory.py | 72 ++++++++++------- tests/unit/primitives/test_pool.py | 49 ++++++++++++ tests/unit/primitives/test_system.py | 32 ++++++-- tests/unit/test_allocate.py | 69 +++++++++++++++- tests/unit/test_io.py | 14 ++++ tests/unit/test_validate.py | 65 +++++++++++++-- tests/unit/test_visualize.py | 31 ++++--- uv.lock | 11 +++ 33 files changed, 478 insertions(+), 163 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a7aef3c..532bc29 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ all = [ dev = [ "ipykernel>=6.30.0", "ipywidgets>=8.1.0", + "nanobind>=2.0", "nbconvert>=7.16.0", "pre-commit>=4.0.0", "pytest>=8.0", diff --git a/src/python/omnimalloc/_allocate.py b/src/python/omnimalloc/_allocate.py index 11846a4..1b2fecb 100644 --- a/src/python/omnimalloc/_allocate.py +++ b/src/python/omnimalloc/_allocate.py @@ -3,37 +3,59 @@ # -from typing import TypeVar, cast +from collections.abc import Sequence +from typing import TypeAlias, TypeVar, cast, overload from .allocators import DEFAULT_ALLOCATOR, BaseAllocator -from .primitives import Memory, Pool, System +from .primitives import Allocation, Memory, Pool, System from .validate import validate_allocation +AllocatorLike: TypeAlias = BaseAllocator | type[BaseAllocator] | str | None + T = TypeVar("T", System, Memory, Pool) +@overload def allocate( entity: T, - allocator: BaseAllocator | type[BaseAllocator] | str | None = None, - *, + allocator: AllocatorLike = None, + validate: bool = False, +) -> T: ... + + +@overload +def allocate( + entity: Sequence[Allocation], + allocator: AllocatorLike = None, validate: bool = False, -) -> T: - """Return the entity (System, Memory, or Pool) with offsets assigned. +) -> tuple[Allocation, ...]: ... + - `allocator` accepts an instance, a class, a registry name, or `None` - (the default allocator); `validate=True` additionally runs - `validate_allocation` on the result. +def allocate( + entity: System | Memory | Pool | Sequence[Allocation], + allocator: AllocatorLike = None, + validate: bool = False, +) -> System | Memory | Pool | tuple[Allocation, ...]: + """Return the entity with offsets assigned. + + Accepts a System, Memory, or Pool (returned as the same type) or a raw + sequence of Allocations (returned as a tuple in input order). `allocator` + accepts an instance, a class, a registry name, or `None` (the default + allocator); `validate=True` additionally runs `validate_allocation` on + the result. """ if allocator is None: allocator = DEFAULT_ALLOCATOR - resolved = BaseAllocator.resolve(allocator) + resolved = cast("BaseAllocator", BaseAllocator.resolve(allocator)) - # ty doesn't understand that TypeVar T (System|Memory|Pool) all have allocate method - allocated = entity.allocate(resolved) # type: ignore[invalid-argument-type] + if isinstance(entity, System | Memory | Pool): + allocated = entity.allocate(resolved) + else: + allocated = Pool.from_allocations(entity).allocate(resolved).allocations if validate: validate_allocation(allocated) - return cast("T", allocated) + return allocated diff --git a/src/python/omnimalloc/allocators/genetic.py b/src/python/omnimalloc/allocators/genetic.py index 70d93d2..f4d51fa 100644 --- a/src/python/omnimalloc/allocators/genetic.py +++ b/src/python/omnimalloc/allocators/genetic.py @@ -44,7 +44,6 @@ class GeneticAllocator(GreedyAllocator): def __init__( self, - *, seed: int = DEFAULT_SEED, population_size: int = 100, max_generations: int = 50, diff --git a/src/python/omnimalloc/allocators/greedy.py b/src/python/omnimalloc/allocators/greedy.py index 6950343..084840c 100644 --- a/src/python/omnimalloc/allocators/greedy.py +++ b/src/python/omnimalloc/allocators/greedy.py @@ -77,7 +77,7 @@ class GreedyByAllAllocator(GreedyAllocator): `num_threads=None` uses all cores. """ - def __init__(self, *, num_threads: int | None = None) -> None: + def __init__(self, num_threads: int | None = None) -> None: ensure_valid_num_threads(num_threads) self._num_threads = num_threads diff --git a/src/python/omnimalloc/allocators/greedy_base.py b/src/python/omnimalloc/allocators/greedy_base.py index e00f2b8..2dfddde 100644 --- a/src/python/omnimalloc/allocators/greedy_base.py +++ b/src/python/omnimalloc/allocators/greedy_base.py @@ -72,7 +72,6 @@ def _allocate( def allocate_parallel( allocations: tuple[Allocation, ...], variants: tuple[BaseAllocator, ...], - *, num_threads: int | None = None, ) -> tuple[Allocation, ...]: """Run each variant and return the lowest peak memory results. diff --git a/src/python/omnimalloc/allocators/hillclimb.py b/src/python/omnimalloc/allocators/hillclimb.py index 2eb9f1c..9f1b4a2 100644 --- a/src/python/omnimalloc/allocators/hillclimb.py +++ b/src/python/omnimalloc/allocators/hillclimb.py @@ -33,7 +33,6 @@ class HillClimbAllocator(GreedyAllocator): def __init__( self, - *, max_iterations: int = 100, seed: int = DEFAULT_SEED, acceptance_temperature: float = 2.0, diff --git a/src/python/omnimalloc/allocators/minimalloc.py b/src/python/omnimalloc/allocators/minimalloc.py index b2256af..e430c0f 100644 --- a/src/python/omnimalloc/allocators/minimalloc.py +++ b/src/python/omnimalloc/allocators/minimalloc.py @@ -51,7 +51,7 @@ class MinimallocAllocator(BaseAllocator): supports_vector_time = False def __init__( - self, *, timeout: float | None = DEFAULT_TIMEOUT, capacity: int = 1 * TB + self, timeout: float | None = DEFAULT_TIMEOUT, capacity: int = 1 * TB ) -> None: _require_minimalloc() ensure_valid_timeout(timeout) diff --git a/src/python/omnimalloc/allocators/omni.py b/src/python/omnimalloc/allocators/omni.py index aeb6817..cfe703b 100644 --- a/src/python/omnimalloc/allocators/omni.py +++ b/src/python/omnimalloc/allocators/omni.py @@ -23,7 +23,7 @@ class OmniAllocator(BaseAllocator): supports_vector_time = True - def __init__(self, *, linearize_budget: int | None = DEFAULT_WORK_BUDGET) -> None: + def __init__(self, linearize_budget: int | None = DEFAULT_WORK_BUDGET) -> None: ensure_valid_budget(linearize_budget, name="linearize_budget") self._linearize_budget = linearize_budget diff --git a/src/python/omnimalloc/allocators/random.py b/src/python/omnimalloc/allocators/random.py index 3b8f43d..ab111a0 100644 --- a/src/python/omnimalloc/allocators/random.py +++ b/src/python/omnimalloc/allocators/random.py @@ -15,7 +15,7 @@ class RandomAllocator(GreedyAllocator): """Randomized allocator that tries multiple random orders and picks the best.""" - def __init__(self, *, num_trials: int = 100, seed: int = DEFAULT_SEED) -> None: + def __init__(self, num_trials: int = 100, seed: int = DEFAULT_SEED) -> None: ensure_non_negative(num_trials, "num_trials") self._seed = seed self._num_trials = num_trials diff --git a/src/python/omnimalloc/allocators/simulated_annealing.py b/src/python/omnimalloc/allocators/simulated_annealing.py index 63f6234..b66ad64 100644 --- a/src/python/omnimalloc/allocators/simulated_annealing.py +++ b/src/python/omnimalloc/allocators/simulated_annealing.py @@ -31,7 +31,6 @@ class SimulatedAnnealingAllocator(BaseAllocator): def __init__( self, - *, seed: int = DEFAULT_SEED, max_iterations: int = 3000, initial_temperature: float = 3.0, diff --git a/src/python/omnimalloc/allocators/supermalloc.py b/src/python/omnimalloc/allocators/supermalloc.py index 99f226b..31a6775 100644 --- a/src/python/omnimalloc/allocators/supermalloc.py +++ b/src/python/omnimalloc/allocators/supermalloc.py @@ -133,7 +133,6 @@ class SupermallocAllocator(BaseAllocator): def __init__( self, - *, timeout: float | None = DEFAULT_TIMEOUT, heuristics: tuple[Heuristic, ...] = DEFAULT_HEURISTICS, num_threads: int | None = None, diff --git a/src/python/omnimalloc/allocators/tabu_search.py b/src/python/omnimalloc/allocators/tabu_search.py index 208084a..d9ff1c3 100644 --- a/src/python/omnimalloc/allocators/tabu_search.py +++ b/src/python/omnimalloc/allocators/tabu_search.py @@ -30,7 +30,6 @@ class TabuSearchAllocator(BaseAllocator): def __init__( self, - *, seed: int = DEFAULT_SEED, max_iterations: int = 500, neighborhood_size: int = 20, diff --git a/src/python/omnimalloc/allocators/telamalloc.py b/src/python/omnimalloc/allocators/telamalloc.py index fab890b..fcbcdaa 100644 --- a/src/python/omnimalloc/allocators/telamalloc.py +++ b/src/python/omnimalloc/allocators/telamalloc.py @@ -36,7 +36,6 @@ class TelamallocAllocator(BaseAllocator): def __init__( self, - *, seed: int = DEFAULT_SEED, max_backtracks: int = 10000, timeout: float | None = DEFAULT_TIMEOUT, diff --git a/src/python/omnimalloc/analysis/__init__.py b/src/python/omnimalloc/analysis/__init__.py index b2f6cf0..f2f9975 100644 --- a/src/python/omnimalloc/analysis/__init__.py +++ b/src/python/omnimalloc/analysis/__init__.py @@ -4,6 +4,10 @@ from ._conflicts import conflict_degrees as conflict_degrees from ._conflicts import conflicts as conflicts +from ._pressure import antichain_pressure as antichain_pressure +from ._pressure import ( + antichain_pressure_per_allocation as antichain_pressure_per_allocation, +) from ._pressure import closure_pressure as closure_pressure from ._pressure import ( closure_pressure_per_allocation as closure_pressure_per_allocation, diff --git a/src/python/omnimalloc/analysis/_pressure.py b/src/python/omnimalloc/analysis/_pressure.py index 247dde1..41a3fa8 100644 --- a/src/python/omnimalloc/analysis/_pressure.py +++ b/src/python/omnimalloc/analysis/_pressure.py @@ -4,7 +4,7 @@ from collections.abc import Sequence -from omnimalloc._cpp import antichain_pressure +from omnimalloc._cpp import antichain_pressure as _antichain_pressure from omnimalloc._cpp import ( antichain_pressure_per_allocation as _antichain_pressure_per_allocation, ) @@ -21,7 +21,7 @@ from omnimalloc.primitives.utils import ensure_unique_ids -def pressure( +def antichain_pressure( allocations: Sequence[Allocation], work_budget: int | None = DEFAULT_WORK_BUDGET ) -> int: """Peak memory pressure: exact max-weight antichain of the happens-before order. @@ -36,7 +36,7 @@ def pressure( the flow work exceeds `work_budget`; pass `None` to always compute. """ ensure_valid_budget(work_budget) - return antichain_pressure(allocations, work_budget) + return _antichain_pressure(allocations, work_budget) def closure_pressure( @@ -46,7 +46,7 @@ def closure_pressure( C++ enumeration of the join-closure of the birth clocks. Pairwise- concurrent allocations need not share a cut, so this can sit strictly - below `pressure`; both soundly lower-bound any placement's peak. Raises + below `antichain_pressure`; both soundly lower-bound any placement's peak. Raises `RuntimeError` once the closure exceeds `closure_cap`; pass `None` to always enumerate — memory then grows with the closure itself. """ @@ -58,7 +58,7 @@ def placement_pressure(allocations: Sequence[Allocation]) -> int: """Peak of a placement: the highest occupied address, max(offset + size). Simply the pressure the placement realizes after allocation — an upper - bound on `pressure` (and so on `closure_pressure`), equal to the max + bound on `antichain_pressure` (and so on `closure_pressure`), equal to the max entry of `placement_pressure_per_allocation`. Raises `ValueError` on unplaced input. """ @@ -71,14 +71,14 @@ def placement_pressure(allocations: Sequence[Allocation]) -> int: return max(heights, default=0) -def pressure_per_allocation( +def antichain_pressure_per_allocation( allocations: Sequence[Allocation], work_budget: int | None = DEFAULT_WORK_BUDGET ) -> dict[IdType, int]: """Peak pressure over each allocation's own lifetime, keyed by id. The max-weight antichain through each allocation: the tightest order-derived lower bound on the pressure any placement can exhibit - while that allocation is live; the max entry equals `pressure`. + while that allocation is live; the max entry equals `antichain_pressure`. Linearizable instances take one O(N log N) window sweep, genuinely partial orders solve one pinned antichain (min flow over the conflict neighborhood) per distinct lifetime — exact, but built to certify @@ -98,7 +98,7 @@ def closure_pressure_per_allocation( """Exact realizable peak while each allocation is live, keyed by id. The max total size at any join-closure cut where the allocation is - live. Can sit elementwise strictly below `pressure_per_allocation`, + live. Can sit elementwise strictly below `antichain_pressure_per_allocation`, since pairwise-concurrent allocations need not share a cut; the max entry equals `closure_pressure`. Raises `RuntimeError` once the closure exceeds `closure_cap`; pass `None` to always enumerate — @@ -128,6 +128,11 @@ def placement_pressure_per_allocation( return _keyed_by_id(allocations, peaks) +# The antichain bound is the canonical/default pressure metric +pressure = antichain_pressure +pressure_per_allocation = antichain_pressure_per_allocation + + def _keyed_by_id( allocations: Sequence[Allocation], peaks: list[int] ) -> dict[IdType, int]: diff --git a/src/python/omnimalloc/benchmark/timer.py b/src/python/omnimalloc/benchmark/timer.py index 90fc2c9..930457d 100644 --- a/src/python/omnimalloc/benchmark/timer.py +++ b/src/python/omnimalloc/benchmark/timer.py @@ -20,7 +20,7 @@ class Timer: multiple threads may result in race conditions and inconsistent state. """ - def __init__(self, *, auto_start: bool = False) -> None: + def __init__(self, auto_start: bool = False) -> None: self._start_ns: int | None = None self._stop_ns: int | None = None self._is_running: bool = False diff --git a/src/python/omnimalloc/io.py b/src/python/omnimalloc/io.py index ef0fc02..4d35b41 100644 --- a/src/python/omnimalloc/io.py +++ b/src/python/omnimalloc/io.py @@ -3,6 +3,7 @@ # import csv +from collections.abc import Sequence from pathlib import Path from .analysis.clock import time_components @@ -39,7 +40,7 @@ def _write_pool(pool: Pool, path: Path) -> Path: # Any placed allocation brings in the offset column (minimalloc's # solution format; unplaced rows leave the cell blank), so save/load # round-trips placements instead of silently stripping them. - with_offsets = any(alloc.is_allocated for alloc in pool.allocations) + with_offsets = pool.any_allocated fields = ("id", "lower", "upper", "size") if with_offsets: fields = (*fields, "offset") @@ -56,11 +57,12 @@ def _write_pool(pool: Pool, path: Path) -> Path: def save_allocation( - entity: System | Memory | Pool, path: str | Path + entity: System | Memory | Pool | Sequence[Allocation], path: str | Path ) -> tuple[Path, ...]: """Save the entity's pools to disk as minimalloc-format CSV files. - Saving a `Pool` writes exactly `path`; a `Memory` or `System` fans out + Saving a `Pool` (or a raw sequence of Allocations, saved as one pool) + writes exactly `path`; a `Memory` or `System` fans out to one `_.csv` per pool (the CSV is pool-level minimalloc interchange). Returns the tuple of paths actually written. Pools with any placed allocation include an `offset` column (minimalloc's solution @@ -72,10 +74,10 @@ def save_allocation( """ path_ = Path(path) path_.parent.mkdir(parents=True, exist_ok=True) + if not isinstance(entity, System | Memory | Pool): + entity = Pool.from_allocations(entity) if isinstance(entity, Pool): return (_write_pool(entity, path_),) - if not isinstance(entity, (Memory, System)): - raise TypeError(f"Unsupported entity type: {type(entity)!r}") return tuple( _write_pool(pool, path_.with_name(f"{path_.stem}_{name}.csv")) for name, pool in _collect_pools(entity).items() diff --git a/src/python/omnimalloc/primitives/memory.py b/src/python/omnimalloc/primitives/memory.py index 6c4abb6..b4fa25d 100644 --- a/src/python/omnimalloc/primitives/memory.py +++ b/src/python/omnimalloc/primitives/memory.py @@ -19,19 +19,19 @@ class Memory: """A physical memory unit containing one or more pools. - `capacity` is the declared memory limit (an input); the computed extent - of a placement is `used_size`. `capacity=None` means unbounded. + `size` is the declared memory size (an input); the computed extent of a + placement is `used_size`. `size=None` means unbounded. """ id: IdType pools: tuple[Pool, ...] - capacity: int | None = None + size: int | None = None def __post_init__(self) -> None: if len({pool.id for pool in self.pools}) != len(self.pools): raise ValueError("pool ids must be unique") - if self.capacity is not None: - ensure_non_negative(self.capacity, "capacity") + if self.size is not None: + ensure_non_negative(self.size, "size") @cached_property def used_size(self) -> int: @@ -43,9 +43,14 @@ def is_allocated(self) -> bool: """True if all pools have been allocated.""" return all(pool.is_allocated for pool in self.pools) + @cached_property + def any_allocated(self) -> bool: + """True if any pool has a placed allocation.""" + return any(pool.any_allocated for pool in self.pools) + def with_pools(self, pools: tuple[Pool, ...]) -> "Memory": """Return new Memory with specified pools.""" - return Memory(id=self.id, capacity=self.capacity, pools=pools) + return Memory(id=self.id, size=self.size, pools=pools) def allocate(self, allocator: "BaseAllocator") -> "Memory": """Apply allocator to all pools.""" diff --git a/src/python/omnimalloc/primitives/pool.py b/src/python/omnimalloc/primitives/pool.py index 5ae52c1..a2f898e 100644 --- a/src/python/omnimalloc/primitives/pool.py +++ b/src/python/omnimalloc/primitives/pool.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # +from collections.abc import Sequence from dataclasses import dataclass from functools import cached_property from typing import TYPE_CHECKING @@ -13,6 +14,7 @@ from omnimalloc.common.validation import ensure_non_negative from .allocation import Allocation, IdType +from .utils import ensure_allocations @dataclass(frozen=True) @@ -29,13 +31,14 @@ def __post_init__(self) -> None: if self.offset is not None: ensure_non_negative(self.offset, "offset") + @classmethod + def from_allocations(cls, allocations: Sequence[Allocation]) -> "Pool": + """Wrap a raw sequence of allocations in an anonymous pool.""" + return cls(id=0, allocations=ensure_allocations(allocations)) + @cached_property def size(self) -> int: - """Memory extent from the pool base (offset 0) to the highest allocated end. - - Counts any gap below the lowest allocation, so it is consistent with - how `overlaps`, pool stacking, and visualization interpret pool extent. - """ + """Memory extent from the pool base (offset 0) to the highest allocated end.""" if not self.is_allocated: raise ValueError("cannot compute size of unallocated pool") ends = [ @@ -47,11 +50,7 @@ def size(self) -> int: @cached_property def pressure(self) -> int: - """Peak memory pressure (max cut through all buffer lifetimes). - - Exact: scalar sweep for linearizable lifetimes, max-weight antichain - for non-linearizable vector-clock instances (see `analysis.pressure`). - """ + """Peak memory pressure (max cut through all buffer lifetimes).""" return _pressure(self.allocations) @cached_property @@ -68,6 +67,11 @@ def is_allocated(self) -> bool: """True if all allocations have been assigned memory offsets.""" return all(alloc.offset is not None for alloc in self.allocations) + @cached_property + def any_allocated(self) -> bool: + """True if any allocation has been assigned a memory offset.""" + return any(alloc.offset is not None for alloc in self.allocations) + def overlaps(self, other: "Pool") -> bool: """True if pools overlap in memory space.""" if self.offset is None or other.offset is None: @@ -82,10 +86,17 @@ def with_allocations(self, allocations: tuple[Allocation, ...]) -> "Pool": return Pool(id=self.id, offset=self.offset, allocations=allocations) def allocate(self, allocator: "BaseAllocator") -> "Pool": - """Apply allocator to assign memory offsets to all allocations.""" + """Apply allocator to assign memory offsets, preserving allocation order.""" allocated = allocator.allocate(self.allocations) - if {a.id for a in allocated} != {a.id for a in self.allocations}: + if len(allocated) != len(self.allocations) or {a.id for a in allocated} != { + a.id for a in self.allocations + }: raise ValueError( f"allocator {allocator!s} returned a different allocation set" ) - return self.with_allocations(allocated) + # Allocators may reorder internally; restore the pool's input order so + # positions in the returned tuple keep corresponding to the request + placed_by_id = {a.id: a for a in allocated} + return self.with_allocations( + tuple(placed_by_id[a.id] for a in self.allocations) + ) diff --git a/src/python/omnimalloc/primitives/system.py b/src/python/omnimalloc/primitives/system.py index 7cd704a..2439d72 100644 --- a/src/python/omnimalloc/primitives/system.py +++ b/src/python/omnimalloc/primitives/system.py @@ -29,6 +29,11 @@ def is_allocated(self) -> bool: """True if all memories have been allocated.""" return all(memory.is_allocated for memory in self.memories) + @property + def any_allocated(self) -> bool: + """True if any memory has a placed allocation.""" + return any(memory.any_allocated for memory in self.memories) + def with_memories(self, memories: tuple[Memory, ...]) -> "System": """Return new System with specified memories.""" return System(id=self.id, memories=memories) diff --git a/src/python/omnimalloc/primitives/utils.py b/src/python/omnimalloc/primitives/utils.py index cca746e..d116ae2 100644 --- a/src/python/omnimalloc/primitives/utils.py +++ b/src/python/omnimalloc/primitives/utils.py @@ -11,3 +11,15 @@ def ensure_unique_ids(allocations: Sequence[Allocation]) -> None: """Raise if any allocation id repeats; id-keyed placement assumes uniqueness.""" if len({alloc.id for alloc in allocations}) != len(allocations): raise ValueError("allocation ids must be unique") + + +def ensure_allocations(allocations: object) -> tuple[Allocation, ...]: + """Coerce a raw sequence to a tuple, requiring every element be an Allocation.""" + if isinstance(allocations, str | bytes) or not isinstance(allocations, Sequence): + raise TypeError(f"Unsupported entity type: {type(allocations)!r}") + checked: list[Allocation] = [] + for alloc in allocations: + if not isinstance(alloc, Allocation): + raise TypeError(f"Expected Allocation, got {type(alloc)!r}") + checked.append(alloc) + return tuple(checked) diff --git a/src/python/omnimalloc/validate.py b/src/python/omnimalloc/validate.py index 1e66b6f..92588ee 100644 --- a/src/python/omnimalloc/validate.py +++ b/src/python/omnimalloc/validate.py @@ -2,8 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 # +from collections.abc import Sequence + from .analysis.clock import uniform_dim from .primitives import Allocation, IdType, Memory, Pool, System +from .primitives.utils import ensure_allocations def _check_unique_ids(entities: tuple[Memory | Pool | Allocation, ...]) -> None: @@ -51,12 +54,12 @@ def _validate_pools(pools: tuple[Pool, ...]) -> None: raise ValueError(f"in pool {pool.id!r}, {e}") from e -def _check_capacity(memory: Memory) -> None: - if memory.capacity is None or not memory.is_allocated: +def _check_size(memory: Memory) -> None: + if memory.size is None or not memory.is_allocated: return - if memory.used_size > memory.capacity: + if memory.used_size > memory.size: raise ValueError( - f"used size {memory.used_size} exceeds capacity {memory.capacity}" + f"used size {memory.used_size} exceeds memory size {memory.size}" ) @@ -65,29 +68,35 @@ def _validate_memories(memories: tuple[Memory, ...]) -> None: for memory in memories: try: _validate_pools(memory.pools) - _check_capacity(memory) + _check_size(memory) except ValueError as e: raise ValueError(f"in memory {memory.id!r}, {e}") from e -def validate_allocation(entity: System | Memory | Pool) -> None: +def validate_allocation(entity: System | Memory | Pool | Sequence[Allocation]) -> None: """Raise ValueError unless the entity is fully placed with no collisions. + Accepts a System, Memory, or Pool, or a raw sequence of Allocations. Checks unique ids, that every allocation and pool carries an offset, that no placed rectangles collide, and that each memory's used size - stays within its declared capacity. + stays within its declared size. """ + if isinstance(entity, System | Memory | Pool): + described = f"{type(entity).__name__} {entity.id!r}" + elif isinstance(entity, Sequence) and not isinstance(entity, str | bytes): + described = f"{len(entity)} allocations" + else: + raise TypeError(f"Unsupported entity type: {type(entity)!r}") + try: if isinstance(entity, System): _validate_memories(entity.memories) elif isinstance(entity, Memory): _validate_pools(entity.pools) - _check_capacity(entity) + _check_size(entity) elif isinstance(entity, Pool): _validate_allocations(entity.allocations) else: - raise TypeError(f"Unsupported entity type: {type(entity)!r}") + _validate_allocations(ensure_allocations(entity)) except ValueError as e: - raise ValueError( - f"Validation of {type(entity).__name__} {entity.id!r} failed, {e}." - ) from e + raise ValueError(f"Validation of {described} failed, {e}.") from e diff --git a/src/python/omnimalloc/visualize.py b/src/python/omnimalloc/visualize.py index 69b8b5f..9d88540 100644 --- a/src/python/omnimalloc/visualize.py +++ b/src/python/omnimalloc/visualize.py @@ -4,6 +4,7 @@ import logging from collections import defaultdict +from collections.abc import Sequence from pathlib import Path from typing import Final, Literal, NamedTuple @@ -219,23 +220,23 @@ def _select_lanes( def _get_y_limits(system: System) -> dict[Memory, tuple[int, int]]: limits: dict[Memory, tuple[int, int]] = {} for memory in system.memories: - capacity = memory.capacity + size = memory.size used = memory.used_size - if capacity is None: - # No capacity declared, scale to 1.2x used + if size is None: + # No size declared, scale to 1.2x used y_limit = used * 1.2 - elif used > capacity: - # Usage exceeds capacity, scale to 1.2x usage + elif used > size: + # Usage exceeds the declared size, scale to 1.2x usage y_limit = used * 1.2 - elif used >= capacity * 0.5: - # Usage is 50-100% of capacity, use capacity as limit - y_limit = capacity + elif used >= size * 0.5: + # Usage is 50-100% of the declared size, use the size as limit + y_limit = size else: - # Usage below 50% of capacity, scale to 2x usage + # Usage below 50% of the declared size, scale to 2x usage y_limit = used * 2 # Clamp to at least 1 so downstream tick spacing stays positive @@ -307,7 +308,7 @@ def _draw_pool_background( def _draw_limit_lines(ax: Axes, limits: dict[str, int]) -> None: - """Draw annotated horizontal lines for used size, capacity, and extras.""" + """Draw annotated horizontal lines for used size, declared size, and extras.""" _, x_max = ax.get_xlim() for name, value in limits.items(): ax.axhline(value, color="black", linestyle="--", linewidth=1, alpha=0.8) @@ -343,11 +344,9 @@ def _set_axes_ticks(ax: Axes, y_limit: int, num_ticks: int = 8) -> None: def _memory_title(memory: Memory, threads: str) -> str: - capacity = memory.capacity - capacity_str = ( - _format_bytes(capacity) if capacity is not None else "Unbounded Capacity" - ) - return f"{memory.id} ({capacity_str}, {len(memory.pools)} pools{threads})" + size = memory.size + size_str = _format_bytes(size) if size is not None else "Unbounded Size" + return f"{memory.id} ({size_str}, {len(memory.pools)} pools{threads})" def _lane_panels( @@ -412,16 +411,31 @@ def _projection_panels(system: System) -> tuple[list[_Panel], str | None]: return panels, caveat +def _add_legend(fig: Figure) -> None: + """Add figure legend for allocation kinds.""" + handles = [ + Patch(color=color, label=kind.name, alpha=0.8) + for kind, color in KIND_COLOR_MAP.items() + ] + fig.legend( + handles=handles, + loc="outside lower center", + ncol=len(handles), + fontsize=8, + title="Allocation Kinds", + ) + + def _set_axes_limits( ax: Axes, x_limits: tuple[int, int], y_limits: tuple[int, int], - capacity: int | None, + size: int | None, ) -> None: """Set axis limits and add scaling notice if needed.""" ax.set_xlim(x_limits) ax.set_ylim(y_limits) - if capacity is not None and y_limits[1] < capacity: + if size is not None and y_limits[1] < size: ax.text( 0.02, 0.98, @@ -433,21 +447,6 @@ def _set_axes_limits( ) -def _add_legend(fig: Figure) -> None: - """Add figure legend for allocation kinds.""" - handles = [ - Patch(color=color, label=kind.name, alpha=0.8) - for kind, color in KIND_COLOR_MAP.items() - ] - fig.legend( - handles=handles, - loc="outside lower center", - ncol=len(handles), - fontsize=8, - title="Allocation Kinds", - ) - - # TODO(fpedd): Add a pools size descriptor on the right side of each pool @@ -462,7 +461,7 @@ def _draw_panel( if panel.title is not None: ax.set_title(panel.title) ax.set_xlabel(panel.xlabel) - _set_axes_limits(ax, panel.x_limits, y_limits, memory.capacity) + _set_axes_limits(ax, panel.x_limits, y_limits, memory.size) _set_axes_ticks(ax, y_limits[1]) if panel.note is not None: ax.text( @@ -488,10 +487,10 @@ def _draw_panel( _draw_allocation(ax, alloc, y_offset, color, extent) _draw_pool_background(ax, y_offset, pool.size, colors) - # Draw used-size, capacity, and extra capacity lines + # Draw used-size, declared-size, and extra capacity lines limits: dict[str, int] = {"used": memory.used_size} - if memory.capacity is not None: - limits["capacity"] = memory.capacity + if memory.size is not None: + limits["size"] = memory.size for label, per_memory_capacity in capacities.items(): if memory.id in per_memory_capacity: limits[label] = per_memory_capacity[memory.id] @@ -551,9 +550,8 @@ def _visualize_system( def plot_allocation( - entity: System | Memory | Pool, + entity: System | Memory | Pool | Sequence[Allocation], path: Path | str | None = None, - *, capacities: dict[str, dict[IdType, int]] | None = None, view: Literal["panel", "lanes"] = "panel", max_lanes: int | None = None, @@ -561,7 +559,8 @@ def plot_allocation( """Plot an allocated entity: `path=None` displays the figure, `path=...` saves it. Args: - entity: The entity to plot (System, Memory, or Pool). + entity: The entity to plot (System, Memory, Pool, or a raw sequence + of Allocations, plotted as a single pool). path: Where to save the figure; `None` displays it instead. capacities: Extra capacity lines to draw, keyed by label then by memory id (byte limits, drawn as horizontal lines). @@ -588,6 +587,9 @@ def plot_allocation( if not HAS_MATPLOTLIB: require_optional("matplotlib", "visualization") + if not isinstance(entity, System | Memory | Pool): + entity = Pool.from_allocations(entity) + if isinstance(entity, Pool): entity = Memory(id=f"memory_{entity.id}", pools=(entity,)) diff --git a/tests/integration/test_omni_torture.py b/tests/integration/test_omni_torture.py index 70762ef..e4d2bba 100644 --- a/tests/integration/test_omni_torture.py +++ b/tests/integration/test_omni_torture.py @@ -61,7 +61,9 @@ def _certify( peak = _peak(placed) assert pressure(allocations, work_budget=None) <= peak assert peak <= sum(a.size for a in allocations) - assert max(placement_pressure_per_allocation(placed).values()) == peak + per_allocation = placement_pressure_per_allocation(placed) + assert per_allocation == placement_pressure_per_allocation(placed, None) + assert max(per_allocation.values()) == peak _assert_conflicting_pairs_disjoint(placed) return peak diff --git a/tests/unit/analysis/test_pressure.py b/tests/unit/analysis/test_pressure.py index e2b1afd..43f6250 100644 --- a/tests/unit/analysis/test_pressure.py +++ b/tests/unit/analysis/test_pressure.py @@ -8,6 +8,8 @@ import pytest from omnimalloc.allocators.omni import OmniAllocator from omnimalloc.analysis import ( + antichain_pressure, + antichain_pressure_per_allocation, closure_pressure, closure_pressure_per_allocation, placement_pressure, @@ -18,6 +20,11 @@ from omnimalloc.primitives import Allocation +def test_default_names_alias_antichain() -> None: + assert pressure is antichain_pressure + assert pressure_per_allocation is antichain_pressure_per_allocation + + def test_pressure_empty_is_zero() -> None: assert pressure(()) == 0 diff --git a/tests/unit/primitives/test_memory.py b/tests/unit/primitives/test_memory.py index 35d392f..0a888a4 100644 --- a/tests/unit/primitives/test_memory.py +++ b/tests/unit/primitives/test_memory.py @@ -16,7 +16,7 @@ def test_basic_creation_with_int_id_simple() -> None: memory = Memory(id=301, pools=(pool,)) assert memory.id == 301 assert len(memory.pools) == 1 - assert memory.capacity is None + assert memory.size is None def test_basic_creation_with_int_id() -> None: @@ -34,15 +34,15 @@ def test_basic_creation_with_str_id() -> None: memory = Memory(id="mem_ddr", pools=(pool,)) assert memory.id == "mem_ddr" assert len(memory.pools) == 1 - assert memory.capacity is None + assert memory.size is None -def test_creation_with_capacity() -> None: - """Test creating a memory with declared capacity.""" +def test_creation_with_size() -> None: + """Test creating a memory with declared size.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - memory = Memory(id=301, pools=(pool,), capacity=1000) - assert memory.capacity == 1000 + memory = Memory(id=301, pools=(pool,), size=1000) + assert memory.size == 1000 def test_creation_with_multiple_pools() -> None: @@ -63,20 +63,20 @@ def test_empty_memory() -> None: assert memory.is_allocated is True -def test_negative_capacity() -> None: - """Test that negative capacity raises ValueError.""" +def test_negative_size() -> None: + """Test that negative size raises ValueError.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - with pytest.raises(ValueError, match="capacity must be non-negative"): - Memory(id=301, pools=(pool,), capacity=-1) + with pytest.raises(ValueError, match="size must be non-negative"): + Memory(id=301, pools=(pool,), size=-1) -def test_zero_capacity() -> None: - """Test that zero capacity is valid.""" +def test_zero_size() -> None: + """Test that zero size is valid.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - memory = Memory(id=301, pools=(pool,), capacity=0) - assert memory.capacity == 0 + memory = Memory(id=301, pools=(pool,), size=0) + assert memory.size == 0 def test_duplicate_pool_ids() -> None: @@ -113,11 +113,11 @@ def test_used_size_empty_memory() -> None: assert memory.used_size == 0 -def test_used_size_can_exceed_capacity() -> None: - """Capacity is a declared limit, not a constructor constraint.""" +def test_used_size_can_exceed_size() -> None: + """Size is a declared limit, not a constructor constraint.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - memory = Memory(id=301, pools=(pool,), capacity=10) + memory = Memory(id=301, pools=(pool,), size=10) assert memory.used_size == 100 @@ -163,12 +163,12 @@ def test_with_pools_replace() -> None: alloc2 = Allocation(id=102, size=200, start=0, end=10, offset=0) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2,)) - memory = Memory(id=301, pools=(pool1,), capacity=1000) + memory = Memory(id=301, pools=(pool1,), size=1000) new_memory = memory.with_pools((pool2,)) assert len(new_memory.pools) == 1 assert new_memory.pools[0].id == 212 assert new_memory.id == memory.id - assert new_memory.capacity == memory.capacity + assert new_memory.size == memory.size def test_with_pools_immutability() -> None: @@ -211,11 +211,11 @@ def test_cannot_modify_pools() -> None: memory.pools = () # type: ignore[misc] -def test_cannot_modify_capacity() -> None: - """Test that capacity cannot be modified.""" - memory = Memory(id=301, pools=(), capacity=1000) +def test_cannot_modify_size() -> None: + """Test that size cannot be modified.""" + memory = Memory(id=301, pools=(), size=1000) with pytest.raises(AttributeError): - memory.capacity = 2000 # type: ignore[misc] + memory.size = 2000 # type: ignore[misc] def test_large_values() -> None: @@ -224,9 +224,25 @@ def test_large_values() -> None: alloc2 = Allocation(id=102, size=10**11, start=0, end=100, offset=0) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2,)) - memory = Memory(id=999, pools=(pool1, pool2), capacity=10**15) + memory = Memory(id=999, pools=(pool1, pool2), size=10**15) assert memory.used_size == 10**12 + 10**11 - assert memory.capacity == 10**15 + assert memory.size == 10**15 + + +def test_any_allocated_across_pools() -> None: + placed = Pool( + id=1, allocations=(Allocation(id=1, size=10, start=0, end=5, offset=0),) + ) + unplaced = Pool(id=2, allocations=(Allocation(id=2, size=10, start=0, end=5),)) + memory = Memory(id=1, pools=(placed, unplaced)) + assert memory.any_allocated is True + assert memory.is_allocated is False + + +def test_any_allocated_empty_memory() -> None: + memory = Memory(id=1, pools=()) + assert memory.any_allocated is False + assert memory.is_allocated is True def test_complex_memory_structure() -> None: @@ -236,7 +252,7 @@ def test_complex_memory_structure() -> None: alloc3 = Allocation(id=103, size=75, start=10, end=20, offset=0) pool1 = Pool(id=211, allocations=(alloc1, alloc2)) pool2 = Pool(id=212, allocations=(alloc3,)) - memory = Memory(id=400, pools=(pool1, pool2), capacity=500) + memory = Memory(id=400, pools=(pool1, pool2), size=500) assert memory.used_size == pool1.size + pool2.size assert memory.is_allocated is True @@ -247,7 +263,7 @@ def test_allocate_with_allocator() -> None: alloc2 = Allocation(id=102, size=50, start=5, end=15) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2,)) - memory = Memory(id=301, pools=(pool1, pool2), capacity=1000) + memory = Memory(id=301, pools=(pool1, pool2), size=1000) assert memory.is_allocated is False allocator = NaiveAllocator() @@ -255,7 +271,7 @@ def test_allocate_with_allocator() -> None: assert allocated_memory.is_allocated is True assert allocated_memory.id == memory.id - assert allocated_memory.capacity == memory.capacity + assert allocated_memory.size == memory.size assert len(allocated_memory.pools) == 2 # Original memory should be unchanged assert memory.is_allocated is False diff --git a/tests/unit/primitives/test_pool.py b/tests/unit/primitives/test_pool.py index e8c7f39..1a74da6 100644 --- a/tests/unit/primitives/test_pool.py +++ b/tests/unit/primitives/test_pool.py @@ -3,6 +3,7 @@ # import pytest +from omnimalloc.allocators.greedy import GreedyBySizeAllocator from omnimalloc.allocators.naive import NaiveAllocator from omnimalloc.primitives import Allocation from omnimalloc.primitives.pool import Pool @@ -385,3 +386,51 @@ def test_size_counts_gap_below_lowest_allocation() -> None: alloc = Allocation(id=1, size=100, start=0, end=10, offset=1000) pool = Pool(id="p", allocations=(alloc,)) assert pool.size == 1100 + + +def test_any_allocated_empty_pool() -> None: + pool = Pool(id=1, allocations=()) + assert pool.any_allocated is False + assert pool.is_allocated is True + + +def test_any_allocated_none_placed() -> None: + pool = Pool(id=1, allocations=(Allocation(id=1, size=10, start=0, end=5),)) + assert pool.any_allocated is False + + +def test_any_allocated_partially_placed() -> None: + alloc1 = Allocation(id=1, size=10, start=0, end=5, offset=0) + alloc2 = Allocation(id=2, size=10, start=0, end=5) + pool = Pool(id=1, allocations=(alloc1, alloc2)) + assert pool.any_allocated is True + assert pool.is_allocated is False + + +def test_from_allocations_wraps_sequence() -> None: + allocations = [Allocation(id=1, size=10, start=0, end=5)] + pool = Pool.from_allocations(allocations) + assert pool.allocations == tuple(allocations) + + +def test_from_allocations_rejects_non_allocation_elements() -> None: + with pytest.raises(TypeError, match="Expected Allocation"): + Pool.from_allocations((1, 2)) + + +def test_from_allocations_rejects_non_sequence() -> None: + with pytest.raises(TypeError, match="Unsupported entity type"): + Pool.from_allocations("abc") + + +def test_allocate_preserves_allocation_order() -> None: + allocations = tuple( + Allocation(id=i, size=10 * (i + 1), start=0, end=5) for i in range(8) + ) + pool = Pool(id=1, allocations=allocations) + allocated = pool.allocate(GreedyBySizeAllocator()) + assert allocated.is_allocated is True + assert tuple(a.id for a in allocated.allocations) == tuple(range(8)) + assert tuple(a.size for a in allocated.allocations) == tuple( + a.size for a in allocations + ) diff --git a/tests/unit/primitives/test_system.py b/tests/unit/primitives/test_system.py index b929482..f468fd0 100644 --- a/tests/unit/primitives/test_system.py +++ b/tests/unit/primitives/test_system.py @@ -168,8 +168,8 @@ def test_complex_system_structure() -> None: pool1 = Pool(id=211, allocations=(alloc1, alloc2)) pool2 = Pool(id=212, allocations=(alloc3,)) pool3 = Pool(id=213, allocations=(alloc4,)) - memory1 = Memory(id=311, pools=(pool1, pool2), capacity=500) - memory2 = Memory(id=312, pools=(pool3,), capacity=300) + memory1 = Memory(id=311, pools=(pool1, pool2), size=500) + memory2 = Memory(id=312, pools=(pool3,), size=300) system = System(id=401, memories=(memory1, memory2)) assert len(system.memories) == 2 assert system.is_allocated is True @@ -193,7 +193,7 @@ def test_single_memory_system() -> None: """Test system with single memory.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - memory = Memory(id=311, pools=(pool,), capacity=1000) + memory = Memory(id=311, pools=(pool,), size=1000) system = System(id=401, memories=(memory,)) assert len(system.memories) == 1 assert system.is_allocated is True @@ -215,14 +215,32 @@ def test_multiple_pools_per_memory() -> None: assert system.is_allocated is True +def test_any_allocated_across_memories() -> None: + alloc = Allocation(id=1, size=10, start=0, end=5, offset=0) + placed = Memory(id=1, pools=(Pool(id=1, allocations=(alloc,)),)) + unplaced = Memory( + id=2, + pools=(Pool(id=2, allocations=(Allocation(id=2, size=10, start=0, end=5),)),), + ) + system = System(id=1, memories=(placed, unplaced)) + assert system.any_allocated is True + assert system.is_allocated is False + + +def test_any_allocated_empty_system() -> None: + system = System(id=1, memories=()) + assert system.any_allocated is False + assert system.is_allocated is True + + def test_large_system() -> None: """Test system with large values.""" alloc1 = Allocation(id=101, size=10**12, start=0, end=100, offset=0) alloc2 = Allocation(id=102, size=10**11, start=0, end=100, offset=0) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2,)) - memory1 = Memory(id=311, pools=(pool1,), capacity=10**15) - memory2 = Memory(id=312, pools=(pool2,), capacity=10**14) + memory1 = Memory(id=311, pools=(pool1,), size=10**15) + memory2 = Memory(id=312, pools=(pool2,), size=10**14) system = System(id=999, memories=(memory1, memory2)) assert memory1.used_size == 10**12 assert memory2.used_size == 10**11 @@ -237,8 +255,8 @@ def test_allocate_with_allocator() -> None: alloc3 = Allocation(id=103, size=75, start=10, end=20) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2, alloc3)) - memory1 = Memory(id=311, pools=(pool1,), capacity=1000) - memory2 = Memory(id=312, pools=(pool2,), capacity=500) + memory1 = Memory(id=311, pools=(pool1,), size=1000) + memory2 = Memory(id=312, pools=(pool2,), size=500) system = System(id=401, memories=(memory1, memory2)) assert system.is_allocated is False diff --git a/tests/unit/test_allocate.py b/tests/unit/test_allocate.py index da2849d..dbdea6e 100644 --- a/tests/unit/test_allocate.py +++ b/tests/unit/test_allocate.py @@ -2,8 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 # +import pytest from omnimalloc import allocate -from omnimalloc.allocators.greedy import GreedyAllocator +from omnimalloc.allocators.greedy import GreedyAllocator, GreedyBySizeAllocator from omnimalloc.allocators.naive import NaiveAllocator from omnimalloc.primitives import Allocation, AllocationKind, Memory, Pool, System @@ -140,14 +141,14 @@ def test_allocate_complex_hierarchy() -> None: pool1 = Pool(id=1, allocations=tuple(allocations_pool1)) pool2 = Pool(id=2, allocations=tuple(allocations_pool2)) - memory1 = Memory(id=1, pools=(pool1, pool2), capacity=2048) + memory1 = Memory(id=1, pools=(pool1, pool2), size=2048) # Second memory allocations_pool3 = [ Allocation(id=i + 10, size=75, start=i * 5, end=(i + 1) * 5) for i in range(4) ] pool3 = Pool(id=3, allocations=tuple(allocations_pool3)) - memory2 = Memory(id=2, pools=(pool3,), capacity=1024) + memory2 = Memory(id=2, pools=(pool3,), size=1024) system = System(id=1, memories=(memory1, memory2)) @@ -246,10 +247,70 @@ def test_allocate_memory_calculates_used_size() -> None: alloc2 = Allocation(id=2, size=150, start=0, end=10) pool1 = Pool(id=1, allocations=(alloc1,)) pool2 = Pool(id=2, allocations=(alloc2,)) - memory = Memory(id=1, pools=(pool1, pool2), capacity=1000) + memory = Memory(id=1, pools=(pool1, pool2), size=1000) allocator = NaiveAllocator() allocated_memory = allocate(memory, allocator) # Each pool gets its allocations placed assert allocated_memory.used_size == 250 # 100 + 150 + + +def test_allocate_raw_allocations_returns_tuple() -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=5), + Allocation(id=2, size=150, start=5, end=10), + ) + allocated = allocate(allocations, NaiveAllocator(), validate=True) + assert isinstance(allocated, tuple) + assert all(a.offset is not None for a in allocated) + assert {a.id for a in allocated} == {1, 2} + + +def test_allocate_raw_allocations_accepts_list_and_registry_name() -> None: + allocations = [ + Allocation(id="a", size=100, start=0, end=5), + Allocation(id="b", size=150, start=0, end=5), + ] + allocated = allocate(allocations, "naive") + assert isinstance(allocated, tuple) + assert all(a.is_allocated for a in allocated) + + +def test_allocate_raw_allocations_does_not_mutate_input() -> None: + allocations = (Allocation(id=1, size=100, start=0, end=5),) + allocate(allocations, NaiveAllocator()) + assert allocations[0].offset is None + + +def test_allocate_raw_allocations_duplicate_ids_raise() -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=5), + Allocation(id=1, size=150, start=5, end=10), + ) + with pytest.raises(ValueError, match="allocation ids must be unique"): + allocate(allocations, NaiveAllocator()) + + +def test_allocate_raw_allocations_empty() -> None: + assert allocate((), NaiveAllocator()) == () + + +def test_allocate_raw_allocations_rejects_non_allocation_elements() -> None: + with pytest.raises(TypeError, match="Expected Allocation"): + allocate((1, 2, 3), NaiveAllocator()) + + +def test_allocate_rejects_unsupported_entity() -> None: + with pytest.raises(TypeError, match="Unsupported entity type"): + allocate("naive") + + +def test_allocate_raw_allocations_preserves_input_order() -> None: + allocations = tuple( + Allocation(id=i, size=10 * (i + 1), start=0, end=5) for i in range(10) + ) + allocated = allocate(allocations, GreedyBySizeAllocator(), validate=True) + assert tuple(a.id for a in allocated) == tuple(a.id for a in allocations) + assert all(a.offset is not None for a in allocated) + assert tuple(a.size for a in allocated) == tuple(a.size for a in allocations) diff --git a/tests/unit/test_io.py b/tests/unit/test_io.py index 8f77ef5..1335800 100644 --- a/tests/unit/test_io.py +++ b/tests/unit/test_io.py @@ -32,6 +32,20 @@ def test_save_pool_writes_exactly_path(tmp_path: Path) -> None: assert file_path.read_text() == "id,lower,upper,size\na,0,3,4\nb,2,9,8\n" +def test_save_raw_allocations_writes_one_pool(tmp_path: Path) -> None: + file_path = tmp_path / "problem.csv" + written = save_allocation(list(make_pool().allocations), file_path) + assert written == (file_path,) + assert file_path.read_text() == "id,lower,upper,size\na,0,3,4\nb,2,9,8\n" + + +def test_save_raw_allocations_rejects_non_allocation_elements( + tmp_path: Path, +) -> None: + with pytest.raises(TypeError, match="Expected Allocation"): + save_allocation((1, 2), tmp_path / "problem.csv") + + def test_save_creates_missing_parent_directories(tmp_path: Path) -> None: file_path = tmp_path / "nested" / "dir" / "problem.csv" written = save_allocation(make_pool(), file_path) diff --git a/tests/unit/test_validate.py b/tests/unit/test_validate.py index da7d05f..fd7e921 100644 --- a/tests/unit/test_validate.py +++ b/tests/unit/test_validate.py @@ -202,20 +202,20 @@ def test_validate_complex_hierarchy_with_error_raises() -> None: validate_allocation(system) -def test_validate_memory_over_capacity_fails() -> None: +def test_validate_memory_over_size_fails() -> None: pool = Pool( id="p", allocations=(Allocation(id=1, size=100, start=0, end=5, offset=0),) ) - memory = Memory(id="m", capacity=50, pools=(pool,)) - with pytest.raises(ValueError, match="exceeds capacity"): + memory = Memory(id="m", size=50, pools=(pool,)) + with pytest.raises(ValueError, match="exceeds memory size"): validate_allocation(memory) -def test_validate_memory_within_capacity_passes() -> None: +def test_validate_memory_within_size_passes() -> None: pool = Pool( id="p", allocations=(Allocation(id=1, size=100, start=0, end=5, offset=0),) ) - memory = Memory(id="m", capacity=100, pools=(pool,)) + memory = Memory(id="m", size=100, pools=(pool,)) validate_allocation(memory) @@ -240,3 +240,58 @@ def test_validate_pool_mixed_dimensions_rejected() -> None: pool = Pool(id=1, allocations=(alloc1, alloc2)) with pytest.raises(ValueError, match="share one clock dimension"): validate_allocation(pool) + + +def test_validate_raw_allocations_valid() -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=5, offset=0), + Allocation(id=2, size=100, start=5, end=10, offset=0), + ) + validate_allocation(allocations) + + +def test_validate_raw_allocations_list() -> None: + validate_allocation([Allocation(id=1, size=10, start=0, end=5, offset=0)]) + + +def test_validate_raw_allocations_empty() -> None: + validate_allocation(()) + + +def test_validate_raw_allocations_overlapping_raises() -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=10, offset=0), + Allocation(id=2, size=100, start=5, end=15, offset=50), + ) + with pytest.raises(ValueError, match=r"Validation of 2 allocations failed"): + validate_allocation(allocations) + + +def test_validate_raw_allocations_unallocated_raises() -> None: + with pytest.raises(ValueError, match="is not allocated"): + validate_allocation((Allocation(id=1, size=10, start=0, end=5),)) + + +def test_validate_raw_allocations_duplicate_ids_raise() -> None: + allocations = ( + Allocation(id=1, size=10, start=0, end=5, offset=0), + Allocation(id=1, size=10, start=5, end=10, offset=0), + ) + with pytest.raises(ValueError, match="duplicate id"): + validate_allocation(allocations) + + +def test_validate_raw_allocations_rejects_non_allocation_elements() -> None: + with pytest.raises(TypeError, match="Expected Allocation"): + validate_allocation([1, 2, 3]) + + +def test_validate_raw_allocations_complex_mix() -> None: + valid = tuple( + Allocation(id=i, size=50, start=5 * i, end=5 * (i + 1), offset=0) + for i in range(10) + ) + validate_allocation(valid) + overlapping = (*valid, Allocation(id=99, size=50, start=12, end=18, offset=25)) + with pytest.raises(ValueError, match=r"Validation of 11 allocations failed"): + validate_allocation(overlapping) diff --git a/tests/unit/test_visualize.py b/tests/unit/test_visualize.py index 6819ee8..216f2d4 100644 --- a/tests/unit/test_visualize.py +++ b/tests/unit/test_visualize.py @@ -77,7 +77,7 @@ def test_visualize_memory_with_multiple_pools(artifacts_dir: Path) -> None: pool2 = Pool(id=2, allocations=(alloc2,), offset=200) pool3 = Pool(id=3, allocations=(alloc3,), offset=500) - memory = Memory(id="mem_1", pools=(pool1, pool2, pool3), capacity=1000) + memory = Memory(id="mem_1", pools=(pool1, pool2, pool3), size=1000) output_path = artifacts_dir / "test_memory_pools.pdf" plot_allocation(memory, output_path) @@ -89,14 +89,14 @@ def test_visualize_system_with_multiple_memories(artifacts_dir: Path) -> None: # Memory 1: Simple pool alloc1 = Allocation(id=1, size=100, start=0, end=5, offset=0) pool1 = Pool(id=1, allocations=(alloc1,), offset=0) - memory1 = Memory(id="ddr4_1", pools=(pool1,), capacity=500) + memory1 = Memory(id="ddr4_1", pools=(pool1,), size=500) # Memory 2: Multiple pools alloc2 = Allocation(id=2, size=150, start=0, end=10, offset=0) alloc3 = Allocation(id=3, size=75, start=5, end=15, offset=0) pool2 = Pool(id=2, allocations=(alloc2,), offset=0) pool3 = Pool(id=3, allocations=(alloc3,), offset=200) - memory2 = Memory(id="ddr4_2", pools=(pool2, pool3), capacity=1000) + memory2 = Memory(id="ddr4_2", pools=(pool2, pool3), size=1000) system = System(id="test_system", memories=(memory1, memory2)) @@ -127,7 +127,7 @@ def test_visualize_complex_hierarchy(artifacts_dir: Path) -> None: pool1 = Pool(id="pool_1", allocations=tuple(allocations_mem1_pool1), offset=0) pool2 = Pool(id="pool_2", allocations=tuple(allocations_mem1_pool2), offset=300) - memory1 = Memory(id="main_memory", pools=(pool1, pool2), capacity=2048) + memory1 = Memory(id="main_memory", pools=(pool1, pool2), size=2048) # Second memory with different allocation patterns allocations_mem2 = [ @@ -142,7 +142,7 @@ def test_visualize_complex_hierarchy(artifacts_dir: Path) -> None: for i in range(0, 20, 5) ] pool3 = Pool(id="pool_3", allocations=tuple(allocations_mem2), offset=0) - memory2 = Memory(id="cache_memory", pools=(pool3,), capacity=1024) + memory2 = Memory(id="cache_memory", pools=(pool3,), size=1024) system = System(id="complex_system", memories=(memory1, memory2)) @@ -159,7 +159,7 @@ def test_visualize_with_string_ids(artifacts_dir: Path) -> None: alloc1 = Allocation(id="workspace_buf", size=100, start=0, end=5, offset=0) alloc2 = Allocation(id="temp_buf", size=150, start=5, end=10, offset=0) pool = Pool(id="tensor_pool", allocations=(alloc1, alloc2), offset=0) - memory = Memory(id="ddr_ram", pools=(pool,), capacity=512) + memory = Memory(id="ddr_ram", pools=(pool,), size=512) output_path = artifacts_dir / "test_string_ids.pdf" plot_allocation(memory, output_path) @@ -193,19 +193,30 @@ def test_visualize_memory_converts_to_system(artifacts_dir: Path) -> None: """Test that visualizing a Memory creates appropriate System wrapper.""" alloc = Allocation(id=1, size=100, start=0, end=10, offset=0) pool = Pool(id=1, allocations=(alloc,), offset=0) - memory = Memory(id=1, pools=(pool,), capacity=500) + memory = Memory(id=1, pools=(pool,), size=500) output_path = artifacts_dir / "test_memory_wrapper.pdf" plot_allocation(memory, output_path) assert output_path.exists() +def test_visualize_raw_allocations_converts_to_pool(artifacts_dir: Path) -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=5, offset=0), + Allocation(id=2, size=150, start=5, end=10, offset=0), + ) + + output_path = artifacts_dir / "test_raw_allocations.pdf" + plot_allocation(allocations, output_path) + assert output_path.exists() + + def test_visualize_with_capacities(artifacts_dir: Path) -> None: """Test visualization with extra capacity lines.""" alloc1 = Allocation(id=1, size=100, start=0, end=5, offset=0) alloc2 = Allocation(id=2, size=150, start=5, end=10, offset=100) pool = Pool(id=1, allocations=(alloc1, alloc2), offset=0) - memory = Memory(id="ddr_mem", pools=(pool,), capacity=1000) + memory = Memory(id="ddr_mem", pools=(pool,), size=1000) system = System(id="test_sys", memories=(memory,)) custom_limits = { @@ -283,7 +294,7 @@ def test_visualize_mixed_dimension_pools(artifacts_dir: Path) -> None: allocations=(Allocation(id=2, size=50, start=(0, 1), end=(2, 3), offset=0),), offset=200, ) - memory = Memory(id="mem", pools=(scalar_pool, vector_pool), capacity=1000) + memory = Memory(id="mem", pools=(scalar_pool, vector_pool), size=1000) output_path = artifacts_dir / "test_mixed_dim_pools.pdf" plot_allocation(memory, output_path) @@ -298,7 +309,7 @@ def test_visualize_empty_pool(artifacts_dir: Path) -> None: allocations=(Allocation(id=1, size=100, start=0, end=4, offset=0),), offset=100, ) - memory = Memory(id="mem", pools=(empty, filled), capacity=500) + memory = Memory(id="mem", pools=(empty, filled), size=500) output_path = artifacts_dir / "test_empty_pool.pdf" plot_allocation(memory, output_path) diff --git a/uv.lock b/uv.lock index e264d07..664bcd4 100644 --- a/uv.lock +++ b/uv.lock @@ -1334,6 +1334,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ad/3f/3d42e9a78fe5edf792a83c074b13b9b770092a4fbf3462872f4303135f09/ml_dtypes-0.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d", size = 168825, upload-time = "2025-11-17T22:32:23.766Z" }, ] +[[package]] +name = "nanobind" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/07/e5/a6fc2f024eaef17c68b94b5f8f56098d31859a7840a11a3dd45da2d3908a/nanobind-2.13.0.tar.gz", hash = "sha256:c7b04d6a6a4cd57985571e605539399b51331ae455d7fce576a5e2fcb89b1dcf", size = 1052101, upload-time = "2026-06-18T06:03:49.491Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/8d/cec69a0a804711fd89f47fad92cd2a438a1ceaafd9dabaf3b8b403b5cbba/nanobind-2.13.0-py3-none-any.whl", hash = "sha256:9268d7201eaf5519ae61b7a83175f2af77088cb4b99dd3ce5b00550f697da7b7", size = 269387, upload-time = "2026-06-18T06:03:47.508Z" }, +] + [[package]] name = "nbclient" version = "0.10.2" @@ -1580,6 +1589,7 @@ all = [ dev = [ { name = "ipykernel" }, { name = "ipywidgets" }, + { name = "nanobind" }, { name = "nbconvert" }, { name = "pre-commit" }, { name = "pytest" }, @@ -1605,6 +1615,7 @@ provides-extras = ["all"] dev = [ { name = "ipykernel", specifier = ">=6.30.0" }, { name = "ipywidgets", specifier = ">=8.1.0" }, + { name = "nanobind", specifier = ">=2.0" }, { name = "nbconvert", specifier = ">=7.16.0" }, { name = "pre-commit", specifier = ">=4.0.0" }, { name = "pytest", specifier = ">=8.0" },