Skip to content

Cut import time and hot path overhead, dedup the first-fit kernels #70

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
fpedd wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Cut import time and hot path overhead, dedup the first-fit kernels #70

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions src/cpp/allocators/best_fit.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,7 @@ namespace omnimalloc {

namespace {

int64_t find_best_fit_offset(
int64_t size, const std::vector<std::pair<int64_t, int64_t>>& spans) {
int64_t find_best_fit_offset(int64_t size, const std::vector<Interval>& spans) {
// Scan every gap between the sorted placed spans, keep the smallest that fits
int64_t cursor = 0;
int64_t best_offset = 0;
Expand All @@ -37,7 +36,7 @@ std::vector<Allocation> best_fit_place(
const std::vector<Allocation>& allocations) {
// Lambda rather than the function pointer so the placement loop inlines
// the offset scan instead of an indirect call per allocation
return place_indexed(allocations, compute_conflict_indices(allocations),
return place_indexed(allocations, build_conflict_adjacency(allocations),
[](int64_t size, const auto& spans) {
return find_best_fit_offset(size, spans);
});
Expand Down
105 changes: 42 additions & 63 deletions src/cpp/allocators/first_fit.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,6 @@

namespace omnimalloc {

namespace {

// Occupied (offset, end) span of a placed allocation, matching the span
// shape that `first_fit_offset` consumes
using Interval = std::pair<int64_t, int64_t>;

// LSD radix sort by offset (the end rides along as payload; equal-offset order
// is irrelevant to the gap scan). Replaces the comparison sort that dominated
// first-fit at scale; pass count scales with the actual offset magnitude.
void sort_intervals_by_lo(std::vector<Interval>& intervals,
std::vector<Interval>& scratch) {
const size_t m = intervals.size();
Expand Down Expand Up @@ -68,14 +59,15 @@ void sort_intervals_by_lo(std::vector<Interval>& intervals,
}
}

namespace {

// First-fit offsets for the allocations taken in `order`, gathering each
// allocation's placed CSR neighbors and reusing the shared gap scan. A
// non-negative `pins[i]` fixes i there, an obstacle before the first scan.
std::vector<int64_t> place_order(const CsrAdjacency& adj,
const std::vector<int64_t>& sizes,
const std::vector<int64_t>& pins,
const std::vector<int32_t>& order) {
constexpr Interval kUnplaced{-1, -1};
std::vector<int64_t> offsets(sizes.size(), -1);
std::vector<Interval> placed(sizes.size(), kUnplaced);
for (size_t i = 0; i < sizes.size(); ++i) {
Expand All @@ -87,21 +79,14 @@ std::vector<int64_t> place_order(const CsrAdjacency& adj,
std::vector<Interval> intervals;
std::vector<Interval> scratch;
for (const int32_t idx : order) {
if (pins[static_cast<size_t>(idx)] >= 0) {
const auto i = static_cast<size_t>(idx);
if (pins[i] >= 0) {
continue;
}
intervals.clear();
for (int64_t e = adj.offsets[idx]; e < adj.offsets[idx + 1]; ++e) {
const Interval span =
placed[static_cast<size_t>(adj.neighbors[static_cast<size_t>(e)])];
if (span.first >= 0) {
intervals.push_back(span);
}
}
sort_intervals_by_lo(intervals, scratch);
const int64_t best = first_fit_offset(sizes[idx], intervals);
offsets[idx] = best;
placed[static_cast<size_t>(idx)] = {best, best + sizes[idx]};
gather_placed_spans(adj.row(i), placed, intervals, scratch);
const int64_t best = first_fit_offset(sizes[i], intervals);
offsets[i] = best;
placed[i] = {best, best + sizes[i]};
}
return offsets;
}
Expand Down Expand Up @@ -301,21 +286,7 @@ PortfolioPlacement place_portfolio(const std::vector<Allocation>& allocations,
return best;
}

void gather_spans(const std::vector<size_t>& neighbors,
const std::vector<std::optional<int64_t>>& offsets,
const std::vector<Allocation>& allocations,
std::vector<std::pair<int64_t, int64_t>>& spans) {
spans.clear();
for (size_t j : neighbors) {
if (offsets[j].has_value()) {
spans.emplace_back(*offsets[j], *offsets[j] + allocations[j].size());
}
}
std::sort(spans.begin(), spans.end());
}

int64_t first_fit_offset(
int64_t size, const std::vector<std::pair<int64_t, int64_t>>& spans) {
int64_t first_fit_offset(int64_t size, const std::vector<Interval>& spans) {
int64_t best_offset = 0;
for (const auto& [offset, end] : spans) {
if (offset - best_offset >= size) {
Expand All @@ -327,26 +298,31 @@ int64_t first_fit_offset(
}

std::vector<Allocation> first_fit_place_indexed(
const std::vector<Allocation>& allocations,
const ConflictIndices& indices) {
const std::vector<Allocation>& allocations, const CsrAdjacency& adj) {
// Lambda rather than the function pointer so the placement loop inlines
// the offset scan instead of an indirect call per allocation
return place_indexed(allocations, indices,
[](int64_t size, const auto& spans) {
return first_fit_offset(size, spans);
});
return place_indexed(allocations, adj, [](int64_t size, const auto& spans) {
return first_fit_offset(size, spans);
});
}

std::vector<Allocation> first_fit_place(
const std::vector<Allocation>& allocations) {
return first_fit_place_indexed(allocations,
compute_conflict_indices(allocations));
build_conflict_adjacency(allocations));
}

FirstFitPlacer::FirstFitPlacer(std::vector<Allocation> allocations)
: allocations_(std::move(allocations)),
indices_(compute_conflict_indices(allocations_)) {
adj_(build_conflict_adjacency(allocations_)) {
check_total_size(allocations_);
const size_t n = allocations_.size();
sizes_.resize(n);
pins_.resize(n);
std::ranges::transform(allocations_, sizes_.begin(), &Allocation::size);
std::ranges::transform(allocations_, pins_.begin(), [](const Allocation& a) {
return a.offset().value_or(-1); // -1 marks a free allocation
});
}

void FirstFitPlacer::check_order(const std::vector<size_t>& order) const {
Expand All @@ -365,43 +341,46 @@ void FirstFitPlacer::check_order(const std::vector<size_t>& order) const {
}
}

std::vector<std::optional<int64_t>> FirstFitPlacer::place_offsets(
std::vector<Interval> FirstFitPlacer::place_spans(
const std::vector<size_t>& order) const {
// Pre-set offsets are pins: obstacles from the first scan, never re-placed
std::vector<std::optional<int64_t>> offsets(allocations_.size());
for (size_t i = 0; i < allocations_.size(); ++i) {
offsets[i] = allocations_[i].offset();
std::vector<Interval> placed(sizes_.size(), kUnplaced);
for (size_t i = 0; i < sizes_.size(); ++i) {
if (pins_[i] >= 0) {
placed[i] = {pins_[i], pins_[i] + sizes_[i]};
}
}
std::vector<std::pair<int64_t, int64_t>> spans;
for (size_t idx : order) {
const Allocation& alloc = allocations_[idx];
if (alloc.offset().has_value()) {
std::vector<Interval> spans;
std::vector<Interval> scratch;
for (const size_t idx : order) {
if (pins_[idx] >= 0) {
continue;
}
gather_spans(indices_[idx], offsets, allocations_, spans);
offsets[idx] = first_fit_offset(alloc.size(), spans);
gather_placed_spans(adj_.row(idx), placed, spans, scratch);
const int64_t offset = first_fit_offset(sizes_[idx], spans);
placed[idx] = {offset, offset + sizes_[idx]};
}
return offsets;
return placed;
}

std::vector<Allocation> FirstFitPlacer::place(
const std::vector<size_t>& order) const {
check_order(order);
const auto offsets = place_offsets(order);
const std::vector<Interval> spans = place_spans(order);
std::vector<Allocation> placed;
placed.reserve(order.size());
for (size_t idx : order) {
placed.push_back(allocations_[idx].with_offset(*offsets[idx]));
for (const size_t idx : order) {
placed.push_back(allocations_[idx].with_offset(spans[idx].first));
}
return placed;
}

int64_t FirstFitPlacer::peak(const std::vector<size_t>& order) const {
check_order(order);
const auto offsets = place_offsets(order);
const std::vector<Interval> spans = place_spans(order);
int64_t peak = 0;
for (size_t idx : order) {
peak = std::max(peak, *offsets[idx] + allocations_[idx].size());
for (const size_t idx : order) {
peak = std::max(peak, spans[idx].second);
}
return peak;
}
Expand Down
97 changes: 63 additions & 34 deletions src/cpp/allocators/first_fit.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,38 @@

namespace omnimalloc {

// Occupied (offset, end) spans of the already-placed neighbors of one
// allocation, sorted by offset so the gap scans can go left-to-right
void gather_spans(const std::vector<size_t>& neighbors,
const std::vector<std::optional<int64_t>>& offsets,
const std::vector<Allocation>& allocations,
std::vector<std::pair<int64_t, int64_t>>& spans);
// Occupied (offset, end) span of a placed allocation, matching the span
// shape that `first_fit_offset` consumes; kUnplaced marks an unplaced one
// (offsets are validated non-negative).
using Interval = std::pair<int64_t, int64_t>;
inline constexpr Interval kUnplaced{-1, -1};

// LSD radix sort by offset (the end rides along as payload; equal-offset order
// is irrelevant to the gap scans). Replaces the comparison sort that dominated
// first-fit at scale; pass count scales with the actual offset magnitude.
void sort_intervals_by_lo(std::vector<Interval>& intervals,
std::vector<Interval>& scratch);

// Occupied spans of the already-placed neighbors of one allocation, sorted by
// offset so the gap scans can go left-to-right; `scratch` backs the sort.
template <typename Neighbors>
void gather_placed_spans(const Neighbors& neighbors,
const std::vector<Interval>& placed,
std::vector<Interval>& spans,
std::vector<Interval>& scratch) {
spans.clear();
for (const auto neighbor : neighbors) {
const Interval span = placed[static_cast<size_t>(neighbor)];
if (span.first >= 0) {
spans.push_back(span);
}
}
sort_intervals_by_lo(spans, scratch);
}

// First-fit: lowest offset where `size` fits between the sorted spans
[[nodiscard]] int64_t first_fit_offset(
int64_t size, const std::vector<std::pair<int64_t, int64_t>>& spans);
[[nodiscard]] int64_t first_fit_offset(int64_t size,
const std::vector<Interval>& spans);

// Offsets (aligned with `allocations`) and peak of the winning placement.
struct PortfolioPlacement {
Expand All @@ -40,43 +62,50 @@ struct PortfolioPlacement {

// Greedily place allocations in input order using first-fit; computes the
// conflict relation natively (unbudgeted by design: placement kernels never
// give up mid-run). Map reuse across many orders is FirstFitPlacer's job.
// give up mid-run). Adjacency reuse across many orders is FirstFitPlacer's job.
[[nodiscard]] std::vector<Allocation> first_fit_place(
const std::vector<Allocation>& allocations);

// Greedily place allocations in order using first-fit over an index-based
// Greedily place allocations in order using first-fit over the CSR conflict
// adjacency (the fast path: each step only visits the allocation's neighbors)
[[nodiscard]] std::vector<Allocation> first_fit_place_indexed(
const std::vector<Allocation>& allocations, const ConflictIndices& indices);
const std::vector<Allocation>& allocations, const CsrAdjacency& adj);

// Shared placement skeleton of the first-fit and best-fit placers: place in
// index order, choosing each offset with `choose_offset` over the sorted spans
// of already-placed neighbors. Seeding pins makes them obstacles throughout.
template <typename OffsetFn>
[[nodiscard]] std::vector<Allocation> place_indexed(
const std::vector<Allocation>& allocations, const ConflictIndices& indices,
const std::vector<Allocation>& allocations, const CsrAdjacency& adj,
OffsetFn choose_offset) {
check_total_size(allocations);
std::vector<std::optional<int64_t>> offsets(allocations.size());
for (size_t i = 0; i < allocations.size(); ++i) {
offsets[i] = allocations[i].offset();
const size_t n = allocations.size();
std::vector<Interval> placed(n, kUnplaced);
for (size_t i = 0; i < n; ++i) {
if (const std::optional<int64_t> pin = allocations[i].offset()) {
placed[i] = {*pin, *pin + allocations[i].size()};
}
}
std::vector<std::pair<int64_t, int64_t>> spans;
std::vector<Allocation> placed;
placed.reserve(allocations.size());
for (size_t i = 0; i < allocations.size(); ++i) {
if (!allocations[i].offset().has_value()) {
gather_spans(indices[i], offsets, allocations, spans);
offsets[i] = choose_offset(allocations[i].size(), spans);
std::vector<Interval> spans;
std::vector<Interval> scratch;
std::vector<Allocation> result;
result.reserve(n);
for (size_t i = 0; i < n; ++i) {
if (placed[i].first < 0) {
gather_placed_spans(adj.row(i), placed, spans, scratch);
const int64_t size = allocations[i].size();
const int64_t offset = choose_offset(size, spans);
placed[i] = {offset, offset + size};
}
placed.push_back(allocations[i].with_offset(*offsets[i]));
result.push_back(allocations[i].with_offset(placed[i].first));
}
return placed;
return result;
}

// Resident first-fit placer for the order-search allocators (genetic, random,
// hill-climb): owns the allocations and their conflict maps, so placing many
// candidate orders passes only an index permutation across the Python boundary.
// hill-climb): owns the allocations, their CSR adjacency, and flat size/pin
// snapshots, so placing many candidate orders passes only an index permutation
// across the Python boundary.
class FirstFitPlacer {
public:
explicit FirstFitPlacer(std::vector<Allocation> allocations);
Expand All @@ -90,23 +119,23 @@ class FirstFitPlacer {
[[nodiscard]] std::vector<Allocation> place(
const std::vector<size_t>& order) const;

// The resident index adjacency, for the local searches' inner loops.
[[nodiscard]] const ConflictIndices& indices() const noexcept {
return indices_;
}
// The resident CSR adjacency, for the local searches' inner loops.
[[nodiscard]] const CsrAdjacency& adjacency() const noexcept { return adj_; }

private:
// Throw std::invalid_argument unless every index in `order` is in range
// and no index repeats.
void check_order(const std::vector<size_t>& order) const;

// Offsets (indexed like allocations_) of a first-fit placement in `order`;
// assumes `order` has been checked.
[[nodiscard]] std::vector<std::optional<int64_t>> place_offsets(
// Placed spans (indexed like allocations_) of a first-fit placement in
// `order`; assumes `order` has been checked.
[[nodiscard]] std::vector<Interval> place_spans(
const std::vector<size_t>& order) const;

std::vector<Allocation> allocations_;
ConflictIndices indices_;
CsrAdjacency adj_;
std::vector<int64_t> sizes_;
std::vector<int64_t> pins_; // -1 marks a free allocation
};

} // namespace omnimalloc
10 changes: 5 additions & 5 deletions src/cpp/allocators/local_search.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,13 +43,13 @@ std::vector<size_t> initial_order(const std::vector<Allocation>& allocations) {

std::vector<size_t> earlier_neighbors(const std::vector<size_t>& order,
size_t target_pos,
const ConflictIndices& indices) {
const CsrAdjacency& adj) {
// Mark the target's conflicts, then keep the earlier positions holding
// one; the marks cost a pass over the order, the alternative a hash
// lookup per earlier position
std::vector<char> conflicting(order.size(), 0);
for (size_t other : indices[order[target_pos]]) {
conflicting[other] = 1;
for (const int32_t other : adj.row(order[target_pos])) {
conflicting[static_cast<size_t>(other)] = 1;
}
std::vector<size_t> neighbors;
for (size_t pos = 0; pos < target_pos; ++pos) {
Expand All @@ -66,12 +66,12 @@ std::vector<size_t> earlier_neighbors(const std::vector<size_t>& order,

std::optional<std::pair<size_t, size_t>> propose_peak_swap(
const std::vector<size_t>& peaks, const std::vector<size_t>& order,
const ConflictIndices& indices, std::mt19937_64& rng) {
const CsrAdjacency& adj, std::mt19937_64& rng) {
assert(!peaks.empty()); // full placements always attain their peak
std::uniform_int_distribution<size_t> pick_peak(0, peaks.size() - 1);
const size_t target_pos = peaks[pick_peak(rng)];
const std::vector<size_t> neighbors =
earlier_neighbors(order, target_pos, indices);
earlier_neighbors(order, target_pos, adj);
if (neighbors.empty()) {
return std::nullopt;
}
Expand Down
Loading
Loading