diff --git a/src/cpp/allocators/best_fit.cpp b/src/cpp/allocators/best_fit.cpp index 4472c99..8818497 100644 --- a/src/cpp/allocators/best_fit.cpp +++ b/src/cpp/allocators/best_fit.cpp @@ -12,8 +12,7 @@ namespace omnimalloc { namespace { -int64_t find_best_fit_offset( - int64_t size, const std::vector>& spans) { +int64_t find_best_fit_offset(int64_t size, const std::vector& spans) { // Scan every gap between the sorted placed spans, keep the smallest that fits int64_t cursor = 0; int64_t best_offset = 0; @@ -37,7 +36,7 @@ std::vector best_fit_place( const std::vector& 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); }); diff --git a/src/cpp/allocators/first_fit.cpp b/src/cpp/allocators/first_fit.cpp index d6b987d..d5a20c7 100644 --- a/src/cpp/allocators/first_fit.cpp +++ b/src/cpp/allocators/first_fit.cpp @@ -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; - -// 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& intervals, std::vector& scratch) { const size_t m = intervals.size(); @@ -68,6 +59,8 @@ void sort_intervals_by_lo(std::vector& 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. @@ -75,7 +68,6 @@ std::vector place_order(const CsrAdjacency& adj, const std::vector& sizes, const std::vector& pins, const std::vector& order) { - constexpr Interval kUnplaced{-1, -1}; std::vector offsets(sizes.size(), -1); std::vector placed(sizes.size(), kUnplaced); for (size_t i = 0; i < sizes.size(); ++i) { @@ -87,21 +79,14 @@ std::vector place_order(const CsrAdjacency& adj, std::vector intervals; std::vector scratch; for (const int32_t idx : order) { - if (pins[static_cast(idx)] >= 0) { + const auto i = static_cast(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(adj.neighbors[static_cast(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(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; } @@ -301,21 +286,7 @@ PortfolioPlacement place_portfolio(const std::vector& allocations, return best; } -void gather_spans(const std::vector& neighbors, - const std::vector>& offsets, - const std::vector& allocations, - std::vector>& 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>& spans) { +int64_t first_fit_offset(int64_t size, const std::vector& spans) { int64_t best_offset = 0; for (const auto& [offset, end] : spans) { if (offset - best_offset >= size) { @@ -327,26 +298,31 @@ int64_t first_fit_offset( } std::vector first_fit_place_indexed( - const std::vector& allocations, - const ConflictIndices& indices) { + const std::vector& 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 first_fit_place( const std::vector& allocations) { return first_fit_place_indexed(allocations, - compute_conflict_indices(allocations)); + build_conflict_adjacency(allocations)); } FirstFitPlacer::FirstFitPlacer(std::vector 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& order) const { @@ -365,43 +341,46 @@ void FirstFitPlacer::check_order(const std::vector& order) const { } } -std::vector> FirstFitPlacer::place_offsets( +std::vector FirstFitPlacer::place_spans( const std::vector& order) const { // Pre-set offsets are pins: obstacles from the first scan, never re-placed - std::vector> offsets(allocations_.size()); - for (size_t i = 0; i < allocations_.size(); ++i) { - offsets[i] = allocations_[i].offset(); + std::vector 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> spans; - for (size_t idx : order) { - const Allocation& alloc = allocations_[idx]; - if (alloc.offset().has_value()) { + std::vector spans; + std::vector 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 FirstFitPlacer::place( const std::vector& order) const { check_order(order); - const auto offsets = place_offsets(order); + const std::vector spans = place_spans(order); std::vector 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& order) const { check_order(order); - const auto offsets = place_offsets(order); + const std::vector 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; } diff --git a/src/cpp/allocators/first_fit.hpp b/src/cpp/allocators/first_fit.hpp index e861375..8931e3b 100644 --- a/src/cpp/allocators/first_fit.hpp +++ b/src/cpp/allocators/first_fit.hpp @@ -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& neighbors, - const std::vector>& offsets, - const std::vector& allocations, - std::vector>& 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; +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& intervals, + std::vector& 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 +void gather_placed_spans(const Neighbors& neighbors, + const std::vector& placed, + std::vector& spans, + std::vector& scratch) { + spans.clear(); + for (const auto neighbor : neighbors) { + const Interval span = placed[static_cast(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>& spans); +[[nodiscard]] int64_t first_fit_offset(int64_t size, + const std::vector& spans); // Offsets (aligned with `allocations`) and peak of the winning placement. struct PortfolioPlacement { @@ -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 first_fit_place( const std::vector& 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 first_fit_place_indexed( - const std::vector& allocations, const ConflictIndices& indices); + const std::vector& 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 [[nodiscard]] std::vector place_indexed( - const std::vector& allocations, const ConflictIndices& indices, + const std::vector& allocations, const CsrAdjacency& adj, OffsetFn choose_offset) { check_total_size(allocations); - std::vector> 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 placed(n, kUnplaced); + for (size_t i = 0; i < n; ++i) { + if (const std::optional pin = allocations[i].offset()) { + placed[i] = {*pin, *pin + allocations[i].size()}; + } } - std::vector> spans; - std::vector 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 spans; + std::vector scratch; + std::vector 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 allocations); @@ -90,23 +119,23 @@ class FirstFitPlacer { [[nodiscard]] std::vector place( const std::vector& 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& order) const; - // Offsets (indexed like allocations_) of a first-fit placement in `order`; - // assumes `order` has been checked. - [[nodiscard]] std::vector> place_offsets( + // Placed spans (indexed like allocations_) of a first-fit placement in + // `order`; assumes `order` has been checked. + [[nodiscard]] std::vector place_spans( const std::vector& order) const; std::vector allocations_; - ConflictIndices indices_; + CsrAdjacency adj_; + std::vector sizes_; + std::vector pins_; // -1 marks a free allocation }; } // namespace omnimalloc diff --git a/src/cpp/allocators/local_search.cpp b/src/cpp/allocators/local_search.cpp index e28cd87..d7e3944 100644 --- a/src/cpp/allocators/local_search.cpp +++ b/src/cpp/allocators/local_search.cpp @@ -43,13 +43,13 @@ std::vector initial_order(const std::vector& allocations) { std::vector earlier_neighbors(const std::vector& 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 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(other)] = 1; } std::vector neighbors; for (size_t pos = 0; pos < target_pos; ++pos) { @@ -66,12 +66,12 @@ std::vector earlier_neighbors(const std::vector& order, std::optional> propose_peak_swap( const std::vector& peaks, const std::vector& 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 pick_peak(0, peaks.size() - 1); const size_t target_pos = peaks[pick_peak(rng)]; const std::vector neighbors = - earlier_neighbors(order, target_pos, indices); + earlier_neighbors(order, target_pos, adj); if (neighbors.empty()) { return std::nullopt; } diff --git a/src/cpp/allocators/local_search.hpp b/src/cpp/allocators/local_search.hpp index bfc0b1c..c0ba25c 100644 --- a/src/cpp/allocators/local_search.hpp +++ b/src/cpp/allocators/local_search.hpp @@ -33,13 +33,13 @@ namespace omnimalloc { // adjacency: this runs per sampled move, where hashing ids would dominate. [[nodiscard]] std::vector earlier_neighbors( const std::vector& order, size_t target_pos, - const ConflictIndices& indices); + const CsrAdjacency& adj); // One random peak-lowering move for the local searches: a random position // among `peaks` paired with a random earlier temporal neighbor, or nullopt // when the chosen target has no earlier position to swap with. [[nodiscard]] std::optional> propose_peak_swap( const std::vector& peaks, const std::vector& order, - const ConflictIndices& indices, std::mt19937_64& rng); + const CsrAdjacency& adj, std::mt19937_64& rng); } // namespace omnimalloc diff --git a/src/cpp/allocators/simulated_annealing.cpp b/src/cpp/allocators/simulated_annealing.cpp index 9a8384f..e652eec 100644 --- a/src/cpp/allocators/simulated_annealing.cpp +++ b/src/cpp/allocators/simulated_annealing.cpp @@ -43,7 +43,7 @@ std::vector simulated_annealing_place( peak_positions(current_placed, current_peak); const auto proposal = - propose_peak_swap(peaks, order, placer.indices(), rng); + propose_peak_swap(peaks, order, placer.adjacency(), rng); if (!proposal) { temperature *= config.cooling_rate; continue; diff --git a/src/cpp/allocators/supermalloc.cpp b/src/cpp/allocators/supermalloc.cpp index 07bd9b3..3de92b0 100644 --- a/src/cpp/allocators/supermalloc.cpp +++ b/src/cpp/allocators/supermalloc.cpp @@ -473,29 +473,25 @@ void Partition::order_indices(std::vector& indices, indices = std::move(sorted); } -Solution Partition::first_fit(const std::vector& order) const { - const int n = static_cast(data_->allocations.size()); - std::vector offsets(static_cast(n), -1); +Packing Partition::first_fit(const std::vector& order) const { + const size_t n = data_->allocations.size(); + std::vector offsets(n, -1); + std::vector placed(n, kUnplaced); int64_t height = 0; - std::vector> intervals; - for (int i : order) { - intervals.clear(); - for (int j : data_->overlaps[i]) { - if (offsets[j] >= 0) { - intervals.emplace_back(offsets[j], offsets[j] + data_->alloc_sizes[j]); - } - } - std::sort(intervals.begin(), intervals.end()); - + std::vector intervals; + std::vector scratch; + for (const int i : order) { + gather_placed_spans(data_->overlaps[i], placed, intervals, scratch); const int64_t size = data_->alloc_sizes[i]; const int64_t offset = first_fit_offset(size, intervals); - offsets[i] = offset; + offsets[static_cast(i)] = offset; + placed[static_cast(i)] = {offset, offset + size}; height = std::max(height, offset + size); } - return Solution{apply_offsets(data_->allocations, offsets), height}; + return Packing{std::move(offsets), height}; } -Solution Partition::greedy_pack(const std::string& heuristic) const { +Packing Partition::greedy_pack(const std::string& heuristic) const { std::vector order(data_->allocations.size()); std::iota(order.begin(), order.end(), 0); order_indices(order, heuristic, 0, static_cast(data_->sections.size())); @@ -948,7 +944,7 @@ Solution greedy_pack_portfolio(const Partition& partition, validate_heuristics(heuristics); const auto deadline = compute_deadline(timeout); - std::vector> results(heuristics.size()); + std::vector> results(heuristics.size()); std::atomic next{0}; // The first heuristic is packed regardless of the deadline so that at @@ -963,11 +959,13 @@ Solution greedy_pack_portfolio(const Partition& partition, }, num_threads, heuristics.size()); - std::optional best; + // Only the winner is materialized into placed Allocation copies. + std::optional best; for (auto& r : results) { - if (r && (!best || r->peak < best->peak)) best = std::move(*r); + if (r && (!best || r->height < best->height)) best = std::move(*r); } - return std::move(*best); + return Solution{apply_offsets(partition.allocations(), best->offsets), + best->height}; } std::optional try_solve_many(const std::vector& partitions, diff --git a/src/cpp/allocators/supermalloc.hpp b/src/cpp/allocators/supermalloc.hpp index 7231345..f16edb4 100644 --- a/src/cpp/allocators/supermalloc.hpp +++ b/src/cpp/allocators/supermalloc.hpp @@ -35,6 +35,13 @@ struct Solution { int64_t peak; }; +// One greedy packing as index-aligned offsets plus its peak height; +// materialized into a Solution only for a portfolio's winner. +struct Packing { + std::vector offsets; + int64_t height; +}; + // A temporal allocation problem: immutable structure (sections, overlaps, // spans) shared across copies, plus mutable search state (offsets, floors, // totals, best_height) that the hot loop updates via `apply_at`/`revert`. @@ -114,7 +121,7 @@ class Partition { // First-fit packing in `heuristic` order (empty keeps the input order): // each buffer takes the lowest gap among its already-placed overlaps. Cheap // incumbent for a fresh (fully unplaced) partition. - [[nodiscard]] Solution greedy_pack(const std::string& heuristic) const; + [[nodiscard]] Packing greedy_pack(const std::string& heuristic) const; // Reorder allocations by `heuristic`: each character is a descending sort // key (one of A, C, L, O, T, U, W, Z; throws otherwise), original index @@ -168,7 +175,7 @@ class Partition { int start, int end) const; // First-fit packing in `order`; body of `greedy_pack`. - [[nodiscard]] Solution first_fit(const std::vector& order) const; + [[nodiscard]] Packing first_fit(const std::vector& order) const; // Derive the incremental search state (candidates, tops, cuts) from // `offsets_`, `min_offsets_`, and the section spans. diff --git a/src/cpp/allocators/tabu_search.cpp b/src/cpp/allocators/tabu_search.cpp index 2211fea..2c3402f 100644 --- a/src/cpp/allocators/tabu_search.cpp +++ b/src/cpp/allocators/tabu_search.cpp @@ -58,20 +58,20 @@ std::vector tabu_search_place( // Sample a neighborhood of candidate swaps and keep the best admissible // one: non-tabu, or tabu but beating the best-ever solution (aspiration). + // Only the peak decides, so samples pay a peak scan, not a placement. 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) { - // Each sample costs a full placement, so the budget is read here too: - // checking only per iteration lets a whole neighborhood run past it + // Each sample costs a full first-fit pass, so the budget is read here + // too: checking only per iteration lets a whole neighborhood run past it if (deadline_expired(deadline)) { break; } const auto proposal = - propose_peak_swap(peaks, order, placer.indices(), rng); + propose_peak_swap(peaks, order, placer.adjacency(), rng); if (!proposal) { continue; } @@ -82,8 +82,7 @@ std::vector tabu_search_place( 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); + const int64_t candidate_peak = placer.peak(order); std::swap(order[target_pos], order[other_pos]); // undo; reapplied if chosen @@ -91,7 +90,6 @@ std::vector tabu_search_place( 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; @@ -102,8 +100,9 @@ std::vector tabu_search_place( continue; // every sampled move was tabu without meeting aspiration } + // One placement per iteration: peak_positions needs the winner realized std::swap(order[best_p1], order[best_p2]); - current_placed = std::move(best_candidate_placed); + current_placed = placer.place(order); current_peak = best_candidate_peak; if (!best_is_tabu) { tabu_until[tabu_key(order[best_p1], order[best_p2], n)] = diff --git a/src/cpp/allocators/telamalloc.cpp b/src/cpp/allocators/telamalloc.cpp index 246e5d8..75ba1f2 100644 --- a/src/cpp/allocators/telamalloc.cpp +++ b/src/cpp/allocators/telamalloc.cpp @@ -30,8 +30,8 @@ constexpr int64_t kUnbounded = std::numeric_limits::max() / 4; // 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 ConflictIndices& neighbors) { - const int n = static_cast(neighbors.size()); +std::vector> build_phases(const CsrAdjacency& adj) { + const int n = static_cast(adj.size()); std::vector> phases; std::vector visited(n, 0); for (int seed = 0; seed < n; ++seed) { @@ -45,10 +45,10 @@ std::vector> build_phases(const ConflictIndices& neighbors) { int idx = stack.back(); stack.pop_back(); phase.push_back(idx); - for (size_t other : neighbors[idx]) { - if (!visited[other]) { - visited[other] = 1; - stack.push_back(static_cast(other)); + for (const int32_t other : adj.row(static_cast(idx))) { + if (!visited[static_cast(other)]) { + visited[static_cast(other)] = 1; + stack.push_back(other); } } } @@ -88,10 +88,9 @@ QueueKey queue_key(const Allocation& alloc, int idx, int evictions, // fitting gap, evicting the cheapest blocking set on conflict and requeueing it // at raised priority; a spent eviction budget wipes and re-packs the phase. std::optional> pack_phase( - const std::vector& allocations, - const ConflictIndices& neighbors, const std::vector& phase, - int64_t capacity, int max_backtracks, const Deadline& deadline, - bool size_major, uint64_t seed) { + const std::vector& allocations, const CsrAdjacency& adj, + 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); @@ -106,6 +105,7 @@ std::optional> pack_phase( // Scratch buffers for the placement loop, reused across queue pops. std::vector> occupied; // (offset, neighbor) std::vector> spans; // (offset, end) + std::vector candidates; // repair offsets while (true) { for (int idx : phase) { @@ -130,9 +130,9 @@ std::optional> pack_phase( // spans inherit occupied's offset order (ties differ only in end // order, which cannot change a gap scan's result), so one sort does. occupied.clear(); - for (size_t other : neighbors[idx]) { - if (offsets[other] >= 0) { - occupied.emplace_back(offsets[other], static_cast(other)); + for (const int32_t other : adj.row(static_cast(idx))) { + if (offsets[static_cast(other)] >= 0) { + occupied.emplace_back(offsets[static_cast(other)], other); } } std::sort(occupied.begin(), occupied.end()); @@ -153,7 +153,8 @@ std::optional> pack_phase( // 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}; + candidates.clear(); + candidates.push_back(0); for (const auto& [offset, other] : occupied) { const int64_t end = offset + allocations[other].size(); if (end + size <= capacity) { @@ -232,16 +233,15 @@ int64_t phase_peak(const std::vector& allocations, // search on capacity down toward the load lower bound. Budget-exhausted // attempts count as infeasible, keeping the search anytime rather than exact. void solve_phase(const std::vector& allocations, - const ConflictIndices& neighbors, - const std::vector& phase, int64_t lower_bound, - const TelamallocConfig& config, const Deadline& deadline, - std::vector& result) { + const CsrAdjacency& adj, 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, + auto by_duration = pack_phase(allocations, adj, phase, kUnbounded, 0, std::nullopt, false, config.seed); - auto by_size = pack_phase(allocations, neighbors, phase, kUnbounded, 0, + auto by_size = pack_phase(allocations, adj, 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); @@ -258,7 +258,7 @@ void solve_phase(const std::vector& allocations, } const int64_t mid = low + (high - low) / 2; auto attempt = - pack_phase(allocations, neighbors, phase, mid, config.max_backtracks, + pack_phase(allocations, adj, phase, mid, config.max_backtracks, deadline, size_major, config.seed); if (attempt) { best = std::move(*attempt); @@ -284,13 +284,13 @@ std::vector telamalloc_place( // Bound by kUnbounded so the unbounded-capacity pack_phase incumbents in // solve_phase can never fail and the cursor arithmetic cannot overflow. check_total_size(allocations, kUnbounded); - const ConflictIndices neighbors = compute_conflict_indices(allocations); + const CsrAdjacency adj = build_conflict_adjacency(allocations); if (allocations.size() < 2) { - return first_fit_place_indexed(allocations, neighbors); + return first_fit_place_indexed(allocations, adj); } const Deadline deadline = make_deadline(config.timeout); - const auto phases = build_phases(neighbors); + const auto phases = build_phases(adj); // 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. @@ -305,8 +305,8 @@ std::vector telamalloc_place( std::vector result(allocations.size(), -1); for (const auto& [lower_bound, p] : order) { - solve_phase(allocations, neighbors, phases[p], lower_bound, config, - deadline, result); + solve_phase(allocations, adj, phases[p], lower_bound, config, deadline, + result); } return apply_offsets(allocations, result); } diff --git a/src/cpp/analysis/clock.hpp b/src/cpp/analysis/clock.hpp index f4e558b..be170a9 100644 --- a/src/cpp/analysis/clock.hpp +++ b/src/cpp/analysis/clock.hpp @@ -318,17 +318,22 @@ inline std::vector interval_peaks( const std::vector>& times, const std::vector& weights) { const std::vector bounds = slot_bounds(times); + // Compressed once, then reused by both the delta pass and the queries + std::vector> windows(times.size()); + for (size_t i = 0; i < times.size(); ++i) { + windows[i] = {slot_index(bounds, times[i].first), + slot_index(bounds, times[i].second)}; + } std::vector pressure(bounds.size(), 0); for (size_t i = 0; i < times.size(); ++i) { - pressure[slot_index(bounds, times[i].first)] += weights[i]; - pressure[slot_index(bounds, times[i].second)] -= weights[i]; + pressure[windows[i].first] += weights[i]; + pressure[windows[i].second] -= weights[i]; } std::partial_sum(pressure.begin(), pressure.end(), pressure.begin()); const MaxSegtree live(pressure); std::vector peaks(times.size()); for (size_t i = 0; i < times.size(); ++i) { - peaks[i] = live.max(slot_index(bounds, times[i].first), - slot_index(bounds, times[i].second)); + peaks[i] = live.max(windows[i].first, windows[i].second); } return peaks; } @@ -339,6 +344,12 @@ inline std::vector interval_peaks( struct CsrAdjacency { std::vector offsets; std::vector neighbors; + + [[nodiscard]] size_t size() const noexcept { return offsets.size() - 1; } + [[nodiscard]] std::span row(size_t index) const noexcept { + return {neighbors.data() + offsets[index], + static_cast(offsets[index + 1] - offsets[index])}; + } }; // Ceiling on the neighbor entries one adjacency may materialize, 4 bytes each. diff --git a/src/cpp/analysis/conflicts.cpp b/src/cpp/analysis/conflicts.cpp index 1386534..df1db0c 100644 --- a/src/cpp/analysis/conflicts.cpp +++ b/src/cpp/analysis/conflicts.cpp @@ -5,6 +5,7 @@ #include "conflicts.hpp" #include +#include #include #include #include @@ -37,16 +38,6 @@ void check_sweep_budget(const ConflictSweep& sweep, } } -// Pairwise happens-before adjacency unpacked from CSR form -ConflictIndices indices_from_adjacency(size_t n, const CsrAdjacency& adj) { - ConflictIndices indices(n); - for (size_t i = 0; i < n; ++i) { - indices[i].assign(adj.neighbors.begin() + adj.offsets[i], - adj.neighbors.begin() + adj.offsets[i + 1]); - } - return indices; -} - } // namespace CsrAdjacency build_conflict_adjacency( @@ -55,12 +46,6 @@ CsrAdjacency build_conflict_adjacency( return sweep.adjacency(parallel_threads(sweep.count())); } -ConflictIndices compute_conflict_indices( - const std::vector& allocations) { - return indices_from_adjacency(allocations.size(), - build_conflict_adjacency(allocations)); -} - ConflictGraph::ConflictGraph(const std::vector& allocations, std::optional work_budget, std::optional max_entries) { @@ -93,8 +78,8 @@ int64_t ConflictGraph::degree(size_t index) const { std::vector ConflictGraph::neighbors(size_t index) const { check_index(index); - return {adj_.neighbors.begin() + adj_.offsets[index], - adj_.neighbors.begin() + adj_.offsets[index + 1]}; + const std::span row = adj_.row(index); + return {row.begin(), row.end()}; } std::vector conflict_degrees( diff --git a/src/cpp/analysis/conflicts.hpp b/src/cpp/analysis/conflicts.hpp index 046c9f0..2be9508 100644 --- a/src/cpp/analysis/conflicts.hpp +++ b/src/cpp/analysis/conflicts.hpp @@ -13,13 +13,6 @@ namespace omnimalloc { -// Index-based conflict adjacency: position i -> positions conflicting with i -using ConflictIndices = std::vector>; - -// Map each allocation index to the indices of conflicting allocations -[[nodiscard]] ConflictIndices compute_conflict_indices( - const std::vector& allocations); - // Per-allocation count of conflicting allocations, aligned with `allocations` // and counted with multiplicity. Scalar timelines count in O(N log N) without // enumerating pairs; on vector clocks `work_budget` bounds the sweep. @@ -44,7 +37,7 @@ class ConflictGraph { std::optional work_budget, std::optional max_entries = std::nullopt); - [[nodiscard]] size_t size() const noexcept { return adj_.offsets.size() - 1; } + [[nodiscard]] size_t size() const noexcept { return adj_.size(); } // Conflicting pairs in the relation, each counted once [[nodiscard]] uint64_t pair_count() const noexcept { return adj_.neighbors.size() / 2; diff --git a/src/cpp/analysis/linearize.cpp b/src/cpp/analysis/linearize.cpp index bcd57ed..45a0261 100644 --- a/src/cpp/analysis/linearize.cpp +++ b/src/cpp/analysis/linearize.cpp @@ -169,21 +169,18 @@ std::optional>> linearize_times( } // Ranks: distinct counts ascending. On a chain, equal counts mean equal - // predecessor sets, so they share a rank and any member represents it. + // predecessor sets, so they share a rank (the last pushed one) and any + // member represents it. std::vector unique_counts; std::vector representative; + std::vector start_rank(k); for (size_t pos = 0; pos < k; ++pos) { const auto si = static_cast(by_count[pos]); if (unique_counts.empty() || counts[si] != unique_counts.back()) { unique_counts.push_back(counts[si]); representative.push_back(by_count[pos]); } - } - std::vector start_rank(k); - for (size_t si = 0; si < k; ++si) { - start_rank[si] = std::lower_bound(unique_counts.begin(), - unique_counts.end(), counts[si]) - - unique_counts.begin(); + start_rank[si] = static_cast(unique_counts.size()) - 1; } // End rank: smallest rank whose representative start dominates the end, diff --git a/src/python/omnimalloc/__init__.py b/src/python/omnimalloc/__init__.py index bc6f1b6..7dae607 100644 --- a/src/python/omnimalloc/__init__.py +++ b/src/python/omnimalloc/__init__.py @@ -4,12 +4,8 @@ """Static memory allocation: place buffers with temporal bounds, minimize peak.""" -from importlib.metadata import version as _version +import typing -__version__ = _version("omnimalloc") - -from ._allocate import allocate as allocate -from .allocators import available_allocators as available_allocators from .common.parallel import max_threads as max_threads from .common.parallel import set_max_threads as set_max_threads from .primitives import Allocation as Allocation @@ -21,4 +17,30 @@ from .primitives import TimePoint as TimePoint from .primitives import VectorClock as VectorClock from .validate import validate_allocation as validate_allocation -from .visualize import plot_allocation as plot_allocation + +if typing.TYPE_CHECKING: + from ._allocate import allocate as allocate + from .allocators import available_allocators as available_allocators + from .visualize import plot_allocation as plot_allocation + +# Heavy submodules load on first attribute access to keep `import omnimalloc` fast. +_LAZY_EXPORTS: dict[str, str] = { + "allocate": "._allocate", + "available_allocators": ".allocators", + "plot_allocation": ".visualize", +} + + +def __getattr__(name: str) -> object: + if name == "__version__": + from importlib.metadata import version + + value: object = version("omnimalloc") + elif name in _LAZY_EXPORTS: + from importlib import import_module + + value = getattr(import_module(_LAZY_EXPORTS[name], __name__), name) + else: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + globals()[name] = value + return value diff --git a/src/python/omnimalloc/allocators/base.py b/src/python/omnimalloc/allocators/base.py index 4333df7..63d8044 100644 --- a/src/python/omnimalloc/allocators/base.py +++ b/src/python/omnimalloc/allocators/base.py @@ -54,12 +54,15 @@ def allocate( def _ensure_preconditions( self, allocations: tuple["Allocation", ...] - ) -> dict["IdType", int | None]: + ) -> dict["IdType", int]: """Shared entry contract; returns the pinned offsets keyed by id.""" ensure_unique_ids(allocations, "allocation") uniform_dim(allocations) self.ensure_supported(allocations) - pins = {alloc.id: alloc.offset for alloc in allocations if alloc.is_allocated} + pins: dict[IdType, int] = {} + for alloc in allocations: + if alloc.offset is not None: + pins[alloc.id] = alloc.offset if pins: self._ensure_pins_placeable(allocations) return pins @@ -68,7 +71,7 @@ def _ensure_postconditions( self, allocations: tuple["Allocation", ...], placed: tuple["Allocation", ...], - pins: dict["IdType", int | None], + pins: dict["IdType", int], ) -> None: """Shared exit contract: same set, fully placed, pins untouched.""" self._ensure_same_set(allocations, placed) @@ -128,7 +131,7 @@ def _ensure_same_set( raise ValueError(f"{self.name()} returned a different allocation set") def _ensure_placed( - self, placed: tuple["Allocation", ...], pins: dict["IdType", int | None] + self, placed: tuple["Allocation", ...], pins: dict["IdType", int] ) -> None: """Every allocation comes back placed, and every pin comes back put.""" unplaced = next((alloc for alloc in placed if alloc.offset is None), None) diff --git a/src/python/omnimalloc/allocators/genetic.py b/src/python/omnimalloc/allocators/genetic.py index 190368b..338c24b 100644 --- a/src/python/omnimalloc/allocators/genetic.py +++ b/src/python/omnimalloc/allocators/genetic.py @@ -19,9 +19,11 @@ from .greedy import ( GreedyAllocator, + _conflict_load, + _conflict_size_load, + _order_by_load, + _sort_degrees, order_by_area, - order_by_conflict, - order_by_conflict_size, order_by_duration, order_by_size, order_by_start, @@ -41,6 +43,13 @@ _GLOBAL_RNG_LOCK = threading.Lock() +def _evaluate_permutation( + permutation: list[int], placer: FirstFitPlacer +) -> tuple[float]: + """Evaluate a permutation by computing its greedy peak memory usage.""" + return (float(placer.peak(permutation)),) + + class GeneticAllocator(GreedyAllocator): """Genetic algorithm allocator that evolves greedy placement orders. @@ -91,27 +100,22 @@ def __init__( fitness=creator.OmnimallocFitnessMin, # ty: ignore[unresolved-attribute] ) - def _evaluate_permutation( - self, permutation: list[int], placer: FirstFitPlacer - ) -> tuple[float]: - """Evaluate a permutation by computing its greedy peak memory usage.""" - return (float(placer.peak(permutation)),) - def _heuristic_permutations( self, allocations: tuple[Allocation, ...] ) -> list[list[int]]: """Create seed permutations mirroring the greedy sort heuristics.""" - orders = ( - order_by_size, - order_by_duration, - order_by_area, - order_by_conflict, - order_by_conflict_size, - order_by_start, + degrees = _sort_degrees(allocations) + orderings = ( + order_by_size(allocations), + order_by_duration(allocations), + order_by_area(allocations), + _order_by_load(allocations, degrees, _conflict_load), + _order_by_load(allocations, degrees, _conflict_size_load), + order_by_start(allocations), ) positions = {alloc.id: i for i, alloc in enumerate(allocations)} permutations = [ - [positions[alloc.id] for alloc in order(allocations)] for order in orders + [positions[alloc.id] for alloc in ordering] for ordering in orderings ] return permutations[: self._population_size] @@ -146,7 +150,7 @@ def _evolve(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...] creator.OmnimallocIndividual, # ty: ignore[unresolved-attribute] toolbox.indices, # ty: ignore[unresolved-attribute] ) - toolbox.register("evaluate", self._evaluate_permutation, placer=placer) + toolbox.register("evaluate", _evaluate_permutation, placer=placer) toolbox.register("mate", tools.cxOrdered) toolbox.register("mutate", tools.mutShuffleIndexes, indpb=0.05) # TODO(fpedd): Try larger tournsize and selNSGA2 diff --git a/src/python/omnimalloc/allocators/greedy.py b/src/python/omnimalloc/allocators/greedy.py index 10fb6c5..531cade 100644 --- a/src/python/omnimalloc/allocators/greedy.py +++ b/src/python/omnimalloc/allocators/greedy.py @@ -3,9 +3,10 @@ # import logging +from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor -from omnimalloc._cpp import first_fit_place +from omnimalloc._cpp import FirstFitPlacer, first_fit_place from omnimalloc.analysis import conflict_degrees, placement_pressure from omnimalloc.analysis._clock import time_components, uniform_dim from omnimalloc.common.constants import DEFAULT_WORK_BUDGET @@ -38,28 +39,38 @@ def order_by_area(allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...] return tuple(sorted(allocations, key=lambda a: a.area, reverse=True)) -def order_by_conflict(allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: - """Order by conflict degree (most conflicted first).""" - degrees = _sort_degrees(allocations) +def _conflict_load(_alloc: Allocation, degree: int) -> int: + return degree + + +def _conflict_size_load(alloc: Allocation, degree: int) -> int: + return degree * alloc.size + + +def _order_by_load( + allocations: tuple[Allocation, ...], + degrees: list[int], + load: Callable[[Allocation, int], int], +) -> tuple[Allocation, ...]: + """Order by `load` over precomputed conflict degrees (largest first, size ties).""" paired = sorted( zip(allocations, degrees, strict=True), - key=lambda pair: (pair[1], pair[0].size), + key=lambda pair: (load(pair[0], pair[1]), pair[0].size), reverse=True, ) return tuple(alloc for alloc, _ in paired) +def order_by_conflict(allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + """Order by conflict degree (most conflicted first).""" + return _order_by_load(allocations, _sort_degrees(allocations), _conflict_load) + + def order_by_conflict_size( allocations: tuple[Allocation, ...], ) -> tuple[Allocation, ...]: """Order by conflict degree times size (largest first).""" - degrees = _sort_degrees(allocations) - paired = sorted( - zip(allocations, degrees, strict=True), - key=lambda pair: (pair[1] * pair[0].size, pair[0].size), - reverse=True, - ) - return tuple(alloc for alloc, _ in paired) + return _order_by_load(allocations, _sort_degrees(allocations), _conflict_size_load) def order_by_start(allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: @@ -168,6 +179,25 @@ def _allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, .. return super()._allocate(order_by_size(allocations)) +def _order_by_input(allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: + return allocations + + +_OrderFn = Callable[[tuple[Allocation, ...]], tuple[Allocation, ...]] + +# The portfolio's variants in their tie-breaking order: equal peaks go to the +# earlier entry, mirroring the variant tuple this table replaced. +_PORTFOLIO_ORDERS: tuple[tuple[str, _OrderFn], ...] = ( + ("greedy", _order_by_input), + ("greedy_by_size", order_by_size), + ("greedy_by_duration", order_by_duration), + ("greedy_by_area", order_by_area), + ("greedy_by_conflict", order_by_conflict), + ("greedy_by_conflict_size", order_by_conflict_size), + ("greedy_by_start", order_by_start), +) + + class GreedyByAllAllocator(GreedyAllocator): """Greedy allocator that runs every variant and keeps the best result.""" @@ -176,13 +206,33 @@ def __init__(self, num_threads: int | None = None) -> None: self._num_threads = num_threads def _allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, ...]: - variants: tuple[BaseAllocator, ...] = ( - GreedyAllocator(), - GreedyBySizeAllocator(), - GreedyByDurationAllocator(), - GreedyByAreaAllocator(), - GreedyByConflictAllocator(), - GreedyByConflictSizeAllocator(), - GreedyByStartAllocator(), - ) - return allocate_parallel(allocations, variants, num_threads=self._num_threads) + # One resident placer serves every variant: the conflict adjacency is + # built once and each order crosses the C++ boundary as a permutation. + # A workload the placer refuses would have failed all seven variants. + try: + placer = FirstFitPlacer(allocations) + except ValueError as e: + raise RuntimeError("Every allocator variant failed") from e + + positions = {alloc.id: i for i, alloc in enumerate(allocations)} + + def score(order: _OrderFn) -> tuple[int, list[int]]: + permutation = [positions[alloc.id] for alloc in order(allocations)] + return placer.peak(permutation), permutation + + workers = min(resolve_num_threads(self._num_threads), len(_PORTFOLIO_ORDERS)) + scored = [] + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = [pool.submit(score, order) for _, order in _PORTFOLIO_ORDERS] + for (name, _), future in zip(_PORTFOLIO_ORDERS, futures, strict=True): + try: + scored.append(future.result()) + except Exception: + logger.warning( + "Variant %s failed; skipping it", name, exc_info=True + ) + + if not scored: + raise RuntimeError("Every allocator variant failed") + _, best_permutation = min(scored, key=lambda item: item[0]) + return tuple(placer.place(best_permutation)) diff --git a/src/python/omnimalloc/allocators/random.py b/src/python/omnimalloc/allocators/random.py index 76a5c09..f200b66 100644 --- a/src/python/omnimalloc/allocators/random.py +++ b/src/python/omnimalloc/allocators/random.py @@ -25,13 +25,13 @@ def _allocate(self, allocations: tuple[Allocation, ...]) -> tuple[Allocation, .. rng = random.Random(self._seed) placer = FirstFitPlacer(allocations) order = list(range(len(allocations))) - rng.shuffle(order) - best_order, best_peak = list(order), placer.peak(order) + best_order: list[int] = [] + best_peak: int | None = None - for _ in range(self._num_trials - 1): + for _ in range(self._num_trials): rng.shuffle(order) peak = placer.peak(order) - if peak < best_peak: + if best_peak is None or peak < best_peak: best_order, best_peak = list(order), peak return tuple(placer.place(best_order)) diff --git a/src/python/omnimalloc/benchmark/converters/onnx.py b/src/python/omnimalloc/benchmark/converters/onnx.py index 84cbdc7..ac0ba16 100644 --- a/src/python/omnimalloc/benchmark/converters/onnx.py +++ b/src/python/omnimalloc/benchmark/converters/onnx.py @@ -3,36 +3,26 @@ # import logging +from importlib.util import find_spec from pathlib import Path +from typing import TYPE_CHECKING from omnimalloc.common.optional import require_optional from omnimalloc.primitives import AllocationKind from .model import Buffer, Model, Op -try: +if TYPE_CHECKING: import onnx - HAS_ONNX = True -except ImportError: - from types import SimpleNamespace - - HAS_ONNX = False - onnx = SimpleNamespace( # ty: ignore[invalid-assignment] - checker=SimpleNamespace(check_model=None), - shape_inference=SimpleNamespace(infer_shapes=None), - load_model=None, - helper=SimpleNamespace(tensor_dtype_to_np_dtype=None), - ModelProto=None, - TensorProto=None, - ValueInfoProto=None, - NodeProto=None, - ) +HAS_ONNX = find_spec("onnx") is not None logger = logging.getLogger(__name__) -def _from_onnx_model(onnx_model: onnx.ModelProto) -> Model: +def _from_onnx_model(onnx_model: "onnx.ModelProto") -> Model: + import onnx + onnx.checker.check_model(onnx_model, full_check=True) onnx_model = onnx.shape_inference.infer_shapes( onnx_model, @@ -76,7 +66,9 @@ def _add_buffer(buffer: Buffer) -> None: return Model(id=name, ops=ops, buffers=buffers) -def _tensor_proto_to_buffer(tensor: onnx.TensorProto) -> Buffer: +def _tensor_proto_to_buffer(tensor: "onnx.TensorProto") -> Buffer: + import onnx + original_shape = tuple(tensor.dims) shape = tuple(dim for dim in original_shape if dim > 0) if len(shape) != len(original_shape): @@ -93,8 +85,10 @@ def _tensor_proto_to_buffer(tensor: onnx.TensorProto) -> Buffer: def _value_info_to_buffer( - value_info: onnx.ValueInfoProto, kind: AllocationKind + value_info: "onnx.ValueInfoProto", kind: AllocationKind ) -> Buffer: + import onnx + tt = value_info.type.tensor_type original_shape = tuple(int(dim.dim_value) for dim in tt.shape.dim) shape = tuple(dim for dim in original_shape if dim > 0) @@ -112,7 +106,7 @@ def _value_info_to_buffer( def _node_to_op( - node: onnx.NodeProto, buffers: dict[str | int, Buffer], op_id: str + node: "onnx.NodeProto", buffers: dict[str | int, Buffer], op_id: str ) -> Op: input_buffers = [] for name in node.input: @@ -136,10 +130,11 @@ def _node_to_op( ) -def from_onnx(onnx_input: onnx.ModelProto | str | Path) -> Model: +def from_onnx(onnx_input: "onnx.ModelProto | str | Path") -> Model: """Convert ONNX model or file path to Model.""" if not HAS_ONNX: require_optional("onnx", "ONNX model conversion") + import onnx if isinstance(onnx_input, (str, Path)): return _from_onnx_model(onnx.load_model(onnx_input)) diff --git a/src/python/omnimalloc/benchmark/results/export.py b/src/python/omnimalloc/benchmark/results/export.py index afd0801..c6512b3 100644 --- a/src/python/omnimalloc/benchmark/results/export.py +++ b/src/python/omnimalloc/benchmark/results/export.py @@ -123,8 +123,7 @@ def _write_allocator_reports( allocator_dir = source_dir / allocator_name allocator_dir.mkdir(parents=True, exist_ok=True) - for variant_label in sorted(variant_dict.keys()): - reports = variant_dict[variant_label] + for variant_label, reports in variant_dict.items(): variant_dir = allocator_dir / variant_label variant_dir.mkdir(parents=True, exist_ok=True) @@ -152,11 +151,11 @@ def _write_source_reports( source_dir = base_dir / "sources" / source_name / "allocators" source_dir.mkdir(parents=True, exist_ok=True) - for allocator_name in sorted(allocator_dict.keys()): + for allocator_name, variant_dict in allocator_dict.items(): _write_allocator_reports( source_dir, allocator_name, - allocator_dict[allocator_name], + variant_dict, visualize_iterations, pbar, ) @@ -165,6 +164,7 @@ def _write_source_reports( def _write_nested_reports( base_dir: Path, campaign: BenchmarkCampaign, visualize_iterations: bool ) -> None: + # The grouping property already sorts every level, so iteration order holds reports_by_source = campaign.reports_by_source_allocator_variant total_iterations = ( @@ -181,11 +181,11 @@ def _write_nested_reports( unit=unit, leave=False, ) as pbar: - for source_name in sorted(reports_by_source.keys()): + for source_name, allocator_dict in reports_by_source.items(): _write_source_reports( base_dir, source_name, - reports_by_source[source_name], + allocator_dict, visualize_iterations, pbar, ) diff --git a/src/python/omnimalloc/benchmark/results/report.py b/src/python/omnimalloc/benchmark/results/report.py index b94f290..adf145f 100644 --- a/src/python/omnimalloc/benchmark/results/report.py +++ b/src/python/omnimalloc/benchmark/results/report.py @@ -66,8 +66,6 @@ def variant_label(self) -> str: """Human-readable label for this variant.""" if self.variant_id is None: return f"{self.num_allocations}" - if isinstance(self.variant_id, str): - return self.variant_id return f"{self.variant_id}" @property @@ -81,7 +79,8 @@ def num_allocations(self) -> int: @property def total_num_allocations(self) -> int: - return sum(r.num_allocations for r in self.results) + # __post_init__ enforces a uniform count across the iterations + return self.num_results * self.num_allocations @property def num_results(self) -> int: diff --git a/src/python/omnimalloc/benchmark/results/visualize.py b/src/python/omnimalloc/benchmark/results/visualize.py index e0c4bed..f587373 100644 --- a/src/python/omnimalloc/benchmark/results/visualize.py +++ b/src/python/omnimalloc/benchmark/results/visualize.py @@ -3,8 +3,9 @@ # import logging +from importlib.util import find_spec from pathlib import Path -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from omnimalloc.common.optional import require_optional @@ -12,26 +13,11 @@ from .report import BenchmarkReport from .result import BenchmarkResult -try: - import matplotlib.pyplot as plt +if TYPE_CHECKING: from matplotlib.axes import Axes from matplotlib.figure import Figure - from matplotlib.lines import Line2D - - HAS_MATPLOTLIB = True -except ImportError: - from types import SimpleNamespace - HAS_MATPLOTLIB = False - plt = SimpleNamespace( # ty: ignore[invalid-assignment] - subplots=None, - savefig=None, - show=None, - close=None, - ) - Line2D = None # ty: ignore[invalid-assignment] - Axes = None # ty: ignore[invalid-assignment] - Figure = None # ty: ignore[invalid-assignment] +HAS_MATPLOTLIB = find_spec("matplotlib") is not None logger = logging.getLogger(__name__) @@ -72,8 +58,8 @@ def _get_sorted_reports( def _draw_graphs( - ax: Axes, - ax2: Axes, + ax: "Axes", + ax2: "Axes", name: str, color: str, is_categorical: bool, @@ -160,11 +146,13 @@ def _draw_graphs( def _draw_subplot( - ax: Axes, + ax: "Axes", source_name: str, source_data: dict[str, dict[str, tuple[BenchmarkReport, ...]]], allocator_names: tuple[str, ...], ) -> None: + import matplotlib.pyplot as plt + ax2 = ax.twinx() is_categorical = _is_categorical(source_data) @@ -202,7 +190,7 @@ def _draw_subplot( ax.set_title(f"Source: {source_name}", fontsize=12, fontweight="bold", pad=10) -def _add_footer(campaign: BenchmarkCampaign, fig: Figure) -> None: +def _add_footer(campaign: BenchmarkCampaign, fig: "Figure") -> None: metadata_text = _format_metadata(campaign.metadata) txt = fig.text( 0.5, @@ -217,7 +205,9 @@ def _add_footer(campaign: BenchmarkCampaign, fig: Figure) -> None: txt._get_wrap_line_width = lambda: fig.bbox.width * 0.90 # ty: ignore[unresolved-attribute] # noqa: SLF001 -def _add_legend(fig: Figure, allocator_names: tuple[str, ...]) -> None: +def _add_legend(fig: "Figure", allocator_names: tuple[str, ...]) -> None: + from matplotlib.lines import Line2D + handles = [ Line2D( [], @@ -241,7 +231,9 @@ def _add_legend(fig: Figure, allocator_names: tuple[str, ...]) -> None: ) -def _create_figure(num_sources: int) -> tuple[Figure, list[Axes]]: +def _create_figure(num_sources: int) -> tuple["Figure", list["Axes"]]: + import matplotlib.pyplot as plt + fig, axs = plt.subplots( nrows=num_sources, ncols=1, @@ -254,6 +246,8 @@ def _visualize_campaign( campaign: BenchmarkCampaign, path: Path | str | None, ) -> None: + import matplotlib.pyplot as plt + source_names = campaign.source_names allocator_names = campaign.allocator_names reports_by_source = campaign.reports_by_source_allocator_variant diff --git a/src/python/omnimalloc/benchmark/sources/adversarial.py b/src/python/omnimalloc/benchmark/sources/adversarial.py index 8cd366c..d8bd09e 100644 --- a/src/python/omnimalloc/benchmark/sources/adversarial.py +++ b/src/python/omnimalloc/benchmark/sources/adversarial.py @@ -10,6 +10,7 @@ from .base import BaseSource from .sizes import SizeDistribution, sample_sizes +from .validation import ensure_duration_range, ensure_size_range class SkewedSource(BaseSource): @@ -40,14 +41,8 @@ def __init__( duration_max: int = 64, seed: int | None = DEFAULT_SEED, ) -> None: - if size_min <= 0: - raise ValueError("size_min must be positive") - if size_max < size_min: - raise ValueError("size_max must be >= size_min") - if duration_min <= 0: - raise ValueError("duration_min must be positive") - if duration_max < duration_min: - raise ValueError("duration_max must be >= duration_min") + ensure_size_range(size_min, size_max) + ensure_duration_range(duration_min, duration_max) if time_max <= duration_max: raise ValueError("time_max must be > duration_max") super().__init__(num_allocations=num_allocations) @@ -62,8 +57,8 @@ def __init__( def get_allocations( self, num_allocations: int | None = None, skip: int = 0 ) -> tuple[Allocation, ...]: - num = num_allocations if num_allocations is not None else self.num_allocations - rng = random.Random(None if self.seed is None else self.seed + skip) + num = self._resolve_count(num_allocations) + rng = self._variant_rng(skip) sizes = sample_sizes(rng, num, self.distribution, self.size_min, self.size_max) allocations = [] @@ -96,10 +91,7 @@ def __init__( ) -> None: if not 0.0 <= noise < 1.0: raise ValueError("noise must be in [0, 1)") - if size_min <= 0: - raise ValueError("size_min must be positive") - if size_max < size_min: - raise ValueError("size_max must be >= size_min") + ensure_size_range(size_min, size_max) super().__init__(num_allocations=num_allocations) self.noise = noise self.size_min = size_min @@ -109,12 +101,12 @@ def __init__( def get_allocations( self, num_allocations: int | None = None, skip: int = 0 ) -> tuple[Allocation, ...]: - num = num_allocations if num_allocations is not None else self.num_allocations + num = self._resolve_count(num_allocations) if num < self._GROUP: raise ValueError( f"TwoPlusTwoSource needs at least {self._GROUP} allocations, got {num}" ) - rng = random.Random(None if self.seed is None else self.seed + skip) + rng = self._variant_rng(skip) obstructions = max(1, round(num * (1.0 - self.noise)) // self._GROUP) allocations = [] diff --git a/src/python/omnimalloc/benchmark/sources/base.py b/src/python/omnimalloc/benchmark/sources/base.py index e5baec3..4656a47 100644 --- a/src/python/omnimalloc/benchmark/sources/base.py +++ b/src/python/omnimalloc/benchmark/sources/base.py @@ -3,6 +3,7 @@ # import inspect +import random from abc import abstractmethod from typing import ClassVar @@ -21,6 +22,9 @@ class BaseSource(Registered): _strip_suffix: ClassVar[str] = "Source" _label_fields: ClassVar[tuple[str, ...]] = () + # Declared for _variant_rng; only seeded subclasses set and use it + seed: int | None + def __init__( self, num_allocations: int = 100, @@ -68,6 +72,14 @@ def is_parameterizable(self) -> bool: """Whether this source can generate arbitrary allocation counts.""" return True + def _resolve_count(self, num_allocations: int | None) -> int: + """The requested allocation count, defaulting to the configured one.""" + return num_allocations if num_allocations is not None else self.num_allocations + + def _variant_rng(self, skip: int) -> random.Random: + """Deterministic per-variant stream: same seed and skip, same problem.""" + return random.Random(None if self.seed is None else self.seed + skip) + def get_known_optimum(self, variant_id: IdType | None = None) -> int | None: """Provably achievable peak size for a variant, or None if unknown. diff --git a/src/python/omnimalloc/benchmark/sources/generator.py b/src/python/omnimalloc/benchmark/sources/generator.py index 0598ba4..fe28216 100644 --- a/src/python/omnimalloc/benchmark/sources/generator.py +++ b/src/python/omnimalloc/benchmark/sources/generator.py @@ -10,17 +10,16 @@ from omnimalloc.primitives import Allocation, AllocationKind from .base import BaseSource +from .validation import ensure_duration_range, ensure_size_range class _GeneratorSource(BaseSource): """Seeded source drawing each allocation from a per-class generation hook.""" - seed: int | None - def get_allocations( self, num_allocations: int | None = None, skip: int = 0 ) -> tuple[Allocation, ...]: - total = num_allocations if num_allocations is not None else self.num_allocations + total = self._resolve_count(num_allocations) rng = random.Random(self.seed) for _ in range(skip): @@ -61,18 +60,12 @@ def __init__( seed: int | None = DEFAULT_SEED, ) -> None: super().__init__(num_allocations=num_allocations) - if size_min <= 0: - raise ValueError("size_min must be positive") - if size_max < size_min: - raise ValueError("size_max must be >= size_min") + ensure_size_range(size_min, size_max) if time_min < 0: raise ValueError("time_min must be non-negative") if time_max <= time_min: raise ValueError("time_max must be > time_min") - if duration_min <= 0: - raise ValueError("duration_min must be positive") - if duration_max < duration_min: - raise ValueError("duration_max must be >= duration_min") + ensure_duration_range(duration_min, duration_max) if duration_max > (time_max - time_min): raise ValueError("duration_max must fit within time bounds") if kinds and kind_weights and len(kinds) != len(kind_weights): @@ -181,10 +174,7 @@ def __init__( raise ValueError("size_exponent_max must be >= size_exponent_min") if time_max <= 0: raise ValueError("time_max must be positive") - if duration_min <= 0: - raise ValueError("duration_min must be positive") - if duration_max < duration_min: - raise ValueError("duration_max must be >= duration_min") + ensure_duration_range(duration_min, duration_max) self.size_exponent_min = size_exponent_min self.size_exponent_max = size_exponent_max @@ -225,10 +215,7 @@ def __init__( seed: int | None = DEFAULT_SEED, ) -> None: super().__init__(num_allocations=num_allocations) - if size_min <= 0: - raise ValueError("size_min must be positive") - if size_max < size_min: - raise ValueError("size_max must be >= size_min") + ensure_size_range(size_min, size_max) if time_window < 2: raise ValueError("time_window must be at least 2") @@ -271,14 +258,8 @@ def __init__( seed: int | None = DEFAULT_SEED, ) -> None: super().__init__(num_allocations=num_allocations) - if size_min <= 0: - raise ValueError("size_min must be positive") - if size_max < size_min: - raise ValueError("size_max must be >= size_min") - if duration_min <= 0: - raise ValueError("duration_min must be positive") - if duration_max < duration_min: - raise ValueError("duration_max must be >= duration_min") + ensure_size_range(size_min, size_max) + ensure_duration_range(duration_min, duration_max) self.size_min = size_min self.size_max = size_max @@ -289,7 +270,7 @@ def __init__( def get_allocations( self, num_allocations: int | None = None, skip: int = 0 ) -> tuple[Allocation, ...]: - total = num_allocations if num_allocations is not None else self.num_allocations + total = self._resolve_count(num_allocations) rng = random.Random(self.seed) current_time = 0 diff --git a/src/python/omnimalloc/benchmark/sources/huggingface.py b/src/python/omnimalloc/benchmark/sources/huggingface.py index 4e8a411..865a3d3 100644 --- a/src/python/omnimalloc/benchmark/sources/huggingface.py +++ b/src/python/omnimalloc/benchmark/sources/huggingface.py @@ -6,7 +6,7 @@ from collections import defaultdict from importlib.util import find_spec from pathlib import Path -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar from omnimalloc.benchmark.converters.model import model_to_allocations from omnimalloc.benchmark.converters.onnx import from_onnx @@ -16,26 +16,24 @@ from ..utils import tqdm # noqa: TID252 from .base import BaseSource -try: +if TYPE_CHECKING: from huggingface_hub import HfApi, ModelInfo - HAS_HUGGINGFACE_HUB = True -except ImportError: - HAS_HUGGINGFACE_HUB = False - HfApi = None # ty: ignore[invalid-assignment] - ModelInfo = None # ty: ignore[invalid-assignment] - +HAS_HUGGINGFACE_HUB = find_spec("huggingface_hub") is not None HAS_ONNX = find_spec("onnx") is not None -def _get_hf_api() -> HfApi: +def _get_hf_api() -> "HfApi": """Get HfApi instance, checking that dependency is available.""" if not HAS_HUGGINGFACE_HUB: require_optional("huggingface-hub", "HuggingfaceSource") + # Deferred so importing this module never pays for the hub client + from huggingface_hub import HfApi + return HfApi() -def _list_onnx_models(limit: int = 10) -> list[ModelInfo]: +def _list_onnx_models(limit: int = 10) -> list["ModelInfo"]: """Return ONNX models from Hugging Face Hub, excluding the legacy repository.""" hf_api = _get_hf_api() models = hf_api.list_models(author="onnxmodelzoo", limit=limit + 1) @@ -43,8 +41,8 @@ def _list_onnx_models(limit: int = 10) -> list[ModelInfo]: def _filter_onnx_opsets( - model_infos: list[ModelInfo], min_opset: int = 16 -) -> list[ModelInfo]: + model_infos: list["ModelInfo"], min_opset: int = 16 +) -> list["ModelInfo"]: """Filter ONNX models to only include the highest opset per base model name.""" model_groups = defaultdict(list) @@ -62,8 +60,9 @@ def _filter_onnx_opsets( def _gather_download_info( - model_infos: list[ModelInfo], + model_infos: list["ModelInfo"], filename_filter: str, + num_models: int, max_file_size_mb: float | None = 200, ) -> dict[str, str]: """Gather information about which models to download, filtering by size.""" @@ -71,6 +70,9 @@ def _gather_download_info( id_file_map = {} for model_info in model_infos: + # Each candidate costs one HTTP round-trip, so stop once enough passed + if len(id_file_map) == num_models: + break repo_files = hf_api.list_repo_tree(model_info.id, recursive=True) onnx_files = [ f @@ -122,9 +124,8 @@ def _download_onnx_models( """Download ONNX models and return local file paths.""" models = _list_onnx_models(limit=num_models * 5) filtered = _filter_onnx_opsets(models) - id_file_map = _gather_download_info(filtered, ".onnx", max_file_size_mb=200) - id_file_map_limited = dict(list(id_file_map.items())[:num_models]) - return _download_files(id_file_map_limited, output_dir, ".onnx") + id_file_map = _gather_download_info(filtered, ".onnx", num_models) + return _download_files(id_file_map, output_dir, ".onnx") class HuggingfaceSource(BaseSource): diff --git a/src/python/omnimalloc/benchmark/sources/minimalloc.py b/src/python/omnimalloc/benchmark/sources/minimalloc.py index 2f964b7..34b161b 100644 --- a/src/python/omnimalloc/benchmark/sources/minimalloc.py +++ b/src/python/omnimalloc/benchmark/sources/minimalloc.py @@ -56,6 +56,25 @@ def _prefix_ids(pool: Pool) -> Pool: ) +def _load_pools(subset: MinimallocSubset, csv_dir: Path | None) -> list[Pool]: + """Load and id-qualify every CSV pool in the subset's directory.""" + resolved = csv_dir if csv_dir is not None else _checkout_csv_dir(subset) + # Sort for a filesystem-independent, reproducible variant order + files = sorted(resolved.glob("*.csv")) if resolved is not None else [] + if resolved is None: + logger.warning( + f"Not running from a source checkout, so the " + f"{subset.value!r} subset has no datasets; pass " + "csv_dir to read them from an install." + ) + elif not files: + logger.warning( + f"No Minimalloc CSVs found in {resolved}; the " + f"{subset.value!r} subset yields no variants." + ) + return [_prefix_ids(load_allocation(f)) for f in files] + + class MinimallocSource(BaseSource): """Fixed source loading pools from a directory of Minimalloc CSV files. @@ -73,36 +92,12 @@ def __init__( self.subset = MinimallocSubset(subset) # The label must carry an explicit csv_dir but not the checkout default self.csv_dir = Path(csv_dir) if csv_dir is not None else None - self._cached_pools: list[Pool] | None = None + self._pools = _load_pools(self.subset, self.csv_dir) # An empty dataset keeps the base invariant; the accessors raise instead num_allocs = sum(len(p.allocations) for p in self._pools) super().__init__(num_allocations=max(num_allocs, 1)) - @property - def _pools(self) -> list[Pool]: - if self._cached_pools is None: - csv_dir = ( - self.csv_dir - if self.csv_dir is not None - else _checkout_csv_dir(self.subset) - ) - # Sort for a filesystem-independent, reproducible variant order - files = sorted(csv_dir.glob("*.csv")) if csv_dir is not None else [] - if csv_dir is None: - logger.warning( - f"Not running from a source checkout, so the " - f"{self.subset.value!r} subset has no datasets; pass " - "csv_dir to read them from an install." - ) - elif not files: - logger.warning( - f"No Minimalloc CSVs found in {csv_dir}; the " - f"{self.subset.value!r} subset yields no variants." - ) - self._cached_pools = [_prefix_ids(load_allocation(f)) for f in files] - return self._cached_pools - def _all_allocations(self) -> tuple[Allocation, ...]: return tuple(alloc for pool in self._pools for alloc in pool.allocations) diff --git a/src/python/omnimalloc/benchmark/sources/sizes.py b/src/python/omnimalloc/benchmark/sources/sizes.py index 9cefb2c..e29d0ad 100644 --- a/src/python/omnimalloc/benchmark/sources/sizes.py +++ b/src/python/omnimalloc/benchmark/sources/sizes.py @@ -7,6 +7,8 @@ from enum import Enum from typing import Final +from .validation import ensure_size_range + class SizeDistribution(str, Enum): """Size distribution families, ordered from flattest to most skewed. @@ -46,10 +48,7 @@ def sample_sizes( `bimodal` the 90/10 accelerator mix, `dominant` one buffer at 90%. """ distribution = SizeDistribution(distribution) - if size_min <= 0: - raise ValueError("size_min must be positive") - if size_max < size_min: - raise ValueError("size_max must be >= size_min") + ensure_size_range(size_min, size_max) if count <= 0: return [] diff --git a/src/python/omnimalloc/benchmark/sources/sync_patterns.py b/src/python/omnimalloc/benchmark/sources/sync_patterns.py index e3eebdd..aa2ffa8 100644 --- a/src/python/omnimalloc/benchmark/sources/sync_patterns.py +++ b/src/python/omnimalloc/benchmark/sources/sync_patterns.py @@ -13,6 +13,7 @@ from .base import BaseSource from .sizes import SizeDistribution, sample_sizes +from .validation import ensure_size_range class SyncPattern(str, Enum): @@ -82,10 +83,7 @@ def __init__( raise ValueError("sync_period must be positive") if group_size is not None and group_size <= 0: raise ValueError("group_size must be positive") - if size_min <= 0: - raise ValueError("size_min must be positive") - if size_max < size_min: - raise ValueError("size_max must be >= size_min") + ensure_size_range(size_min, size_max) if max_lifetime is not None and max_lifetime <= 0: raise ValueError("max_lifetime must be positive") super().__init__(num_allocations=num_allocations) @@ -104,8 +102,8 @@ def __init__( def get_allocations( self, num_allocations: int | None = None, skip: int = 0 ) -> tuple[Allocation, ...]: - num = num_allocations if num_allocations is not None else self.num_allocations - rng = random.Random(None if self.seed is None else self.seed + skip) + num = self._resolve_count(num_allocations) + rng = self._variant_rng(skip) steps = self.steps or max(4 * self.speed_skew, 2 * num // self.num_threads) max_lifetime = self.max_lifetime or max(1, steps // 4) snapshots = self._simulate(steps, rng) diff --git a/src/python/omnimalloc/benchmark/sources/tiling_base.py b/src/python/omnimalloc/benchmark/sources/tiling_base.py index 6696c33..97b85ff 100644 --- a/src/python/omnimalloc/benchmark/sources/tiling_base.py +++ b/src/python/omnimalloc/benchmark/sources/tiling_base.py @@ -11,6 +11,7 @@ from omnimalloc.primitives import Allocation, IdType, Pool, TimePoint from .base import BaseSource +from .validation import ensure_duration_range, ensure_size_range TimeT_co = TypeVar("TimeT_co", bound=TimePoint, covariant=True) @@ -54,10 +55,8 @@ def __init__( seed: int | None, ) -> None: super().__init__(num_allocations=num_allocations) - if size_min <= 0: - raise ValueError("size_min must be positive") - if duration_min <= 0: - raise ValueError("duration_min must be positive") + ensure_size_range(size_min) + ensure_duration_range(duration_min) self.capacity = capacity self.makespan = makespan self.size_min = size_min @@ -103,10 +102,6 @@ def _build_tiles( ) return tiles - def _variant_rng(self, skip: int) -> random.Random: - """Deterministic per-variant stream: same seed and skip, same problem.""" - return random.Random(None if self.seed is None else self.seed + skip) - def _placed_tiles(self, num: int, skip: int) -> Sequence[_Tile[TimePoint]]: """One placed tile per allocation to generate; the subclass hook.""" return self._build_tiles(num, self._variant_rng(skip)) @@ -114,7 +109,7 @@ def _placed_tiles(self, num: int, skip: int) -> Sequence[_Tile[TimePoint]]: def _tile_allocations( self, num_allocations: int | None, skip: int, with_offsets: bool ) -> tuple[Allocation, ...]: - num = num_allocations if num_allocations is not None else self.num_allocations + num = self._resolve_count(num_allocations) return tuple( Allocation( id=skip + i, diff --git a/src/python/omnimalloc/benchmark/sources/validation.py b/src/python/omnimalloc/benchmark/sources/validation.py new file mode 100644 index 0000000..6451187 --- /dev/null +++ b/src/python/omnimalloc/benchmark/sources/validation.py @@ -0,0 +1,19 @@ +# +# SPDX-License-Identifier: Apache-2.0 +# + + +def ensure_size_range(size_min: int, size_max: int | None = None) -> None: + """Raise ValueError unless size_min is positive and size_max at least matches.""" + if size_min <= 0: + raise ValueError("size_min must be positive") + if size_max is not None and size_max < size_min: + raise ValueError("size_max must be >= size_min") + + +def ensure_duration_range(duration_min: int, duration_max: int | None = None) -> None: + """Raise ValueError unless duration_min is positive and duration_max matches.""" + if duration_min <= 0: + raise ValueError("duration_min must be positive") + if duration_max is not None and duration_max < duration_min: + raise ValueError("duration_max must be >= duration_min") diff --git a/src/python/omnimalloc/benchmark/utils.py b/src/python/omnimalloc/benchmark/utils.py index ee8513e..54bd6b0 100644 --- a/src/python/omnimalloc/benchmark/utils.py +++ b/src/python/omnimalloc/benchmark/utils.py @@ -4,25 +4,27 @@ from typing import Any -try: - from tqdm.auto import tqdm -except ImportError: - class _DummyProgressBar: - """No-op stand-in for tqdm's total=... progress bar.""" +class _DummyProgressBar: + """No-op stand-in for tqdm's total=... progress bar.""" - def __enter__(self) -> "_DummyProgressBar": - return self + def __enter__(self) -> "_DummyProgressBar": + return self - def __exit__(self, *args: object) -> None: - pass + def __exit__(self, *args: object) -> None: + pass - def update(self, n: int = 1) -> None: - pass + def update(self, n: int = 1) -> None: + pass - def tqdm(iterable: Any = None, **kwargs: Any) -> Any: # noqa: ARG001, ANN401 - """No-op tqdm fallback when tqdm is not installed.""" + +def tqdm(iterable: Any = None, **kwargs: Any) -> Any: # noqa: ANN401 + """Lazily imported tqdm.auto progress bar; a no-op when tqdm is not installed.""" + try: + from tqdm.auto import tqdm as tqdm_auto + except ImportError: if iterable is None: # When called with total= instead of an iterable return _DummyProgressBar() return iterable + return tqdm_auto(iterable, **kwargs) diff --git a/src/python/omnimalloc/common/constants.py b/src/python/omnimalloc/common/constants.py index 1e584ee..5a1b733 100644 --- a/src/python/omnimalloc/common/constants.py +++ b/src/python/omnimalloc/common/constants.py @@ -36,9 +36,3 @@ MB: Final[int] = 1_024 * KB GB: Final[int] = 1_024 * MB TB: Final[int] = 1_024 * GB - -# Frequency units in hertz -HZ: Final[int] = 1 -KHZ: Final[int] = 1_000 * HZ -MHZ: Final[int] = 1_000 * KHZ -GHZ: Final[int] = 1_000 * MHZ diff --git a/src/python/omnimalloc/visualize.py b/src/python/omnimalloc/visualize.py index e791295..cf86657 100644 --- a/src/python/omnimalloc/visualize.py +++ b/src/python/omnimalloc/visualize.py @@ -3,8 +3,9 @@ # from collections.abc import Sequence +from importlib.util import find_spec from pathlib import Path -from typing import Final, Literal, NamedTuple +from typing import TYPE_CHECKING, Final, Literal, NamedTuple from omnimalloc.analysis import antichain_pressure, conflict_degrees, try_linearize from omnimalloc.analysis._clock import time_components, uniform_dim @@ -19,31 +20,11 @@ System, ) -try: - import matplotlib.pyplot as plt +if TYPE_CHECKING: from matplotlib.axes import Axes from matplotlib.figure import Figure - from matplotlib.patches import Patch, Rectangle - from matplotlib.ticker import FuncFormatter, MaxNLocator, MultipleLocator - - HAS_MATPLOTLIB = True - -except ImportError: - from types import SimpleNamespace - HAS_MATPLOTLIB = False - - plt = SimpleNamespace( # ty: ignore[invalid-assignment] - subplots=None, - show=None, - ) - Axes = None # ty: ignore[invalid-assignment] - Figure = None # ty: ignore[invalid-assignment] - Patch = None # ty: ignore[invalid-assignment] - Rectangle = None # ty: ignore[invalid-assignment] - FuncFormatter = None # ty: ignore[invalid-assignment] - MaxNLocator = None # ty: ignore[invalid-assignment] - MultipleLocator = None # ty: ignore[invalid-assignment] +HAS_MATPLOTLIB = find_spec("matplotlib") is not None # Rendering guard: plot annotation gives up within tens of milliseconds # rather than stall on a second-scale conflict sweep. @@ -269,13 +250,15 @@ def _get_y_offsets(system: System) -> dict[Memory, dict[Pool, int]]: def _draw_allocation( - ax: Axes, + ax: "Axes", alloc: Allocation, offset: int, color: str, extent: tuple[int, int], ) -> None: """Draw a single allocation rectangle over its projected lifetime.""" + from matplotlib.patches import Rectangle + assert alloc.offset is not None start, end = extent y_pos = offset + alloc.offset @@ -299,9 +282,11 @@ def _draw_allocation( def _draw_pool_background( - ax: Axes, y_offset: int, pool_size: int, colors: set[str] + ax: "Axes", y_offset: int, pool_size: int, colors: set[str] ) -> None: """Draw background rectangle for allocation pool (gray for mixed/empty kinds).""" + from matplotlib.patches import Rectangle + color = next(iter(colors)) if len(colors) == 1 else "gray" x_min, x_max = ax.get_xlim() rect = Rectangle( @@ -315,7 +300,7 @@ def _draw_pool_background( ax.add_patch(rect) -def _draw_limit_lines(ax: Axes, limits: dict[str, int]) -> None: +def _draw_limit_lines(ax: "Axes", limits: dict[str, int]) -> None: """Draw annotated horizontal lines for used size, declared size, and extras.""" _, x_max = ax.get_xlim() for name, value in limits.items(): @@ -339,8 +324,10 @@ def _draw_limit_lines(ax: Axes, limits: dict[str, int]) -> None: ) -def _set_axes_ticks(ax: Axes, y_limit: int, num_ticks: int = 8) -> None: +def _set_axes_ticks(ax: "Axes", y_limit: int, num_ticks: int = 8) -> None: """Configure axis ticks and formatters.""" + from matplotlib.ticker import FuncFormatter, MaxNLocator, MultipleLocator + tick_size = y_limit / num_ticks divisor, unit = _byte_unit(y_limit) ax.yaxis.set_major_locator(MultipleLocator(tick_size)) @@ -419,8 +406,10 @@ def _projection_panels(system: System) -> tuple[list[_Panel], str | None]: return panels, caveat -def _add_legend(fig: Figure) -> None: +def _add_legend(fig: "Figure") -> None: """Add figure legend for allocation kinds.""" + from matplotlib.patches import Patch + handles = [ Patch(color=color, label=kind.name, alpha=0.8) for kind, color in KIND_COLOR_MAP.items() @@ -435,7 +424,7 @@ def _add_legend(fig: Figure) -> None: def _set_axes_limits( - ax: Axes, + ax: "Axes", x_limits: tuple[int, int], y_limits: tuple[int, int], size: int | None, @@ -459,7 +448,7 @@ def _set_axes_limits( def _draw_panel( - ax: Axes, + ax: "Axes", panel: _Panel, y_limits: tuple[int, int], y_offsets: dict[Pool, int], @@ -511,6 +500,8 @@ def _visualize_system( view: Literal["panel", "lanes"], max_lanes: int | None, ) -> None: + import matplotlib.pyplot as plt + if view == "lanes": panels, caveat = _lane_panels(system, max_lanes) else: diff --git a/tests/unit/test_public_api.py b/tests/unit/test_public_api.py index 4cca94c..40a6fea 100644 --- a/tests/unit/test_public_api.py +++ b/tests/unit/test_public_api.py @@ -81,6 +81,8 @@ def _public_names(module: ModuleType) -> set[str]: def test_top_level_api_is_pinned() -> None: + for name in TOP_LEVEL_API: + getattr(omnimalloc, name) assert _public_names(omnimalloc) == TOP_LEVEL_API