diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ac349e..3ec4881 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -39,9 +39,13 @@ find_package(nanobind CONFIG REQUIRED) nanobind_add_module( _cpp NB_STATIC + src/cpp/allocators/best_fit.cpp src/cpp/allocators/greedy.cpp src/cpp/allocators/greedy_base.cpp + src/cpp/allocators/simulated_annealing.cpp src/cpp/allocators/supermalloc/partition.cpp + src/cpp/allocators/tabu_search.cpp + src/cpp/allocators/telamalloc.cpp src/cpp/primitives/allocation.cpp src/cpp/bindings.cpp) diff --git a/src/cpp/allocators/best_fit.cpp b/src/cpp/allocators/best_fit.cpp new file mode 100644 index 0000000..e997733 --- /dev/null +++ b/src/cpp/allocators/best_fit.cpp @@ -0,0 +1,60 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#include "best_fit.hpp" + +#include + +namespace omnimalloc { + +namespace { + +int64_t find_best_fit_offset(const Allocation& current_alloc, + const std::vector& placed_allocations, + const TemporalOverlaps& overlaps) { + // Scan every gap between placed overlaps and keep the smallest that fits + int64_t cursor = 0; + int64_t best_offset = 0; + int64_t best_gap = -1; // negative sentinel: no finite fitting gap yet + for (const auto* placed : + placed_overlapping(current_alloc, placed_allocations, overlaps)) { + int64_t gap = placed->offset().value() - cursor; + if (gap >= current_alloc.size() && (best_gap < 0 || gap < best_gap)) { + best_gap = gap; + best_offset = cursor; + } + cursor = std::max(cursor, placed->offset().value() + placed->size()); + } + + // No finite gap fit: place after the last overlapping allocation + return best_gap < 0 ? cursor : best_offset; +} + +} // namespace + +std::vector BestFitAllocator::allocate( + const std::vector& allocations) const { + const TemporalOverlaps overlaps = compute_temporal_overlaps(allocations); + std::vector placed_allocations; + placed_allocations.reserve(allocations.size()); + for (const auto& alloc : allocations) { + int64_t best_offset = + find_best_fit_offset(alloc, placed_allocations, overlaps); + placed_allocations.push_back(alloc.with_offset(best_offset)); + } + + return placed_allocations; +} + +} // namespace omnimalloc + +namespace std { + +size_t hash::operator()( + const omnimalloc::BestFitAllocator&) const noexcept { + // Stateless class - all instances are equal, use constant hash + return 0x517cc1b7; // arbitrary constant +} + +} // namespace std diff --git a/src/cpp/allocators/best_fit.hpp b/src/cpp/allocators/best_fit.hpp new file mode 100644 index 0000000..cf387b5 --- /dev/null +++ b/src/cpp/allocators/best_fit.hpp @@ -0,0 +1,36 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include + +#include "greedy_base.hpp" +#include "primitives/allocation.hpp" + +namespace omnimalloc { + +// Best-fit placement: like first-fit, but among the gaps left by +// already-placed overlapping allocations it picks the smallest one that fits +// (ties broken by lowest offset) rather than the first one found. Falls back +// to placing after the last overlapping allocation when no finite gap fits. +class BestFitAllocator { + public: + BestFitAllocator() = default; + + // Allocate the given allocations using a best-fit greedy strategy. + std::vector allocate( + const std::vector& allocations) const; + + bool operator==(const BestFitAllocator&) const noexcept = default; +}; + +} // namespace omnimalloc + +namespace std { +template <> +struct hash { + size_t operator()(const omnimalloc::BestFitAllocator&) const noexcept; +}; +} // namespace std diff --git a/src/cpp/allocators/greedy_base.cpp b/src/cpp/allocators/greedy_base.cpp index 75e1134..ded3758 100644 --- a/src/cpp/allocators/greedy_base.cpp +++ b/src/cpp/allocators/greedy_base.cpp @@ -5,6 +5,7 @@ #include "greedy_base.hpp" #include +#include #include #include @@ -40,32 +41,78 @@ TemporalOverlaps compute_temporal_overlaps( return overlaps; } -namespace { +int64_t peak_of(const std::vector& placed) { + int64_t peak = 0; + for (const auto& alloc : placed) { + if (const auto height = alloc.height()) { + peak = std::max(peak, *height); + } + } + return peak; +} -int64_t find_best_offset(const Allocation& current_alloc, - const std::vector& placed_allocations, - const TemporalOverlaps& overlaps) { +std::vector placed_overlapping( + const Allocation& alloc, const std::vector& placed, + const TemporalOverlaps& overlaps) { // Collect overlapping allocations that have been placed std::vector overlapping; - auto it = overlaps.find(current_alloc.id()); + auto it = overlaps.find(alloc.id()); if (it != overlaps.end()) { overlapping.reserve(it->second.size()); - for (const auto& placed : placed_allocations) { - if (it->second.count(placed.id())) { - overlapping.push_back(&placed); + for (const auto& candidate : placed) { + if (it->second.count(candidate.id())) { + overlapping.push_back(&candidate); } } } - // Sort by offset to enable first-fit algorithm + // Sort by offset so callers can scan the free gaps left-to-right std::sort(overlapping.begin(), overlapping.end(), [](const Allocation* a, const Allocation* b) { return a->offset().value() < b->offset().value(); }); + return overlapping; +} + +std::vector initial_order(const std::vector& allocations) { + std::vector order(allocations.size()); + std::iota(order.begin(), order.end(), size_t{0}); + std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { + return allocations[a].size() > allocations[b].size(); + }); + return order; +} + +std::vector earlier_neighbors( + const std::vector& order, size_t target_pos, + const std::vector& allocations, + const TemporalOverlaps& overlaps) { + std::vector neighbors; + auto it = overlaps.find(allocations[order[target_pos]].id()); + if (it != overlaps.end()) { + for (size_t pos = 0; pos < target_pos; ++pos) { + if (it->second.count(allocations[order[pos]].id())) { + neighbors.push_back(pos); + } + } + } + if (neighbors.empty()) { + neighbors.resize(target_pos); + std::iota(neighbors.begin(), neighbors.end(), size_t{0}); + } + return neighbors; +} + +namespace { + +int64_t find_best_offset(const Allocation& current_alloc, + const std::vector& placed_allocations, + const TemporalOverlaps& overlaps) { // Find best offset using first-fit: scan for first gap that fits int64_t best_offset = 0; - for (const auto* placed : overlapping) { + for (const auto* placed : + placed_overlapping(current_alloc, placed_allocations, overlaps)) { int64_t gap = placed->offset().value() - best_offset; if (gap >= current_alloc.size()) { break; // Found a fitting gap @@ -107,13 +154,7 @@ std::vector FirstFitPlacer::place( } int64_t FirstFitPlacer::evaluate(const std::vector& order) const { - int64_t peak = 0; - for (const auto& alloc : place(order)) { - if (const auto height = alloc.height()) { - peak = std::max(peak, *height); - } - } - return peak; + return peak_of(place(order)); } } // namespace omnimalloc diff --git a/src/cpp/allocators/greedy_base.hpp b/src/cpp/allocators/greedy_base.hpp index 6d2c807..d29c8d0 100644 --- a/src/cpp/allocators/greedy_base.hpp +++ b/src/cpp/allocators/greedy_base.hpp @@ -29,6 +29,27 @@ using TemporalOverlaps = const std::vector& allocations, const TemporalOverlaps& overlaps); +// Peak memory (highest end offset) across the placed allocations +[[nodiscard]] int64_t peak_of(const std::vector& placed); + +// Already-placed allocations that temporally overlap `alloc`, sorted by +// offset so callers can scan the free gaps left-to-right +[[nodiscard]] std::vector placed_overlapping( + const Allocation& alloc, const std::vector& placed, + const TemporalOverlaps& overlaps); + +// Indices into `allocations`, sorted largest-size-first: a decent, cheap +// starting order for the order-search allocators. +[[nodiscard]] std::vector initial_order( + const std::vector& allocations); + +// Positions before `target_pos` in `order` whose allocation temporally +// overlaps the one at `target_pos`, or every earlier position if none do. +[[nodiscard]] std::vector earlier_neighbors( + const std::vector& order, size_t target_pos, + const std::vector& allocations, + const TemporalOverlaps& overlaps); + // Resident first-fit placer for the order-search allocators (genetic, random, // hill-climb): owns the allocations and their overlap map (computed once) so // that placing many candidate orderings only passes an index permutation across diff --git a/src/cpp/allocators/simulated_annealing.cpp b/src/cpp/allocators/simulated_annealing.cpp new file mode 100644 index 0000000..c35fd4d --- /dev/null +++ b/src/cpp/allocators/simulated_annealing.cpp @@ -0,0 +1,99 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#include "simulated_annealing.hpp" + +#include +#include +#include +#include + +#include "greedy_base.hpp" + +namespace omnimalloc { + +SimulatedAnnealingAllocator::SimulatedAnnealingAllocator( + SimulatedAnnealingConfig config) + : config_(config) {} + +std::vector SimulatedAnnealingAllocator::allocate( + const std::vector& allocations) const { + const FirstFitPlacer placer(allocations); + std::vector order = initial_order(allocations); + if (allocations.size() < 2) { + return placer.place(order); + } + + std::vector current_placed = placer.place(order); + int64_t current_peak = peak_of(current_placed); + std::vector best_placed = current_placed; + int64_t best_peak = current_peak; + + std::mt19937_64 rng(config_.seed); + std::uniform_real_distribution unit(0.0, 1.0); + double temperature = config_.initial_temperature; + + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::duration(config_.max_seconds); + + for (int iteration = 0; iteration < config_.max_iterations; ++iteration) { + if (config_.max_seconds > 0.0 && + std::chrono::steady_clock::now() >= deadline) { + break; + } + + std::vector peak_positions; + for (size_t pos = 0; pos < current_placed.size(); ++pos) { + const auto height = current_placed[pos].height(); + if (height && *height == current_peak) { + peak_positions.push_back(pos); + } + } + if (peak_positions.empty()) { + break; // no placed allocation reaches the peak: nothing to perturb + } + + std::uniform_int_distribution pick_peak(0, + peak_positions.size() - 1); + size_t target_pos = peak_positions[pick_peak(rng)]; + std::vector neighbors = + earlier_neighbors(order, target_pos, allocations, placer.overlaps()); + if (neighbors.empty()) { + temperature *= config_.cooling_rate; + continue; + } + std::uniform_int_distribution pick_neighbor(0, + neighbors.size() - 1); + size_t other_pos = neighbors[pick_neighbor(rng)]; + + std::swap(order[target_pos], order[other_pos]); + std::vector candidate_placed = placer.place(order); + int64_t candidate_peak = peak_of(candidate_placed); + + bool accept = candidate_peak <= current_peak; + if (!accept && current_peak > 0 && temperature > 0.0) { + double worsening_percent = + 100.0 * static_cast(candidate_peak - current_peak) / + static_cast(current_peak); + accept = unit(rng) < std::exp(-worsening_percent / temperature); + } + + if (accept) { + current_placed = std::move(candidate_placed); + current_peak = candidate_peak; + if (current_peak < best_peak) { + best_placed = current_placed; + best_peak = current_peak; + } + } else { + std::swap(order[target_pos], order[other_pos]); // undo the rejected swap + } + + temperature *= config_.cooling_rate; + } + + return best_placed; +} + +} // namespace omnimalloc diff --git a/src/cpp/allocators/simulated_annealing.hpp b/src/cpp/allocators/simulated_annealing.hpp new file mode 100644 index 0000000..33d343e --- /dev/null +++ b/src/cpp/allocators/simulated_annealing.hpp @@ -0,0 +1,47 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include +#include + +#include "primitives/allocation.hpp" + +namespace omnimalloc { + +// Cooling schedule and iteration budget for `SimulatedAnnealingAllocator`. +struct SimulatedAnnealingConfig { + uint64_t seed = 42; + int max_iterations = 3000; + // Percent memory worsening accepted with probability 1/e at iteration 0; + // decays geometrically by `cooling_rate` every iteration. + double initial_temperature = 3.0; + double cooling_rate = 0.998; + // Wall-clock budget checked once per iteration; 0 disables it. Each + // iteration re-evaluates a full O(n) placement, so `max_iterations` alone + // does not bound runtime as `allocations` grows - this does. + double max_seconds = 2.0; +}; + +// Simulated annealing over first-fit placement orders. Each iteration swaps a +// currently-peak allocation with an earlier temporal neighbor, accepting the +// swap outright when it does not worsen the peak and otherwise with a +// Metropolis probability that anneals to zero over `max_iterations`. The +// entire search runs natively (no Python round trip per candidate), so it can +// evaluate far more candidate placements per second than an +// equivalent Python-orchestrated local search. +class SimulatedAnnealingAllocator { + public: + explicit SimulatedAnnealingAllocator( + SimulatedAnnealingConfig config = SimulatedAnnealingConfig{}); + + [[nodiscard]] std::vector allocate( + const std::vector& allocations) const; + + private: + SimulatedAnnealingConfig config_; +}; + +} // namespace omnimalloc diff --git a/src/cpp/allocators/tabu_search.cpp b/src/cpp/allocators/tabu_search.cpp new file mode 100644 index 0000000..aeb531d --- /dev/null +++ b/src/cpp/allocators/tabu_search.cpp @@ -0,0 +1,133 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#include "tabu_search.hpp" + +#include +#include +#include +#include + +#include "greedy_base.hpp" + +namespace omnimalloc { + +namespace { + +// Order-independent key for the pair of original allocation indices a swap +// exchanges, used to record/check tabu status. +int64_t tabu_key(size_t a, size_t b, size_t num_allocations) { + if (a > b) { + std::swap(a, b); + } + return static_cast(a * num_allocations + b); +} + +} // namespace + +TabuSearchAllocator::TabuSearchAllocator(TabuSearchConfig config) + : config_(config) {} + +std::vector TabuSearchAllocator::allocate( + const std::vector& allocations) const { + const FirstFitPlacer placer(allocations); + std::vector order = initial_order(allocations); + if (allocations.size() < 2) { + return placer.place(order); + } + + const size_t n = allocations.size(); + + std::vector current_placed = placer.place(order); + int64_t current_peak = peak_of(current_placed); + std::vector best_placed = current_placed; + int64_t best_peak = current_peak; + + std::mt19937_64 rng(config_.seed); + std::unordered_map tabu_until; + + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::duration(config_.max_seconds); + + for (int iteration = 0; iteration < config_.max_iterations; ++iteration) { + if (config_.max_seconds > 0.0 && + std::chrono::steady_clock::now() >= deadline) { + break; + } + + std::vector peak_positions; + for (size_t pos = 0; pos < current_placed.size(); ++pos) { + const auto height = current_placed[pos].height(); + if (height && *height == current_peak) { + peak_positions.push_back(pos); + } + } + if (peak_positions.empty()) { + break; // no placed allocation reaches the peak: nothing to perturb + } + std::uniform_int_distribution pick_peak(0, + peak_positions.size() - 1); + + // Sample a neighborhood of candidate swaps and keep the best admissible + // one: non-tabu, or tabu but beating the best-ever solution (aspiration). + size_t best_p1 = 0; + size_t best_p2 = 0; + int64_t best_candidate_peak = -1; + std::vector best_candidate_placed; + bool best_is_tabu = false; + + for (int sample = 0; sample < config_.neighborhood_size; ++sample) { + size_t target_pos = peak_positions[pick_peak(rng)]; + std::vector neighbors = + earlier_neighbors(order, target_pos, allocations, placer.overlaps()); + if (neighbors.empty()) { + continue; + } + std::uniform_int_distribution pick_neighbor(0, + neighbors.size() - 1); + size_t other_pos = neighbors[pick_neighbor(rng)]; + + auto tabu_it = + tabu_until.find(tabu_key(order[target_pos], order[other_pos], n)); + bool is_tabu = tabu_it != tabu_until.end() && tabu_it->second > iteration; + + std::swap(order[target_pos], order[other_pos]); + std::vector candidate_placed = placer.place(order); + int64_t candidate_peak = peak_of(candidate_placed); + std::swap(order[target_pos], + order[other_pos]); // undo; reapplied if chosen + + bool aspires = candidate_peak < best_peak; + if ((!is_tabu || aspires) && + (best_candidate_peak < 0 || candidate_peak < best_candidate_peak)) { + best_candidate_peak = candidate_peak; + best_candidate_placed = std::move(candidate_placed); + best_p1 = target_pos; + best_p2 = other_pos; + best_is_tabu = is_tabu; + } + } + + if (best_candidate_peak < 0) { + continue; // every sampled move was tabu without meeting aspiration + } + + std::swap(order[best_p1], order[best_p2]); + current_placed = std::move(best_candidate_placed); + current_peak = best_candidate_peak; + if (!best_is_tabu) { + tabu_until[tabu_key(order[best_p1], order[best_p2], n)] = + iteration + config_.tabu_tenure; + } + + if (current_peak < best_peak) { + best_placed = current_placed; + best_peak = current_peak; + } + } + + return best_placed; +} + +} // namespace omnimalloc diff --git a/src/cpp/allocators/tabu_search.hpp b/src/cpp/allocators/tabu_search.hpp new file mode 100644 index 0000000..1ee897d --- /dev/null +++ b/src/cpp/allocators/tabu_search.hpp @@ -0,0 +1,48 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include +#include + +#include "primitives/allocation.hpp" + +namespace omnimalloc { + +// Neighborhood size, iteration budget, and tabu memory for +// `TabuSearchAllocator`. +struct TabuSearchConfig { + uint64_t seed = 42; + int max_iterations = 500; + int neighborhood_size = 20; // candidate swaps sampled per iteration + int tabu_tenure = 15; // iterations a reversed swap stays forbidden + // Wall-clock budget checked once per iteration; 0 disables it. Each + // iteration evaluates `neighborhood_size` full O(n) placements, so + // `max_iterations` alone does not bound runtime as `allocations` grows - + // this does. + double max_seconds = 2.0; +}; + +// Tabu search over first-fit placement orders. Each iteration samples +// `neighborhood_size` candidate swaps between a currently-peak allocation and +// an earlier temporal neighbor, and moves to the best-scoring candidate that +// is not tabu (or, per the aspiration criterion, a tabu candidate that beats +// the best solution found so far). The swap just made is then forbidden from +// being immediately reversed for `tabu_tenure` iterations, which lets the +// search climb out of local optima without cycling between the same two +// orders. Runs entirely in C++ for the same reason as +// `SimulatedAnnealingAllocator`: no Python round trip per candidate. +class TabuSearchAllocator { + public: + explicit TabuSearchAllocator(TabuSearchConfig config = TabuSearchConfig{}); + + [[nodiscard]] std::vector allocate( + const std::vector& allocations) const; + + private: + TabuSearchConfig config_; +}; + +} // namespace omnimalloc diff --git a/src/cpp/allocators/telamalloc.cpp b/src/cpp/allocators/telamalloc.cpp new file mode 100644 index 0000000..2457ccf --- /dev/null +++ b/src/cpp/allocators/telamalloc.cpp @@ -0,0 +1,362 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#include "telamalloc.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "greedy_base.hpp" + +namespace omnimalloc { + +namespace { + +using Clock = std::chrono::steady_clock; +using Deadline = std::optional; + +// Effectively-unbounded capacity: makes a pack attempt plain first-fit. +constexpr int64_t kUnbounded = std::numeric_limits::max() / 4; + +// Index-based temporal adjacency via a sweep over start/end events; ends +// sort before starts at equal times, so touching [start, end) intervals do +// not overlap. +std::vector> build_neighbors( + const std::vector& allocations) { + const int n = static_cast(allocations.size()); + std::vector> events; + events.reserve(2 * static_cast(n)); + for (int i = 0; i < n; ++i) { + events.emplace_back(allocations[i].start(), true, i); + events.emplace_back(allocations[i].end(), false, i); + } + std::sort(events.begin(), events.end()); + + std::vector> neighbors(n); + std::unordered_set active; + for (const auto& [time, is_start, idx] : events) { + if (is_start) { + for (int other : active) { + neighbors[idx].push_back(other); + neighbors[other].push_back(idx); + } + active.insert(idx); + } else { + active.erase(idx); + } + } + return neighbors; +} + +// Connected components of the overlap graph: the paper's "phases". Buffers +// in different components never interact, so each packs independently. +std::vector> build_phases( + const std::vector>& neighbors) { + const int n = static_cast(neighbors.size()); + std::vector> phases; + std::vector visited(n, 0); + for (int seed = 0; seed < n; ++seed) { + if (visited[seed]) { + continue; + } + std::vector phase; + std::vector stack{seed}; + visited[seed] = 1; + while (!stack.empty()) { + int idx = stack.back(); + stack.pop_back(); + phase.push_back(idx); + for (int other : neighbors[idx]) { + if (!visited[other]) { + visited[other] = 1; + stack.push_back(other); + } + } + } + phases.push_back(std::move(phase)); + } + return phases; +} + +// Peak simultaneous load of `phase`: a lower bound on its achievable peak. +int64_t load_lower_bound(const std::vector& allocations, + const std::vector& phase) { + std::vector> events; + events.reserve(2 * phase.size()); + for (int idx : phase) { + events.emplace_back(allocations[idx].start(), true, idx); + events.emplace_back(allocations[idx].end(), false, idx); + } + std::sort(events.begin(), events.end()); + + int64_t load = 0; + int64_t peak = 0; + for (const auto& [time, is_start, idx] : events) { + load += is_start ? allocations[idx].size() : -allocations[idx].size(); + peak = std::max(peak, load); + } + return peak; +} + +// Queue order: most-evicted first (squeaky wheel), then the paper's tiers +// (longest lifetime, largest size — or size-major when `size_major`), then +// lowest index for determinism. +using QueueKey = std::tuple; + +QueueKey queue_key(const Allocation& alloc, int idx, int evictions, + bool size_major) { + if (size_major) { + return {-evictions, -alloc.size(), -alloc.duration(), idx}; + } + return {-evictions, -alloc.duration(), -alloc.size(), idx}; +} + +// One capacity attempt over one phase: place buffers in queue order at the +// lowest fitting gap among placed temporal neighbors; on conflict (minor +// backtracking), evict the cheapest blocking set (fewest bytes, then fewest +// buffers, then lowest offset) and place anyway. Evicted buffers re-enter +// the queue with raised priority. When a restart's share of the eviction +// budget runs out (major backtracking), every placement is wiped and the +// phase re-packs from scratch; eviction counts survive the wipe, so +// contentious buffers get placed first in the next round (squeaky-wheel +// reordering). Returns nullopt when the total budget or deadline runs out; +// `capacity` at least the phase's max buffer size guarantees offset 0 is +// always a repair candidate, so every iteration makes progress. +std::optional> pack_phase( + const std::vector& allocations, + const std::vector>& neighbors, + const std::vector& phase, int64_t capacity, int max_backtracks, + const Deadline& deadline, bool size_major, uint64_t seed) { + std::vector offsets(allocations.size(), -1); + std::vector evictions(allocations.size(), 0); + std::mt19937_64 rng(seed); + + // A restart's eviction share scales with the phase, not the budget: a + // larger budget then buys more squeaky-wheel restarts (the diversification + // mechanism) rather than deeper local repair within one. + const int per_restart = + std::min(max_backtracks, 4 * static_cast(phase.size())); + int total_backtracks = 0; + + while (true) { + for (int idx : phase) { + offsets[idx] = -1; + } + std::set pending; + for (int idx : phase) { + pending.insert( + queue_key(allocations[idx], idx, evictions[idx], size_major)); + } + int restart_backtracks = 0; + + while (!pending.empty()) { + if (deadline && Clock::now() >= *deadline) { + return std::nullopt; + } + const int idx = std::get<3>(*pending.begin()); + pending.erase(pending.begin()); + const int64_t size = allocations[idx].size(); + + // Spatial intervals of the already-placed temporal neighbors. + std::vector> occupied; // (offset, neighbor) + for (int other : neighbors[idx]) { + if (offsets[other] >= 0) { + occupied.emplace_back(offsets[other], other); + } + } + std::sort(occupied.begin(), occupied.end()); + + // First-fit: lowest gap among the placed neighbors that fits. + int64_t cursor = 0; + bool in_gap = false; + for (const auto& [offset, other] : occupied) { + if (offset - cursor >= size) { + in_gap = true; + break; + } + cursor = std::max(cursor, offset + allocations[other].size()); + } + if (in_gap || cursor + size <= capacity) { + offsets[idx] = cursor; + continue; + } + + // Conflict: no gap fits below `capacity`. Score the candidate + // placements (offset 0, flush above each blocker, flush below each + // blocker) by the blockers each would displace; take the cheapest. + std::vector candidates{0}; + for (const auto& [offset, other] : occupied) { + const int64_t end = offset + allocations[other].size(); + if (end + size <= capacity) { + candidates.push_back(end); + } + if (offset > size) { + candidates.push_back(offset - size); + } + } + int64_t best_bytes = -1; + int best_count = 0; + int64_t best_offset = 0; + for (int64_t candidate : candidates) { + int64_t bytes = 0; + int count = 0; + for (const auto& [offset, other] : occupied) { + const int64_t end = offset + allocations[other].size(); + if (offset < candidate + size && end > candidate) { + bytes += allocations[other].size(); + ++count; + } + } + if (best_bytes < 0 || + std::tie(bytes, count, candidate) < + std::tie(best_bytes, best_count, best_offset)) { + best_bytes = bytes; + best_count = count; + best_offset = candidate; + } + } + + // Random walk (WalkSAT-style): occasionally take a random repair + // instead of the cheapest one, breaking the cycles a purely + // deterministic min-conflict search falls into. + if (rng() % 4 == 0) { + best_offset = candidates[rng() % candidates.size()]; + } + + for (const auto& [offset, other] : occupied) { + const int64_t end = offset + allocations[other].size(); + if (offset < best_offset + size && end > best_offset) { + offsets[other] = -1; + ++evictions[other]; + pending.insert(queue_key(allocations[other], other, evictions[other], + size_major)); + ++restart_backtracks; + ++total_backtracks; + } + } + offsets[idx] = best_offset; + + if (total_backtracks > max_backtracks) { + return std::nullopt; + } + if (restart_backtracks > per_restart) { + break; // major backtrack: wipe and re-pack in squeaky-wheel order + } + } + if (pending.empty()) { + return offsets; + } + } +} + +int64_t phase_peak(const std::vector& allocations, + const std::vector& phase, + const std::vector& offsets) { + int64_t peak = 0; + for (int idx : phase) { + peak = std::max(peak, offsets[idx] + allocations[idx].size()); + } + return peak; +} + +// Best-effort minimal peak for one phase: first-fit incumbent, then binary +// search on capacity down toward the load lower bound. Failed attempts are +// treated as infeasible even though they are only budget-exhausted, keeping +// the search anytime rather than exact. +void solve_phase(const std::vector& allocations, + const std::vector>& neighbors, + const std::vector& phase, int64_t lower_bound, + const TelamallocConfig& config, const Deadline& deadline, + std::vector& result) { + // Unbounded capacity never conflicts, so these incumbents are plain + // first-fit in each tiered order and cannot fail. The winner's order also + // steers the capacity search below. + auto by_duration = pack_phase(allocations, neighbors, phase, kUnbounded, 0, + std::nullopt, false, config.seed); + auto by_size = pack_phase(allocations, neighbors, phase, kUnbounded, 0, + std::nullopt, true, config.seed); + const int64_t duration_peak = phase_peak(allocations, phase, *by_duration); + const int64_t size_peak = phase_peak(allocations, phase, *by_size); + const bool size_major = size_peak < duration_peak; + + std::vector best = + size_major ? std::move(*by_size) : std::move(*by_duration); + int64_t high = std::min(duration_peak, size_peak); + int64_t low = lower_bound; + + while (low < high) { + if (deadline && Clock::now() >= *deadline) { + break; + } + const int64_t mid = low + (high - low) / 2; + auto attempt = + pack_phase(allocations, neighbors, phase, mid, config.max_backtracks, + deadline, size_major, config.seed); + if (attempt) { + best = std::move(*attempt); + high = phase_peak(allocations, phase, best); + } else { + low = mid + 1; + } + } + + for (int idx : phase) { + result[idx] = best[idx]; + } +} + +} // namespace + +TelamallocAllocator::TelamallocAllocator(TelamallocConfig config) + : config_(config) {} + +std::vector TelamallocAllocator::allocate( + const std::vector& allocations) const { + if (allocations.size() < 2) { + return first_fit_place(allocations, compute_temporal_overlaps(allocations)); + } + + Deadline deadline; + if (config_.max_seconds > 0.0) { + deadline = + Clock::now() + std::chrono::duration_cast( + std::chrono::duration(config_.max_seconds)); + } + + const auto neighbors = build_neighbors(allocations); + const auto phases = build_phases(neighbors); + + // Solve phases in descending load order: the global peak is the max over + // phases, so the dominant phase should get the wall-clock budget first. + std::vector> order; // (lower bound, phase) + order.reserve(phases.size()); + for (size_t p = 0; p < phases.size(); ++p) { + order.emplace_back(load_lower_bound(allocations, phases[p]), p); + } + std::sort(order.begin(), order.end(), [](const auto& a, const auto& b) { + return a.first != b.first ? a.first > b.first : a.second < b.second; + }); + + std::vector result(allocations.size(), -1); + for (const auto& [lower_bound, p] : order) { + solve_phase(allocations, neighbors, phases[p], lower_bound, config_, + deadline, result); + } + + std::vector placed; + placed.reserve(allocations.size()); + for (size_t i = 0; i < allocations.size(); ++i) { + placed.push_back(allocations[i].with_offset(result[i])); + } + return placed; +} + +} // namespace omnimalloc diff --git a/src/cpp/allocators/telamalloc.hpp b/src/cpp/allocators/telamalloc.hpp new file mode 100644 index 0000000..8eabc3e --- /dev/null +++ b/src/cpp/allocators/telamalloc.hpp @@ -0,0 +1,52 @@ +// +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include +#include + +#include "primitives/allocation.hpp" + +namespace omnimalloc { + +// Search budgets for `TelamallocAllocator`. +struct TelamallocConfig { + // Seeds the random-walk step of the conflict repair (see `pack_phase`); + // results are deterministic for a fixed seed. + uint64_t seed = 42; + // Eviction (backtrack) budget per capacity attempt; an attempt that + // exhausts it reports the capacity as unreachable. + int max_backtracks = 10000; + // Wall-clock budget for the whole allocate(); 0 disables it, leaving the + // per-attempt `max_backtracks` as the only bound. + double max_seconds = 2.0; +}; + +// TelaMalloc-style allocator after Maas et al., ASPLOS 2023 ("TelaMalloc: +// Efficient On-Chip Memory Allocation for Production Machine Learning +// Accelerators"). The paper packs buffers below a fixed on-chip capacity by +// interleaving a tiered buffer-ordering heuristic with constraint-solver +// feedback and conflict-directed backtracking. This adaptation minimizes +// peak memory instead: it splits the problem into phases (connected +// components of the temporal-overlap graph), packs each phase in the +// paper's tiered order (longest lifetime, then largest size) with +// min-conflict eviction as the backtracking mechanism, and binary-searches +// each phase's capacity between its load lower bound and a first-fit +// incumbent. Evicted buffers re-enter the queue with raised priority +// (squeaky-wheel), standing in for the paper's minor/major backtracking +// levels; an occasional seeded random-walk repair breaks the cycles a purely +// deterministic min-conflict search can fall into. +class TelamallocAllocator { + public: + explicit TelamallocAllocator(TelamallocConfig config = TelamallocConfig{}); + + [[nodiscard]] std::vector allocate( + const std::vector& allocations) const; + + private: + TelamallocConfig config_; +}; + +} // namespace omnimalloc diff --git a/src/cpp/bindings.cpp b/src/cpp/bindings.cpp index 94fe0bb..ef413aa 100644 --- a/src/cpp/bindings.cpp +++ b/src/cpp/bindings.cpp @@ -13,9 +13,13 @@ #include +#include "allocators/best_fit.hpp" #include "allocators/greedy.hpp" #include "allocators/greedy_base.hpp" +#include "allocators/simulated_annealing.hpp" #include "allocators/supermalloc/partition.hpp" +#include "allocators/tabu_search.hpp" +#include "allocators/telamalloc.hpp" #include "primitives/allocation.hpp" #include "primitives/buffer_kind.hpp" #include "primitives/id_type.hpp" @@ -121,6 +125,75 @@ NB_MODULE(_cpp, m) { .def("__eq__", &GreedyAllocator::operator==) .def("__hash__", std::hash{}); + // BestFitAllocator class + nb::class_(m, "BestFitAllocatorCpp") + .def(nb::init<>()) + .def("allocate", &BestFitAllocator::allocate, "allocations"_a, + nb::rv_policy::move) + .def("__str__", + [](const BestFitAllocator&) { return "BestFitAllocator()"; }) + .def("__repr__", + [](const BestFitAllocator&) { return "BestFitAllocator()"; }) + .def("__eq__", &BestFitAllocator::operator==) + .def("__hash__", std::hash{}); + + // SimulatedAnnealingConfig / SimulatedAnnealingAllocator classes + constexpr SimulatedAnnealingConfig kDefaultSaConfig{}; + nb::class_(m, "SimulatedAnnealingConfig") + .def(nb::init(), + "seed"_a = kDefaultSaConfig.seed, + "max_iterations"_a = kDefaultSaConfig.max_iterations, + "initial_temperature"_a = kDefaultSaConfig.initial_temperature, + "cooling_rate"_a = kDefaultSaConfig.cooling_rate, + "max_seconds"_a = kDefaultSaConfig.max_seconds) + .def_rw("seed", &SimulatedAnnealingConfig::seed) + .def_rw("max_iterations", &SimulatedAnnealingConfig::max_iterations) + .def_rw("initial_temperature", + &SimulatedAnnealingConfig::initial_temperature) + .def_rw("cooling_rate", &SimulatedAnnealingConfig::cooling_rate) + .def_rw("max_seconds", &SimulatedAnnealingConfig::max_seconds); + + nb::class_(m, "SimulatedAnnealingAllocatorCpp") + .def(nb::init(), "config"_a = kDefaultSaConfig) + .def("allocate", &SimulatedAnnealingAllocator::allocate, "allocations"_a, + nb::call_guard(), nb::rv_policy::move); + + // TabuSearchConfig / TabuSearchAllocator classes + constexpr TabuSearchConfig kDefaultTabuConfig{}; + nb::class_(m, "TabuSearchConfig") + .def(nb::init(), + "seed"_a = kDefaultTabuConfig.seed, + "max_iterations"_a = kDefaultTabuConfig.max_iterations, + "neighborhood_size"_a = kDefaultTabuConfig.neighborhood_size, + "tabu_tenure"_a = kDefaultTabuConfig.tabu_tenure, + "max_seconds"_a = kDefaultTabuConfig.max_seconds) + .def_rw("seed", &TabuSearchConfig::seed) + .def_rw("max_iterations", &TabuSearchConfig::max_iterations) + .def_rw("neighborhood_size", &TabuSearchConfig::neighborhood_size) + .def_rw("tabu_tenure", &TabuSearchConfig::tabu_tenure) + .def_rw("max_seconds", &TabuSearchConfig::max_seconds); + + nb::class_(m, "TabuSearchAllocatorCpp") + .def(nb::init(), "config"_a = kDefaultTabuConfig) + .def("allocate", &TabuSearchAllocator::allocate, "allocations"_a, + nb::call_guard(), nb::rv_policy::move); + + // TelamallocConfig / TelamallocAllocator classes + constexpr TelamallocConfig kDefaultTelaConfig{}; + nb::class_(m, "TelamallocConfig") + .def(nb::init(), + "seed"_a = kDefaultTelaConfig.seed, + "max_backtracks"_a = kDefaultTelaConfig.max_backtracks, + "max_seconds"_a = kDefaultTelaConfig.max_seconds) + .def_rw("seed", &TelamallocConfig::seed) + .def_rw("max_backtracks", &TelamallocConfig::max_backtracks) + .def_rw("max_seconds", &TelamallocConfig::max_seconds); + + nb::class_(m, "TelamallocAllocatorCpp") + .def(nb::init(), "config"_a = kDefaultTelaConfig) + .def("allocate", &TelamallocAllocator::allocate, "allocations"_a, + nb::call_guard(), nb::rv_policy::move); + // SearchOptions class constexpr SearchOptions kDefaultOptions{}; nb::class_(m, "SearchOptions") diff --git a/src/python/omnimalloc/_cpp.pyi b/src/python/omnimalloc/_cpp.pyi index 22408bd..c7ec2e6 100644 --- a/src/python/omnimalloc/_cpp.pyi +++ b/src/python/omnimalloc/_cpp.pyi @@ -102,6 +102,121 @@ class GreedyAllocatorCpp: def __hash__(self) -> int: ... +class BestFitAllocatorCpp: + def __init__(self) -> None: ... + + def allocate(self, allocations: Sequence[Allocation]) -> list[Allocation]: ... + + def __str__(self) -> str: ... + + def __repr__(self) -> str: ... + + def __eq__(self, arg: BestFitAllocatorCpp, /) -> bool: ... + + def __hash__(self) -> int: ... + +class SimulatedAnnealingConfig: + def __init__(self, seed: int = 42, max_iterations: int = 3000, initial_temperature: float = 3.0, cooling_rate: float = 0.998, max_seconds: float = 2.0) -> None: ... + + @property + def seed(self) -> int: ... + + @seed.setter + def seed(self, arg: int, /) -> None: ... + + @property + def max_iterations(self) -> int: ... + + @max_iterations.setter + def max_iterations(self, arg: int, /) -> None: ... + + @property + def initial_temperature(self) -> float: ... + + @initial_temperature.setter + def initial_temperature(self, arg: float, /) -> None: ... + + @property + def cooling_rate(self) -> float: ... + + @cooling_rate.setter + def cooling_rate(self, arg: float, /) -> None: ... + + @property + def max_seconds(self) -> float: ... + + @max_seconds.setter + def max_seconds(self, arg: float, /) -> None: ... + +class SimulatedAnnealingAllocatorCpp: + def __init__(self, config: SimulatedAnnealingConfig = ...) -> None: ... + + def allocate(self, allocations: Sequence[Allocation]) -> list[Allocation]: ... + +class TabuSearchConfig: + def __init__(self, seed: int = 42, max_iterations: int = 500, neighborhood_size: int = 20, tabu_tenure: int = 15, max_seconds: float = 2.0) -> None: ... + + @property + def seed(self) -> int: ... + + @seed.setter + def seed(self, arg: int, /) -> None: ... + + @property + def max_iterations(self) -> int: ... + + @max_iterations.setter + def max_iterations(self, arg: int, /) -> None: ... + + @property + def neighborhood_size(self) -> int: ... + + @neighborhood_size.setter + def neighborhood_size(self, arg: int, /) -> None: ... + + @property + def tabu_tenure(self) -> int: ... + + @tabu_tenure.setter + def tabu_tenure(self, arg: int, /) -> None: ... + + @property + def max_seconds(self) -> float: ... + + @max_seconds.setter + def max_seconds(self, arg: float, /) -> None: ... + +class TabuSearchAllocatorCpp: + def __init__(self, config: TabuSearchConfig = ...) -> None: ... + + def allocate(self, allocations: Sequence[Allocation]) -> list[Allocation]: ... + +class TelamallocConfig: + def __init__(self, seed: int = 42, max_backtracks: int = 10000, max_seconds: float = 2.0) -> None: ... + + @property + def seed(self) -> int: ... + + @seed.setter + def seed(self, arg: int, /) -> None: ... + + @property + def max_backtracks(self) -> int: ... + + @max_backtracks.setter + def max_backtracks(self, arg: int, /) -> None: ... + + @property + def max_seconds(self) -> float: ... + + @max_seconds.setter + def max_seconds(self, arg: float, /) -> None: ... + +class TelamallocAllocatorCpp: + def __init__(self, config: TelamallocConfig = ...) -> None: ... + + def allocate(self, allocations: Sequence[Allocation]) -> list[Allocation]: ... + class SearchOptions: def __init__(self, canonical: bool = True, dominance: bool = True, floor_inference: bool = True, monotonic_floor: bool = True, decompose: bool = True) -> None: ... diff --git a/src/python/omnimalloc/allocators/__init__.py b/src/python/omnimalloc/allocators/__init__.py index c0e6936..a1dea65 100644 --- a/src/python/omnimalloc/allocators/__init__.py +++ b/src/python/omnimalloc/allocators/__init__.py @@ -3,6 +3,7 @@ # from .base import BaseAllocator as BaseAllocator +from .best_fit import BestFitAllocator as BestFitAllocator from .genetic import GeneticAllocator as GeneticAllocator from .greedy import GreedyAllocator as GreedyAllocator from .greedy import GreedyByAllAllocator as GreedyByAllAllocator @@ -26,8 +27,16 @@ from .minimalloc import MinimallocAllocator as MinimallocAllocator from .naive import NaiveAllocator as NaiveAllocator from .random import RandomAllocator as RandomAllocator +from .simulated_annealing import ( + SimulatedAnnealingAllocator as SimulatedAnnealingAllocator, +) +from .simulated_annealing import SimulatedAnnealingConfig as SimulatedAnnealingConfig from .supermalloc import SupermallocAllocator as SupermallocAllocator from .supermalloc import SupermallocConfig as SupermallocConfig +from .tabu_search import TabuSearchAllocator as TabuSearchAllocator +from .tabu_search import TabuSearchConfig as TabuSearchConfig +from .telamalloc import TelamallocAllocator as TelamallocAllocator +from .telamalloc import TelamallocConfig as TelamallocConfig from .utils import AVAILABLE_ALLOCATORS as AVAILABLE_ALLOCATORS from .utils import DEFAULT_ALLOCATOR as DEFAULT_ALLOCATOR from .utils import get_allocator_by_name as get_allocator_by_name diff --git a/src/python/omnimalloc/allocators/best_fit.py b/src/python/omnimalloc/allocators/best_fit.py new file mode 100644 index 0000000..64fbee5 --- /dev/null +++ b/src/python/omnimalloc/allocators/best_fit.py @@ -0,0 +1,27 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +from functools import cached_property + +from omnimalloc._cpp import BestFitAllocatorCpp as _BestFitAllocatorCpp +from omnimalloc.primitives import Allocation + +from .base import BaseAllocator + + +class BestFitAllocator(BaseAllocator): + """Greedy allocator that places each buffer in the smallest sufficient gap. + + Unlike first-fit (which takes the first gap wide enough among overlapping + placements), best-fit scans every such gap and picks the tightest one, a + classic bin-packing strategy that tends to leave larger, more broadly + useful gaps free for later, bigger allocations. + """ + + @cached_property + def _cpp_allocator(self) -> _BestFitAllocatorCpp: + return _BestFitAllocatorCpp() + + def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + return tuple(self._cpp_allocator.allocate(list(allocations))) diff --git a/src/python/omnimalloc/allocators/simulated_annealing.py b/src/python/omnimalloc/allocators/simulated_annealing.py new file mode 100644 index 0000000..609fee6 --- /dev/null +++ b/src/python/omnimalloc/allocators/simulated_annealing.py @@ -0,0 +1,73 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +from dataclasses import dataclass + +from omnimalloc._cpp import ( + SimulatedAnnealingAllocatorCpp as _SimulatedAnnealingAllocatorCpp, +) +from omnimalloc._cpp import SimulatedAnnealingConfig as _SimulatedAnnealingConfig +from omnimalloc.primitives import Allocation + +from .base import BaseAllocator + + +@dataclass(frozen=True) +class SimulatedAnnealingConfig: + """Cooling schedule and iteration budget for `SimulatedAnnealingAllocator`.""" + + seed: int = 42 + max_iterations: int = 3000 + initial_temperature: float = 3.0 + cooling_rate: float = 0.998 + max_seconds: float = 2.0 + + def __post_init__(self) -> None: + if self.max_iterations <= 0: + raise ValueError( + f"max_iterations must be positive, got {self.max_iterations}" + ) + if self.initial_temperature < 0: + raise ValueError( + f"initial_temperature must be non-negative, " + f"got {self.initial_temperature}" + ) + if not 0.0 < self.cooling_rate <= 1.0: + raise ValueError(f"cooling_rate must be in (0, 1], got {self.cooling_rate}") + if self.max_seconds < 0: + raise ValueError( + f"max_seconds must be non-negative, got {self.max_seconds}" + ) + + def to_cpp_config(self) -> _SimulatedAnnealingConfig: + return _SimulatedAnnealingConfig( + seed=self.seed, + max_iterations=self.max_iterations, + initial_temperature=self.initial_temperature, + cooling_rate=self.cooling_rate, + max_seconds=self.max_seconds, + ) + + +class SimulatedAnnealingAllocator(BaseAllocator): + """Simulated annealing over first-fit placement orders, run entirely in C++. + + Repeatedly swaps a peak allocation with an earlier temporal neighbor, + accepting improving swaps outright and worsening ones with a probability + that anneals to zero over `max_iterations`. Because the whole search loop + (including every candidate placement) runs natively, it evaluates far more + candidates per second than an equivalent Python-orchestrated local search + such as `HillClimbAllocator`. Each iteration re-evaluates a full placement + of every allocation, so `max_seconds` (default 2s) bounds wall-clock time + as the input grows, independent of `max_iterations`. + """ + + def __init__(self, config: SimulatedAnnealingConfig | None = None) -> None: + self._config = config or SimulatedAnnealingConfig() + + def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + if not allocations: + return allocations + cpp_allocator = _SimulatedAnnealingAllocatorCpp(self._config.to_cpp_config()) + return tuple(cpp_allocator.allocate(list(allocations))) diff --git a/src/python/omnimalloc/allocators/tabu_search.py b/src/python/omnimalloc/allocators/tabu_search.py new file mode 100644 index 0000000..9a92c06 --- /dev/null +++ b/src/python/omnimalloc/allocators/tabu_search.py @@ -0,0 +1,71 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +from dataclasses import dataclass + +from omnimalloc._cpp import TabuSearchAllocatorCpp as _TabuSearchAllocatorCpp +from omnimalloc._cpp import TabuSearchConfig as _TabuSearchConfig +from omnimalloc.primitives import Allocation + +from .base import BaseAllocator + + +@dataclass(frozen=True) +class TabuSearchConfig: + """Neighborhood size, iteration budget, and tabu memory for TabuSearchAllocator.""" + + seed: int = 42 + max_iterations: int = 500 + neighborhood_size: int = 20 + tabu_tenure: int = 15 + max_seconds: float = 2.0 + + def __post_init__(self) -> None: + if self.max_iterations <= 0: + raise ValueError( + f"max_iterations must be positive, got {self.max_iterations}" + ) + if self.neighborhood_size <= 0: + raise ValueError( + f"neighborhood_size must be positive, got {self.neighborhood_size}" + ) + if self.tabu_tenure <= 0: + raise ValueError(f"tabu_tenure must be positive, got {self.tabu_tenure}") + if self.max_seconds < 0: + raise ValueError( + f"max_seconds must be non-negative, got {self.max_seconds}" + ) + + def to_cpp_config(self) -> _TabuSearchConfig: + return _TabuSearchConfig( + seed=self.seed, + max_iterations=self.max_iterations, + neighborhood_size=self.neighborhood_size, + tabu_tenure=self.tabu_tenure, + max_seconds=self.max_seconds, + ) + + +class TabuSearchAllocator(BaseAllocator): + """Tabu search over first-fit placement orders, run entirely in C++. + + Each iteration samples a neighborhood of candidate swaps between a + currently-peak allocation and an earlier temporal neighbor, and moves to + the best-scoring candidate that is not tabu (or, per the aspiration + criterion, a tabu one that beats the best solution found so far). The + swap just made is then forbidden from being immediately reversed for + `tabu_tenure` iterations, which helps the search escape local optima + without cycling between the same two orders. Each iteration evaluates + `neighborhood_size` full placements, so `max_seconds` (default 2s) bounds + wall-clock time as the input grows, independent of `max_iterations`. + """ + + def __init__(self, config: TabuSearchConfig | None = None) -> None: + self._config = config or TabuSearchConfig() + + def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + if not allocations: + return allocations + cpp_allocator = _TabuSearchAllocatorCpp(self._config.to_cpp_config()) + return tuple(cpp_allocator.allocate(list(allocations))) diff --git a/src/python/omnimalloc/allocators/telamalloc.py b/src/python/omnimalloc/allocators/telamalloc.py new file mode 100644 index 0000000..026b304 --- /dev/null +++ b/src/python/omnimalloc/allocators/telamalloc.py @@ -0,0 +1,66 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +from dataclasses import dataclass + +from omnimalloc._cpp import TelamallocAllocatorCpp as _TelamallocAllocatorCpp +from omnimalloc._cpp import TelamallocConfig as _TelamallocConfig +from omnimalloc.primitives import Allocation + +from .base import BaseAllocator + + +@dataclass(frozen=True) +class TelamallocConfig: + """Search budgets for TelamallocAllocator.""" + + seed: int = 42 + max_backtracks: int = 10000 + max_seconds: float = 2.0 + + def __post_init__(self) -> None: + if self.max_backtracks < 0: + raise ValueError( + f"max_backtracks must be non-negative, got {self.max_backtracks}" + ) + if self.max_seconds < 0: + raise ValueError( + f"max_seconds must be non-negative, got {self.max_seconds}" + ) + + def to_cpp_config(self) -> _TelamallocConfig: + return _TelamallocConfig( + seed=self.seed, + max_backtracks=self.max_backtracks, + max_seconds=self.max_seconds, + ) + + +class TelamallocAllocator(BaseAllocator): + """TelaMalloc-style allocator (Maas et al., ASPLOS 2023), run entirely in C++. + + The paper packs buffers below a fixed on-chip capacity by interleaving a + tiered buffer-ordering heuristic with constraint-solver feedback and + conflict-directed backtracking; it ships in Google's TPUv4 and Pixel 6. + No reference implementation is public, so this is an adaptation that + minimizes peak memory instead of testing satisfiability: independent + phases (connected components of the temporal-overlap graph) are each + packed in the paper's tiered order (longest lifetime, then largest size) + with min-conflict eviction as minor backtracking and squeaky-wheel + restarts as major backtracking, while a binary search drives each phase's + capacity from a first-fit incumbent down toward its load lower bound. An + occasional seeded random-walk repair breaks min-conflict cycles; results + are deterministic for a fixed `seed`, and setting `max_seconds` (default + 2s) to 0 makes them reproducible across machines via `max_backtracks` + alone. + """ + + def __init__(self, config: TelamallocConfig | None = None) -> None: + self._config = config or TelamallocConfig() + + def allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + if not allocations: + return allocations + cpp_allocator = _TelamallocAllocatorCpp(self._config.to_cpp_config()) + return tuple(cpp_allocator.allocate(list(allocations))) diff --git a/tests/unit/allocators/test_best_fit.py b/tests/unit/allocators/test_best_fit.py new file mode 100644 index 0000000..0d7cd8b --- /dev/null +++ b/tests/unit/allocators/test_best_fit.py @@ -0,0 +1,101 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +from omnimalloc.allocators.best_fit import BestFitAllocator +from omnimalloc.allocators.greedy import GreedyAllocator +from omnimalloc.allocators.greedy_base import peak_memory +from omnimalloc.primitives import Allocation, Pool +from omnimalloc.validate import validate_allocation + + +def _is_valid(result: tuple[Allocation, ...]) -> bool: + return validate_allocation(Pool(id="test_pool", allocations=result)) + + +def test_best_fit_empty() -> None: + allocator = BestFitAllocator() + result = allocator.allocate(()) + assert len(result) == 0 + + +def test_best_fit_single() -> None: + allocator = BestFitAllocator() + alloc = Allocation(id=1, size=100, start=0, end=10) + result = allocator.allocate((alloc,)) + assert len(result) == 1 + assert result[0].offset == 0 + + +def test_best_fit_no_temporal_overlap_shares_offset() -> None: + allocator = BestFitAllocator() + alloc1 = Allocation(id=1, size=100, start=0, end=10) + alloc2 = Allocation(id=2, size=200, start=10, end=20) + result = allocator.allocate((alloc1, alloc2)) + by_id = {a.id: a for a in result} + assert by_id[1].offset == 0 + assert by_id[2].offset == 0 + + +def test_best_fit_all_overlap_stacks_sequentially() -> None: + allocator = BestFitAllocator() + allocs = tuple(Allocation(id=i, size=100, start=0, end=10) for i in range(5)) + result = allocator.allocate(allocs) + assert _is_valid(result) + assert peak_memory(result) == 500 + + +def test_best_fit_preserves_allocations() -> None: + allocator = BestFitAllocator() + allocs = ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=50, start=5, end=15), + ) + result = allocator.allocate(allocs) + assert len(result) == len(allocs) + assert {a.id for a in result} == {1, 2} + assert all(a.offset is not None for a in result) + + +def test_best_fit_deterministic() -> None: + allocs = tuple( + Allocation(id=i, size=(i % 5 + 1) * 100, start=i % 3, end=i % 3 + i % 7 + 1) + for i in range(20) + ) + result1 = BestFitAllocator().allocate(allocs) + result2 = BestFitAllocator().allocate(allocs) + assert all(r1.offset == r2.offset for r1, r2 in zip(result1, result2, strict=True)) + + +def test_best_fit_chooses_tighter_gap_than_first_fit() -> None: + allocs = ( + Allocation(id="a", size=10, start=0, end=100), + Allocation(id="filler1", size=50, start=1, end=2), + Allocation(id="b", size=10, start=1, end=55), + Allocation(id="filler2", size=25, start=3, end=4), + Allocation(id="c", size=10, start=3, end=55), + Allocation(id="t", size=15, start=50, end=55), + ) + + first_fit = {a.id: a.offset for a in GreedyAllocator().allocate(allocs)} + best_fit = {a.id: a.offset for a in BestFitAllocator().allocate(allocs)} + + assert first_fit["t"] == 10 + assert best_fit["t"] == 45 + assert _is_valid(BestFitAllocator().allocate(allocs)) + assert peak_memory(tuple(BestFitAllocator().allocate(allocs))) == peak_memory( + tuple(GreedyAllocator().allocate(allocs)) + ) + + +def test_best_fit_complex_overlap_produces_valid_allocation() -> None: + allocator = BestFitAllocator() + allocs = ( + Allocation(id=1, size=100, start=0, end=5), + Allocation(id=2, size=100, start=3, end=8), + Allocation(id=3, size=100, start=6, end=10), + Allocation(id=4, size=50, start=0, end=10), + Allocation(id=5, size=300, start=2, end=4), + ) + result = allocator.allocate(allocs) + assert _is_valid(result) diff --git a/tests/unit/allocators/test_simulated_annealing.py b/tests/unit/allocators/test_simulated_annealing.py new file mode 100644 index 0000000..fd50b8b --- /dev/null +++ b/tests/unit/allocators/test_simulated_annealing.py @@ -0,0 +1,124 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +import pytest +from omnimalloc.allocators.greedy import GreedyBySizeAllocator +from omnimalloc.allocators.greedy_base import peak_memory +from omnimalloc.allocators.simulated_annealing import ( + SimulatedAnnealingAllocator, + SimulatedAnnealingConfig, +) +from omnimalloc.primitives import Allocation, Pool +from omnimalloc.validate import validate_allocation + + +def _is_valid(result: tuple[Allocation, ...]) -> bool: + return validate_allocation(Pool(id="test_pool", allocations=result)) + + +def test_simulated_annealing_empty() -> None: + allocator = SimulatedAnnealingAllocator() + result = allocator.allocate(()) + assert len(result) == 0 + + +def test_simulated_annealing_single() -> None: + allocator = SimulatedAnnealingAllocator() + alloc = Allocation(id=1, size=100, start=0, end=10) + result = allocator.allocate((alloc,)) + assert len(result) == 1 + assert result[0].offset == 0 + + +def test_simulated_annealing_rejects_non_positive_iterations() -> None: + with pytest.raises(ValueError, match="max_iterations must be positive"): + SimulatedAnnealingConfig(max_iterations=0) + + +def test_simulated_annealing_rejects_negative_temperature() -> None: + with pytest.raises(ValueError, match="initial_temperature must be non-negative"): + SimulatedAnnealingConfig(initial_temperature=-1.0) + + +def test_simulated_annealing_rejects_invalid_cooling_rate() -> None: + with pytest.raises(ValueError, match="cooling_rate must be in"): + SimulatedAnnealingConfig(cooling_rate=0.0) + with pytest.raises(ValueError, match="cooling_rate must be in"): + SimulatedAnnealingConfig(cooling_rate=1.5) + + +def test_simulated_annealing_preserves_allocations() -> None: + allocator = SimulatedAnnealingAllocator() + allocs = ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=50, start=5, end=15), + ) + result = allocator.allocate(allocs) + assert len(result) == len(allocs) + assert {a.id for a in result} == {1, 2} + assert all(a.offset is not None for a in result) + + +def test_simulated_annealing_no_temporal_overlap_shares_offset() -> None: + allocator = SimulatedAnnealingAllocator() + alloc1 = Allocation(id=1, size=100, start=0, end=10) + alloc2 = Allocation(id=2, size=200, start=10, end=20) + result = allocator.allocate((alloc1, alloc2)) + by_id = {a.id: a for a in result} + assert by_id[1].offset == 0 + assert by_id[2].offset == 0 + + +def test_simulated_annealing_all_overlap_stacks_sequentially() -> None: + allocator = SimulatedAnnealingAllocator() + allocs = tuple(Allocation(id=i, size=100, start=0, end=10) for i in range(5)) + result = allocator.allocate(allocs) + assert _is_valid(result) + assert peak_memory(result) == 500 + + +def test_simulated_annealing_deterministic() -> None: + allocs = tuple( + Allocation(id=i, size=(i % 5 + 1) * 100, start=i % 3, end=i % 3 + i % 7 + 1) + for i in range(20) + ) + config = SimulatedAnnealingConfig(max_iterations=200) + result1 = SimulatedAnnealingAllocator(config).allocate(allocs) + result2 = SimulatedAnnealingAllocator(config).allocate(allocs) + assert all(r1.offset == r2.offset for r1, r2 in zip(result1, result2, strict=True)) + + +def test_simulated_annealing_produces_valid_allocation_on_dense_overlap() -> None: + config = SimulatedAnnealingConfig(max_iterations=300) + allocator = SimulatedAnnealingAllocator(config) + allocs = tuple( + Allocation( + id=i, + size=(i % 7 + 1) * 10, + start=i % 4, + end=i % 4 + (i % 3 + 1) * 3, + ) + for i in range(30) + ) + result = allocator.allocate(allocs) + assert _is_valid(result) + assert {a.id for a in result} == {a.id for a in allocs} + + +def test_simulated_annealing_matches_or_beats_single_pass_greedy() -> None: + allocs = tuple( + Allocation( + id=i, + size=(i * 37 % 50 + 1) * 10, + start=i % 6, + end=i % 6 + (i * 13 % 5 + 1), + ) + for i in range(40) + ) + greedy_peak = peak_memory(GreedyBySizeAllocator().allocate(allocs)) + annealed = SimulatedAnnealingAllocator( + SimulatedAnnealingConfig(max_iterations=2000) + ).allocate(allocs) + assert _is_valid(annealed) + assert peak_memory(annealed) <= greedy_peak diff --git a/tests/unit/allocators/test_tabu_search.py b/tests/unit/allocators/test_tabu_search.py new file mode 100644 index 0000000..364952e --- /dev/null +++ b/tests/unit/allocators/test_tabu_search.py @@ -0,0 +1,117 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +import pytest +from omnimalloc.allocators.greedy import GreedyBySizeAllocator +from omnimalloc.allocators.greedy_base import peak_memory +from omnimalloc.allocators.tabu_search import TabuSearchAllocator, TabuSearchConfig +from omnimalloc.primitives import Allocation, Pool +from omnimalloc.validate import validate_allocation + + +def _is_valid(result: tuple[Allocation, ...]) -> bool: + return validate_allocation(Pool(id="test_pool", allocations=result)) + + +def test_tabu_search_empty() -> None: + allocator = TabuSearchAllocator() + result = allocator.allocate(()) + assert len(result) == 0 + + +def test_tabu_search_single() -> None: + allocator = TabuSearchAllocator() + alloc = Allocation(id=1, size=100, start=0, end=10) + result = allocator.allocate((alloc,)) + assert len(result) == 1 + assert result[0].offset == 0 + + +def test_tabu_search_rejects_non_positive_iterations() -> None: + with pytest.raises(ValueError, match="max_iterations must be positive"): + TabuSearchConfig(max_iterations=0) + + +def test_tabu_search_rejects_non_positive_neighborhood_size() -> None: + with pytest.raises(ValueError, match="neighborhood_size must be positive"): + TabuSearchConfig(neighborhood_size=0) + + +def test_tabu_search_rejects_non_positive_tabu_tenure() -> None: + with pytest.raises(ValueError, match="tabu_tenure must be positive"): + TabuSearchConfig(tabu_tenure=0) + + +def test_tabu_search_preserves_allocations() -> None: + allocator = TabuSearchAllocator() + allocs = ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=50, start=5, end=15), + ) + result = allocator.allocate(allocs) + assert len(result) == len(allocs) + assert {a.id for a in result} == {1, 2} + assert all(a.offset is not None for a in result) + + +def test_tabu_search_no_temporal_overlap_shares_offset() -> None: + allocator = TabuSearchAllocator() + alloc1 = Allocation(id=1, size=100, start=0, end=10) + alloc2 = Allocation(id=2, size=200, start=10, end=20) + result = allocator.allocate((alloc1, alloc2)) + by_id = {a.id: a for a in result} + assert by_id[1].offset == 0 + assert by_id[2].offset == 0 + + +def test_tabu_search_all_overlap_stacks_sequentially() -> None: + allocator = TabuSearchAllocator() + allocs = tuple(Allocation(id=i, size=100, start=0, end=10) for i in range(5)) + result = allocator.allocate(allocs) + assert _is_valid(result) + assert peak_memory(result) == 500 + + +def test_tabu_search_deterministic() -> None: + allocs = tuple( + Allocation(id=i, size=(i % 5 + 1) * 100, start=i % 3, end=i % 3 + i % 7 + 1) + for i in range(20) + ) + config = TabuSearchConfig(max_iterations=50) + result1 = TabuSearchAllocator(config).allocate(allocs) + result2 = TabuSearchAllocator(config).allocate(allocs) + assert all(r1.offset == r2.offset for r1, r2 in zip(result1, result2, strict=True)) + + +def test_tabu_search_produces_valid_allocation_on_dense_overlap() -> None: + allocator = TabuSearchAllocator(TabuSearchConfig(max_iterations=80)) + allocs = tuple( + Allocation( + id=i, + size=(i % 7 + 1) * 10, + start=i % 4, + end=i % 4 + (i % 3 + 1) * 3, + ) + for i in range(30) + ) + result = allocator.allocate(allocs) + assert _is_valid(result) + assert {a.id for a in result} == {a.id for a in allocs} + + +def test_tabu_search_matches_or_beats_single_pass_greedy() -> None: + allocs = tuple( + Allocation( + id=i, + size=(i * 37 % 50 + 1) * 10, + start=i % 6, + end=i % 6 + (i * 13 % 5 + 1), + ) + for i in range(40) + ) + greedy_peak = peak_memory(GreedyBySizeAllocator().allocate(allocs)) + config = TabuSearchConfig(max_iterations=200) + searched = TabuSearchAllocator(config).allocate(allocs) + assert _is_valid(searched) + assert peak_memory(searched) <= greedy_peak diff --git a/tests/unit/allocators/test_telamalloc.py b/tests/unit/allocators/test_telamalloc.py new file mode 100644 index 0000000..2c654b3 --- /dev/null +++ b/tests/unit/allocators/test_telamalloc.py @@ -0,0 +1,145 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + +import pytest +from omnimalloc.allocators.greedy import GreedyBySizeAllocator +from omnimalloc.allocators.greedy_base import peak_memory +from omnimalloc.allocators.telamalloc import TelamallocAllocator, TelamallocConfig +from omnimalloc.primitives import Allocation, Pool +from omnimalloc.validate import validate_allocation + + +def _is_valid(result: tuple[Allocation, ...]) -> bool: + return validate_allocation(Pool(id="test_pool", allocations=result)) + + +def test_telamalloc_empty() -> None: + allocator = TelamallocAllocator() + result = allocator.allocate(()) + assert len(result) == 0 + + +def test_telamalloc_single() -> None: + allocator = TelamallocAllocator() + alloc = Allocation(id=1, size=100, start=0, end=10) + result = allocator.allocate((alloc,)) + assert len(result) == 1 + assert result[0].offset == 0 + + +def test_telamalloc_rejects_negative_backtracks() -> None: + with pytest.raises(ValueError, match="max_backtracks must be non-negative"): + TelamallocConfig(max_backtracks=-1) + + +def test_telamalloc_rejects_negative_seconds() -> None: + with pytest.raises(ValueError, match="max_seconds must be non-negative"): + TelamallocConfig(max_seconds=-1.0) + + +def test_telamalloc_no_temporal_overlap_shares_offset() -> None: + allocator = TelamallocAllocator() + alloc1 = Allocation(id=1, size=100, start=0, end=10) + alloc2 = Allocation(id=2, size=200, start=10, end=20) + result = allocator.allocate((alloc1, alloc2)) + by_id = {a.id: a for a in result} + assert by_id[1].offset == 0 + assert by_id[2].offset == 0 + + +def test_telamalloc_preserves_allocations() -> None: + allocator = TelamallocAllocator() + allocs = ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=50, start=5, end=15), + ) + result = allocator.allocate(allocs) + assert len(result) == len(allocs) + assert {a.id for a in result} == {1, 2} + assert all(a.offset is not None for a in result) + + +def test_telamalloc_preserves_input_order() -> None: + allocator = TelamallocAllocator() + allocs = tuple( + Allocation(id=i, size=(i % 3 + 1) * 10, start=0, end=10) for i in range(6) + ) + result = allocator.allocate(allocs) + assert [a.id for a in result] == [a.id for a in allocs] + + +def test_telamalloc_all_overlap_stacks_sequentially() -> None: + allocator = TelamallocAllocator() + allocs = tuple(Allocation(id=i, size=100, start=0, end=10) for i in range(5)) + result = allocator.allocate(allocs) + assert _is_valid(result) + assert peak_memory(result) == 500 + + +def test_telamalloc_overlapping_reach_lower_bound() -> None: + config = TelamallocConfig(max_seconds=0) + result = TelamallocAllocator(config).allocate( + ( + Allocation(id=1, size=100, start=0, end=10), + Allocation(id=2, size=50, start=5, end=15), + Allocation(id=3, size=25, start=0, end=15), + ) + ) + assert _is_valid(result) + assert peak_memory(result) == 175 + + +def test_telamalloc_independent_phases_share_address_space() -> None: + allocator = TelamallocAllocator() + early = tuple(Allocation(id=i, size=100, start=0, end=10) for i in range(3)) + late = tuple(Allocation(id=10 + i, size=100, start=20, end=30) for i in range(3)) + result = allocator.allocate(early + late) + assert _is_valid(result) + assert peak_memory(result) == 300 + offsets = sorted(a.offset for a in result if a.id >= 10) + assert offsets == [0, 100, 200] + + +def test_telamalloc_deterministic_without_deadline() -> None: + allocs = tuple( + Allocation(id=i, size=(i % 5 + 1) * 100, start=i % 3, end=i % 3 + i % 7 + 1) + for i in range(20) + ) + config = TelamallocConfig(max_seconds=0) + result1 = TelamallocAllocator(config).allocate(allocs) + result2 = TelamallocAllocator(config).allocate(allocs) + assert all(r1.offset == r2.offset for r1, r2 in zip(result1, result2, strict=True)) + + +def test_telamalloc_produces_valid_allocation_on_dense_overlap() -> None: + allocator = TelamallocAllocator() + allocs = tuple( + Allocation( + id=i, + size=(i % 7 + 1) * 10, + start=i % 4, + end=i % 4 + (i % 3 + 1) * 3, + ) + for i in range(30) + ) + result = allocator.allocate(allocs) + assert _is_valid(result) + assert {a.id for a in result} == {a.id for a in allocs} + + +def test_telamalloc_matches_or_beats_single_pass_greedy() -> None: + allocs = tuple( + Allocation( + id=i, + size=(i * 37 % 50 + 1) * 10, + start=i % 6, + end=i % 6 + (i * 13 % 5 + 1), + ) + for i in range(40) + ) + greedy_peak = peak_memory(GreedyBySizeAllocator().allocate(allocs)) + config = TelamallocConfig(max_seconds=0) + result = TelamallocAllocator(config).allocate(allocs) + assert _is_valid(result) + assert peak_memory(result) <= greedy_peak