omnimalloc 0.3.0__cp313-cp313-win_amd64.whl

omnimalloc/__init__.py

@@ -0,0 +1,19 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from importlib.metadata import version
+
+__version__ = version("omnimalloc")
+
+from .allocate import run_allocation as run_allocation
+from .allocators import get_available_allocators as get_available_allocators
+from .allocators import get_default_allocator as get_default_allocator
+from .primitives import Allocation as Allocation
+from .primitives import BufferKind as BufferKind
+from .primitives import IdType as IdType
+from .primitives import Memory as Memory
+from .primitives import Pool as Pool
+from .primitives import System as System
+from .validate import validate_allocation as validate_allocation
+from .visualize import plot_allocation as plot_allocation

omnimalloc/_cpp.pyi

@@ -0,0 +1,85 @@
+from collections.abc import Sequence
+import enum
+
+
+class BufferKind(enum.Enum):
+    def __str__(self) -> str: ...
+
+    def __repr__(self) -> str: ...
+
+    WORKSPACE = 0
+
+    CONSTANT = 1
+
+    INPUT = 2
+
+    OUTPUT = 3
+
+    @property
+    def is_io(self) -> bool: ...
+
+    def __hash__(self) -> int: ...
+
+class Allocation:
+    def __init__(self, id: int | str, size: int, start: int, end: int, offset: int | None = None, kind: BufferKind | None = None) -> None: ...
+
+    @property
+    def id(self) -> int | str: ...
+
+    @property
+    def size(self) -> int: ...
+
+    @property
+    def start(self) -> int: ...
+
+    @property
+    def end(self) -> int: ...
+
+    @property
+    def offset(self) -> int | None: ...
+
+    @property
+    def kind(self) -> BufferKind | None: ...
+
+    @property
+    def is_allocated(self) -> bool: ...
+
+    @property
+    def duration(self) -> int: ...
+
+    @property
+    def height(self) -> int | None: ...
+
+    @property
+    def area(self) -> int: ...
+
+    def overlaps_temporally(self, other: Allocation) -> bool: ...
+
+    def overlaps_spatially(self, other: Allocation) -> bool: ...
+
+    def overlaps(self, other: Allocation) -> bool: ...
+
+    def with_offset(self, offset: int) -> Allocation: ...
+
+    def with_kind(self, kind: BufferKind) -> Allocation: ...
+
+    def __str__(self) -> str: ...
+
+    def __repr__(self) -> str: ...
+
+    def __eq__(self, arg: Allocation, /) -> bool: ...
+
+    def __hash__(self) -> int: ...
+
+class GreedyAllocatorCpp:
+    def __init__(self) -> None: ...
+
+    def allocate(self, allocations: Sequence[Allocation]) -> list[Allocation]: ...
+
+    def __str__(self) -> str: ...
+
+    def __repr__(self) -> str: ...
+
+    def __eq__(self, arg: GreedyAllocatorCpp, /) -> bool: ...
+
+    def __hash__(self) -> int: ...

omnimalloc/allocate.py

@@ -0,0 +1,49 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+
+from typing import TypeVar, cast
+
+from .allocators import BaseAllocator, get_default_allocator
+from .primitives import Memory, Pool, System
+from .validate import validate_allocation
+
+T = TypeVar("T", System, Memory, Pool)
+
+
+def run_allocation(
+    entity: T,
+    allocator: BaseAllocator | type[BaseAllocator] | str | None = None,
+    validate: bool = False,
+) -> T:
+    """Run allocation on the given entity using the provided allocator.
+
+    Args:
+        entity: The entity to allocate (System, Memory, or Pool).
+        allocator: The allocator to use (instance, class, or name).
+        validate: Whether to validate the allocated entity.
+
+    Returns:
+        The allocated entity with the same type as the input.
+    """
+
+    if allocator is None:
+        allocator = get_default_allocator()
+
+    if isinstance(allocator, str):
+        allocator = BaseAllocator.get(allocator)
+
+    if isinstance(allocator, type):
+        allocator = allocator()
+
+    # At this point, allocator is guaranteed to be a BaseAllocator instance
+    assert isinstance(allocator, BaseAllocator)
+
+    # ty doesn't understand that TypeVar T (System|Memory|Pool) all have allocate method
+    allocated = entity.allocate(allocator)  # type: ignore[invalid-argument-type]
+
+    if validate:
+        validate_allocation(allocated)
+
+    return cast("T", allocated)

omnimalloc/allocators/__init__.py

@@ -0,0 +1,25 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from .base import BaseAllocator as BaseAllocator
+from .genetic import GeneticAllocator as GeneticAllocator
+from .greedy import GreedyAllocator as GreedyAllocator
+from .greedy import GreedyByAreaAllocator as GreedyByAreaAllocator
+from .greedy import GreedyByConflictAllocator as GreedyByConflictAllocator
+from .greedy import GreedyByDurationAllocator as GreedyByDurationAllocator
+from .greedy import GreedyBySizeAllocator as GreedyBySizeAllocator
+from .greedy_cpp import GreedyAllocatorCpp as GreedyAllocatorCpp
+from .greedy_cpp import GreedyByAreaAllocatorCpp as GreedyByAreaAllocatorCpp
+from .greedy_cpp import GreedyByConflictAllocatorCpp as GreedyByConflictAllocatorCpp
+from .greedy_cpp import GreedyByDurationAllocatorCpp as GreedyByDurationAllocatorCpp
+from .greedy_cpp import GreedyBySizeAllocatorCpp as GreedyBySizeAllocatorCpp
+from .hillclimb import HillClimbAllocator as HillClimbAllocator
+from .minimalloc import MinimallocAllocator as MinimallocAllocator
+from .naive import NaiveAllocator as NaiveAllocator
+from .random import RandomAllocator as RandomAllocator
+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
+from .utils import get_available_allocators as get_available_allocators
+from .utils import get_default_allocator as get_default_allocator

omnimalloc/allocators/base.py

@@ -0,0 +1,26 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+from omnimalloc.common.registry import Registered
+
+if TYPE_CHECKING:
+    from omnimalloc.primitives import Allocation
+
+
+class BaseAllocator(Registered):
+    """Base class for allocators with automatic registry."""
+
+    @abstractmethod
+    def allocate(
+        self, allocations: tuple["Allocation", ...]
+    ) -> tuple["Allocation", ...]:
+        """Run the allocator on the given allocations."""
+        ...
+
+    def reset(self) -> None:
+        """Optional: reset allocator state/config. Override if needed."""
+        ...

omnimalloc/allocators/genetic.py

@@ -0,0 +1,187 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+import random
+
+import numpy as np
+
+from omnimalloc.common.optional import require_optional
+from omnimalloc.primitives import Allocation
+
+from .greedy import GreedyAllocator
+
+try:
+    from deap import algorithms, base, creator, tools
+
+    HAS_DEAP = True
+
+except ImportError:
+    from types import SimpleNamespace
+    from typing import Any
+
+    HAS_DEAP = False
+
+    algorithms: Any = SimpleNamespace(
+        eaSimple=None,
+    )
+    base: Any = SimpleNamespace(
+        Fitness=None,
+        Toolbox=None,
+    )
+    creator: Any = SimpleNamespace(
+        create=None,
+    )
+    tools: Any = SimpleNamespace(
+        initIterate=None,
+        selTournament=None,
+        cxOrdered=None,
+        mutShuffleIndexes=None,
+    )
+
+
+class GeneticAllocator(GreedyAllocator):
+    """Genetic algorithm allocator that evolves permutation orders."""
+
+    def __init__(
+        self,
+        seed: int = 42,
+        population_size: int = 100,
+        num_generations: int = 50,
+        crossover_prob: float = 0.7,
+        mutation_prob: float = 0.2,
+        tournament_size: int = 3,
+    ) -> None:
+        """Initialize the genetic allocator."""
+        if not HAS_DEAP:
+            require_optional("deap", "GeneticAllocator")
+
+        self.seed = seed
+        self.population_size = population_size
+        self.num_generations = num_generations
+        self.crossover_prob = crossover_prob
+        self.mutation_prob = mutation_prob
+        self.tournament_size = tournament_size
+
+        # Setup DEAP creators (only once per class)
+        if not hasattr(creator, "FitnessMin"):
+            creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
+        if not hasattr(creator, "Individual"):
+            # FitnessMin is dynamically created by DEAP
+            creator.create("Individual", list, fitness=creator.FitnessMin)  # type: ignore[possibly-missing-attribute]
+
+    def _evaluate_permutation(
+        self, permutation: list[int], allocations: tuple[Allocation, ...]
+    ) -> tuple[float]:
+        """Evaluate a permutation by computing peak memory usage."""
+
+        # Apply permutation
+        permuted_allocs = tuple(allocations[i] for i in permutation)
+
+        # Run greedy allocation
+        result = super().allocate(permuted_allocs)
+
+        # Calculate peak memory usage
+        if not result:
+            return (0.0,)
+
+        peak_memory = max(alloc.height for alloc in result if alloc.height is not None)
+        return (float(peak_memory),)
+
+    def _create_heuristic_permutations(
+        self, allocations: tuple[Allocation, ...]
+    ) -> list[list[int]]:
+        """Create permutations based on greedy heuristics."""
+
+        permutations = []
+
+        # Create index mapping for original order
+        indexed_allocs = list(enumerate(allocations))
+
+        # 1. Greedy by size (largest first)
+        sorted_by_size = sorted(indexed_allocs, key=lambda x: x[1].size, reverse=True)
+        permutations.append([idx for idx, _ in sorted_by_size])
+
+        return permutations
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        """Evolve permutations using genetic algorithm to find best allocation."""
+
+        # Set random seeds for deterministic behavior
+        random.seed(self.seed)
+        np.random.default_rng(self.seed)
+
+        if not allocations:
+            return allocations
+
+        if len(allocations) == 1:
+            return super().allocate(allocations)
+
+        n = len(allocations)
+
+        # Setup toolbox
+        toolbox = base.Toolbox()
+
+        # Register individual and population generators
+        toolbox.register("indices", random.sample, range(n), n)  # Random permutation
+        # Individual and indices are dynamically created by DEAP
+        toolbox.register(
+            "individual",
+            tools.initIterate,
+            creator.Individual,  # type: ignore[possibly-missing-attribute]
+            toolbox.indices,  # type: ignore[possibly-missing-attribute]
+        )
+
+        # Register genetic operators
+        toolbox.register(
+            "evaluate", self._evaluate_permutation, allocations=allocations
+        )
+        toolbox.register("mate", tools.cxOrdered)
+        toolbox.register("mutate", tools.mutShuffleIndexes, indpb=0.05)
+        # TODO(fpedd): Try larger tournsize and selNSGA2
+        toolbox.register("select", tools.selTournament, tournsize=self.tournament_size)
+
+        # Create initial population with heuristic seeding
+        population = []
+
+        # Add heuristic-based individuals (5 heuristics)
+        heuristic_perms = self._create_heuristic_permutations(allocations)
+        for perm in heuristic_perms:
+            # Individual is dynamically created by DEAP
+            individual = creator.Individual(perm)  # type: ignore[possibly-missing-attribute]
+            population.append(individual)
+
+        # Fill rest with random individuals
+        for _ in range(self.population_size - len(heuristic_perms)):
+            # individual() is dynamically registered on toolbox
+            individual = toolbox.individual()  # type: ignore[possibly-missing-attribute]
+            population.append(individual)
+
+        # Track best individual
+        hall_of_fame = tools.HallOfFame(maxsize=1)
+
+        # Setup statistics
+        stats = tools.Statistics(key=lambda ind: ind.fitness.values)
+        stats.register("min", np.min)
+        stats.register("avg", np.mean)
+
+        # Run genetic algorithm
+        # TODO(fpedd): Try eaMuPlusLambda and eaMuCommaLambda
+        population, _ = algorithms.eaSimple(
+            population=population,
+            toolbox=toolbox,
+            cxpb=self.crossover_prob,
+            mutpb=self.mutation_prob,
+            ngen=self.num_generations,
+            stats=stats,
+            halloffame=hall_of_fame,
+            verbose=False,
+        )
+
+        # Get best permutation
+        best_individual = hall_of_fame[0]
+        best_permutation = list(best_individual)
+
+        # Apply best permutation and allocate
+        permuted_allocs = tuple(allocations[i] for i in best_permutation)
+        return super().allocate(permuted_allocs)

omnimalloc/allocators/greedy.py

@@ -0,0 +1,82 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from omnimalloc.primitives import Allocation
+
+from .base import BaseAllocator
+
+
+class GreedyAllocator(BaseAllocator):
+    """Base greedy allocator using first-fit strategy."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        placed_allocations: list[Allocation] = []
+
+        for current_alloc in allocations:
+            # Collect overlapping allocations sorted by offset
+            overlapping = [
+                placed
+                for placed in placed_allocations
+                if current_alloc.overlaps_temporally(placed)
+            ]
+            overlapping.sort(key=lambda a: a.offset or 0, reverse=False)
+
+            # Find offset using first-fit (outperforms best-fit in practice)
+            best_offset = 0
+            for placed in overlapping:
+                assert placed.offset is not None
+                gap = placed.offset - best_offset
+                if gap >= current_alloc.size:
+                    break
+                best_offset = max(best_offset, placed.offset + placed.size)
+
+            new_alloc = current_alloc.with_offset(best_offset)
+            placed_allocations.append(new_alloc)
+
+        return tuple(placed_allocations)
+
+
+class GreedyByDurationAllocator(GreedyAllocator):
+    """Greedy allocator sorting by duration (longest first)."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        sorted_allocs = sorted(allocations, key=lambda a: a.duration, reverse=True)
+        return super().allocate(tuple(sorted_allocs))
+
+
+class GreedyByConflictAllocator(GreedyAllocator):
+    """Greedy allocator sorting by conflict degree (most conflicted first)."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        conflict_degrees = {}
+        for alloc in allocations:
+            conflicts = sum(
+                1
+                for other in allocations
+                if other != alloc and alloc.overlaps_temporally(other)
+            )
+            conflict_degrees[alloc] = conflicts
+
+        sorted_allocs = sorted(
+            allocations, key=lambda a: (conflict_degrees[a], a.size), reverse=True
+        )
+        return super().allocate(tuple(sorted_allocs))
+
+
+class GreedyByAreaAllocator(GreedyAllocator):
+    """Greedy allocator sorting by area (size * duration, largest first)."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        sorted_allocs = sorted(
+            allocations, key=lambda a: a.size * a.duration, reverse=True
+        )
+        return super().allocate(tuple(sorted_allocs))
+
+
+class GreedyBySizeAllocator(GreedyAllocator):
+    """Greedy allocator sorting by size (largest first)."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        sorted_allocs = sorted(allocations, key=lambda a: a.size, reverse=True)
+        return super().allocate(tuple(sorted_allocs))

omnimalloc/allocators/greedy_cpp.py

@@ -0,0 +1,65 @@
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+
+from omnimalloc._cpp import GreedyAllocatorCpp as _GreedyAllocatorCpp
+from omnimalloc.primitives import Allocation
+
+from .base import BaseAllocator
+
+
+class GreedyAllocatorCpp(BaseAllocator):
+    """C++ implementation of the base greedy allocator using first-fit strategy."""
+
+    def __init__(self) -> None:
+        self._cpp_allocator = _GreedyAllocatorCpp()
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        allocs_list = list(allocations)
+        result = self._cpp_allocator.allocate(allocs_list)
+        return tuple(result)
+
+
+class GreedyByDurationAllocatorCpp(GreedyAllocatorCpp):
+    """C++ greedy allocator sorting by duration (longest first)."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        sorted_allocs = sorted(allocations, key=lambda a: a.duration, reverse=True)
+        return super().allocate(tuple(sorted_allocs))
+
+
+class GreedyByConflictAllocatorCpp(GreedyAllocatorCpp):
+    """C++ greedy allocator sorting by conflict degree (most conflicted first)."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        conflict_degrees = {}
+        for alloc in allocations:
+            conflicts = sum(
+                1
+                for other in allocations
+                if other != alloc and alloc.overlaps_temporally(other)
+            )
+            conflict_degrees[alloc] = conflicts
+
+        sorted_allocs = sorted(
+            allocations, key=lambda a: (conflict_degrees[a], a.size), reverse=True
+        )
+        return super().allocate(tuple(sorted_allocs))
+
+
+class GreedyByAreaAllocatorCpp(GreedyAllocatorCpp):
+    """C++ greedy allocator sorting by area (size * duration, largest first)."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        sorted_allocs = sorted(
+            allocations, key=lambda a: a.size * a.duration, reverse=True
+        )
+        return super().allocate(tuple(sorted_allocs))
+
+
+class GreedyBySizeAllocatorCpp(GreedyAllocatorCpp):
+    """C++ greedy allocator sorting by size (largest first)."""
+
+    def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]:
+        sorted_allocs = sorted(allocations, key=lambda a: a.size, reverse=True)
+        return super().allocate(tuple(sorted_allocs))