diff --git a/.gitignore b/.gitignore index 3b59b5a..1903991 100644 --- a/.gitignore +++ b/.gitignore @@ -61,6 +61,7 @@ Thumbs.db *.log .env *.pdf +*.pkl CLAUDE.md .claude *.json diff --git a/CMakeLists.txt b/CMakeLists.txt index 3ec4881..cdf433e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -44,6 +44,7 @@ nanobind_add_module( src/cpp/allocators/greedy_base.cpp src/cpp/allocators/simulated_annealing.cpp src/cpp/allocators/supermalloc/partition.cpp + src/cpp/allocators/sweep.cpp src/cpp/allocators/tabu_search.cpp src/cpp/allocators/telamalloc.cpp src/cpp/primitives/allocation.cpp diff --git a/pyproject.toml b/pyproject.toml index 88c22d6..1dcd357 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,7 +87,8 @@ archs = ["arm64"] [tool.ruff] line-length = 88 target-version = "py310" -exclude = ["*.ipynb", "*.pyi"] +# research/ holds throwaway prototype and benchmark scripts, not library code +exclude = ["*.ipynb", "*.pyi", "research/"] [tool.ruff.lint] select = ["ALL"] @@ -149,4 +150,4 @@ log_cli_date_format = "%Y-%m-%d %H:%M:%S" python-version = "3.10" [tool.ty.src] -exclude = ["tests/", "*.pyi"] +exclude = ["tests/", "*.pyi", "research/"] diff --git a/research/bench.py b/research/bench.py new file mode 100644 index 0000000..cfa6427 --- /dev/null +++ b/research/bench.py @@ -0,0 +1,167 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# +"""Quality benchmark: candidates vs GreedyByAll (best-of-7 C++ first-fit).""" + +import csv +import math +import pickle +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) + +from candidates import ( # ty: ignore[unresolved-import] + CANDIDATES, + lower_bound, + peak_of, + validate, +) +from omnimalloc.allocators.greedy_cpp import ( + GreedyAllocatorCpp, + GreedyByAreaAllocatorCpp, + GreedyByConflictAllocatorCpp, + GreedyByConflictSizeAllocatorCpp, + GreedyByDurationAllocatorCpp, + GreedyBySizeAllocatorCpp, + GreedyByStartAllocatorCpp, +) +from omnimalloc.benchmark.sources.generator import ( + HighContentionSource, + PowerOf2Source, + RandomSource, + SequentialSource, +) +from omnimalloc.primitives import Allocation + +EXTERNAL = Path(__file__).parent.parent / "external" / "minimalloc" + +GREEDY_VARIANTS = ( + GreedyAllocatorCpp, + GreedyBySizeAllocatorCpp, + GreedyByDurationAllocatorCpp, + GreedyByAreaAllocatorCpp, + GreedyByConflictAllocatorCpp, + GreedyByConflictSizeAllocatorCpp, + GreedyByStartAllocatorCpp, +) + + +def load_problems(): + problems = {} + for csv_path in sorted(EXTERNAL.glob("*/*.csv")): + with open(csv_path) as f: + rows = list(csv.DictReader(f)) + allocs = [(int(r["size"]), int(r["lower"]), int(r["upper"])) for r in rows] + problems[f"mm_{csv_path.parent.name[0]}_{csv_path.stem.split('.')[0]}"] = allocs + + sources = { + "rand_100": RandomSource(num_allocations=100), + "rand_1000": RandomSource(num_allocations=1000), + "rand_5000": RandomSource(num_allocations=5000, time_max=100000), + "contention_1000": HighContentionSource(num_allocations=1000), + "pow2_1000": PowerOf2Source(num_allocations=1000, time_max=1000), + "seq_1000": SequentialSource(num_allocations=1000), + } + for name, source in sources.items(): + problems[name] = [(a.size, a.start, a.end) for a in source.get_allocations()] + problems.update(load_onnx_problems()) + return problems + + +ONNX_MODELS = ( + "resnet18_Opset18_timm", + "alexnet_Opset17", + "adv_inception_v3_Opset18", + "beit_large_patch16_224_Opset17", + "cait_m48_448_Opset18", + "convnext_large_384_in22ft1k_Opset18", + "convmixer_1536_20_Opset16", + "coat_mini_Opset18", +) + + +def load_onnx_problems(): + cache = Path(__file__).parent / ".onnx_cache.pkl" + if cache.exists(): + return pickle.loads(cache.read_bytes()) + + from omnimalloc.benchmark.converters.model import model_to_allocations + from omnimalloc.benchmark.converters.onnx import from_onnx + + hub = Path.home() / ".cache" / "huggingface" / "hub" + problems = {} + for model_name in ONNX_MODELS: + paths = list(hub.glob(f"models--onnxmodelzoo--{model_name}/snapshots/*/*.onnx")) + if not paths: + continue + model = from_onnx(paths[0]) + allocs = model_to_allocations(model) + problems[f"onnx_{model_name.split('_Opset')[0]}"] = [ + (a.size, a.start, a.end) for a in allocs + ] + cache.write_bytes(pickle.dumps(problems)) + return problems + + +def greedy_by_all(allocs): + entities = tuple( + Allocation(id=i, size=s, start=b, end=e) for i, (s, b, e) in enumerate(allocs) + ) + best = None + for variant in GREEDY_VARIANTS: + placed = variant().allocate(entities) + peak = max(a.offset + a.size for a in placed if a.offset is not None) + if best is None or peak < best: + best = peak + return best + + +def main(): + names = sys.argv[1:] or list(CANDIDATES) + problems = load_problems() + ratios_lb = {n: [] for n in names} + ratios_gba = {n: [] for n in names} + times = dict.fromkeys(names, 0.0) + + header = f"{'problem':<18}{'N':>6}{'LB':>12}{'gba/LB':>8}{'gba_t':>8}" + for n in names: + header += f"{n:>{max(len(n) + 2, 9)}}" + print(header) + + gba_time_total = 0.0 + for pname, allocs in problems.items(): + lb = lower_bound(allocs) + t0 = time.perf_counter() + gba = greedy_by_all(allocs) + gba_t = time.perf_counter() - t0 + gba_time_total += gba_t + row = f"{pname:<18}{len(allocs):>6}{lb:>12}{gba / lb:>8.4f}{gba_t:>8.2f}" + for n in names: + t0 = time.perf_counter() + offsets = CANDIDATES[n](allocs) + times[n] += time.perf_counter() - t0 + assert validate(allocs, offsets), f"{n} invalid on {pname}" + peak = peak_of(allocs, offsets) + ratios_lb[n].append(peak / lb) + ratios_gba[n].append(peak / gba) + row += f"{peak / gba:>{max(len(n) + 2, 9)}.4f}" + print(row) + + print("\n--- geomean peak/greedy_by_all (lower is better, <1 beats baseline) ---") + for n in names: + gm_gba = math.exp(sum(map(math.log, ratios_gba[n])) / len(ratios_gba[n])) + gm_lb = math.exp(sum(map(math.log, ratios_lb[n])) / len(ratios_lb[n])) + wins = sum(r < 0.9999 for r in ratios_gba[n]) + ties = sum(abs(r - 1) < 0.0001 for r in ratios_gba[n]) + print( + f"{n:<22} vs_gba={gm_gba:.4f} vs_lb={gm_lb:.4f} " + f"wins={wins} ties={ties} losses={len(ratios_gba[n]) - wins - ties} " + f"total_t={times[n]:.2f}s" + ) + print(f"greedy_by_all total_t={gba_time_total:.2f}s") + + +if __name__ == "__main__": + main() diff --git a/research/candidates.py b/research/candidates.py new file mode 100644 index 0000000..6124508 --- /dev/null +++ b/research/candidates.py @@ -0,0 +1,644 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# +"""Prototype O(N log N) allocator candidates. + +All candidates take a list of (size, start, end) tuples and return a list of +offsets aligned with the input order. Pure Python, optimized for clarity. +""" + +import heapq +import random +from bisect import bisect_left, bisect_right + +# --------------------------------------------------------------------------- +# Building blocks +# --------------------------------------------------------------------------- + + +class SegTree: + """Segment tree over time slots: range max query + range assign update. + + Range assign is only valid when the assigned value is >= the current max + over the range (true for skyline placement: offset+size >= watermark). + """ + + def __init__(self, n: int) -> None: + self.n = n + self.max = [0] * (4 * n) + self.lazy = [None] * (4 * n) + + def _push(self, node: int) -> None: + val = self.lazy[node] + if val is not None: + for child in (2 * node, 2 * node + 1): + self.max[child] = val + self.lazy[child] = val + self.lazy[node] = None + + def query( + self, lo: int, hi: int, node: int = 1, nlo: int = 0, nhi: int | None = None + ) -> int: + if nhi is None: + nhi = self.n + if hi <= nlo or nhi <= lo: + return 0 + if lo <= nlo and nhi <= hi: + return self.max[node] + self._push(node) + mid = (nlo + nhi) // 2 + return max( + self.query(lo, hi, 2 * node, nlo, mid), + self.query(lo, hi, 2 * node + 1, mid, nhi), + ) + + def assign( + self, + lo: int, + hi: int, + val: int, + node: int = 1, + nlo: int = 0, + nhi: int | None = None, + ) -> None: + if nhi is None: + nhi = self.n + if hi <= nlo or nhi <= lo: + return + if lo <= nlo and nhi <= hi: + self.max[node] = val + self.lazy[node] = val + return + self._push(node) + mid = (nlo + nhi) // 2 + self.assign(lo, hi, val, 2 * node, nlo, mid) + self.assign(lo, hi, val, 2 * node + 1, mid, nhi) + self.max[node] = max(self.max[2 * node], self.max[2 * node + 1]) + + +def compress_times(allocs): + """Map starts/ends to slot indices. Returns (slot_starts, slot_ends, n_slots).""" + times = sorted({t for size, start, end in allocs for t in (start, end)}) + index = {t: i for i, t in enumerate(times)} + starts = [index[a[1]] for a in allocs] + ends = [index[a[2]] for a in allocs] + return starts, ends, len(times) + + +def skyline_place(allocs, order): + """Place allocations in the given index order at the top of the skyline.""" + starts, ends, n_slots = compress_times(allocs) + tree = SegTree(max(n_slots, 1)) + offsets = [0] * len(allocs) + for i in order: + offset = tree.query(starts[i], ends[i]) + offsets[i] = offset + tree.assign(starts[i], ends[i], offset + allocs[i][0]) + return offsets + + +def compute_conflicts(allocs): + """Number of temporally overlapping allocations, via sweep. O(N log N).""" + starts = sorted(a[1] for a in allocs) + ends = sorted(a[2] for a in allocs) + return [ + bisect_left(starts, end) - bisect_right(ends, start) - 1 + for size, start, end in allocs + ] + + +def lower_bound(allocs): + """Max over time of the sum of active sizes.""" + events = [] + for size, start, end in allocs: + events.append((start, size)) + events.append((end, -size)) + events.sort() + peak = active = 0 + for _, delta in events: + active += delta + peak = max(peak, active) + return peak + + +def peak_of(allocs, offsets): + return max((o + a[0] for a, o in zip(allocs, offsets, strict=False)), default=0) + + +# --------------------------------------------------------------------------- +# Orders +# --------------------------------------------------------------------------- + + +def order_indices(allocs, key, reverse=True): + return sorted(range(len(allocs)), key=key, reverse=reverse) + + +ORDERS = { + "size": lambda allocs: order_indices(allocs, lambda i: allocs[i][0]), + "duration": lambda allocs: order_indices( + allocs, lambda i: allocs[i][2] - allocs[i][1] + ), + "area": lambda allocs: order_indices( + allocs, lambda i: allocs[i][0] * (allocs[i][2] - allocs[i][1]) + ), + "start": lambda allocs: sorted( + range(len(allocs)), key=lambda i: (allocs[i][1], -allocs[i][0]) + ), +} + + +def conflict_orders(allocs): + conflicts = compute_conflicts(allocs) + return { + "conflict": order_indices(allocs, lambda i: (conflicts[i], allocs[i][0])), + "conflict_size": order_indices( + allocs, lambda i: (conflicts[i] * allocs[i][0], allocs[i][0]) + ), + } + + +def all_orders(allocs): + orders = {name: fn(allocs) for name, fn in ORDERS.items()} + orders.update(conflict_orders(allocs)) + orders["input"] = list(range(len(allocs))) + return orders + + +# --------------------------------------------------------------------------- +# Compaction (monotone improvement pass) +# --------------------------------------------------------------------------- + + +def compact(allocs, offsets, max_passes=10): + """Re-place in ascending-offset order; peak never increases. Iterate.""" + for _ in range(max_passes): + order = sorted(range(len(allocs)), key=lambda i: (offsets[i], allocs[i][1])) + new_offsets = skyline_place(allocs, order) + if new_offsets == offsets: + break + offsets = new_offsets + return offsets + + +# --------------------------------------------------------------------------- +# Candidate allocators +# --------------------------------------------------------------------------- + + +def skyline_by(name): + def place(allocs): + order = all_orders(allocs)[name] + return skyline_place(allocs, order) + + place.__name__ = f"skyline_{name}" + return place + + +def skyline_by_all(allocs): + best = None + for order in all_orders(allocs).values(): + offsets = skyline_place(allocs, order) + if best is None or peak_of(allocs, offsets) < peak_of(allocs, best): + best = offsets + return best + + +def skyline_by_all_compact(allocs): + best = None + for order in all_orders(allocs).values(): + offsets = compact(allocs, skyline_place(allocs, order)) + if best is None or peak_of(allocs, offsets) < peak_of(allocs, best): + best = offsets + return best + + +def sweep_fit(allocs, best_fit=False): + """Chronological sweep with an address-ordered free list (true gap reuse). + + At each start event (larger sizes first on ties) allocate from the free + list (first-fit or best-fit); at each end event free and coalesce. + """ + n = len(allocs) + events = [] + for i, (size, start, end) in enumerate(allocs): + events.append((start, 0, -size, i)) + events.append((end, -1, 0, i)) + # Ends processed before starts at the same time (half-open intervals). + events.sort() + + INF = float("inf") + free_offsets = [0] # address-ordered gap start offsets + free_sizes = [INF] # parallel gap lengths; last gap is unbounded + offsets = [0] * n + + for _, kind, _, i in events: + size = allocs[i][0] + if kind == -1: # free event + offset = offsets[i] + pos = bisect_left(free_offsets, offset) + # Coalesce with successor gap + if pos < len(free_offsets) and free_offsets[pos] == offset + size: + size += free_sizes[pos] + del free_offsets[pos], free_sizes[pos] + # Coalesce with predecessor gap + if pos > 0 and free_offsets[pos - 1] + free_sizes[pos - 1] == offset: + free_sizes[pos - 1] += size + else: + free_offsets.insert(pos, offset) + free_sizes.insert(pos, size) + else: # allocation event + if best_fit: + pos = min( + (p for p in range(len(free_offsets)) if free_sizes[p] >= size), + key=lambda p: free_sizes[p], + ) + else: + pos = next(p for p in range(len(free_offsets)) if free_sizes[p] >= size) + offsets[i] = free_offsets[pos] + if free_sizes[pos] == size: + del free_offsets[pos], free_sizes[pos] + else: + free_offsets[pos] += size + free_sizes[pos] -= size + return offsets + + +def sweep_first_fit(allocs): + return sweep_fit(allocs, best_fit=False) + + +def sweep_best_fit(allocs): + return sweep_fit(allocs, best_fit=True) + + +def validate(allocs, offsets): + """Check no temporal+spatial overlap via sweep over active offset intervals.""" + events = [] + for i, (size, start, end) in enumerate(allocs): + events.append((end, 0, i)) + events.append((start, 1, i)) + events.sort() + active_offsets = [] # sorted offsets of active allocations + active_ids = [] + for _, is_start, i in events: + offset, size = offsets[i], allocs[i][0] + if is_start: + pos = bisect_left(active_offsets, offset) + if pos > 0: + j = active_ids[pos - 1] + if offsets[j] + allocs[j][0] > offset: + return False + if pos < len(active_offsets) and offset + size > active_offsets[pos]: + return False + active_offsets.insert(pos, offset) + active_ids.insert(pos, i) + else: + pos = bisect_left(active_offsets, offset) + while active_ids[pos] != i: + pos += 1 + del active_offsets[pos], active_ids[pos] + return True + + +# --------------------------------------------------------------------------- +# Two-ended sweep: large sizes from the bottom, small from the top of a band +# --------------------------------------------------------------------------- + + +def sweep_two_ended(allocs, quantile=0.5): + """Chronological sweep; small allocs are placed best-fit, large first-fit.""" + sizes = sorted(a[0] for a in allocs) + threshold = sizes[int(len(sizes) * quantile)] if sizes else 0 + + n = len(allocs) + events = [] + for i, (size, start, end) in enumerate(allocs): + events.append((start, 0, -size, i)) + events.append((end, -1, 0, i)) + events.sort() + + INF = float("inf") + free_offsets = [0] + free_sizes = [INF] + offsets = [0] * n + + for _, kind, _, i in events: + size = allocs[i][0] + if kind == -1: + offset = offsets[i] + pos = bisect_left(free_offsets, offset) + if pos < len(free_offsets) and free_offsets[pos] == offset + size: + size += free_sizes[pos] + del free_offsets[pos], free_sizes[pos] + if pos > 0 and free_offsets[pos - 1] + free_sizes[pos - 1] == offset: + free_sizes[pos - 1] += size + else: + free_offsets.insert(pos, offset) + free_sizes.insert(pos, size) + else: + fits = [p for p in range(len(free_offsets)) if free_sizes[p] >= size] + if allocs[i][0] >= threshold: + pos = fits[0] + else: + pos = min(fits, key=lambda p: free_sizes[p]) + offsets[i] = free_offsets[pos] + if free_sizes[pos] == size: + del free_offsets[pos], free_sizes[pos] + else: + free_offsets[pos] += size + free_sizes[pos] -= size + return offsets + + +# --------------------------------------------------------------------------- +# Hybrid: exact first-fit on the K largest allocs, sweep the rest around them +# --------------------------------------------------------------------------- + + +def first_fit_exact(allocs, order, best_fit=False): + """Quadratic reference first-fit (used only on small K subsets). + + With best_fit=True, picks the smallest fitting gap instead of the lowest + (Pisarchyk & Lee 2020, Algorithm 3), falling back to the watermark top. + """ + placed = [] # (offset, top, start, end) + offsets = [0] * len(allocs) + for i in order: + size, start, end = allocs[i] + blockers = sorted( + (p for p in placed if p[2] < end and start < p[3]), + key=lambda p: p[0], + ) + cursor = 0 + best = None # (gap_size, offset) + for b_off, b_top, _, _ in blockers: + gap = b_off - cursor + if gap >= size: + if not best_fit: + best = (gap, cursor) + break + if best is None or gap < best[0]: + best = (gap, cursor) + cursor = max(cursor, b_top) + if best is None: + best = (0, cursor) + offsets[i] = best[1] + placed.append((best[1], best[1] + size, start, end)) + return offsets + + +def hybrid_obstacles( + allocs, k=None, obstacle_order="conflict_size", select="size", best_fit=False +): + """Top-K (by size or area) placed exactly first; rest swept around them.""" + n = len(allocs) + if k is None: + k = min(n, max(16, int(4 * n**0.5))) + select_key = { + "size": lambda i: allocs[i][0], + "area": lambda i: allocs[i][0] * (allocs[i][2] - allocs[i][1]), + }[select] + # Keep selected ids in input order so ordering tie-breaks match the + # stable sorts of the exact greedy variants. + big_ids = sorted(sorted(range(n), key=select_key, reverse=True)[:k]) + big = [allocs[i] for i in big_ids] + big_offsets = first_fit_exact(big, all_orders(big)[obstacle_order], best_fit) + return _sweep_around(allocs, big_ids, big, big_offsets) + + +def _sweep_around(allocs, big_ids_list, big, big_offsets): + """Sweep-place all non-obstacle allocs around fixed obstacle placements.""" + n = len(allocs) + big_ids = set(big_ids_list) + obstacles = [ + (off, off + a[0], a[1], a[2]) for a, off in zip(big, big_offsets, strict=False) + ] + + offsets = [0] * n + for i, off in zip(big_ids_list, big_offsets, strict=False): + offsets[i] = off + + small_ids = [i for i in range(n) if i not in big_ids] + events = [] + for i in small_ids: + size, start, end = allocs[i] + events.append((start, 0, -size, i)) + events.append((end, -1, 0, i)) + events.sort() + + INF = float("inf") + free_offsets = [0] + free_sizes = [INF] + + for _, kind, _, i in events: + size, start, end = allocs[i] + if kind == -1: + offset = offsets[i] + pos = bisect_left(free_offsets, offset) + if pos < len(free_offsets) and free_offsets[pos] == offset + size: + size += free_sizes[pos] + del free_offsets[pos], free_sizes[pos] + if pos > 0 and free_offsets[pos - 1] + free_sizes[pos - 1] == offset: + free_sizes[pos - 1] += size + else: + free_offsets.insert(pos, offset) + free_sizes.insert(pos, size) + else: + # Forbidden offset bands: obstacles overlapping [start, end) + forbidden = sorted( + (o[0], o[1]) for o in obstacles if o[2] < end and start < o[3] + ) + merged = [] + for lo, hi in forbidden: + if merged and lo <= merged[-1][1]: + merged[-1] = (merged[-1][0], max(merged[-1][1], hi)) + else: + merged.append((lo, hi)) + offsets[i] = _first_fit_with_forbidden( + free_offsets, free_sizes, merged, size + ) + return offsets + + +def _first_fit_with_forbidden(free_offsets, free_sizes, forbidden, size): + """Lowest offset inside a free-list gap avoiding forbidden bands; splices.""" + f_idx = 0 + for pos in range(len(free_offsets)): + gap_lo = free_offsets[pos] + gap_hi = gap_lo + free_sizes[pos] + candidate = gap_lo + while f_idx > 0 and forbidden[f_idx - 1][1] > candidate: + f_idx -= 1 + while True: + while f_idx < len(forbidden) and forbidden[f_idx][1] <= candidate: + f_idx += 1 + if f_idx < len(forbidden) and forbidden[f_idx][0] < candidate + size: + candidate = forbidden[f_idx][1] + if candidate + size > gap_hi: + break + continue + if candidate + size <= gap_hi: + # Splice [candidate, candidate+size) out of this gap + del free_offsets[pos], free_sizes[pos] + if candidate > gap_lo: + free_offsets.insert(pos, gap_lo) + free_sizes.insert(pos, candidate - gap_lo) + pos += 1 + if gap_hi > candidate + size: + free_offsets.insert(pos, candidate + size) + free_sizes.insert(pos, gap_hi - (candidate + size)) + return candidate + break + raise AssertionError("Unbounded top gap must always fit") + + +def hybrid_size(allocs): + return hybrid_obstacles(allocs, obstacle_order="size") + + +def hybrid_conflict_size(allocs): + return hybrid_obstacles(allocs, obstacle_order="conflict_size") + + +def first_fit_exact_noisy(allocs, rng, sigma=0.25): + """Exact first-fit with a log-normally perturbed size ordering.""" + order = sorted( + range(len(allocs)), + key=lambda i: allocs[i][0] * rng.lognormvariate(0, sigma), + reverse=True, + ) + return first_fit_exact(allocs, order) + + +def noisy_restarts(allocs, restarts=8, k=512, seed=0): + """Best of hybrid runs with perturbed obstacle orderings.""" + rng = random.Random(seed) + n = len(allocs) + k = min(n, k) + select_key = lambda i: allocs[i][0] # noqa: E731 + by_size = sorted(range(n), key=select_key, reverse=True) + big = [allocs[i] for i in by_size[:k]] + + best = None + for _ in range(restarts): + big_offsets = first_fit_exact_noisy(big, rng) + offsets = _sweep_around(allocs, by_size[:k], big, big_offsets) + if best is None or peak_of(allocs, offsets) < peak_of(allocs, best): + best = offsets + return best + + +def sweep_ff_compact(allocs): + return compact(allocs, sweep_first_fit(allocs)) + + +def sweep_bf_compact(allocs): + return compact(allocs, sweep_best_fit(allocs)) + + +# --------------------------------------------------------------------------- +# Boxing-lite: power-of-2 size classes, optimal channel assignment per class +# --------------------------------------------------------------------------- + + +def boxed_channels(allocs): + """Round sizes to powers of 2; per class, greedy channel assignment + (optimal interval-graph coloring); stack class bands; caller compacts. + """ + n = len(allocs) + groups = {} + for i, (size, start, end) in enumerate(allocs): + groups.setdefault((size - 1).bit_length(), []).append(i) + + offsets = [0] * n + base = 0 + for c in sorted(groups, reverse=True): + width = 1 << c + expiry = [] # (end, channel) of live allocs + free = [] # released channel numbers + max_ch = 0 + for i in sorted(groups[c], key=lambda i: allocs[i][1]): + start = allocs[i][1] + while expiry and expiry[0][0] <= start: + heapq.heappush(free, heapq.heappop(expiry)[1]) + if free: + ch = heapq.heappop(free) + else: + ch = max_ch + max_ch += 1 + offsets[i] = base + ch * width + heapq.heappush(expiry, (allocs[i][2], ch)) + base += max_ch * width + return offsets + + +def boxed(allocs): + return boxed_channels(allocs) + + +def boxed_compact(allocs): + return compact(allocs, boxed_channels(allocs)) + + +def fast_by_all(allocs): + candidates = ( + sweep_first_fit, + sweep_best_fit, + sweep_two_ended, + hybrid_size, + hybrid_conflict_size, + ) + best = None + for fn in candidates: + offsets = compact(allocs, fn(allocs)) + if best is None or peak_of(allocs, offsets) < peak_of(allocs, best): + best = offsets + return best + + +def fast_by_all_v2(allocs, k=1024): + members = [ + lambda a: sweep_first_fit(a), + lambda a: sweep_best_fit(a), + lambda a: sweep_two_ended(a), + lambda a: hybrid_obstacles(a, k=k, obstacle_order="size"), + lambda a: hybrid_obstacles(a, k=k, obstacle_order="conflict_size"), + lambda a: hybrid_obstacles(a, k=k, obstacle_order="duration"), + lambda a: hybrid_obstacles(a, k=k, obstacle_order="conflict"), + lambda a: hybrid_obstacles(a, k=k, obstacle_order="area"), + lambda a: hybrid_obstacles(a, k=k, obstacle_order="area", select="area"), + lambda a: hybrid_obstacles(a, k=k, obstacle_order="size", best_fit=True), + lambda a: hybrid_obstacles( + a, k=k, obstacle_order="conflict_size", best_fit=True + ), + ] + best = None + for fn in members: + offsets = fn(allocs) + if best is None or peak_of(allocs, offsets) < peak_of(allocs, best): + best = offsets + return best + + +CANDIDATES = { + "skyline_size": skyline_by("size"), + "skyline_area": skyline_by("area"), + "skyline_duration": skyline_by("duration"), + "skyline_start": skyline_by("start"), + "skyline_input": skyline_by("input"), + "skyline_all": skyline_by_all, + "skyline_all_compact": skyline_by_all_compact, + "sweep_first_fit": sweep_first_fit, + "sweep_best_fit": sweep_best_fit, + "sweep_ff_compact": sweep_ff_compact, + "sweep_bf_compact": sweep_bf_compact, + "sweep_two_ended": sweep_two_ended, + "hybrid_size": hybrid_size, + "hybrid_conflict_size": hybrid_conflict_size, + "fast_by_all": fast_by_all, + "fast_by_all_v2": fast_by_all_v2, + "noisy_restarts": noisy_restarts, + "boxed": boxed, + "boxed_compact": boxed_compact, +} diff --git a/research/fast-allocators.md b/research/fast-allocators.md new file mode 100644 index 0000000..1d71363 --- /dev/null +++ b/research/fast-allocators.md @@ -0,0 +1,224 @@ +# Fast (O(N log N)) allocators — research notes + +Goal: allocators with O(N) or O(N log N) runtime that match or beat +`GreedyByAll` (best of 7 sort orders + first-fit, currently O(N²)–O(N² log N)) +in peak memory. Simplicity is king. + +## Why the current greedy is quadratic + +- `compute_temporal_overlaps` (src/cpp/allocators/greedy_base.cpp) materializes + every temporally-overlapping *pair* → O(N²) time/space when lifetimes overlap + densely (the common case in NN workloads). +- `find_best_offset` scans + sorts all placed overlapping allocations per + placement → O(N² log N) total worst case. + +## Key insight for O(N log N) + +First-fit needs 2D (time × offset) gap knowledge. But *skyline placement* +(place each allocation at the max of the per-time watermark over its lifetime, +then raise the watermark over that interval to offset+size) only needs a +segment tree over compressed time with range-max query + range-assign update: +O(log N) per allocation. The assign is valid because the placement offset is +by definition ≥ the watermark everywhere in the interval. + +Cost vs first-fit: skyline cannot tuck allocations into holes *below* the +watermark. Hypothesis: with good orderings + compaction passes the gap is +small. + +**Compaction pass (monotone, O(N log N))**: given any valid solution, re-run +skyline placement with allocations ordered by ascending current offset. By +induction each allocation's new offset ≤ old offset, so peak never increases. +Iterate a few times / until fixpoint. Cheap post-improvement for *any* +candidate. + +## Idea list + +| # | Idea | Complexity | Status | +|---|------|-----------|--------| +| 1 | Skyline placement, size-desc order (segment tree) | O(N log N) | todo | +| 2 | Skyline with other orders (area, duration, conflict, start) | O(N log N) | todo | +| 3 | Compaction passes (skyline re-place in ascending-offset order), iterate | O(kN log N) | todo | +| 4 | Time-sweep first-fit: chronological events, address-ordered free list (true gap reuse) | ~O(N log N) | todo | +| 5 | Time-sweep best-fit variant | ~O(N log N) | todo | +| 6 | Conflict order computed via sweep (O(N log N)) feeding skyline | O(N log N) | todo | +| 7 | Ensemble "fast by all": best of 1/2/4/5 each + compaction | O(N log N) | todo | +| 8 | Randomized restarts: perturb sort key (size × U[a,b]), keep best | O(kN log N) | todo | +| 9 | Gergov 3-approx (lit review pending — implementability?) | O(N log N)? | todo | +| 10 | Buchsbaum boxing / size-class rounding (idealloc-style) | O(N log N) | todo | +| 11 | Best-of-two-offsets skyline: try top-of-watermark vs known hole list | ? | todo | +| 12 | True first-fit in size order via offset-bucketed structures | O(N log² N)? | probably too complex | + +## Benchmark setup + +- Problems: 13 minimalloc CSVs (`external/minimalloc/{challenging,small,examples}`), + synthetic sources from `omnimalloc.benchmark.sources` (random, high-contention, + power-of-2, sequential) at N = 100 / 1000 / 5000. +- Baseline: `GreedyByAll` (C++ variants, sequential) peak. +- Lower bound: max over time of Σ active sizes (sweep) — normalizes quality. +- Quality metric: peak / lower_bound, and peak / greedy_by_all_peak. +- Runtime scaling: candidates timed at N = 1k / 10k / 50k (Python protos; + algorithmic scaling is what matters, final impl in C++). + +## Results + +Metric: geomean(peak / greedy_by_all_peak) over 25 problems (13 minimalloc, +6 synthetic, small examples). <1.0 beats the quadratic baseline. + +| Candidate | vs gba | Notes | +|---|---|---| +| skyline_size | 1.44 | no hole reuse → bad | +| skyline_all (6 orders) | 1.38 | still bad; compaction provably no-op on skyline | +| sweep_first_fit | 1.13 | chronological + address-ordered free list; *ties Start-order greedy exactly* (F/G/H) | +| sweep_best_fit | 1.11 | | +| sweep_two_ended | 1.13 | large→first-fit, small→best-fit | +| hybrid_size (K=4√N) | 1.11 | top-K exact first-fit + sweep-around-obstacles | +| hybrid_conflict_size | 1.09 | | +| fast_by_all (5 members + compact) | **1.02** | 3 wins / 10 ties / 12 losses | + +Findings: + +1. **Gap reuse is essential.** Skyline (watermark-only) placement loses ~40%. + All good candidates use a real free list. +2. **Which greedy variant wins**: Size/ConflictSize on most minimalloc + instances and randoms; Start on F/G/H (sweep_ff ties those exactly — + it *is* greedy-by-start with true chronological gap reuse); Duration on D; + Conflict on high-contention. +3. **Compaction (re-place ascending-offset via skyline) never helps sweeps** — + provable: chronological first-fit placements are temporally blocked from + below at placement time and stay blocked. It occasionally helps hybrids. +4. **K-sensitivity of hybrid is noisy, not monotone.** Going K=64→∞ moves + geomean only 1.065→1.032 (size order). Ensemble diversity beats exactness: + single exact orders at K=∞ still lose to gba (best-of-7 effect). + +## Round 2 results + +| Candidate | vs gba | Notes | +|---|---|---| +| noisy_restarts (8× perturbed order) | 1.02 | good anytime knob, wins on a few | +| boxed (pow2 classes + optimal channels) | 1.86 | rounding waste kills it | +| boxed_compact | 1.21 | compaction recovers a lot, still bad | +| hybrid K sweep 64→∞ | 1.065→1.032 | noisy, not monotone; ensemble > exactness | +| **fast_by_all_v2 (9 members, K=1024)** | **0.9955** | **2 wins / 22 ties / 1 loss (rand_5000 +1.7%)** | + +fast_by_all_v2 members: sweep_ff, sweep_bf, sweep_two_ended, hybrid with +obstacle orders {size, conflict_size, conflict, duration, area, +area-selected+area-ordered}, best-of. Tie-break fix (keep obstacle set in +input order so stable sorts match the exact greedy variants) fixed mm_c_D. + +More findings: + +5. **Compaction contributes nothing to the ensemble** (helps individual weak + members, never the final best). Dropped → no segment tree needed at all. +6. Scaling (Python protos vs C++ quadratic greedy, single variant): + N=20k: sweep_ff 0.04 s, hybrid(K=512) 0.43 s, C++ greedy_by_size 5.9 s. + N=50k: sweep_ff 0.11 s, hybrid 1.1 s (C++ quadratic not feasible). +7. Whole 25-problem suite: v2 in pure Python 3.5 s vs C++ greedy_by_all 3.6 s. + +## Winning recipe (to be ported to C++) + +1. `sweep(fit)`: sort 2N events by (time, end-before-start, size-desc); + maintain address-ordered free list (map offset→len) with coalescing; + place at first-fit / best-fit / two-ended (large→ff, small→bf) gap. +2. `hybrid(order, select, K)`: top-K by size (or area), exact quadratic + first-fit on those K only (K fixed, default 1024 → O(1) asymptotically), + then sweep the rest treating the K as fixed obstacles (forbidden offset + bands during temporally-overlapping placements). +3. `sweep_by_all`: best of the 9 members above. + +Complexity: O(N log N + N·m + K²), m = avg temporally-overlapping obstacles +per small alloc (≤ K, typically tiny). K fixed → O(N log N). + +## Final implementation (shipped) + +`src/cpp/allocators/sweep.{hpp,cpp}` + bindings, Python wrappers in +`allocators/sweep.py`, tests in `tests/unit/allocators/test_sweep.py`. + +Allocators (all registered): + +- `SweepAllocator` / `SweepBestFitAllocator` / `SweepTwoEndedAllocator` — + chronological sweep, address-ordered coalescing free list (C++ `std::map`). +- `HybridSweepAllocator` (conflict_size order) + `BySize` / `ByDuration` / + `ByArea` variants — exact first-fit on the top-1024 obstacles + (`max_obstacles` configurable), sweep-around-obstacles for the rest. + Ordering/selection stays in Python (repo pattern), placement in C++ + (`sweep_place` / `hybrid_sweep_place`). +- `SweepByAllAllocator` — best of the 7, via `allocate_parallel`. + +Leave-one-out pruning dropped the plain-area and pure-conflict hybrid members +(zero ensemble contribution); 7 members mirror GreedyByAll's portfolio size. + +### Final numbers (33 problems: 13 minimalloc, 6 synthetic, 8 ONNX models, 6 small) + +- C++ ↔ prototype parity: **0 mismatches** across 33 × 7 runs; all valid. +- Quality: **sweep_by_all/greedy_by_all geomean 0.9967 — 2 wins, 29 ties, + 2 losses (worst +1.7%)**, i.e. matches or beats the quadratic baseline. +- Suite runtime: sweep_by_all 1.0 s sequential vs greedy_by_all 6.6 s. +- Scaling (RandomSource): + +| N | sweep | hybrid | sweep_by_all | greedy_by_size (quadratic) | +|---|---|---|---|---| +| 10k | 0.01 s | 0.03 s | 0.14 s | 1.55 s | +| 50k | 0.03 s | 0.10 s | 0.57 s | ~39 s (extrap.) | +| 200k | 0.13 s | 0.39 s | 2.38 s | — | +| 1M | 0.83 s | 2.58 s | 24.4 s | ~4.3 h (extrap.) | + +(sweep_by_all at 1M is dominated by pickling 8M allocations through the +process pool, not by the algorithms; single members stay near-linear.) + +### Ideas tried (scoreboard) + +1. Skyline/watermark seg-tree, 6 orders — 1.38–1.44, rejected (≡ TVM USMP rule) +2. Skyline ensemble + compaction — 1.38, rejected +3. Compaction passes (monotone re-place) — no ensemble contribution, dropped +4. Time-sweep first-fit — 1.11, **kept** +5. Time-sweep best-fit — 1.10, **kept** +6. Two-ended sweep (median split ff/bf) — 1.11, **kept** +7. Hybrid obstacles, 6 order/select combos — 1.05–1.17 solo, **4 kept** +8. Best-fit-among-gaps exact rule (TFLite Alg. 3) — no ensemble gain, dropped +9. Noisy-restart perturbation — 1.02, viable anytime knob, not shipped +10. Power-of-2 boxing + optimal channels (+compact) — 1.21, rejected +11. Gergov 3-approx — rejected on implementability (lit review) +12. True sub-quadratic first-fit-by-size — rejected on complexity; the + obstacle budget makes it unnecessary (exact ≡ greedy for N ≤ 1024) + +## Literature notes (from background research agent) + +Theory (all ratios vs LOAD = max-over-time sum of live sizes, our `lower_bound`): + +- DSA is NP-complete (Stockmeyer '76). Best implementable guarantee: + **Gergov 1999, 3·LOAD in O(n log n)** — but it is a 2-page abstract nobody + has ever implemented in production; research-grade reconstruction effort. + Gergov 1996 (5-approx, via "2-allocations") is better documented. +- **Buchsbaum et al. STOC'03**: boxing/size-class recursion gives (2+ε) always + and OPT+o(OPT) when h_max ≪ L; brutal constants (ε⁻⁶). **idealloc** + (arXiv:2504.04874, Rust, 2025) is the first real implementation — a + stochastic iterated box/unbox/first-fit ensemble; 0–2.6% fragmentation at + up to 567k buffers. Notably: its bootstrap/fallback is plain + size-desc-then-lifetime-desc first-fit ("big rocks first"), and even its + first-fit is O(N·degree) via interference lists — *not* N log N. +- No o(OPT) additive result possible in general; gap governed by h_max/L. +- idealloc's evaluation: **no single ordering dominates across benchmarks** — + independent confirmation of the portfolio/best-of-k approach. + +Production planners: + +- TFLite (Pisarchyk & Lee 2020): greedy-by-size, **best-fit-among-gaps** + placement; hits the LB on 5/6 nets. O(N²)-ish. (Tried the best-fit-gaps + rule in our hybrids: no ensemble gain on our suite.) +- TFLM GreedyMemoryPlanner: O(N·live-set) via purging dead records in a + chronological sweep — industry's "quadratic but small live set" answer. +- TVM USMP greedy = pure watermark rule (no gap reuse) — **exactly equivalent + to our skyline segment-tree placement** (which is its O(N log N) impl). +- XLA heap simulator: decreasing-size best-fit with an interval tree over + placed chunks; O(N·overlaps) worst case. +- MiniMalloc's "canonical solutions" justify offset-ordered first-fit + placements as a dominance class containing an optimum. + +Address-ordered first-fit (our sweep): + +- Brent (TOPLAS '89): AO first-fit findable in O(log) per op → chronological + sweep is genuinely O(N log N). +- Wilson/Johnstone ('95/'98): AO-first-fit and best-fit ≈ 1% fragmentation on + real traces; policy (address order + immediate coalescing) is what matters. +- Worst case (online): Θ(min{log h_max, log χ})-competitive (Luby–Naor–Orda + '94) — no constant guarantee, empirically excellent. Matches our data. diff --git a/src/cpp/allocators/sweep.cpp b/src/cpp/allocators/sweep.cpp new file mode 100644 index 0000000..bde65b7 --- /dev/null +++ b/src/cpp/allocators/sweep.cpp @@ -0,0 +1,251 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#include "sweep.hpp" + +#include +#include +#include +#include +#include +#include + +#include "greedy_base.hpp" + +namespace omnimalloc { + +namespace { + +constexpr int64_t kUnbounded = std::numeric_limits::max() / 4; + +// Every placement offset is bounded by the live footprint, so capping the +// total size at kUnbounded guarantees the top free gap fits every request +// and no offset arithmetic overflows. +void check_total_size(const std::vector& allocations) { + int64_t total = 0; + for (const auto& alloc : allocations) { + if (alloc.size() > kUnbounded - total) { + throw std::overflow_error("Total allocation size exceeds sweep range"); + } + total += alloc.size(); + } +} + +// Obstacle occupying offsets [lo, hi) during times [start, end). +struct Band { + int64_t lo; + int64_t hi; + int64_t start; + int64_t end; + + auto operator<=>(const Band&) const = default; +}; + +// Address-ordered free list over offsets with immediate coalescing. The top +// gap is unbounded, so every request succeeds. +class FreeList { + public: + FreeList() { gaps_.emplace(0, kUnbounded); } + + // Lowest offset with `size` free space outside the merged, offset-sorted + // `forbidden` bands; carves the placement out of the free list. + int64_t take_first( + int64_t size, const std::vector>& forbidden) { + auto band = forbidden.begin(); + for (auto it = gaps_.begin(); it != gaps_.end(); ++it) { + const int64_t gap_lo = it->first; + const int64_t gap_hi = gap_lo + it->second; + int64_t candidate = gap_lo; + // A band that bumped the candidate past the previous gap may still + // cover this gap's start, so back up to the first band ending above it. + while (band != forbidden.begin() && std::prev(band)->second > candidate) { + --band; + } + while (true) { + while (band != forbidden.end() && band->second <= candidate) { + ++band; + } + if (band != forbidden.end() && band->first < candidate + size) { + candidate = band->second; + if (candidate + size > gap_hi) { + break; + } + continue; + } + if (candidate + size <= gap_hi) { + carve(it, candidate, size); + return candidate; + } + break; + } + } + return 0; // Unreachable: the unbounded top gap always fits. + } + + // Smallest gap that fits (ties: lowest offset); unaware of forbidden + // bands, so callers must fall back to take_first while bands are active. + int64_t take_best(int64_t size) { + auto best = gaps_.end(); + for (auto it = gaps_.begin(); it != gaps_.end(); ++it) { + if (it->second >= size && + (best == gaps_.end() || it->second < best->second)) { + best = it; + } + } + const int64_t offset = best->first; + carve(best, offset, size); + return offset; + } + + void release(int64_t offset, int64_t size) { + auto next = gaps_.lower_bound(offset); + if (next != gaps_.end() && next->first == offset + size) { + size += next->second; + next = gaps_.erase(next); + } + if (next != gaps_.begin()) { + auto prev = std::prev(next); + if (prev->first + prev->second == offset) { + prev->second += size; + return; + } + } + gaps_.emplace(offset, size); + } + + private: + void carve(std::map::iterator it, int64_t offset, + int64_t size) { + const int64_t gap_lo = it->first; + const int64_t gap_hi = gap_lo + it->second; + gaps_.erase(it); + if (offset > gap_lo) { + gaps_.emplace(gap_lo, offset - gap_lo); + } + if (gap_hi > offset + size) { + gaps_.emplace(offset + size, gap_hi - (offset + size)); + } + } + + std::map gaps_; // offset -> length +}; + +struct SweepEvent { + int64_t time; + bool is_start; // false sorts first: frees precede same-time placements + int64_t neg_size; + size_t idx; + + auto operator<=>(const SweepEvent&) const = default; +}; + +int64_t median_size(const std::vector& allocations, + const std::vector& indices) { + std::vector sizes; + sizes.reserve(indices.size()); + for (size_t idx : indices) { + sizes.push_back(allocations[idx].size()); + } + const auto mid = sizes.begin() + static_cast(sizes.size() / 2); + std::nth_element(sizes.begin(), mid, sizes.end()); + return *mid; +} + +// Sweep-place `indices` around the offset-sorted `obstacles`, writing offsets. +void sweep_indices(const std::vector& allocations, + const std::vector& indices, SweepFit fit, + const std::vector& obstacles, + std::vector& offsets) { + std::vector events; + events.reserve(indices.size() * 2); + for (size_t idx : indices) { + const auto& alloc = allocations[idx]; + events.push_back({alloc.start(), true, -alloc.size(), idx}); + events.push_back({alloc.end(), false, 0, idx}); + } + std::sort(events.begin(), events.end()); + + const int64_t threshold = fit == SweepFit::kTwoEnded && !indices.empty() + ? median_size(allocations, indices) + : 0; + + FreeList free_list; + std::vector> forbidden; + for (const auto& event : events) { + const auto& alloc = allocations[event.idx]; + if (!event.is_start) { + free_list.release(offsets[event.idx], alloc.size()); + continue; + } + forbidden.clear(); + for (const auto& band : obstacles) { + if (band.start < alloc.end() && alloc.start() < band.end) { + if (!forbidden.empty() && band.lo <= forbidden.back().second) { + forbidden.back().second = std::max(forbidden.back().second, band.hi); + } else { + forbidden.emplace_back(band.lo, band.hi); + } + } + } + const bool best_fit = + (fit == SweepFit::kBest || + (fit == SweepFit::kTwoEnded && alloc.size() < threshold)) && + forbidden.empty(); + offsets[event.idx] = best_fit + ? free_list.take_best(alloc.size()) + : free_list.take_first(alloc.size(), forbidden); + } +} + +std::vector with_offsets(const std::vector& allocations, + const std::vector& offsets) { + std::vector placed; + placed.reserve(allocations.size()); + for (size_t i = 0; i < allocations.size(); ++i) { + placed.push_back(allocations[i].with_offset(offsets[i])); + } + return placed; +} + +} // namespace + +std::vector sweep_place(const std::vector& allocations, + SweepFit fit) { + check_total_size(allocations); + std::vector indices(allocations.size()); + std::iota(indices.begin(), indices.end(), size_t{0}); + std::vector offsets(allocations.size(), 0); + sweep_indices(allocations, indices, fit, {}, offsets); + return with_offsets(allocations, offsets); +} + +std::vector hybrid_sweep_place( + const std::vector& allocations, size_t num_obstacles) { + check_total_size(allocations); + num_obstacles = std::min(num_obstacles, allocations.size()); + const std::vector prefix(allocations.begin(), + allocations.begin() + num_obstacles); + auto placed = first_fit_place(prefix, compute_temporal_overlaps(prefix)); + + std::vector obstacles; + obstacles.reserve(num_obstacles); + for (const auto& alloc : placed) { + obstacles.push_back({alloc.offset().value(), alloc.height().value(), + alloc.start(), alloc.end()}); + } + std::sort(obstacles.begin(), obstacles.end()); + + std::vector offsets(allocations.size(), 0); + std::vector rest(allocations.size() - num_obstacles); + std::iota(rest.begin(), rest.end(), num_obstacles); + sweep_indices(allocations, rest, SweepFit::kFirst, obstacles, offsets); + + placed.reserve(allocations.size()); + for (size_t i = num_obstacles; i < allocations.size(); ++i) { + placed.push_back(allocations[i].with_offset(offsets[i])); + } + return placed; +} + +} // namespace omnimalloc diff --git a/src/cpp/allocators/sweep.hpp b/src/cpp/allocators/sweep.hpp new file mode 100644 index 0000000..cbf5253 --- /dev/null +++ b/src/cpp/allocators/sweep.hpp @@ -0,0 +1,36 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include +#include +#include + +#include "primitives/allocation.hpp" + +namespace omnimalloc { + +// Placement policy for the chronological sweep allocator. +enum class SweepFit : std::uint8_t { + kFirst, // lowest-offset gap that fits (address-ordered first-fit) + kBest, // smallest gap that fits (ties broken by lowest offset) + kTwoEnded, // sizes >= median use first-fit, smaller ones best-fit +}; + +// Chronological sweep placement: allocation/free events are processed in time +// order while an address-ordered, coalescing free list tracks available +// offsets. Same-time events free before allocating; same-time allocations are +// placed largest first. O(N (log N + G)), G = concurrent free gaps. +[[nodiscard]] std::vector sweep_place( + const std::vector& allocations, SweepFit fit); + +// Portfolio placement: the first `num_obstacles` allocations (already in the +// desired placement order) are placed with exact first-fit; the remaining +// allocations are swept chronologically around them, treating the obstacle +// placements as forbidden offset bands while temporally overlapping. +[[nodiscard]] std::vector hybrid_sweep_place( + const std::vector& allocations, size_t num_obstacles); + +} // namespace omnimalloc diff --git a/src/cpp/bindings.cpp b/src/cpp/bindings.cpp index ef413aa..77c10d2 100644 --- a/src/cpp/bindings.cpp +++ b/src/cpp/bindings.cpp @@ -18,6 +18,7 @@ #include "allocators/greedy_base.hpp" #include "allocators/simulated_annealing.hpp" #include "allocators/supermalloc/partition.hpp" +#include "allocators/sweep.hpp" #include "allocators/tabu_search.hpp" #include "allocators/telamalloc.hpp" #include "primitives/allocation.hpp" @@ -106,6 +107,19 @@ NB_MODULE(_cpp, m) { m.def("first_fit_place", &first_fit_place, "allocations"_a, "overlaps"_a, nb::rv_policy::move); + // Sweep placement: chronological free-list allocators + nb::enum_(m, "SweepFit") + .value("FIRST", SweepFit::kFirst) + .value("BEST", SweepFit::kBest) + .value("TWO_ENDED", SweepFit::kTwoEnded); + + m.def("sweep_place", &sweep_place, "allocations"_a, + "fit"_a = SweepFit::kFirst, nb::call_guard(), + nb::rv_policy::move); + m.def("hybrid_sweep_place", &hybrid_sweep_place, "allocations"_a, + "num_obstacles"_a, nb::call_guard(), + nb::rv_policy::move); + // FirstFitPlacer class: resident placer for the order-search allocators nb::class_(m, "FirstFitPlacer") .def(nb::init>(), "allocations"_a) diff --git a/src/python/omnimalloc/_cpp.pyi b/src/python/omnimalloc/_cpp.pyi index e8c8fb3..41d6641 100644 --- a/src/python/omnimalloc/_cpp.pyi +++ b/src/python/omnimalloc/_cpp.pyi @@ -79,6 +79,17 @@ def compute_temporal_overlaps(allocations: Sequence[Allocation]) -> dict[int | s def first_fit_place(allocations: Sequence[Allocation], overlaps: Mapping[int | str, Set[int | str]]) -> list[Allocation]: ... +class SweepFit(enum.Enum): + FIRST = 0 + + BEST = 1 + + TWO_ENDED = 2 + +def sweep_place(allocations: Sequence[Allocation], fit: SweepFit = SweepFit.FIRST) -> list[Allocation]: ... + +def hybrid_sweep_place(allocations: Sequence[Allocation], num_obstacles: int) -> list[Allocation]: ... + class FirstFitPlacer: def __init__(self, allocations: Sequence[Allocation]) -> None: ... diff --git a/src/python/omnimalloc/allocators/__init__.py b/src/python/omnimalloc/allocators/__init__.py index 336f32e..993da13 100644 --- a/src/python/omnimalloc/allocators/__init__.py +++ b/src/python/omnimalloc/allocators/__init__.py @@ -34,6 +34,14 @@ from .simulated_annealing import SimulatedAnnealingConfig as SimulatedAnnealingConfig from .supermalloc import SupermallocAllocator as SupermallocAllocator from .supermalloc import SupermallocConfig as SupermallocConfig +from .sweep import HybridSweepAllocator as HybridSweepAllocator +from .sweep import HybridSweepByAreaAllocator as HybridSweepByAreaAllocator +from .sweep import HybridSweepByDurationAllocator as HybridSweepByDurationAllocator +from .sweep import HybridSweepBySizeAllocator as HybridSweepBySizeAllocator +from .sweep import SweepAllocator as SweepAllocator +from .sweep import SweepBestFitAllocator as SweepBestFitAllocator +from .sweep import SweepByAllAllocator as SweepByAllAllocator +from .sweep import SweepTwoEndedAllocator as SweepTwoEndedAllocator from .tabu_search import TabuSearchAllocator as TabuSearchAllocator from .tabu_search import TabuSearchConfig as TabuSearchConfig from .telamalloc import TelamallocAllocator as TelamallocAllocator diff --git a/src/python/omnimalloc/allocators/sweep.py b/src/python/omnimalloc/allocators/sweep.py new file mode 100644 index 0000000..079df53 --- /dev/null +++ b/src/python/omnimalloc/allocators/sweep.py @@ -0,0 +1,127 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +import heapq +from typing import ClassVar + +from omnimalloc._cpp import SweepFit, hybrid_sweep_place, sweep_place +from omnimalloc.primitives import Allocation + +from .base import BaseAllocator +from .greedy_base import ( + allocate_parallel, + order_by_area, + order_by_conflict_size, + order_by_duration, + order_by_size, +) + +# Obstacle budget for the hybrid sweep: the exact quadratic phase runs on at +# most this many allocations, keeping the overall runtime O(N log N). +DEFAULT_MAX_OBSTACLES = 1024 + + +class SweepAllocator(BaseAllocator): + """Chronological sweep with an address-ordered coalescing free list. + + Processes allocation/free events in time order, placing each allocation + into the lowest free gap that fits and reusing freed space. O(N log N), + suitable for very large problems where the greedy allocators are too slow. + """ + + _fit: ClassVar[SweepFit] = SweepFit.FIRST + + def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + return tuple(sweep_place(allocations, self._fit)) + + +class SweepBestFitAllocator(SweepAllocator): + """Sweep allocator placing into the smallest fitting gap.""" + + _fit = SweepFit.BEST + + +class SweepTwoEndedAllocator(SweepAllocator): + """Sweep allocator: above-median sizes use first-fit, smaller best-fit.""" + + _fit = SweepFit.TWO_ENDED + + +class HybridSweepAllocator(BaseAllocator): + """Exact first-fit for the largest allocations, sweep for the rest. + + The top max_obstacles allocations by _rank are placed with the exact + quadratic first-fit in _order; the remaining allocations are swept + chronologically around them, treating the fixed placements as forbidden + bands. O(N log N) with the default budget. + """ + + _order = staticmethod(order_by_conflict_size) + + def __init__(self, max_obstacles: int = DEFAULT_MAX_OBSTACLES) -> None: + if max_obstacles < 0: + raise ValueError(f"max_obstacles must be non-negative, got {max_obstacles}") + self._max_obstacles = max_obstacles + + @staticmethod + def _rank(allocation: Allocation) -> int: + return allocation.size + + def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + if not allocations: + return allocations + big_ids = set( + heapq.nlargest( + self._max_obstacles, + range(len(allocations)), + key=lambda i: self._rank(allocations[i]), + ) + ) + # Split in input order so the stable ordering tie-breaks match the + # exact greedy variants. + big = tuple(a for i, a in enumerate(allocations) if i in big_ids) + rest = tuple(a for i, a in enumerate(allocations) if i not in big_ids) + ordered = self._order(big) + rest + return tuple(hybrid_sweep_place(ordered, len(big))) + + +class HybridSweepBySizeAllocator(HybridSweepAllocator): + """Hybrid sweep ordering the exact phase by size (largest first).""" + + _order = staticmethod(order_by_size) + + +class HybridSweepByDurationAllocator(HybridSweepAllocator): + """Hybrid sweep ordering the exact phase by duration (longest first).""" + + _order = staticmethod(order_by_duration) + + +class HybridSweepByAreaAllocator(HybridSweepAllocator): + """Hybrid sweep selecting and ordering the exact phase by area.""" + + _order = staticmethod(order_by_area) + + @staticmethod + def _rank(allocation: Allocation) -> int: + return allocation.size * allocation.duration + + +class SweepByAllAllocator(BaseAllocator): + """Runs every sweep and hybrid sweep variant and keeps the best result.""" + + def __init__(self, cores: int | None = None) -> None: + self._cores = cores + + def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + variants: tuple[BaseAllocator, ...] = ( + SweepAllocator(), + SweepBestFitAllocator(), + SweepTwoEndedAllocator(), + HybridSweepAllocator(), + HybridSweepBySizeAllocator(), + HybridSweepByDurationAllocator(), + HybridSweepByAreaAllocator(), + ) + return allocate_parallel(variants, allocations, cores=self._cores) diff --git a/tests/unit/allocators/test_sweep.py b/tests/unit/allocators/test_sweep.py new file mode 100644 index 0000000..6ee621f --- /dev/null +++ b/tests/unit/allocators/test_sweep.py @@ -0,0 +1,175 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +import pytest +from omnimalloc.allocators.greedy import ( + GreedyByDurationAllocator, + GreedyBySizeAllocator, +) +from omnimalloc.allocators.greedy_base import peak_memory +from omnimalloc.allocators.sweep import ( + HybridSweepAllocator, + HybridSweepByAreaAllocator, + HybridSweepByDurationAllocator, + HybridSweepBySizeAllocator, + SweepAllocator, + SweepBestFitAllocator, + SweepByAllAllocator, + SweepTwoEndedAllocator, +) +from omnimalloc.benchmark.sources.generator import RandomSource +from omnimalloc.primitives import Allocation, Pool +from omnimalloc.validate import validate_allocation + +ALL_SWEEP_ALLOCATORS = ( + SweepAllocator, + SweepBestFitAllocator, + SweepTwoEndedAllocator, + HybridSweepAllocator, + HybridSweepBySizeAllocator, + HybridSweepByDurationAllocator, + HybridSweepByAreaAllocator, +) + + +def _is_valid(result: tuple[Allocation, ...]) -> bool: + return validate_allocation(Pool(id="test_pool", allocations=result)) + + +@pytest.mark.parametrize("allocator_cls", ALL_SWEEP_ALLOCATORS) +def test_sweep_empty(allocator_cls: type) -> None: + assert allocator_cls().allocate(()) == () + + +@pytest.mark.parametrize("allocator_cls", ALL_SWEEP_ALLOCATORS) +def test_sweep_single(allocator_cls: type) -> None: + result = allocator_cls().allocate((Allocation(id=1, size=100, start=0, end=10),)) + assert len(result) == 1 + assert result[0].offset == 0 + + +def test_sweep_no_temporal_overlap_shares_offset() -> None: + allocs = ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=200, start=10, end=20), + ) + by_id = {a.id: a for a in SweepAllocator().allocate(allocs)} + assert by_id[1].offset == 0 + assert by_id[2].offset == 0 + + +def test_sweep_all_overlap_stacks_sequentially() -> None: + allocs = tuple(Allocation(id=i, size=100, start=0, end=10) for i in range(5)) + result = SweepAllocator().allocate(allocs) + assert _is_valid(result) + assert peak_memory(result) == 500 + + +def test_sweep_reuses_freed_gap() -> None: + allocs = ( + Allocation(id="long", size=50, start=0, end=10), + Allocation(id="early", size=100, start=0, end=5), + Allocation(id="late", size=100, start=5, end=10), + ) + by_id = {a.id: a for a in SweepAllocator().allocate(allocs)} + assert by_id["early"].offset == 0 + assert by_id["long"].offset == 100 + assert by_id["late"].offset == 0 + + +def test_sweep_same_time_places_largest_first() -> None: + allocs = ( + Allocation(id="small", size=10, start=0, end=10), + Allocation(id="large", size=100, start=0, end=10), + ) + by_id = {a.id: a for a in SweepAllocator().allocate(allocs)} + assert by_id["large"].offset == 0 + assert by_id["small"].offset == 100 + + +def test_sweep_best_fit_picks_tighter_gap() -> None: + allocs = ( + Allocation(id="pillar_a", size=200, start=0, end=10), + Allocation(id="hole_wide", size=150, start=0, end=2), + Allocation(id="pillar_b", size=100, start=0, end=10), + Allocation(id="hole_tight", size=50, start=0, end=2), + Allocation(id="pillar_c", size=40, start=0, end=10), + Allocation(id="probe", size=50, start=3, end=10), + ) + first = {a.id: a.offset for a in SweepAllocator().allocate(allocs)} + best = {a.id: a.offset for a in SweepBestFitAllocator().allocate(allocs)} + assert first["probe"] == 200 + assert best["probe"] == 450 + + +def test_sweep_does_not_mutate_input() -> None: + allocs = ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=50, start=5, end=15), + ) + result = SweepAllocator().allocate(allocs) + assert all(a.offset is None for a in allocs) + assert {a.id for a in result} == {1, 2} + assert all(a.offset is not None for a in result) + + +def test_hybrid_sweep_by_size_matches_greedy_by_size_below_budget() -> None: + allocs = RandomSource(num_allocations=100, seed=7).get_allocations() + hybrid = HybridSweepBySizeAllocator().allocate(allocs) + greedy = GreedyBySizeAllocator().allocate(allocs) + assert peak_memory(hybrid) == peak_memory(greedy) + + +def test_hybrid_sweep_by_duration_matches_greedy_by_duration_below_budget() -> None: + allocs = RandomSource(num_allocations=100, seed=7).get_allocations() + hybrid = HybridSweepByDurationAllocator().allocate(allocs) + greedy = GreedyByDurationAllocator().allocate(allocs) + assert peak_memory(hybrid) == peak_memory(greedy) + + +def test_hybrid_sweep_negative_max_obstacles_raises() -> None: + with pytest.raises(ValueError, match="max_obstacles"): + HybridSweepAllocator(max_obstacles=-1) + + +@pytest.mark.parametrize( + "allocator_cls", [SweepAllocator, SweepBestFitAllocator, HybridSweepAllocator] +) +def test_sweep_oversized_total_size_raises(allocator_cls: type) -> None: + with pytest.raises(OverflowError): + allocator_cls().allocate((Allocation(id=1, size=2**62, start=0, end=1),)) + + +def test_hybrid_sweep_above_budget_is_valid() -> None: + allocs = RandomSource(num_allocations=200, seed=3).get_allocations() + result = HybridSweepAllocator(max_obstacles=16).allocate(allocs) + assert _is_valid(result) + assert len(result) == len(allocs) + + +@pytest.mark.parametrize("allocator_cls", ALL_SWEEP_ALLOCATORS) +def test_sweep_deterministic(allocator_cls: type) -> None: + allocs = RandomSource(num_allocations=50, seed=11).get_allocations() + result1 = allocator_cls().allocate(allocs) + result2 = allocator_cls().allocate(allocs) + by_id1 = {a.id: a.offset for a in result1} + by_id2 = {a.id: a.offset for a in result2} + assert by_id1 == by_id2 + + +def test_sweep_by_all_not_worse_than_any_member() -> None: + allocs = RandomSource(num_allocations=150, seed=5).get_allocations() + best = peak_memory(SweepByAllAllocator(cores=1).allocate(allocs)) + for allocator_cls in ALL_SWEEP_ALLOCATORS: + assert best <= peak_memory(allocator_cls().allocate(allocs)) + + +def test_all_sweep_allocators_valid_on_complex_workload() -> None: + allocs = RandomSource( + num_allocations=300, duration_max=2000, time_max=4000, seed=13 + ).get_allocations() + for allocator_cls in ALL_SWEEP_ALLOCATORS: + result = allocator_cls().allocate(allocs) + assert _is_valid(result) + assert len(result) == len(allocs)