Implement the public API audit

CMakeLists.txt

@@ -42,11 +42,10 @@ nanobind_add_module(
   NOMINSIZE
   src/cpp/allocators/best_fit.cpp
   src/cpp/allocators/first_fit.cpp
-  src/cpp/allocators/greedy.cpp
   src/cpp/allocators/local_search.cpp
   src/cpp/allocators/omni.cpp
   src/cpp/allocators/simulated_annealing.cpp
-  src/cpp/allocators/supermalloc/partition.cpp
+  src/cpp/allocators/supermalloc.cpp
   src/cpp/allocators/tabu_search.cpp
   src/cpp/allocators/telamalloc.cpp
   src/cpp/analysis/antichain.cpp

README.md

@@ -38,15 +38,15 @@ pip install git+https://github.com/fpedd/omnimalloc.git
 ## Usage
 
 ```python
-from omnimalloc import Allocation, Pool, run_allocation
+import omnimalloc as om
 
-pool = Pool(id="pool", allocations=(
-    Allocation(id=0, size=64, start=0, end=10),
-    Allocation(id=1, size=64, start=12, end=20),
-    Allocation(id=2, size=32, start=5, end=15),
+pool = om.Pool(id="pool", allocations=(
+    om.Allocation(id=0, size=64, start=0, end=10),
+    om.Allocation(id=1, size=64, start=12, end=20),
+    om.Allocation(id=2, size=32, start=5, end=15),
 ))
 
-pool = run_allocation(pool, allocator="supermalloc_allocator", validate=True)
+pool = om.allocate(pool, allocator="supermalloc", validate=True)
 
 print(pool.size)                                     # 96
 print([alloc.offset for alloc in pool.allocations])  # [0, 0, 64]

examples/01_basic.py

@@ -2,19 +2,19 @@
 # SPDX-License-Identifier: Apache-2.0
 #
 
-from omnimalloc import Allocation, Pool, run_allocation
+import omnimalloc as om
 
 
 def main() -> None:
     # Define allocations with temporal bounds
-    alloc_0 = Allocation(id="alloc_0", size=5, start=0, end=10)
-    alloc_1 = Allocation(id="alloc_1", size=5, start=12, end=20)
-    alloc_2 = Allocation(id="alloc_2", size=4, start=5, end=15)
-    alloc_3 = Allocation(id="alloc_3", size=5, start=15, end=23)
+    alloc_0 = om.Allocation(id="alloc_0", size=5, start=0, end=10)
+    alloc_1 = om.Allocation(id="alloc_1", size=5, start=12, end=20)
+    alloc_2 = om.Allocation(id="alloc_2", size=4, start=5, end=15)
+    alloc_3 = om.Allocation(id="alloc_3", size=5, start=15, end=23)
 
     # Create pool and allocate
-    pool = Pool(id="pool_0", allocations=(alloc_0, alloc_1, alloc_2, alloc_3))
-    pool = run_allocation(pool, validate=True)
+    pool = om.Pool(id="pool_0", allocations=(alloc_0, alloc_1, alloc_2, alloc_3))
+    pool = om.allocate(pool, validate=True)
 
     # View results
     print(f"Pool {pool.id!r} size: {pool.size}")

examples/02_plotting.py

@@ -2,27 +2,27 @@
 # SPDX-License-Identifier: Apache-2.0
 #
 
-from omnimalloc import Allocation, Pool, plot_allocation, run_allocation
+import omnimalloc as om
 
 
 def main() -> None:
     # Define allocations with temporal bounds
-    alloc_0 = Allocation(id="alloc_0", size=5, start=0, end=10)
-    alloc_1 = Allocation(id="alloc_1", size=5, start=12, end=20)
-    alloc_2 = Allocation(id="alloc_2", size=4, start=5, end=15)
-    alloc_3 = Allocation(id="alloc_3", size=5, start=15, end=23)
+    alloc_0 = om.Allocation(id="alloc_0", size=5, start=0, end=10)
+    alloc_1 = om.Allocation(id="alloc_1", size=5, start=12, end=20)
+    alloc_2 = om.Allocation(id="alloc_2", size=4, start=5, end=15)
+    alloc_3 = om.Allocation(id="alloc_3", size=5, start=15, end=23)
 
     # Create pool and allocate
-    pool = Pool(id="pool_0", allocations=(alloc_0, alloc_1, alloc_2, alloc_3))
-    pool = run_allocation(pool, validate=True)
+    pool = om.Pool(id="pool_0", allocations=(alloc_0, alloc_1, alloc_2, alloc_3))
+    pool = om.allocate(pool, validate=True)
 
     # View results
     print(f"Pool {pool.id!r} size: {pool.size}")
     for alloc in pool.allocations:
         print(f"  {alloc.id!r} offset: {alloc.offset}")
 
     # Visualize (requires matplotlib)
-    plot_allocation(pool, "allocation.pdf")
+    om.plot_allocation(pool, "allocation.pdf")
 
 
 if __name__ == "__main__":

examples/03_allocators.py

@@ -4,47 +4,39 @@
 
 from pathlib import Path
 
-from omnimalloc import (
-    Allocation,
-    Pool,
-    plot_allocation,
-    run_allocation,
-)
-from omnimalloc.allocators import get_available_allocators, get_default_allocator
+import omnimalloc as om
+from omnimalloc.allocators import DEFAULT_ALLOCATOR, available_allocators
 from omnimalloc.allocators.minimalloc import HAS_MINIMALLOC
 
 
 def main() -> None:
     example_dir = Path("03_example_output")
 
     # Define allocations with temporal bounds
-    alloc_0 = Allocation(id="alloc_0", size=5, start=0, end=10)
-    alloc_1 = Allocation(id="alloc_1", size=5, start=12, end=20)
-    alloc_2 = Allocation(id="alloc_2", size=4, start=5, end=15)
-    alloc_3 = Allocation(id="alloc_3", size=5, start=15, end=23)
+    alloc_0 = om.Allocation(id="alloc_0", size=5, start=0, end=10)
+    alloc_1 = om.Allocation(id="alloc_1", size=5, start=12, end=20)
+    alloc_2 = om.Allocation(id="alloc_2", size=4, start=5, end=15)
+    alloc_3 = om.Allocation(id="alloc_3", size=5, start=15, end=23)
 
     # Create pool and allocate
-    pool = Pool(id="pool_0", allocations=(alloc_0, alloc_1, alloc_2, alloc_3))
+    pool = om.Pool(id="pool_0", allocations=(alloc_0, alloc_1, alloc_2, alloc_3))
 
     # Get and run the default allocator
-    default_allocator_name = get_default_allocator()
-    print(f"Running allocation with default allocator: {default_allocator_name}")
-    pool = run_allocation(pool, allocator=default_allocator_name, validate=True)
+    print(f"Running allocation with default allocator: {DEFAULT_ALLOCATOR}")
+    pool = om.allocate(pool, allocator=DEFAULT_ALLOCATOR, validate=True)
     print(f"Pool {pool.id!r} size: {pool.size}")
-    plot_allocation(
-        pool, example_dir / f"allocation_{default_allocator_name}_default.pdf"
-    )
+    om.plot_allocation(pool, example_dir / f"{DEFAULT_ALLOCATOR}_default.pdf")
 
     # Run allocation with all available allocators
-    for allocator_name in get_available_allocators():
+    for allocator_name in available_allocators():
         # minimalloc is an optional dependency that only builds on some platforms
         if "minimalloc" in allocator_name and not HAS_MINIMALLOC:
             print(f"Skipping unavailable allocator: {allocator_name}")
             continue
         print(f"Running allocation with allocator: {allocator_name}")
-        pool = run_allocation(pool, allocator_name, validate=True)
+        pool = om.allocate(pool, allocator_name, validate=True)
         print(f"Pool {pool.id!r} size: {pool.size}")
-        plot_allocation(pool, example_dir / f"allocation_{allocator_name}.pdf")
+        om.plot_allocation(pool, example_dir / f"{allocator_name}.pdf")
 
 
 if __name__ == "__main__":

examples/04_sources.py

@@ -4,38 +4,35 @@
 
 from pathlib import Path
 
-from omnimalloc import plot_allocation, run_allocation
+import omnimalloc as om
 from omnimalloc.benchmark.sources import (
-    get_available_sources,
-    get_default_source,
-    get_source_by_name,
+    DEFAULT_SOURCE,
+    BaseSource,
+    available_sources,
 )
 
 
+def allocate_and_plot(source: BaseSource, output: Path) -> None:
+    pool = source.get_pool()
+    pool = om.allocate(pool, validate=True)
+    print(f"Pool {pool.id!r} size: {pool.size}")
+    om.plot_allocation(pool, output)
+
+
 def main() -> None:
     example_dir = Path("04_example_output")
 
     # Get and use the default source
-    default_source_name = get_default_source()
-    default_source_class = get_source_by_name(default_source_name)
-    default_source = default_source_class()
-    print(f"Using default source: {default_source_name}")
-
-    # Get allocations from the default source
-    pool = default_source.get_pool()
-    pool = run_allocation(pool, validate=True)
-    print(f"Pool {pool.id!r} size: {pool.size}")
-    plot_allocation(pool, example_dir / f"source_{default_source_name}_default.pdf")
-
-    for source_name in get_available_sources():
-        source_class = get_source_by_name(source_name)
-        source = source_class()
+    default_source = BaseSource.get(DEFAULT_SOURCE)()
+    print(f"Using default source: {DEFAULT_SOURCE}")
+    allocate_and_plot(
+        default_source, example_dir / f"source_{DEFAULT_SOURCE}_default.pdf"
+    )
+
+    for source_name in available_sources():
+        source = BaseSource.get(source_name)()
         print(f"Using source: {source_name}")
-
-        pool = source.get_pool()
-        pool = run_allocation(pool, validate=True)
-        print(f"Pool {pool.id!r} size: {pool.size}")
-        plot_allocation(pool, example_dir / f"source_{source_name}.pdf")
+        allocate_and_plot(source, example_dir / f"source_{source_name}.pdf")
 
 
 if __name__ == "__main__":

examples/05_benchmark.py

@@ -17,25 +17,25 @@ def main() -> None:
 
     # Define allocators, sources, and variants to benchmark
     allocators = (
-        "greedy_by_size_allocator",
-        "greedy_by_size_allocator_cpp",
-        "greedy_by_all_allocator_cpp",
-        "best_fit_allocator",
-        "telamalloc_allocator",
+        "greedy_by_size",
+        "greedy_by_all",
+        "omni",
+        "best_fit",
+        "telamalloc",
     )
     # minimalloc is an optional dependency that only builds on some platforms
     if HAS_MINIMALLOC:
-        allocators += ("minimalloc_allocator",)
+        allocators += ("minimalloc",)
     sources = (
-        "random_source",
-        "minimalloc_source",
-        "huggingface_source",
+        "random",
+        "minimalloc",
+        "huggingface",
     )
     # Counts for the parameterizable source, "first 5" for the fixed ones
     variants = {
-        "random_source": (10, 50, 100, 250, 500),
-        "minimalloc_source": 5,
-        "huggingface_source": 5,
+        "random": (10, 50, 100, 250, 500),
+        "minimalloc": 5,
+        "huggingface": 5,
     }
 
     # Run benchmark campaign

notebooks/custom_allocator.ipynb

@@ -11,14 +11,124 @@
     "%autoreload 2"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
+   "source": [
+    "# Custom allocator\n",
+    "\n",
+    "Subclass `BaseAllocator`, implement the `_allocate` hook, and the registry\n",
+    "picks the allocator up automatically — usable by name everywhere, including\n",
+    "`om.allocate` and the benchmark harness."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1",
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import omnimalloc as om\n",
+    "from omnimalloc.allocators import BaseAllocator, available_allocators\n",
+    "from omnimalloc.benchmark.sources import RandomSource"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "`_allocate` receives validated, non-empty allocations with unique ids and\n",
+    "returns them with offsets assigned. The one constraint to satisfy:\n",
+    "*conflicting* allocations — those whose lifetimes can coexist, as reported\n",
+    "by `om.conflicts` — must occupy disjoint address ranges.\n",
+    "\n",
+    "This allocator places largest-first, each allocation at the lowest offset\n",
+    "not blocked by an already-placed conflict. Because it consumes only the\n",
+    "conflict relation, it also accepts vector-clock lifetimes\n",
+    "(`supports_vector_time = True`)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LargestFirstAllocator(BaseAllocator):\n",
+    "    \"\"\"Place allocations largest-first, each at the lowest conflict-free offset.\"\"\"\n",
+    "\n",
+    "    supports_vector_time = True\n",
+    "\n",
+    "    def _allocate(\n",
+    "        self, allocations: tuple[om.Allocation, ...]\n",
+    "    ) -> tuple[om.Allocation, ...]:\n",
+    "        conflict_map = om.conflicts(allocations)\n",
+    "        sizes = {alloc.id: alloc.size for alloc in allocations}\n",
+    "        offsets: dict[om.IdType, int] = {}\n",
+    "        for alloc in sorted(allocations, key=lambda a: a.size, reverse=True):\n",
+    "            blockers = sorted(\n",
+    "                (offsets[other], offsets[other] + sizes[other])\n",
+    "                for other in conflict_map[alloc.id]\n",
+    "                if other in offsets\n",
+    "            )\n",
+    "            offset = 0\n",
+    "            for lower, upper in blockers:\n",
+    "                if lower - offset >= alloc.size:\n",
+    "                    break\n",
+    "                offset = max(offset, upper)\n",
+    "            offsets[alloc.id] = offset\n",
+    "        return tuple(alloc.with_offset(offsets[alloc.id]) for alloc in allocations)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "Defining the class registered it: the `Allocator` role token is stripped and\n",
+    "the rest snake_cased, so it answers to `\"largest_first\"`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "assert \"largest_first\" in available_allocators()\n",
+    "\n",
+    "pool = RandomSource(num_allocations=200).get_pool()\n",
+    "pool = om.allocate(pool, \"largest_first\", validate=True)\n",
+    "print(\n",
+    "    f\"peak memory {pool.size}, lower bound {pool.pressure}, \"\n",
+    "    f\"efficiency {pool.efficiency:.3f}\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7",
+   "metadata": {},
+   "source": [
+    "Compare it against a few built-in allocators, then plot its placement."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(\"TODO(fpedd): Implement this\")"
+    "for name in (\"naive\", \"greedy_by_size\", \"omni\", \"largest_first\"):\n",
+    "    print(f\"{name:>14}: peak {om.allocate(pool, name).size}\")\n",
+    "\n",
+    "om.plot_allocation(pool)"
    ]
   }
  ],