diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 666eb45..fd68ce8 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -6,6 +6,8 @@ repos:
rev: v6.0.0
hooks:
- id: trailing-whitespace
+ # matplotlib-generated README figures contain trailing whitespace
+ exclude: '^assets/.*\.svg$'
- id: end-of-file-fixer
- id: mixed-line-ending
- id: check-yaml
diff --git a/README.md b/README.md
index d5727f4..ccce04b 100644
--- a/README.md
+++ b/README.md
@@ -10,19 +10,21 @@
-
-
-
-
-
-OmniMalloc is a Python library for **static memory allocation**: given buffers
-with known sizes and lifetimes, assign offsets so that **peak memory is minimized**.
-This is the memory-planning step at the heart of **ML compilers**, embedded
+
+
+
+
+
+
+
+OmniMalloc is a Python library for static memory allocation: given buffers
+with known sizes and lifetimes, assign offsets so that peak memory is minimized.
+This is the memory-planning step at the heart of ML compilers, embedded
runtimes, and accelerator toolchains.
It ships a collection of allocators and allocation algorithms behind one API,
-implemented with an efficient C++ backend. This includes **SuperMalloc**, a new
-allocator that **outperforms the best open-source alternatives** (see benchmarks
+implemented with an efficient C++ backend. This includes SuperMalloc, a new
+allocator that outperforms the best open-source alternatives (see benchmarks
below). OmniMalloc also provides a rich benchmark harness and visualization
tools to develop and evaluate new allocation strategies.
@@ -50,12 +52,14 @@ print([alloc.offset for alloc in pool.allocations]) # [0, 0, 64]
```
On a real problem, the result looks like this: 308 buffers of an ML workload
-packed with zero wasted memory.
+packed with no wasted memory.
-
-
-
-
+
+
+
+
+
+
See [examples](examples/) for allocator selection, visualization, custom
allocation sources, and benchmarking.
diff --git a/scripts/generate_readme_assets.py b/scripts/generate_readme_assets.py
index e313846..8825ee4 100644
--- a/scripts/generate_readme_assets.py
+++ b/scripts/generate_readme_assets.py
@@ -25,15 +25,14 @@
from __future__ import annotations
import argparse
-import importlib
import json
import random
import shutil
import subprocess
import sys
-import time
from dataclasses import dataclass
from pathlib import Path
+from statistics import mean
from typing import TYPE_CHECKING, Any
import matplotlib as mpl
@@ -43,8 +42,10 @@
from matplotlib.ticker import FuncFormatter, MultipleLocator
from omnimalloc import run_allocation, validate_allocation
from omnimalloc.allocators import BaseAllocator
+from omnimalloc.allocators.minimalloc import MinimallocAllocator
from omnimalloc.allocators.supermalloc import SupermallocAllocator, SupermallocConfig
from omnimalloc.benchmark.sources import BaseSource
+from omnimalloc.benchmark.timer import Timer
from omnimalloc.common.units import MB
if TYPE_CHECKING:
@@ -58,35 +59,82 @@
MINIMALLOC_URL = "git+https://github.com/google/minimalloc.git"
SCALING_SIZES = (10, 32, 100, 316, 1000, 3162, 10000)
SCALING_SIZES_SLOW = SCALING_SIZES[:-1] # hill climbing needs minutes at 10k
+ALLOCATION_PROBLEM = "mm-G"
ASSETS_DIR = Path(__file__).resolve().parent.parent / "assets"
-# Hero points: allocator -> (display label, palette role).
-HERO_ALLOCATORS: dict[str, tuple[str, str]] = {
+# Allocator display metadata: registry name -> (label, palette role).
+ALLOCATORS: dict[str, tuple[str, str]] = {
+ "naive_allocator": ("naive", "baseline"),
"random_allocator": ("random search", "baseline"),
"greedy_by_area_allocator_cpp": ("greedy (area)", "greedy"),
"greedy_by_size_allocator_cpp": ("greedy (size)", "greedy"),
"greedy_by_all_allocator_cpp": ("greedy (all)", "greedy"),
"hill_climb_allocator": ("hill climbing", "search"),
"genetic_allocator": ("genetic", "search"),
- "minimalloc": ("minimalloc", "minimalloc"),
- "supermalloc": ("supermalloc", "exact"),
+ "minimalloc_allocator": ("minimalloc", "minimalloc"),
+ "supermalloc_allocator": ("supermalloc", "exact"),
}
+HERO_ALLOCATORS = (
+ "random_allocator",
+ "greedy_by_area_allocator_cpp",
+ "greedy_by_size_allocator_cpp",
+ "greedy_by_all_allocator_cpp",
+ "hill_climb_allocator",
+ "genetic_allocator",
+ "minimalloc_allocator",
+ "supermalloc_allocator",
+)
QUALITY_ALLOCATORS = (
"greedy_by_size_allocator_cpp",
"greedy_by_all_allocator_cpp",
- "minimalloc",
- "supermalloc",
+ "minimalloc_allocator",
+ "supermalloc_allocator",
)
+SCALING_ALLOCATORS = (
+ "naive_allocator",
+ "greedy_by_size_allocator_cpp",
+ "hill_climb_allocator",
+ "minimalloc_allocator",
+ "supermalloc_allocator",
+)
+
+# Quality re-colors greedy (size) so the two greedy variants stay distinguishable.
+QUALITY_ROLES = {"greedy_by_size_allocator_cpp": "greedy_alt"}
+
QUALITY_PROBLEMS = ("mm-A", "mm-C", "mm-H", "mm-K", "pinwheel", "tiling", "random")
+PROBLEM_LABELS = {
+ "mm-A": "minimalloc A",
+ "mm-C": "minimalloc C",
+ "mm-G": "minimalloc G",
+ "mm-H": "minimalloc H",
+ "mm-K": "minimalloc K",
+ "pinwheel": "pinwheel",
+ "tiling": "tiling",
+ "random": "random (easy)",
+}
-SCALING_ALLOCATORS: dict[str, tuple[str, str]] = {
- "naive_allocator": ("naive", "baseline"),
- "greedy_by_size_allocator_cpp": ("greedy (size)", "greedy"),
- "hill_climb_allocator": ("hill climbing", "search"),
- "minimalloc": ("minimalloc", "minimalloc"),
- "supermalloc": ("supermalloc", "exact"),
+# Direct-label offsets in points, tuned per hero point: (dx, dy, ha).
+HERO_LABEL_OFFSETS: dict[str, tuple[float, float, str]] = {
+ "random_allocator": (0, -11, "center"),
+ "greedy_by_area_allocator_cpp": (0, 8, "center"),
+ "greedy_by_size_allocator_cpp": (-8, 0, "right"),
+ "greedy_by_all_allocator_cpp": (0, -11, "center"),
+ "hill_climb_allocator": (0, 8, "center"),
+ "genetic_allocator": (0, -11, "center"),
+ "minimalloc_allocator": (0, 8, "center"),
+ "supermalloc_allocator": (-10, 0, "right"),
+}
+
+# Direct-label offsets in points: (dx, dy, ha). minimalloc's line ends early
+# (no 10k point), so its label anchors left, away from the 10k label cluster.
+SCALING_LABEL_OFFSETS = {
+ "naive_allocator": (4, -2, "left"),
+ "greedy_by_size_allocator_cpp": (4, -4, "left"),
+ "hill_climb_allocator": (4, 2, "left"),
+ "minimalloc_allocator": (0, 9, "center"),
+ "supermalloc_allocator": (4, 4, "left"),
}
@@ -147,37 +195,31 @@ def _pip_install(spec: str) -> None:
def _ensure_minimalloc() -> None:
- """Install Google's minimalloc on demand (no PyPI wheel) and refresh the wrapper."""
+ """Install Google's minimalloc on demand (no PyPI wheel)."""
try:
import minimalloc # type: ignore # noqa: F401
except ImportError:
- pass
- else:
- return
- print(f"minimalloc not installed — installing from {MINIMALLOC_URL} ...")
- _pip_install(MINIMALLOC_URL)
- importlib.invalidate_caches()
- # The wrapper decides availability at import time; reload it post-install.
- importlib.reload(importlib.import_module("omnimalloc.allocators.minimalloc"))
+ print(f"minimalloc not installed — installing from {MINIMALLOC_URL} ...")
+ _pip_install(MINIMALLOC_URL)
def _allocator(name: str, timeout: float = SUPERMALLOC_TIMEOUT) -> BaseAllocator:
- if name == "supermalloc":
+ if name == "supermalloc_allocator":
return SupermallocAllocator(config=SupermallocConfig(timeout=timeout))
- if name == "minimalloc":
- from omnimalloc.allocators.minimalloc import MinimallocAllocator
-
+ if name == "minimalloc_allocator":
return MinimallocAllocator(timeout=int(timeout))
- return BaseAllocator.get(name)()
+ return BaseAllocator.resolve(name)
-def _solve(pool: Pool, allocator: BaseAllocator) -> tuple[float, float, Pool]:
+def _solve(
+ pool: Pool, allocator: BaseAllocator, *, validate: bool = True
+) -> tuple[float, float, Pool]:
"""Time the solve alone; validation is quadratic and would skew timings."""
- start = time.perf_counter()
- solved = run_allocation(pool, allocator=allocator)
- seconds = time.perf_counter() - start
- validate_allocation(solved)
- return seconds, solved.efficiency, solved
+ with Timer() as timer:
+ solved = run_allocation(pool, allocator=allocator)
+ if validate:
+ validate_allocation(solved)
+ return timer.elapsed_s, solved.efficiency, solved
def _hard_suite() -> dict[str, Pool]:
@@ -196,51 +238,56 @@ def collect_data() -> dict[str, Any]:
_ensure_minimalloc()
random.seed(SEED)
suite = _hard_suite()
- hard = {k: v for k, v in suite.items() if k != "random"}
+ hard = [k for k in suite if k != "random"]
- # Hero + quality: every allocator over every problem.
+ # Hero + quality: every problem, but "random" only where quality needs it.
runs: dict[str, dict[str, tuple[float, float]]] = {}
+ supermalloc_pools: dict[str, Pool] = {}
names = set(HERO_ALLOCATORS) | set(QUALITY_ALLOCATORS)
for name in sorted(names):
allocator = _allocator(name)
+ problems = tuple(suite) if name in QUALITY_ALLOCATORS else hard
runs[name] = {}
- for problem, pool in suite.items():
- seconds, efficiency, _ = _solve(pool, allocator)
+ for problem in problems:
+ seconds, efficiency, solved = _solve(suite[problem], allocator)
+ if name == "supermalloc_allocator":
+ supermalloc_pools[problem] = solved
runs[name][problem] = (seconds, efficiency)
print(f"{name:38s} {problem:10s} {seconds:8.3f}s {efficiency:7.2%}")
hero = {
name: {
- "seconds": sum(runs[name][p][0] for p in hard) / len(hard),
- "efficiency": sum(runs[name][p][1] for p in hard) / len(hard) * 100,
+ "seconds": mean(runs[name][p][0] for p in hard),
+ "efficiency": mean(runs[name][p][1] for p in hard),
}
for name in HERO_ALLOCATORS
}
quality = {
- name: {p: runs[name][p][1] * 100 for p in QUALITY_PROBLEMS}
+ name: {p: runs[name][p][1] for p in QUALITY_PROBLEMS}
for name in QUALITY_ALLOCATORS
}
- # Scaling: solve time vs. problem size on the random source.
+ # Scaling: solve time vs. problem size on the random source. Skip validation
+ # here: it is quadratic in pure Python and would dwarf the fast solves at 10k.
source = BaseSource.get("random_source")()
- scaling: dict[str, dict[str, float]] = {}
+ scaling: dict[str, list[list[float]]] = {}
for name in SCALING_ALLOCATORS:
allocator = _allocator(name, SCALING_TIMEOUT)
# Hill climbing needs minutes at 10k; minimalloc can't solve 10k in budget.
- capped = name in ("hill_climb_allocator", "minimalloc")
+ capped = name in ("hill_climb_allocator", "minimalloc_allocator")
sizes = SCALING_SIZES_SLOW if capped else SCALING_SIZES
- scaling[name] = {}
+ scaling[name] = []
for size in sizes:
- seconds, _, _ = _solve(source.get_variant(size), allocator)
- scaling[name][str(size)] = seconds
+ seconds, _, _ = _solve(source.get_variant(size), allocator, validate=False)
+ scaling[name].append([size, seconds])
print(f"{name:38s} n={size:<6d} {seconds:8.3f}s")
- # Allocation rendering: one real problem solved to proven optimality.
- problem = "mm-G"
- _, efficiency, solved = _solve(suite[problem], _allocator("supermalloc"))
+ # Allocation rendering: a real problem solved to proven optimality. The loop
+ # above already solved it with the same budget, so reuse that pool.
+ solved = supermalloc_pools[ALLOCATION_PROBLEM]
allocation = {
- "problem": QUALITY_ROW_LABELS.get(problem, problem),
- "efficiency": efficiency,
+ "problem": PROBLEM_LABELS.get(ALLOCATION_PROBLEM, ALLOCATION_PROBLEM),
+ "efficiency": runs["supermalloc_allocator"][ALLOCATION_PROBLEM][1],
"size": solved.size,
"rects": [[a.start, a.duration, a.offset, a.size] for a in solved.allocations],
}
@@ -262,22 +309,28 @@ def _style(theme: Theme) -> dict[str, Any]:
"axes.facecolor": "none",
"savefig.facecolor": "none",
"savefig.transparent": True,
- "axes.edgecolor": theme.grid,
+ "axes.spines.left": False,
+ "axes.spines.right": False,
+ "axes.spines.top": False,
+ "axes.spines.bottom": False,
"axes.labelcolor": theme.muted,
- "axes.linewidth": 0.8,
- "xtick.color": theme.grid,
- "ytick.color": theme.grid,
+ "axes.labelsize": 8.5,
+ "grid.color": theme.grid,
+ "grid.linewidth": 0.5,
+ "grid.alpha": 0.45,
"xtick.labelcolor": theme.muted,
"ytick.labelcolor": theme.muted,
"xtick.labelsize": 8.0,
"ytick.labelsize": 8.0,
- "axes.labelsize": 8.5,
+ "xtick.major.size": 0,
+ "ytick.major.size": 0,
+ "xtick.minor.size": 0,
+ "ytick.minor.size": 0,
}
-def _despine(ax: Axes, keep: tuple[str, ...] = ("left", "bottom")) -> None:
- for side, spine in ax.spines.items():
- spine.set_visible(side in keep)
+def _quality_role(name: str) -> str:
+ return QUALITY_ROLES.get(name, ALLOCATORS[name][1])
def _title(fig: Figure, theme: Theme, title: str, subtitle: str) -> None:
@@ -304,6 +357,27 @@ def _series_line(fig: Figure, series: list[tuple[str, str]], y: float) -> None:
x += 0.033 + len(label) * 0.0148
+def _optimal_line(
+ ax: Axes,
+ theme: Theme,
+ value: float,
+ axis: str = "y",
+ linewidth: float = 0.8,
+ alpha: float = 0.8,
+ zorder: int = 1,
+) -> None:
+ """Dashed reference line marking the proven-optimal value."""
+ line = ax.axhline if axis == "y" else ax.axvline
+ line(
+ value,
+ color=theme.optimal,
+ linewidth=linewidth,
+ linestyle=(0, (4, 4)),
+ alpha=alpha,
+ zorder=zorder,
+ )
+
+
def _format_seconds(value: float) -> str:
if value < 1e-3:
return f"{value * 1e6:.0f} µs"
@@ -322,12 +396,8 @@ def _format_steps(value: float, _pos: int) -> str:
def _save(fig: Figure, name: str, theme: Theme, preview: Path | None) -> None:
ASSETS_DIR.mkdir(parents=True, exist_ok=True)
- svg = ASSETS_DIR / f"{name}_{theme.name}.svg"
- fig.savefig(svg, bbox_inches="tight", pad_inches=0.02)
- # matplotlib emits trailing whitespace on some lines; strip it so the
- # committed assets stay clean under the trailing-whitespace pre-commit hook.
- svg.write_text(
- "\n".join(line.rstrip() for line in svg.read_text().splitlines()) + "\n"
+ fig.savefig(
+ ASSETS_DIR / f"{name}_{theme.name}.svg", bbox_inches="tight", pad_inches=0.02
)
if preview is not None:
preview.mkdir(parents=True, exist_ok=True)
@@ -342,35 +412,15 @@ def _save(fig: Figure, name: str, theme: Theme, preview: Path | None) -> None:
plt.close(fig)
-# Direct-label offsets in points, tuned per hero point: (dx, dy, ha).
-HERO_LABEL_OFFSETS: dict[str, tuple[float, float, str]] = {
- "random search": (0, -11, "center"),
- "greedy (area)": (0, 8, "center"),
- "greedy (size)": (-8, 0, "right"),
- "greedy (all)": (0, -11, "center"),
- "hill climbing": (0, 8, "center"),
- "genetic": (0, -11, "center"),
- "minimalloc": (0, 8, "center"),
- "supermalloc": (-10, 0, "right"),
-}
-
-
def render_hero(data: dict[str, Any], theme: Theme, preview: Path | None) -> None:
fig, ax = plt.subplots(figsize=(7.4, 3.5))
points = {
- label: (data[name]["seconds"], data[name]["efficiency"], role)
- for name, (label, role) in HERO_ALLOCATORS.items()
+ name: (data[name]["seconds"], data[name]["efficiency"] * 100)
+ for name in HERO_ALLOCATORS
}
- ax.axhline(
- 100,
- color=theme.optimal,
- linewidth=0.8,
- linestyle=(0, (4, 4)),
- alpha=0.8,
- zorder=1,
- )
+ _optimal_line(ax, theme, 100)
ax.annotate(
"optimal",
xy=(1.0, 100),
@@ -385,7 +435,7 @@ def render_hero(data: dict[str, Any], theme: Theme, preview: Path | None) -> Non
# Pareto frontier: lower time and higher efficiency dominate.
ordered = sorted(points.values(), key=lambda p: (p[0], -p[1]))
front, best = [], float("-inf")
- for seconds, efficiency, _ in ordered:
+ for seconds, efficiency in ordered:
if efficiency > best:
front.append((seconds, efficiency))
best = efficiency
@@ -397,9 +447,10 @@ def render_hero(data: dict[str, Any], theme: Theme, preview: Path | None) -> Non
zorder=2,
)
- for label, (seconds, efficiency, role) in points.items():
+ for name, (seconds, efficiency) in points.items():
+ label, role = ALLOCATORS[name]
color = theme.role[role]
- emphasis = label == "supermalloc"
+ emphasis = name == "supermalloc_allocator"
ax.scatter(
seconds,
efficiency,
@@ -409,7 +460,7 @@ def render_hero(data: dict[str, Any], theme: Theme, preview: Path | None) -> Non
edgecolors=theme.ink if emphasis else "none",
zorder=4,
)
- dx, dy, ha = HERO_LABEL_OFFSETS[label]
+ dx, dy, ha = HERO_LABEL_OFFSETS[name]
ax.annotate(
label,
(seconds, efficiency),
@@ -429,10 +480,7 @@ def render_hero(data: dict[str, Any], theme: Theme, preview: Path | None) -> Non
ticks = (1e-3, 1e-2, 1e-1, 1, 10)
ax.set_xticks(ticks)
ax.set_xticklabels([_format_seconds(t) for t in ticks])
- ax.grid(visible=True, axis="both", color=theme.grid, linewidth=0.5, alpha=0.45)
- _despine(ax, keep=())
- ax.tick_params(length=0, which="both")
- ax.minorticks_off()
+ ax.grid(visible=True, axis="both")
ax.set_xlabel("mean solve time (log scale)")
ax.set_ylabel("mean packing efficiency (%)")
@@ -447,41 +495,16 @@ def render_hero(data: dict[str, Any], theme: Theme, preview: Path | None) -> Non
_save(fig, "hero", theme, preview)
-QUALITY_LABELS = {
- "greedy_by_size_allocator_cpp": ("greedy (size)", "greedy_alt"),
- "greedy_by_all_allocator_cpp": ("greedy (all)", "greedy"),
- "minimalloc": ("minimalloc", "minimalloc"),
- "supermalloc": ("supermalloc", "exact"),
-}
-QUALITY_ROW_LABELS = {
- "mm-A": "minimalloc A",
- "mm-C": "minimalloc C",
- "mm-G": "minimalloc G",
- "mm-H": "minimalloc H",
- "mm-K": "minimalloc K",
- "pinwheel": "pinwheel",
- "tiling": "tiling",
- "random": "random (easy)",
-}
-
-
def render_quality(data: dict[str, Any], theme: Theme, preview: Path | None) -> None:
fig, ax = plt.subplots(figsize=(3.8, 3.35))
rows = list(QUALITY_PROBLEMS)
ys = range(len(rows), 0, -1)
- ax.axvline(
- 100,
- color=theme.optimal,
- linewidth=0.8,
- linestyle=(0, (4, 4)),
- alpha=0.8,
- zorder=1,
- )
+ _optimal_line(ax, theme, 100, axis="x")
for row, y in zip(rows, ys, strict=True):
- values = [data[name][row] for name in QUALITY_ALLOCATORS]
+ values = [data[name][row] * 100 for name in QUALITY_ALLOCATORS]
ax.plot(
[min(values), max(values)],
[y, y],
@@ -491,55 +514,43 @@ def render_quality(data: dict[str, Any], theme: Theme, preview: Path | None) ->
solid_capstyle="round",
)
for name, value in zip(QUALITY_ALLOCATORS, values, strict=True):
- _, role = QUALITY_LABELS[name]
- emphasis = name == "supermalloc"
+ emphasis = name == "supermalloc_allocator"
ax.scatter(
value,
y,
s=52 if emphasis else 34,
- color=theme.role[role],
+ color=theme.role[_quality_role(name)],
zorder=4,
linewidths=1.2 if emphasis else 0,
edgecolors=theme.ink if emphasis else "none",
)
ax.set_yticks(list(ys))
- ax.set_yticklabels([QUALITY_ROW_LABELS[r] for r in rows], fontsize=8.5)
+ ax.set_yticklabels([PROBLEM_LABELS[r] for r in rows], fontsize=8.5)
ax.set_ylim(0.4, len(rows) + 0.6)
ax.set_xlim(60, 103)
ax.set_xticks((60, 70, 80, 90, 100))
- ax.grid(visible=True, axis="x", color=theme.grid, linewidth=0.5, alpha=0.45)
- _despine(ax, keep=())
- ax.tick_params(length=0)
+ ax.grid(visible=True, axis="x")
ax.set_xlabel("packing efficiency (%)")
_title(fig, theme, "Quality per problem", "100% = proven lower bound")
- series = [(label, theme.role[role]) for label, role in QUALITY_LABELS.values()]
+ series = [
+ (ALLOCATORS[name][0], theme.role[_quality_role(name)])
+ for name in QUALITY_ALLOCATORS
+ ]
_series_line(fig, series, y=0.842)
fig.subplots_adjust(top=0.775, bottom=0.125, left=0.265, right=0.97)
_save(fig, "quality", theme, preview)
-# Direct-label offsets in points: (dx, dy, ha). minimalloc's line ends early
-# (no 10k point), so its label anchors left, away from the 10k label cluster.
-SCALING_LABEL_OFFSETS = {
- "naive": (4, -2, "left"),
- "greedy (size)": (4, -4, "left"),
- "hill climbing": (4, 2, "left"),
- "minimalloc": (0, 9, "center"),
- "supermalloc": (4, 4, "left"),
-}
-
-
def render_scaling(data: dict[str, Any], theme: Theme, preview: Path | None) -> None:
fig, ax = plt.subplots(figsize=(3.8, 3.35))
- for name, (label, role) in SCALING_ALLOCATORS.items():
- series = data[name]
- sizes = sorted(int(k) for k in series)
- seconds = [series[str(n)] for n in sizes]
+ for name in SCALING_ALLOCATORS:
+ label, role = ALLOCATORS[name]
+ sizes, seconds = zip(*data[name], strict=True)
color = theme.role[role]
- emphasis = name == "supermalloc"
+ emphasis = name == "supermalloc_allocator"
ax.plot(
sizes,
seconds,
@@ -550,7 +561,7 @@ def render_scaling(data: dict[str, Any], theme: Theme, preview: Path | None) ->
markeredgewidth=0,
zorder=4,
)
- dx, dy, ha = SCALING_LABEL_OFFSETS[label]
+ dx, dy, ha = SCALING_LABEL_OFFSETS[name]
ax.annotate(
label,
(sizes[-1], seconds[-1]),
@@ -572,10 +583,7 @@ def render_scaling(data: dict[str, Any], theme: Theme, preview: Path | None) ->
ax.set_ylim(3e-6, 400)
ax.set_yticks(yticks)
ax.set_yticklabels([_format_seconds(t) for t in yticks])
- ax.grid(visible=True, axis="both", color=theme.grid, linewidth=0.5, alpha=0.45)
- _despine(ax, keep=())
- ax.tick_params(length=0, which="both")
- ax.minorticks_off()
+ ax.grid(visible=True, axis="both")
ax.set_xlabel("number of allocations")
ax.set_ylabel("solve time")
@@ -615,9 +623,7 @@ def render_allocation(data: dict[str, Any], theme: Theme, preview: Path | None)
)
)
- ax.axhline(
- size, color=theme.optimal, linewidth=0.9, linestyle=(0, (4, 4)), zorder=5
- )
+ _optimal_line(ax, theme, size, linewidth=0.9, alpha=1.0, zorder=5)
ax.annotate(
f"peak {size / MB:.2f} MB = lower bound (proven optimal)",
xy=(0.995, size),
@@ -635,9 +641,7 @@ def render_allocation(data: dict[str, Any], theme: Theme, preview: Path | None)
ax.yaxis.set_major_locator(MultipleLocator(0.25 * MB))
ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v / MB:g}"))
ax.xaxis.set_major_formatter(FuncFormatter(_format_steps))
- ax.grid(visible=True, axis="y", color=theme.grid, linewidth=0.5, alpha=0.45)
- _despine(ax, keep=())
- ax.tick_params(length=0)
+ ax.grid(visible=True, axis="y")
ax.set_xlabel("time step")
ax.set_ylabel("offset (MB)")
diff --git a/src/python/omnimalloc/allocators/minimalloc.py b/src/python/omnimalloc/allocators/minimalloc.py
index 26a6847..f69a38f 100644
--- a/src/python/omnimalloc/allocators/minimalloc.py
+++ b/src/python/omnimalloc/allocators/minimalloc.py
@@ -12,12 +12,27 @@
try:
import minimalloc as mm # type: ignore
-
- HAS_MINIMALLOC = True
except ImportError:
- HAS_MINIMALLOC = False
mm = cast("Any", None)
+HAS_MINIMALLOC = mm is not None
+
+
+def _require_minimalloc() -> None:
+ """Re-check availability at use time so installs after package import work."""
+ global mm, HAS_MINIMALLOC
+ if mm is not None:
+ return
+ try:
+ import minimalloc # type: ignore
+ except ImportError:
+ # TODO(fpedd): Make minimalloc more easily installable via PyPI
+ raise OptionalDependencyError(
+ "The MinimallocAllocator feature requires 'minimalloc' which is not "
+ "installed.\nInstall manually: pip install git+https://github.com/google/minimalloc.git"
+ ) from None
+ mm, HAS_MINIMALLOC = minimalloc, True
+
def _to_buffer(allocation: Allocation) -> "mm.Buffer":
return mm.Buffer(
@@ -31,13 +46,7 @@ class MinimallocAllocator(BaseAllocator):
"""Wrapper for Google's minimalloc constraint-based allocator."""
def __init__(self, timeout: int = 10, max_capacity: int = 1 * TB) -> None:
- if not HAS_MINIMALLOC:
- # TODO(fpedd): Make minimalloc more easily installable via PyPI
- raise OptionalDependencyError(
- "The MinimallocAllocator feature requires 'minimalloc' which is not "
- "installed.\nInstall manually: pip install git+https://github.com/google/minimalloc.git"
- )
-
+ _require_minimalloc()
self._timeout = timeout
self._max_capacity = max_capacity
-
-