Add best-fit, simulated annealing, tabu search, and telamalloc allocators

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)
 

src/cpp/allocators/best_fit.cpp

@@ -0,0 +1,60 @@
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "best_fit.hpp"
+
+#include <algorithm>
+
+namespace omnimalloc {
+
+namespace {
+
+int64_t find_best_fit_offset(const Allocation& current_alloc,
+                             const std::vector<Allocation>& 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<Allocation> BestFitAllocator::allocate(
+    const std::vector<Allocation>& allocations) const {
+  const TemporalOverlaps overlaps = compute_temporal_overlaps(allocations);
+  std::vector<Allocation> 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<omnimalloc::BestFitAllocator>::operator()(
+    const omnimalloc::BestFitAllocator&) const noexcept {
+  // Stateless class - all instances are equal, use constant hash
+  return 0x517cc1b7;  // arbitrary constant
+}
+
+}  // namespace std

src/cpp/allocators/best_fit.hpp

@@ -0,0 +1,36 @@
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <vector>
+
+#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<Allocation> allocate(
+      const std::vector<Allocation>& allocations) const;
+
+  bool operator==(const BestFitAllocator&) const noexcept = default;
+};
+
+}  // namespace omnimalloc
+
+namespace std {
+template <>
+struct hash<omnimalloc::BestFitAllocator> {
+  size_t operator()(const omnimalloc::BestFitAllocator&) const noexcept;
+};
+}  // namespace std

src/cpp/allocators/greedy_base.cpp

@@ -5,6 +5,7 @@
 #include "greedy_base.hpp"
 
 #include <algorithm>
+#include <numeric>
 #include <tuple>
 #include <utility>
 
@@ -40,32 +41,78 @@ TemporalOverlaps compute_temporal_overlaps(
   return overlaps;
 }
 
-namespace {
+int64_t peak_of(const std::vector<Allocation>& 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<Allocation>& placed_allocations,
-                         const TemporalOverlaps& overlaps) {
+std::vector<const Allocation*> placed_overlapping(
+    const Allocation& alloc, const std::vector<Allocation>& placed,
+    const TemporalOverlaps& overlaps) {
   // Collect overlapping allocations that have been placed
   std::vector<const Allocation*> 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<size_t> initial_order(const std::vector<Allocation>& allocations) {
+  std::vector<size_t> 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<size_t> earlier_neighbors(
+    const std::vector<size_t>& order, size_t target_pos,
+    const std::vector<Allocation>& allocations,
+    const TemporalOverlaps& overlaps) {
+  std::vector<size_t> 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<Allocation>& 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<Allocation> FirstFitPlacer::place(
 }
 
 int64_t FirstFitPlacer::evaluate(const std::vector<size_t>& 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

src/cpp/allocators/greedy_base.hpp

@@ -29,6 +29,27 @@ using TemporalOverlaps =
     const std::vector<Allocation>& allocations,
     const TemporalOverlaps& overlaps);
 
+// Peak memory (highest end offset) across the placed allocations
+[[nodiscard]] int64_t peak_of(const std::vector<Allocation>& placed);
+
+// Already-placed allocations that temporally overlap `alloc`, sorted by
+// offset so callers can scan the free gaps left-to-right
+[[nodiscard]] std::vector<const Allocation*> placed_overlapping(
+    const Allocation& alloc, const std::vector<Allocation>& placed,
+    const TemporalOverlaps& overlaps);
+
+// Indices into `allocations`, sorted largest-size-first: a decent, cheap
+// starting order for the order-search allocators.
+[[nodiscard]] std::vector<size_t> initial_order(
+    const std::vector<Allocation>& 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<size_t> earlier_neighbors(
+    const std::vector<size_t>& order, size_t target_pos,
+    const std::vector<Allocation>& 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

src/cpp/allocators/simulated_annealing.cpp

@@ -0,0 +1,99 @@
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "simulated_annealing.hpp"
+
+#include <chrono>
+#include <cmath>
+#include <random>
+#include <utility>
+
+#include "greedy_base.hpp"
+
+namespace omnimalloc {
+
+SimulatedAnnealingAllocator::SimulatedAnnealingAllocator(
+    SimulatedAnnealingConfig config)
+    : config_(config) {}
+
+std::vector<Allocation> SimulatedAnnealingAllocator::allocate(
+    const std::vector<Allocation>& allocations) const {
+  const FirstFitPlacer placer(allocations);
+  std::vector<size_t> order = initial_order(allocations);
+  if (allocations.size() < 2) {
+    return placer.place(order);
+  }
+
+  std::vector<Allocation> current_placed = placer.place(order);
+  int64_t current_peak = peak_of(current_placed);
+  std::vector<Allocation> best_placed = current_placed;
+  int64_t best_peak = current_peak;
+
+  std::mt19937_64 rng(config_.seed);
+  std::uniform_real_distribution<double> unit(0.0, 1.0);
+  double temperature = config_.initial_temperature;
+
+  const auto deadline = std::chrono::steady_clock::now() +
+                        std::chrono::duration<double>(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<size_t> 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<size_t> pick_peak(0,
+                                                    peak_positions.size() - 1);
+    size_t target_pos = peak_positions[pick_peak(rng)];
+    std::vector<size_t> neighbors =
+        earlier_neighbors(order, target_pos, allocations, placer.overlaps());
+    if (neighbors.empty()) {
+      temperature *= config_.cooling_rate;
+      continue;
+    }
+    std::uniform_int_distribution<size_t> pick_neighbor(0,
+                                                        neighbors.size() - 1);
+    size_t other_pos = neighbors[pick_neighbor(rng)];
+
+    std::swap(order[target_pos], order[other_pos]);
+    std::vector<Allocation> 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<double>(candidate_peak - current_peak) /
+          static_cast<double>(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

src/cpp/allocators/simulated_annealing.hpp

@@ -0,0 +1,47 @@
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+#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<Allocation> allocate(
+      const std::vector<Allocation>& allocations) const;
+
+ private:
+  SimulatedAnnealingConfig config_;
+};
+
+}  // namespace omnimalloc