diff --git a/.gitignore b/.gitignore index 99bca7d..48539b4 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ # Generated files build/ +build-*/ generated/ # Prerequisites diff --git a/docs/build.md b/docs/build.md index b329859..17f85e9 100644 --- a/docs/build.md +++ b/docs/build.md @@ -118,6 +118,45 @@ Each result reports `unit`, `mean`, and `median` as separate fields. The default memory sweep runs up to 16 MiB; set `INFINI_RT_PERF_ENABLE_LARGE=1` to include the 256 MiB case. +`perf_memory_pool` is an A/B benchmark: it runs each workload twice, once +straight through the backend allocator and once through `MemoryPool`. Both arms +share the workload, the iteration count, and the unit, and emit one JSON row +each differing only in the `allocator` param (`direct` vs `pool`), so a consumer +can divide one by the other. A human-readable speedup table is also written to +stderr at the end of the run; stdout stays pure JSON. + +The paired workloads are `SingleBlock` (allocate one block, free it, repeat), +`WorkingSetChurn` (rotate a window of 8 live blocks, the shape an inference loop +produces), `MixedSizeClasses` (rotate 32 size classes, the pool's least +favorable shape), `ThreadScaling` (1/2/4/8 threads on one device, which exposes +the cost of the pool's single mutex), and `ConcurrentMixedSizes` (the same thread +counts but with 16 size classes per thread, so the threads collide on the mutex +while touching different free lists). Unpaired pool-only rows — `MissPath`, +`ConcurrentMissPath`, `AlignedHit`, `ReleaseCached`, `GetStats`, +`AllocateZeroBytes`, `DeallocateNullptr` — measure costs that exist only for the +pool. `ConcurrentMissPath` drops every cached block after each operation, so it +prices the upstream call under contention rather than the free-list hit. + +Concurrent workloads calibrate their op count at run time to fill a ~40 ms +sample window. A fixed count would let a fast allocator finish a batch in a few +microseconds, where thread startup and scheduler placement dominate the +measurement; `ops_per_thread` is therefore an output in the JSON params rather +than a constant, and it differs between backends and between the two arms. + +The pool is instantiated over the `runtime::` dispatch API, so one binary +measures whichever backend the library was built with. To compare backends, +configure one build dir per backend and run each: + +```bash +cmake -S . -B build-perf-cuda \ + -DCMAKE_BUILD_TYPE=Release \ + -DWITH_CPU=OFF -DWITH_NVIDIA=ON \ + -DINFINI_RT_BUILD_PERFORMANCE_TESTING=ON +cmake --build build-perf-cuda -j +python3 scripts/run_performance_tests.py \ + --build-dir build-perf-cuda --output perf-nvidia.json +``` + ## Documentation Enable the Doxygen documentation target with: diff --git a/include/infini/rt/arena_memory_pool.h b/include/infini/rt/arena_memory_pool.h new file mode 100644 index 0000000..2e318e2 --- /dev/null +++ b/include/infini/rt/arena_memory_pool.h @@ -0,0 +1,1618 @@ +#ifndef INFINI_RT_ARENA_MEMORY_POOL_H_ +#define INFINI_RT_ARENA_MEMORY_POOL_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "infini/rt/detail/node_arena.h" +#include "infini/rt/detail/pointer_table.h" + +namespace infini::rt { + +/// Tunables for `ArenaMemoryPool`. Every threshold is a template parameter +/// rather than a hard-coded constant because the interesting behaviors -- +/// growth past the cap, oversize backings, automatic shrink -- are only +/// reachable at multi-hundred-megabyte scale with the production values, which +/// no CI machine (least of all a GPU-less one) can exercise. Tests instantiate +/// the pool with a kilobyte-scale config and drive the same code paths. +struct DefaultArenaConfig { + /// Capacity of the first backing store, and the base of the doubling ramp. + static constexpr std::size_t kInitialCapacity = 64ull << 20; // 64 MB + + /// Ceiling on the doubling ramp. Once reached, further growth adds more + /// backings of this size rather than larger ones. + static constexpr std::size_t kMaxCapacity = 512ull << 20; // 512 MB + + /// Requests at or below this size count as "small" for the shrink heuristic. + static constexpr std::size_t kSmallThreshold = 1ull << 20; // 1 MB + + /// Every slice starts at a multiple of this, whatever the caller asked for. + /// Device allocators guarantee a healthy natural alignment (256 B for + /// `cudaMalloc`) and callers -- plus vectorized kernels -- rely on it, so + /// slicing must not hand back a worse-aligned pointer than `Upstream::Malloc` + /// would have. + static constexpr std::size_t kMinSliceAlignment = 512; + + /// A split leaving less than this behind is not performed; the remainder goes + /// to the caller as internal waste instead of becoming an unusable sliver. + static constexpr std::size_t kMinSplitRemainder = 512; + + /// Consecutive small allocations that mark the end of a burst. Reaching this + /// count triggers one idle scan. + static constexpr std::size_t kShrinkThreshold = 16; + + /// Consecutive idle scans a non-resident backing must be found empty for + /// before it is destroyed. This is the hysteresis that keeps a "one big op + /// plus twenty small ops" loop from destroying and re-creating a backing + /// every iteration -- `cudaFree` implicitly synchronizes the whole device, so + /// thrashing it from the allocation path is far worse than holding the memory + /// for one more round. + static constexpr std::uint32_t kEmptyScansToDestroy = 2; + + /// Blocks one thread may hold per exact size in its front cache. + /// + /// One, because depth beyond one only pays off when several blocks of the same + /// size are live at once, and the loops the cache exists for -- allocate, use, + /// free, allocate the same size again -- hit at depth one already. Depth is + /// not free: a parked block stays in `allocated_`, so it is invisible to + /// coalescing and cannot be merged into a larger request. Deeper caches + /// scatter more such blocks through the backings, which on a workload whose + /// footprint grows monotonically shows up as extra backing growth. Depth one + /// keeps the hit and pays the least for it. + static constexpr std::size_t kThreadCacheDepth = 1; + + /// Total bytes one thread may retain across every one of its cache lists. + /// This, not the depth, is what bounds the cache on large sizes: eight 512 MB + /// blocks per size class would be absurd, and this cap is what stops it. + static constexpr std::size_t kThreadCacheBytes = 8ull << 20; // 8 MB +}; + +namespace detail { + +// The front cache arrived after `Config` had several implementations in this +// tree, so its two knobs are detected rather than required: a `Config` that +// predates them keeps working. Every other threshold is mandatory, because +// omitting one is almost certainly a mistake; these two are the exception only +// because their absence has a safe reading. +// +// The byte fallback is proportional rather than absolute. An arena configured at +// kilobyte scale (the tests) and one configured at production scale differ by +// five orders of magnitude, and a fixed default would let one thread privatize +// an entire backing in the former case. +template +struct ArenaCacheDepth + : std::integral_constant {}; + +template +struct ArenaCacheDepth> + : std::integral_constant {}; + +template +struct ArenaCacheBytes + : std::integral_constant {}; + +template +struct ArenaCacheBytes> + : std::integral_constant {}; + +} // namespace detail + +/// ## Backend-agnostic arena (sub-allocating) allocator. +/// +/// `MemoryPool` calls `Upstream::Malloc` on every cache miss. Device allocators +/// are synchronous and cost hundreds of microseconds, so a workload whose sizes +/// keep missing pays that price over and over. `ArenaMemoryPool` instead +/// requests large *backing stores* from the upstream allocator and satisfies +/// requests by slicing them, so upstream calls scale with the pool's high-water +/// footprint rather than with the number of allocations. +/// +/// Like `MemoryPool`, this is a pure composition over an `Upstream` allocator: +/// any type providing `Malloc(void**, size_t)`, `Free(void*)`, an `Error` type +/// alias, and a `static constexpr Error kSuccess` satisfies the contract, so +/// every `runtime::Runtime<...>` device specialization qualifies and tests can +/// inject a mock upstream. +/// +/// ### Structure +/// +/// Each backing store is modeled as one doubly-linked list of `Chunk`s sorted +/// by address, tiling its usable span with no gaps. A chunk is either free or +/// handed out; allocation splits a free chunk, and `Deallocate` marks a chunk +/// free. Adjacent free chunks are coalescible, so a drained backing can always +/// be reassembled into one chunk spanning its whole span and serve a contiguous +/// request of its full size again. (A bump-pointer design cannot express this: +/// memory released below the bump cursor is physically adjacent to the +/// untouched tail yet unreachable from it, so a drained 512 MB backing could +/// not produce a 512 MB block.) +/// +/// That merge is *deferred* rather than performed on every free. Coalescing +/// touches both physical neighbors and re-inserts into an ordered container -- +/// several cold cache lines inside the lock -- and a workload that cycles +/// through a handful of sizes never needs it, because the next request for a +/// size pops back exactly the chunk that was just released. So `Deallocate` +/// files a released chunk straight into its exact-size bin and returns, and the +/// merging happens in one pass (`CoalesceAll`) only when a request cannot be +/// served from what is already indexed. The cost is that "this region is free" +/// has more than one representation, so emptiness is a per-backing live-chunk +/// count rather than a list-length test, and `Stats::largest_free_chunk` +/// reports the largest coalescible *run* rather than the largest single chunk. +/// +/// Free chunks across all backings are indexed together, so allocation is a +/// best-fit lookup rather than an exact size-class match: splitting and +/// coalescing mean a released block is reusable at any size, not only at the +/// size class it was allocated with. The index is in two parts. Sizes that are +/// a small multiple of `Config::kMinSliceAlignment` go in *fast bins* -- one +/// intrusive list per exact size, with a bitmap over their occupancy -- so a +/// request that matches an occupied bin is served by popping a list head, in +/// O(1) and with no node allocation. Everything else lives in a tree ordered by +/// (size, address), searched in O(log n). Since the bins are exact, a hit there +/// is already the best fit and the tree is not consulted at all; workloads that +/// cycle through a handful of sizes stay entirely on that path. +/// +/// ### Front cache +/// +/// Everything above happens under one mutex, so throughput is bounded by how +/// long each thread holds it -- and best fit plus splitting is a longer hold +/// than a size-class pool's list pop. In front of it sits a small per-thread +/// cache of exact-size blocks: a thread that cycles through a handful of sizes +/// serves its own allocations by popping one of its own lists, touching no +/// shared state and taking no lock at all. +/// +/// `Deallocate` still takes the mutex, because rejecting a foreign pointer or a +/// double free requires the live-block table, which is shared. What it does +/// while holding it is file the block into the *releasing* thread's cache rather +/// than into the global free index, so the next allocation on that thread finds +/// it without the lock. In an alloc/free loop that halves lock acquisitions and +/// removes the expensive half of the work -- the best-fit lookup and the split. +/// +/// A cached block is still counted live by its backing store, which is what +/// keeps automatic shrink from handing that backing upstream while a cache +/// points into it. The retention is bounded per thread, and any path that needs +/// the memory back -- a request nothing indexed can serve, `ReleaseCached`, +/// destruction -- reclaims every cache first. The cost is that `Stats` cannot be +/// maintained on the lock-free path, so `GetStats` folds each cache's own +/// tallies in as it reads them. +/// +/// ### Growth and shrink +/// +/// Backing capacity follows `Config::kInitialCapacity` doubling up to +/// `Config::kMaxCapacity`; past the cap, additional backings of the cap size +/// are added. A single request larger than the cap gets its own exactly sized +/// *oversize* backing -- multiple capped backings cannot serve it, since +/// slicing never spans two upstream allocations. +/// +/// The first non-oversize backing is *resident*: never destroyed except by +/// `ReleaseCached` or destruction, so a steady small workload keeps a warm +/// arena and never thrashes. Every other backing is subject to automatic +/// shrink: once `Config::kShrinkThreshold` consecutive small allocations +/// indicate the burst is over, drained non-resident backings are returned +/// upstream (after `Config::kEmptyScansToDestroy` scans of hysteresis; oversize +/// backings are exempt and go back on the first scan, since holding gigabytes +/// idle is far more expensive than one extra upstream call). +/// +/// ### Caller responsibilities +/// +/// **Streams.** A block is reusable the instant `Deallocate` returns, and +/// coalescing means the memory may come back as a *differently sized* block at +/// a *different offset* handed to an unrelated caller. If a previously launched +/// kernel still reads the old address, it corrupts a live allocation whose size +/// and bounds bear no relation to the original. The pool has no notion of +/// streams: callers must ensure device-side access to a block has completed +/// before calling `Deallocate` (e.g. by synchronizing the stream, or recording +/// an event and waiting on it). +/// +/// **Devices.** A backing is bound to whichever device was current when it was +/// created. The pool does not track device ids, so multi-device use needs one +/// pool instance per device. +/// +/// The pool is thread-safe: every public method takes an internal mutex, and no +/// upstream call is ever made while it is held. It is neither copyable nor +/// movable. +template +class ArenaMemoryPool { + public: + using Error = typename Upstream::Error; + + /// Runtime statistics. Byte counters are live totals; `peak_*` track + /// high-water marks. The remaining counters are monotonic tallies. + struct Stats { + /// Bytes currently handed out to callers (sum of served chunk sizes, + /// including any remainder folded in by the split threshold). + std::size_t bytes_in_use = 0; + + /// Bytes currently held from the upstream allocator: the sum of every + /// backing store's capacity. + std::size_t bytes_reserved = 0; + + /// High-water mark of `bytes_in_use`. + std::size_t peak_bytes_in_use = 0; + + /// High-water mark of `bytes_reserved`. + std::size_t peak_bytes_reserved = 0; + + /// Number of `Allocate` calls that returned a non-null pointer. + std::size_t alloc_count = 0; + + /// Number of `Deallocate` calls that released a live block. + std::size_t free_count = 0; + + /// Allocations served from existing backing stores. + std::size_t cache_hit_count = 0; + + /// Allocations that required a new backing store. + std::size_t cache_miss_count = 0; + + /// Calls into `Upstream::Malloc` that produced a backing store. + std::size_t upstream_alloc_count = 0; + + /// Calls into `Upstream::Free` (one per backing store). + std::size_t upstream_free_count = 0; + + /// Number of live backing stores. + std::size_t backing_count = 0; + + /// Bytes sitting free inside backing stores. This is the pool's + /// fragmentation: reserved but not in use, and not returnable upstream + /// unless a whole backing drains. + std::size_t bytes_free_in_backings = 0; + + /// Size of the largest request the pool could serve without going upstream: + /// the longest run of adjacent free chunks, since coalescing is deferred and + /// any such run can be merged on demand. Together with + /// `bytes_free_in_backings` this shows whether free space is usable or + /// shattered. + std::size_t largest_free_chunk = 0; + + /// Bytes inside served chunks beyond what the caller asked for: remainders + /// too small to split off. A subset of `bytes_in_use`, counted while the + /// chunks carrying them are live. + std::size_t bytes_internal_waste = 0; + + /// Reserved bytes that are in no chunk at all: the head of each backing + /// store trimmed off to bring the first slice up to + /// `Config::kMinSliceAlignment`. Zero whenever the upstream allocator + /// already returns suitably aligned pointers, which device allocators do. + /// + /// Together these close the books: + /// `bytes_reserved == bytes_in_use + bytes_free_in_backings + + /// bytes_unusable`. + std::size_t bytes_unusable = 0; + + /// Backing stores destroyed by automatic shrink. + std::size_t shrink_count = 0; + }; + + ArenaMemoryPool() + : registry_(std::make_shared()), mutex_(registry_->mutex) { + // No lock: nothing else can reach this registry yet. + registry_->pool = this; + } + + ArenaMemoryPool(const ArenaMemoryPool&) = delete; + ArenaMemoryPool& operator=(const ArenaMemoryPool&) = delete; + + /// Frees every backing store, exactly once each. Blocks still outstanding are + /// slices of those backings, so they must not be freed individually -- only + /// the upstream base pointers are valid arguments to `Upstream::Free`. Any + /// pointer from `Allocate` dangles after destruction. + ~ArenaMemoryPool() { + { + // Clearing `pool` under the lock is the handshake with every thread still + // holding a cache: from here on their exit handlers see a dead pool and + // touch nothing but their own node. Reclaiming first keeps the chunk + // arena's bookkeeping consistent while `DestroyChunks` walks it. + std::lock_guard lock(registry_->mutex); + for (ThreadCache* cache = registry_->head; cache != nullptr; + cache = cache->next) { + ReclaimCacheLocked(cache); + } + registry_->pool = nullptr; + } + + for (const auto& backing : backings_) { + DestroyChunks(backing.get()); + Upstream::Free(backing->base); + } + } + + /// Allocates at least `size` bytes by slicing a backing store, requesting a + /// new one from the upstream allocator only when no existing backing has + /// room. `alignment` of `0` uses the pool's natural slice alignment + /// (`Config::kMinSliceAlignment`); otherwise the returned pointer is aligned + /// up to `alignment`, which must be a power of two. + /// + /// On success writes the pointer to `*ptr` and returns `kSuccess`. A `size` of + /// `0` succeeds with `*ptr == nullptr`. On upstream failure the upstream error + /// is returned and `*ptr` is set to `nullptr`. + Error Allocate(void** ptr, std::size_t size, std::size_t alignment = 0) { + if (ptr == nullptr) { + return InvalidValue(); + } + + *ptr = nullptr; + if (size == 0) { + return Upstream::kSuccess; + } + + const std::size_t align = SliceAlignment(alignment); + const std::size_t rounded = RoundSize(size); + + // Front cache first, for the ordinary case: natural alignment and a size the + // bins cover. A hit takes no lock and touches nothing another thread reads. + // + // Over-aligned requests skip it. A cached block is only known to start at + // `kMinSliceAlignment`, so serving one would mean re-checking alignment and + // falling through on failure -- work on the hot path for a rare request. + if (align == Config::kMinSliceAlignment) { + if (void* cached = TryCacheAllocate(rounded); cached != nullptr) { + *ptr = cached; + return Upstream::kSuccess; + } + } + + // Worst case a chunk must cover. Every chunk starts at a multiple of + // `kMinSliceAlignment` (see `AdoptBacking`), so aligning up to `align` + // costs at most the difference between the two. + const std::size_t needed = rounded + align - Config::kMinSliceAlignment; + + std::unique_lock lock(mutex_); + + { + // Backings the idle scan selects are freed with the lock dropped: + // `cudaFree` implicitly synchronizes the whole device, so it must never + // run inside the critical section, let alone on the allocation path. + std::vector> doomed; + UpdateShrinkState(size, &doomed); + if (!doomed.empty()) { + lock.unlock(); + FreeBackings(doomed); + doomed.clear(); + lock.lock(); + } + } + + if (Chunk* chunk = FindFit(needed); chunk != nullptr) { + ++stats_.cache_hit_count; + *ptr = Serve(chunk, rounded, align); + return Upstream::kSuccess; + } + + // No room anywhere: a new backing store is needed. The upstream allocator + // is a synchronous device call costing hundreds of microseconds, so it runs + // with the lock released -- holding it here would stall every other thread, + // including ones that only need to slice an existing backing. + std::size_t capacity = NextCapacity(needed); + lock.unlock(); + + void* base = nullptr; + Error status = Upstream::Malloc(&base, capacity); + if (status != Upstream::kSuccess) { + status = AllocateFallback(&base, &capacity, needed, status); + } + + lock.lock(); + + if (status != Upstream::kSuccess) { + // Another thread may have released space while the lock was down, which + // turns an upstream failure into a hit. + if (Chunk* chunk = FindFit(needed); chunk != nullptr) { + ++stats_.cache_hit_count; + *ptr = Serve(chunk, rounded, align); + return Upstream::kSuccess; + } + return status; + } + + ++stats_.upstream_alloc_count; + Chunk* chunk = AdoptBacking(base, capacity); + ++stats_.cache_miss_count; + *ptr = Serve(chunk, rounded, align); + return Upstream::kSuccess; + } + + /// Returns a block from `Allocate` to its backing store, where it becomes + /// available at its own size immediately and at any larger size once merging + /// runs (deferred to the next request that needs it -- see the class comment). + /// The memory is not handed back to the upstream allocator here; that happens + /// on automatic shrink, `ReleaseCached`, or destruction. `nullptr` is a no-op. + /// Returns an invalid-value error if `ptr` was not produced by this pool (or + /// was already freed). + Error Deallocate(void* ptr) { + if (ptr == nullptr) { + return Upstream::kSuccess; + } + + // The lock is unavoidable here: rejecting a foreign pointer or a double free + // means consulting the live-block table, which is shared. What the front + // cache saves is not this acquisition but the *next* allocation's -- and the + // best-fit lookup and split that would have come with it. + std::lock_guard lock(mutex_); + + // `Find` rather than `Take`: a chunk already parked in some thread's cache is + // still in the table, and telling that case apart from a live block is what + // makes a double free of a cached block detectable. + Chunk* chunk = nullptr; + if (!allocated_.Find(ptr, &chunk) || chunk->cached) { + return InvalidValue(); + } + + ++stats_.free_count; + + // File it into the releasing thread's cache, so the next allocation of this + // size on this thread needs no lock at all. The table entry deliberately + // stays: the block never became free, it changed hands from the caller to + // the cache. + // + // Resolve the bin before touching thread-local state. Sizes the bins do not + // cover cannot be parked at all, and on a workload built from large blocks + // that is every free -- looking the cache up first would spend a TLS access + // and a scan on every one of them only to be turned away. `IfPresent` + // because registering a cache takes `mutex_`, which is already held: a + // thread that has never allocated a binnable size has no cache here, and a + // free is not a reason to give it one. + const std::size_t bin = BinIndex(chunk->size); + if (bin != kNoBin) { + if (ThreadCache* cache = LocalCacheIfPresent(); + cache != nullptr && TryParkCached(cache, bin, chunk)) { + return Upstream::kSuccess; + } + } + + allocated_.Take(ptr, &chunk); + ReleaseChunk(chunk); + return Upstream::kSuccess; + } + + /// Immediately returns every fully drained backing store to the upstream + /// allocator, including the resident one and ignoring shrink hysteresis. + /// Backings with live blocks are untouched. + /// + /// Note this is a coarser knob than `MemoryPool::ReleaseCached`: free space + /// *inside* a backing that still holds live blocks cannot be handed back, so + /// a fragmented pool may release nothing. `Stats::bytes_free_in_backings` + /// reports how much is retained. + void ReleaseCached() { + // Decide under the lock, free outside it. Stats are settled while the lock + // is held: once detached the backings are no longer the pool's, so a + // concurrent `GetStats` sees them gone even though the frees are in flight. + std::vector> doomed; + { + std::lock_guard lock(mutex_); + // Blocks parked in a cache count live, so a backing holding nothing but + // cached blocks would look occupied. Draining the caches first is what + // makes this release everything a caller has actually returned. + ReclaimAllCaches(); + for (std::size_t i = backings_.size(); i-- > 0;) { + if (IsDrained(backings_[i].get())) { + doomed.push_back(Detach(i)); + } + } + if (backings_.empty()) { + // Nothing is warm any more, so the next allocation should restart the + // growth ramp rather than resume at the burst-time capacity. + next_capacity_ = Config::kInitialCapacity; + } + } + + FreeBackings(doomed); + } + + /// Returns a snapshot of the pool's statistics. + Stats GetStats() const { + std::lock_guard lock(mutex_); + Stats stats = stats_; + stats.backing_count = backings_.size(); + stats.largest_free_chunk = LargestFreeChunk(); + + // Fold in what the lock-free path tallied. Hits are counted per cache + // because counting them in `stats_` would mean a shared write on exactly the + // path that exists to avoid one. + std::size_t cached_bytes = 0; + SumCaches(&cached_bytes, &stats.alloc_count, &stats.cache_hit_count); + + // A parked block is charged to `bytes_in_use` internally -- that is how + // shrink knows not to reclaim the backing under it -- but no caller holds it, + // so reporting it as in use would be wrong: `bytes_in_use` would never reach + // zero after a balanced run. Reported as fragmentation instead, which is what + // it is: reserved, not handed out, not returnable upstream. + stats.bytes_in_use -= cached_bytes; + stats.bytes_free_in_backings += cached_bytes; + return stats; + } + + private: + static_assert( + std::is_invocable_v, + "`Upstream::Malloc` must be callable with `(void**, size_t)`."); + static_assert(std::is_invocable_v, + "`Upstream::Free` must be callable with `(void*)`."); + static_assert( + std::is_same_v, Error>, + "`Upstream` must define `static constexpr Error kSuccess`."); + static_assert( + Config::kMinSliceAlignment != 0 && + (Config::kMinSliceAlignment & (Config::kMinSliceAlignment - 1)) == 0, + "`Config::kMinSliceAlignment` must be a power of two."); + static_assert(Config::kInitialCapacity <= Config::kMaxCapacity, + "`Config::kInitialCapacity` must not exceed " + "`Config::kMaxCapacity`."); + + struct BackingStore; + + // A free chunk is published in exactly one of two containers; a served chunk, + // or one mid-coalesce, is in neither. + enum class Location : std::uint8_t { kNone, kFastBin, kTree }; + + // One extent of a backing store, either free or handed out. Chunks tile their + // backing's usable span in address order with no gaps, so `prev`/`next` are + // exactly the physical neighbors and coalescing is a local operation. + struct Chunk { + BackingStore* owner = nullptr; + Chunk* prev = nullptr; + Chunk* next = nullptr; + void* ptr = nullptr; + std::size_t size = 0; + // Bytes the caller asked for (rounded). Smaller than `size` when a + // remainder was too small to split off; the difference is internal waste. + std::size_t requested = 0; + bool is_free = true; + // Parked in a thread's front cache: not free (its backing still counts it + // live, which is what stops shrink from reclaiming the backing underneath + // the cache) but not held by a caller either, so `Deallocate` must reject it + // as a double free. + bool cached = false; + // Which free container currently holds this chunk, so `EraseFree` can + // unlink it without searching -- and can skip the work entirely for a chunk + // that is in neither, which is every chunk arriving on the `Deallocate` + // coalescing path. + Location location = Location::kNone; + // Links in the exact-size fast bin holding this chunk, when one does. The + // list is doubly linked because coalescing removes a neighbor from the + // middle of a bin, which must stay O(1). + // + // A cached chunk is in no fast bin (`location` is `kNone`, and `EraseFree` + // returns early on that), so `fast_next` does double duty as the front + // cache's list link. The cache is LIFO and never unlinks from the middle, so + // it needs only the forward link. + Chunk* fast_prev = nullptr; + Chunk* fast_next = nullptr; + }; + + // `Chunk`s are owned by their backing store and referenced by raw pointer + // from the live-block table, the free set, and their neighbors. They are + // allocated individually and never moved, so those references stay valid for + // a chunk's whole life. + // + // Backings themselves are held by `unique_ptr` in a vector so that erasing + // one -- which shrink does routinely -- does not move the others. `Chunk` + // stores a raw `owner` pointer, which an index into the vector could not do + // safely: every live chunk's index would shift on erase, and those indices + // live inside the live-block table where they cannot be fixed up. + struct BackingStore { + void* base = nullptr; + std::size_t capacity = 0; + Chunk* head = nullptr; + // Chunks currently handed out. Deferred coalescing means a drained backing + // is no longer recognizable from its list length, so emptiness is this + // counter reaching zero. Maintained by `Serve` and `Deallocate`, the only + // two places a chunk changes hands. + std::size_t live_chunks = 0; + // The resident backing is exempt from automatic shrink so a steady small + // workload keeps a warm arena. + bool resident = false; + // Larger than `Config::kMaxCapacity`, created for one request no capped + // backing could serve. Never resident, and exempt from shrink hysteresis. + bool oversize = false; + std::uint32_t empty_scans = 0; + }; + + // Orders free chunks by size, then by address to break ties. Best fit is + // `lower_bound` on the needed size; the largest free chunk is `rbegin`. + struct BySizeThenAddress { + bool operator()(const Chunk* lhs, const Chunk* rhs) const { + if (lhs->size != rhs->size) { + return lhs->size < rhs->size; + } + return reinterpret_cast(lhs->ptr) < + reinterpret_cast(rhs->ptr); + } + }; + + using FreeSet = + std::set>; + + // Number of exact-size fast bins. Bin `i` holds free chunks of exactly + // `(i + 1) * Config::kMinSliceAlignment` bytes, so the bins cover every + // aligned size up to `kFastBinCount * Config::kMinSliceAlignment` (64 KB + // under the default config). Everything else -- oversize chunks, and the + // non-aligned tails a backing's leading trim leaves behind -- stays in the + // tree. + static constexpr std::size_t kFastBinCount = 128; + static constexpr std::size_t kFastBinWords = (kFastBinCount + 63) / 64; + + // "No such bin": either the size is not bin-eligible, or no bin in the + // scanned range is occupied. + static constexpr std::size_t kNoBin = static_cast(-1); + + static Error InvalidValue() { return static_cast(1); } + + static std::size_t CountTrailingZeros(std::uint64_t word) { +#if defined(__GNUC__) || defined(__clang__) + return static_cast(__builtin_ctzll(word)); +#else + std::size_t count = 0; + while ((word & 1ull) == 0) { + word >>= 1; + ++count; + } + return count; +#endif + } + + // The bin holding chunks of exactly `size`, or `kNoBin` if no bin does. + static std::size_t BinIndex(std::size_t size) { + if (size == 0 || + (size & (Config::kMinSliceAlignment - 1)) != 0) { + return kNoBin; + } + const std::size_t multiples = size / Config::kMinSliceAlignment; + return multiples <= kFastBinCount ? multiples - 1 : kNoBin; + } + + // The lowest-indexed bin whose chunks are large enough for `needed`, or + // `kNoBin` when `needed` outruns the bins entirely. + static std::size_t FirstEligibleBin(std::size_t needed) { + const std::size_t multiples = + (needed + Config::kMinSliceAlignment - 1) / Config::kMinSliceAlignment; + if (multiples == 0) { + return 0; + } + return multiples <= kFastBinCount ? multiples - 1 : kNoBin; + } + + static constexpr std::size_t kCacheDepth = + detail::ArenaCacheDepth::value; + static constexpr std::size_t kCacheBytes = + detail::ArenaCacheBytes::value; + + // Exact-size lists a single thread owns outright. Sizes are indexed the same + // way the fast bins are -- `size / kMinSliceAlignment - 1` -- so a request + // whose rounded size lands in range checks one list head and is done, with no + // lock and no shared line touched. + // + // Only the owning thread reads or writes its lists, so nothing here is atomic. + // What *is* shared is `pool`, `next`, and `orphaned`: the pool reaches a live + // cache through the registry to reclaim it, and a cache reaches the pool on + // thread exit to hand its blocks back. Both directions take `mutex_`. + struct ThreadCache { + // Same bins as the global fast bins, so a size is cacheable exactly when it + // is binnable and the two indexes agree. + Chunk* lists[kFastBinCount] = {}; + std::size_t depths[kFastBinCount] = {}; + std::size_t bytes = 0; + + // Tallies this thread accumulated without the pool's lock. `GetStats` folds + // them in rather than having the hot path write shared counters. Only hits + // are tallied here: every miss and every free already holds `mutex_` and + // writes `stats_` directly. + std::size_t alloc_count = 0; + std::size_t cache_hit_count = 0; + + // Held by the owning thread across a pop or a park, and by any other thread + // reclaiming this cache. Uncontended in the common case -- two atomic RMWs, + // an order of magnitude cheaper than the global mutex it stands in front of. + // A reclaimer always takes it *after* `mutex_` and the owner's fast path + // never holds `mutex_`, so there is one lock order and no cycle. + std::atomic busy{false}; + + // Registry link, guarded by `mutex_`. + ThreadCache* next = nullptr; + }; + + // Shared rendezvous between the pool and the threads holding caches into it. + // Either side may die first: a worker can outlive the pool, and the main + // thread's cache is destroyed at process exit, long after any pool on the + // stack. So the lock lives here rather than in the pool, held alive by a + // `shared_ptr` from each side. A thread exiting locks it and reads `pool`; a + // null `pool` means the pool is gone and took its memory upstream with it. + // + // One lock does both jobs -- guarding pool state and guarding this handshake -- + // precisely so there is no second lock to order against the first. + struct Registry { + std::mutex mutex; + ArenaMemoryPool* pool = nullptr; + ThreadCache* head = nullptr; + }; + + // Owns one thread's `ThreadCache` for one pool. The destructor runs on thread + // exit -- or at process exit for the main thread -- and is what hands a + // departing thread's blocks back, so memory is never stranded in the cache of + // a thread that has gone away. + // + // Holding the registry by `shared_ptr` is what makes the destructor safe in + // either order: if the pool went first it cleared `pool` and already reclaimed + // this cache, and all that is left to do is free the node. + class ThreadCacheHandle { + public: + explicit ThreadCacheHandle(std::shared_ptr registry) + : registry_(std::move(registry)), cache_(new ThreadCache) { + std::lock_guard lock(registry_->mutex); + cache_->next = registry_->head; + registry_->head = cache_; + } + + ~ThreadCacheHandle() { + { + std::lock_guard lock(registry_->mutex); + if (registry_->pool != nullptr) { + registry_->pool->RetireCacheLocked(cache_); + } + Unlink(cache_); + } + delete cache_; + } + + ThreadCacheHandle(const ThreadCacheHandle&) = delete; + ThreadCacheHandle& operator=(const ThreadCacheHandle&) = delete; + + ThreadCache* get() const { return cache_; } + + private: + // Caller must hold `registry_->mutex`. + void Unlink(ThreadCache* cache) { + ThreadCache** link = ®istry_->head; + while (*link != nullptr && *link != cache) { + link = &(*link)->next; + } + if (*link != nullptr) { + *link = cache->next; + } + } + + std::shared_ptr registry_; + ThreadCache* cache_; + }; + + // Spin lock over one `ThreadCache`. Uncontended in the common case -- only a + // reclaim landing on a cache whose owner is mid-operation contends -- and the + // critical sections are a handful of loads, so spinning beats parking. The + // owner's fast path holds this and nothing else; a reclaimer holds `mutex_` + // first, so the one lock order is `mutex_` then `busy`. + class CacheGuard { + public: + explicit CacheGuard(ThreadCache* cache) : cache_(cache) { + while (cache_->busy.exchange(true, std::memory_order_acquire)) { + } + } + ~CacheGuard() { cache_->busy.store(false, std::memory_order_release); } + + CacheGuard(const CacheGuard&) = delete; + CacheGuard& operator=(const CacheGuard&) = delete; + + private: + ThreadCache* cache_; + }; + + // One cache per (thread, pool instance) pair, keyed by registry address. That + // pairing matters because a process may run several pools -- one per device -- + // and a block from one is not servable from another. + // + // A small vector rather than a hash: a thread touches one or two pools, so a + // linear scan over a handful of pointers beats a hash lookup. + struct CacheEntry { + const Registry* key; + std::unique_ptr handle; + }; + + // Destroyed at thread exit in reverse order, which is what returns this + // thread's blocks. Function-local `thread_local` in a template gives one + // instance per `ArenaMemoryPool` specialization, so different `Upstream` types + // never share a vector. + static std::vector& CacheMap() { + static thread_local std::vector caches; + return caches; + } + + ThreadCache* LocalCache() { + if (ThreadCache* cache = LocalCacheIfPresent(); cache != nullptr) { + return cache; + } + // Takes `registry_->mutex` -- which is `mutex_` -- to link the new cache in, + // so this must not run with the pool's lock held. + CacheMap().push_back( + CacheEntry{registry_.get(), + std::make_unique(registry_)}); + return CacheMap().back().handle->get(); + } + + // This thread's cache for this pool if it already has one, else `nullptr`. + // Reads thread-local state and takes no lock, so unlike `LocalCache` it is + // safe to call while holding `mutex_`. + ThreadCache* LocalCacheIfPresent() { + for (const CacheEntry& entry : CacheMap()) { + if (entry.key == registry_.get()) { + return entry.handle->get(); + } + } + return nullptr; + } + + static std::size_t RoundUp(std::size_t size, std::size_t granularity) { + return (size + granularity - 1) / granularity * granularity; + } + + // Requests round to the slice alignment and nothing coarser. The 2 MB + // rounding `MemoryPool` applies to large requests exists to make exact + // size-class matching hit; with splitting and coalescing that motivation is + // gone, and coarse rounding would waste up to 2 MB per large allocation. + static std::size_t RoundSize(std::size_t size) { + return RoundUp(size, Config::kMinSliceAlignment); + } + + static std::size_t SliceAlignment(std::size_t alignment) { + return alignment > Config::kMinSliceAlignment ? alignment + : Config::kMinSliceAlignment; + } + + static void* Offset(void* ptr, std::size_t bytes) { + return static_cast(ptr) + bytes; + } + + static std::size_t AlignPadding(void* ptr, std::size_t alignment) { + const auto address = reinterpret_cast(ptr); + const auto aligned = (address + alignment - 1) & + ~static_cast(alignment - 1); + return static_cast(aligned - address); + } + + // A backing is drained when nothing it contains is handed out. With deferred + // coalescing its chunk list may still hold many adjacent free chunks, so this + // does *not* imply the span is available as one extent -- `Detach` runs + // `CoalesceBacking` to restore that before handing the memory back. + static bool IsDrained(const BackingStore* backing) { + return backing->live_chunks == 0; + } + + Chunk* NewChunk() { + void* storage = chunk_arena_.Allocate(sizeof(Chunk), alignof(Chunk)); + return new (storage) Chunk(); + } + + void DeleteChunk(Chunk* chunk) { + chunk->~Chunk(); + chunk_arena_.Deallocate(chunk, sizeof(Chunk), alignof(Chunk)); + } + + void DestroyChunks(BackingStore* backing) { + Chunk* chunk = backing->head; + while (chunk != nullptr) { + Chunk* next = chunk->next; + DeleteChunk(chunk); + chunk = next; + } + backing->head = nullptr; + } + + // Publishes `chunk` as free, in its exact-size bin when one exists and in the + // tree otherwise. Caller must hold `mutex_`. + void InsertFree(Chunk* chunk) { + const std::size_t bin = BinIndex(chunk->size); + if (bin == kNoBin) { + chunk->location = Location::kTree; + free_chunks_.insert(chunk); + return; + } + + chunk->location = Location::kFastBin; + chunk->fast_prev = nullptr; + chunk->fast_next = fast_bins_[bin]; + if (chunk->fast_next != nullptr) { + chunk->fast_next->fast_prev = chunk; + } + fast_bins_[bin] = chunk; + fast_bitmap_[bin / 64] |= 1ull << (bin % 64); + } + + // Removes `chunk` from whichever container holds it. Must be called before + // any change to `chunk->size`, since the size selects the bin. Caller must + // hold `mutex_`. + void EraseFree(Chunk* chunk) { + if (chunk->location == Location::kNone) { + return; + } + if (chunk->location == Location::kTree) { + chunk->location = Location::kNone; + free_chunks_.erase(chunk); + return; + } + + const std::size_t bin = BinIndex(chunk->size); + if (chunk->fast_prev != nullptr) { + chunk->fast_prev->fast_next = chunk->fast_next; + } else { + fast_bins_[bin] = chunk->fast_next; + } + if (chunk->fast_next != nullptr) { + chunk->fast_next->fast_prev = chunk->fast_prev; + } + chunk->fast_prev = nullptr; + chunk->fast_next = nullptr; + chunk->location = Location::kNone; + + if (fast_bins_[bin] == nullptr) { + fast_bitmap_[bin / 64] &= ~(1ull << (bin % 64)); + } + } + + // Pops an exact-size block from this thread's cache, or returns `nullptr`. + // Takes no pool lock and touches no shared state, which is the whole point. + // + // `bin` selects by exact size, so a hit means the block is exactly `rounded` + // bytes: no split, no alignment padding, and no internal waste to account + // for. Alignment beyond `kMinSliceAlignment` is not served from here at all + // (see `Allocate`), so the chunk's own start alignment is sufficient. + void* TryCacheAllocate(std::size_t rounded) { + const std::size_t bin = BinIndex(rounded); + if (bin == kNoBin) { + return nullptr; + } + + ThreadCache* cache = LocalCache(); + CacheGuard guard(cache); + Chunk* chunk = cache->lists[bin]; + if (chunk == nullptr) { + return nullptr; + } + + cache->lists[bin] = chunk->fast_next; + --cache->depths[bin]; + cache->bytes -= chunk->size; + ++cache->alloc_count; + ++cache->cache_hit_count; + + chunk->fast_next = nullptr; + chunk->cached = false; + // Still in `allocated_` and still counted live by its backing -- parking + // never removed either -- so handing it back needs no shared write at all. + chunk->requested = rounded; + return chunk->ptr; + } + + // Parks `chunk` in `cache`, or returns false if the cache has no room for it. + // Caller must hold `mutex_` and must have established that `chunk` is live. + // + // A parked chunk stays in `allocated_` and stays counted in its backing's + // `live_chunks`, so `stats_.bytes_in_use` still covers it -- which is what + // keeps automatic shrink from handing the backing upstream while the cache + // points into it. `GetStats` reclassifies those bytes as free, since no caller + // holds them. + // + // `bin` is the caller's already-resolved `BinIndex(chunk->size)`, never + // `kNoBin`: the caller has to test that anyway to decide whether looking up a + // cache is worth it, so recomputing it here would be the second time. + bool TryParkCached(ThreadCache* cache, std::size_t bin, Chunk* chunk) { + CacheGuard guard(cache); + if (cache->depths[bin] >= kCacheDepth || + cache->bytes + chunk->size > kCacheBytes) { + return false; + } + + // `requested` becomes the whole extent: a cache hit is an exact-size match, + // so nothing parked here carries internal waste, and the pop can leave the + // waste tally alone. + stats_.bytes_internal_waste -= chunk->size - chunk->requested; + chunk->requested = chunk->size; + chunk->cached = true; + + chunk->fast_next = cache->lists[bin]; + cache->lists[bin] = chunk; + ++cache->depths[bin]; + cache->bytes += chunk->size; + return true; + } + + // Hands every block in `cache` back to the global free index. Caller must hold + // `mutex_`; the cache may belong to another thread, which `CacheGuard` covers. + // + // Leaves the tallies in place: they are monotonic and `GetStats` folds them, so + // clearing them here would lose allocations from the totals. + void ReclaimCacheLocked(ThreadCache* cache) { + CacheGuard guard(cache); + if (cache->bytes == 0) { + return; + } + + for (std::size_t bin = 0; bin < kFastBinCount; ++bin) { + Chunk* chunk = cache->lists[bin]; + cache->lists[bin] = nullptr; + cache->depths[bin] = 0; + while (chunk != nullptr) { + Chunk* next = chunk->fast_next; + chunk->fast_next = nullptr; + chunk->cached = false; + ReleaseChunk(chunk); + chunk = next; + } + } + cache->bytes = 0; + } + + // Drains `cache` and absorbs its tallies, for a cache about to leave the + // registry with its thread. `GetStats` folds live caches' tallies as it reads + // them, so a departing one has to hand its own over or the allocations it + // served would vanish from the totals. Caller must hold `mutex_`. + void RetireCacheLocked(ThreadCache* cache) { + ReclaimCacheLocked(cache); + stats_.alloc_count += cache->alloc_count; + stats_.cache_hit_count += cache->cache_hit_count; + cache->alloc_count = 0; + cache->cache_hit_count = 0; + } + + // Drains every registered cache. Caller must hold `mutex_`. + void ReclaimAllCaches() { + for (ThreadCache* cache = registry_->head; cache != nullptr; + cache = cache->next) { + ReclaimCacheLocked(cache); + } + } + + // Total bytes parked across every cache, and the tallies to fold into `Stats`. + // Caller must hold `mutex_`. + void SumCaches(std::size_t* bytes, std::size_t* allocs, + std::size_t* hits) const { + for (ThreadCache* cache = registry_->head; cache != nullptr; + cache = cache->next) { + CacheGuard guard(cache); + *bytes += cache->bytes; + *allocs += cache->alloc_count; + *hits += cache->cache_hit_count; + } + } + + // Returns a chunk that is no longer held by anyone to the global free index. + // `chunk` must already be out of `allocated_` and must not be cached. Caller + // must hold `mutex_`. + // + // Shared by `Deallocate` and by cache reclaim, which differ only in who was + // holding the block; the accounting from here down is identical. + void ReleaseChunk(Chunk* chunk) { + stats_.bytes_in_use -= chunk->size; + stats_.bytes_internal_waste -= chunk->size - chunk->requested; + + chunk->is_free = true; + chunk->requested = 0; + stats_.bytes_free_in_backings += chunk->size; + + // No coalescing here -- see the class comment. Publishing the chunk at its + // own size is the whole critical section, and for the common case (the next + // request for this size pops this very chunk back out of its bin) merging + // would be pure overhead: two neighbor loads, a tree rebalance, and a second + // rebalance to split the result apart again. + BackingStore* owner = chunk->owner; + // Conservative: this free may have created an adjacent free pair, so the + // next `FindFit` miss has to try coalescing before growing. + coalesce_dirty_ = true; + InsertFree(chunk); + + if (--owner->live_chunks == 0 && !owner->resident) { + // This backing just drained and is shrinkable. Merge it now, even though + // merging is otherwise deferred: the leftover slivers are *bait*. A 64 B + // tail sitting in an exact-size bin is the best possible fit for the next + // 64 B request, so it would pull a fresh live block into the one backing + // that was about to be handed back, and shrink could never reclaim it. + // Merged, the backing presents a single large chunk that best fit passes + // over while any smaller one exists. + // + // The resident backing is deliberately excluded: it is shrink-exempt, so + // it has no bait problem, and it is where a near-empty alloc/free loop + // lives -- coalescing it on every drain would put a list walk and a tree + // round trip back on exactly the hot path this deferral exists to clear. + CoalesceBacking(owner); + ++drained_candidates_; + } + } + + // First chunk in the lowest occupied bin at or above `from`, or `nullptr`. + // The bitmap makes this a couple of word scans rather than a walk over 128 + // list heads. Caller must hold `mutex_`. + Chunk* ScanBins(std::size_t from) const { + if (from >= kFastBinCount) { + return nullptr; + } + std::size_t word = from / 64; + std::uint64_t bits = fast_bitmap_[word] & (~0ull << (from % 64)); + while (bits == 0) { + if (++word >= kFastBinWords) { + return nullptr; + } + bits = fast_bitmap_[word]; + } + return fast_bins_[word * 64 + CountTrailingZeros(bits)]; + } + + // Best fit over what is currently indexed, coalescing and retrying once if + // nothing fits. Deferred merging means a failure here does not mean the pool + // is out of room -- adjacent free chunks may add up to a fit -- so growing a + // new backing must never be decided on `FindFitIndexed` alone. + // + // Caller must hold `mutex_`. + Chunk* FindFit(std::size_t needed) { + if (Chunk* chunk = FindFitIndexed(needed); chunk != nullptr) { + return chunk; + } + if (coalesce_dirty_) { + CoalesceAll(); + if (Chunk* chunk = FindFitIndexed(needed); chunk != nullptr) { + return chunk; + } + } + + // Last resort before growing: the memory may be parked in some thread's + // cache. Reclaiming is what keeps the caches from turning retention into an + // upstream call -- a thread must never hold blocks another thread's request + // cannot get back. + if (!AnyCached()) { + return nullptr; + } + ReclaimAllCaches(); + if (Chunk* chunk = FindFitIndexed(needed); chunk != nullptr) { + return chunk; + } + // Reclaim published chunks at their own sizes, which may have created + // adjacent free pairs of its own. + if (!coalesce_dirty_) { + return nullptr; + } + CoalesceAll(); + return FindFitIndexed(needed); + } + + // Whether any registered cache holds anything. Caller must hold `mutex_`. + bool AnyCached() const { + for (ThreadCache* cache = registry_->head; cache != nullptr; + cache = cache->next) { + CacheGuard guard(cache); + if (cache->bytes != 0) { + return true; + } + } + return false; + } + + // Best fit: the smallest indexed free chunk that can hold `needed`. Caller + // must hold `mutex_`. + Chunk* FindFitIndexed(std::size_t needed) { + const std::size_t first = FirstEligibleBin(needed); + + // An occupied exact-fit bin ends the search: no chunk anywhere can be a + // better fit, so the common case -- a workload cycling through a handful of + // sizes -- never touches the tree at all. This is the whole point of the + // bins; a best-fit tree lookup on every allocation was the pool's dominant + // per-call cost. + if (first != kNoBin && fast_bins_[first] != nullptr) { + return fast_bins_[first]; + } + + // Otherwise both containers are candidates. Bins beyond the exact one hold + // aligned sizes; the tree holds everything too large to bin plus the + // non-aligned tails a backing's leading trim leaves behind, either of which + // may be the tighter fit. + Chunk* binned = first == kNoBin ? nullptr : ScanBins(first + 1); + + Chunk probe{}; + probe.size = needed; + probe.ptr = nullptr; // Sorts before any real chunk of the same size. + auto it = free_chunks_.lower_bound(&probe); + Chunk* treed = it == free_chunks_.end() ? nullptr : *it; + + if (binned == nullptr) { + return treed; + } + if (treed == nullptr) { + return binned; + } + return treed->size < binned->size ? treed : binned; + } + + // Largest request servable without going upstream, for `Stats`. Because + // coalescing is deferred, this is the longest *run* of adjacent free chunks, + // not the largest indexed one -- a drained backing sitting as fifty separate + // free chunks can still serve its full span, and reporting the largest single + // chunk would understate the pool's capability by the deferral. + // + // Walks the chunk lists rather than the free index, so it is linear in the + // chunk count. `GetStats` is a diagnostic call, not on the hot path. + // + // Caller must hold `mutex_`. + std::size_t LargestFreeChunk() const { + std::size_t largest = 0; + for (const auto& backing : backings_) { + std::size_t run = 0; + for (const Chunk* chunk = backing->head; chunk != nullptr; + chunk = chunk->next) { + // A cached chunk counts toward the run: no caller holds it, so any + // request that needs the span reclaims the caches and gets it. Excluding + // it would report a drained backing as shattered. + if (chunk->is_free || chunk->cached) { + run += chunk->size; + if (run > largest) { + largest = run; + } + } else { + run = 0; + } + } + } + return largest; + } + + // Merges every adjacent free pair in `backing`, leaving one chunk per + // contiguous free region. Runs in one address-order walk, so a full pass over + // the pool is linear in the number of chunks rather than in merges performed. + // Caller must hold `mutex_`. + void CoalesceBacking(BackingStore* backing) { + for (Chunk* chunk = backing->head; chunk != nullptr;) { + if (!chunk->is_free) { + chunk = chunk->next; + continue; + } + // Absorb the whole free run to the right in one go. Only the survivor is + // re-published, so a run of N chunks costs one insert, not N. + Chunk* next = chunk->next; + if (next == nullptr || !next->is_free) { + chunk = next; + continue; + } + EraseFree(chunk); + while (next != nullptr && next->is_free) { + EraseFree(next); + chunk->size += next->size; + Chunk* after = next->next; + DeleteChunk(next); + next = after; + } + chunk->next = next; + if (next != nullptr) { + next->prev = chunk; + } + InsertFree(chunk); + chunk = next; + } + } + + // Coalesces every backing. Caller must hold `mutex_`. + void CoalesceAll() { + for (const auto& backing : backings_) { + CoalesceBacking(backing.get()); + } + coalesce_dirty_ = false; + } + + // Splits `chunk` at `offset` and returns the tail. Purely structural: the + // caller owns both halves' free-set membership and accounting, because the + // two call sites want different outcomes for the left part. + Chunk* SplitAt(Chunk* chunk, std::size_t offset) { + Chunk* tail = NewChunk(); + tail->owner = chunk->owner; + tail->ptr = Offset(chunk->ptr, offset); + tail->size = chunk->size - offset; + tail->is_free = chunk->is_free; + tail->prev = chunk; + tail->next = chunk->next; + if (chunk->next != nullptr) { + chunk->next->prev = tail; + } + chunk->next = tail; + chunk->size = offset; + return tail; + } + + // Carves `rounded` bytes at `alignment` out of the free chunk `chunk`, + // registers the result, and returns the pointer the caller sees. `chunk` must + // be large enough for the request including alignment padding, which is what + // `FindFit`'s `needed` guarantees. Caller must hold `mutex_`. + void* Serve(Chunk* chunk, std::size_t rounded, std::size_t alignment) { + EraseFree(chunk); + BackingStore* owner = chunk->owner; + if (owner->live_chunks++ == 0) { + DropCandidate(owner); + } + stats_.bytes_free_in_backings -= chunk->size; + // The whole extent is the pool's to carve now; the pieces handed back below + // are re-marked individually. + chunk->is_free = false; + + // Aligning inside the chunk leaves a leading gap. Split it off as its own + // free chunk rather than folding it into the served block: with coalescing + // it is genuinely reusable, which is also why this pool does not need + // `MemoryPool`'s trick of over-requesting `rounded + alignment` upstream. + const std::size_t padding = AlignPadding(chunk->ptr, alignment); + if (padding != 0) { + Chunk* body = SplitAt(chunk, padding); + chunk->is_free = true; + stats_.bytes_free_in_backings += chunk->size; + InsertFree(chunk); + chunk = body; + } + + if (chunk->size - rounded >= Config::kMinSplitRemainder) { + Chunk* tail = SplitAt(chunk, rounded); + tail->is_free = true; + stats_.bytes_free_in_backings += tail->size; + InsertFree(tail); + } + // Otherwise the remainder stays with the served chunk: splitting it off + // would only create a sliver too small to satisfy anything. It is counted + // as internal waste below. + + chunk->requested = rounded; + allocated_.Insert(chunk->ptr, chunk); + + stats_.bytes_in_use += chunk->size; + stats_.bytes_internal_waste += chunk->size - chunk->requested; + if (stats_.bytes_in_use > stats_.peak_bytes_in_use) { + stats_.peak_bytes_in_use = stats_.bytes_in_use; + } + ++stats_.alloc_count; + return chunk->ptr; + } + + // Capacity to request for a new backing store. The floor carries one extra + // slice alignment because the upstream base may not be `kMinSliceAlignment` + // aligned and `AdoptBacking` trims up to that much off the front. + // + // Caller must hold `mutex_`. + std::size_t NextCapacity(std::size_t needed) const { + const std::size_t floor = needed + Config::kMinSliceAlignment; + return next_capacity_ > floor ? next_capacity_ : floor; + } + + // Takes ownership of a fresh upstream allocation and returns its sole free + // chunk. Caller must hold `mutex_`. + Chunk* AdoptBacking(void* base, std::size_t capacity) { + auto backing = std::make_unique(); + backing->base = base; + backing->capacity = capacity; + backing->oversize = capacity > Config::kMaxCapacity; + // Residency follows the growth ramp, not arrival order: a first request + // that happens to be huge gets an oversize backing, and pinning *that* + // forever would be the opposite of what shrink is for. + backing->resident = !backing->oversize && !HasResident(); + + // Trim the front so every chunk in this backing -- and therefore every + // pointer the pool hands out -- starts at a multiple of the slice + // alignment. `base` itself is only guaranteed whatever the upstream + // allocator promises. `capacity` still records what upstream gave us, since + // that is what `Upstream::Free` releases. + const std::size_t lead = AlignPadding(base, Config::kMinSliceAlignment); + + Chunk* chunk = NewChunk(); + chunk->owner = backing.get(); + chunk->ptr = Offset(base, lead); + chunk->size = capacity - lead; + chunk->is_free = true; + backing->head = chunk; + + stats_.bytes_reserved += capacity; + if (stats_.bytes_reserved > stats_.peak_bytes_reserved) { + stats_.peak_bytes_reserved = stats_.bytes_reserved; + } + stats_.bytes_free_in_backings += chunk->size; + stats_.bytes_unusable += lead; + InsertFree(chunk); + // A fresh backing is drained by definition, and the caller `Serve`s out of + // it immediately -- which decrements. Counting it here keeps the pair + // balanced. + if (!backing->resident) { + ++drained_candidates_; + } + + if (!backing->oversize) { + // Advance the ramp, capped. An oversize backing is a one-off exception + // and must not advance it -- otherwise a single 2 GB request would make + // the next ordinary backing 4 GB. + next_capacity_ = capacity >= Config::kMaxCapacity / 2 + ? Config::kMaxCapacity + : capacity * 2; + } + + backings_.push_back(std::move(backing)); + return chunk; + } + + // `drained_candidates_` counts drained, shrinkable (non-resident) backings: + // exactly what an idle scan could find. Keeping it exact is what lets the + // allocation path skip that scan on a single load in the steady state, so + // every transition into or out of "drained and non-resident" must be paired. + // The resident backing is never counted -- it is shrink-exempt, so counting it + // would defeat the skip whenever the pool is idle. + // + // Caller must hold `mutex_`. + void DropCandidate(const BackingStore* backing) { + if (!backing->resident) { + --drained_candidates_; + } + } + + // Caller must hold `mutex_`. + bool HasResident() const { + for (const auto& backing : backings_) { + if (backing->resident) { + return true; + } + } + return false; + } + + // Removes the backing at `index` from the pool, settling stats and dropping + // its chunk from the free set. The returned owner keeps the upstream pointer + // alive until the caller frees it outside the lock. Caller must hold + // `mutex_`, and the backing must be drained. + std::unique_ptr Detach(std::size_t index) { + std::unique_ptr backing = std::move(backings_[index]); + backings_.erase(backings_.begin() + static_cast(index)); + + // Deferred coalescing may have left the drained span as many chunks. Merge + // them so the accounting below (and the `head`-spans-everything assumption) + // holds, exactly as it did when `Deallocate` merged eagerly. + CoalesceBacking(backing.get()); + DropCandidate(backing.get()); + + EraseFree(backing->head); + stats_.bytes_free_in_backings -= backing->head->size; + stats_.bytes_unusable -= backing->capacity - backing->head->size; + stats_.bytes_reserved -= backing->capacity; + ++stats_.upstream_free_count; + DestroyChunks(backing.get()); + return backing; + } + + static void FreeBackings( + const std::vector>& doomed) { + for (const auto& backing : doomed) { + Upstream::Free(backing->base); + } + } + + // Updates the consecutive-small-allocation run and, when it indicates a burst + // has ended, runs one idle scan. Selected backings are moved to `*doomed` for + // the caller to free once the lock is dropped. Caller must hold `mutex_`. + void UpdateShrinkState(std::size_t size, + std::vector>* doomed) { + // Nothing is shrinkable, so there is nothing for a scan to find and no + // reason to touch the run counter. This is the steady state -- one resident + // backing, or every backing holding live blocks -- and skipping it here + // keeps `Allocate` from dirtying a shared cache line on every call, which + // under contention costs a remote-dirty miss for whichever thread holds the + // lock next. + if (drained_candidates_ == 0) { + return; + } + + if (size > Config::kSmallThreshold) { + // A large request means a burst is starting or ongoing: keep everything. + small_alloc_since_last_trim_ = 0; + return; + } + + if (++small_alloc_since_last_trim_ < Config::kShrinkThreshold) { + return; + } + // Reset unconditionally, so a scan that frees nothing does not repeat on + // every subsequent allocation. + small_alloc_since_last_trim_ = 0; + + // The burst is over, so the blocks it left parked in caches are what is + // standing between its backings and the upstream allocator. Reclaiming here + // is what lets the scan below actually find them drained -- and it is also + // where a thread that has gone quiet stops holding memory: the run of small + // allocations that got us here is the signal that nothing needs it. + ReclaimAllCaches(); + + // Reverse iteration keeps the remaining indices valid across erases. + for (std::size_t i = backings_.size(); i-- > 0;) { + BackingStore* backing = backings_[i].get(); + if (backing->resident) { + continue; + } + if (!IsDrained(backing)) { + backing->empty_scans = 0; + continue; + } + ++backing->empty_scans; + // Oversize backings skip the hysteresis: holding gigabytes idle for + // another round of small allocations costs far more than one upstream + // call. + if (!backing->oversize && + backing->empty_scans < Config::kEmptyScansToDestroy) { + continue; + } + doomed->push_back(Detach(i)); + ++stats_.shrink_count; + } + } + + // OOM fallback chain, called with the lock released after `Upstream::Malloc` + // failed. Frees drained backings (the upstream allocator may be exactly what + // is out of memory), retries at the requested capacity, then retries at the + // smallest capacity that can serve this one request. Returns `failure` -- the + // original error -- if none of that helps. + Error AllocateFallback(void** base, std::size_t* capacity, + std::size_t needed, Error failure) { + std::vector> doomed; + { + std::lock_guard lock(mutex_); + // Upstream is out of memory, so every retained byte is worth having back. + ReclaimAllCaches(); + for (std::size_t i = backings_.size(); i-- > 0;) { + if (IsDrained(backings_[i].get())) { + doomed.push_back(Detach(i)); + } + } + } + + if (!doomed.empty()) { + FreeBackings(doomed); + doomed.clear(); + if (Upstream::Malloc(base, *capacity) == Upstream::kSuccess) { + return Upstream::kSuccess; + } + } + + const std::size_t minimum = needed + Config::kMinSliceAlignment; + if (*capacity > minimum && + Upstream::Malloc(base, minimum) == Upstream::kSuccess) { + *capacity = minimum; + return Upstream::kSuccess; + } + return failure; + } + + // Outlives this pool when a thread holding a cache does. Declared first so it + // is constructed before the reference below binds to it. + std::shared_ptr registry_; + // The pool's one lock, living in the registry so a departing thread can take + // it without having to know whether the pool is still there. Named as a member + // because every critical section in this file locks it directly. + std::mutex& mutex_; + + // Declared before `free_chunks_`: the set's nodes come from `free_set_arena_`, + // so the arena must outlive it. Members are destroyed in reverse declaration + // order, which puts the set first. + detail::NodeArena chunk_arena_; + detail::NodeArena free_set_arena_; + + detail::PointerTable allocated_; + FreeSet free_chunks_{BySizeThenAddress{}, + detail::ArenaAllocator{&free_set_arena_}}; + // Exact-size free lists and an occupancy bitmap over them. Together these + // keep the allocation hot path off the tree: a request whose size matches an + // occupied bin is served by popping a list head. + Chunk* fast_bins_[kFastBinCount] = {}; + std::uint64_t fast_bitmap_[kFastBinWords] = {}; + + std::vector> backings_; + std::size_t next_capacity_ = Config::kInitialCapacity; + std::size_t small_alloc_since_last_trim_ = 0; + // Drained non-resident backings: the exact number of things an idle scan could + // find. Zero is the steady state and lets `UpdateShrinkState` return on one + // load. + std::size_t drained_candidates_ = 0; + // Set by any `Deallocate` that may have created an adjacent free pair, cleared + // by `CoalesceAll`. Lets a `FindFit` miss skip the coalescing pass when + // nothing has been released since the last one. + bool coalesce_dirty_ = false; + Stats stats_; +}; + +} // namespace infini::rt + +#endif diff --git a/include/infini/rt/detail/node_arena.h b/include/infini/rt/detail/node_arena.h new file mode 100644 index 0000000..8cc6f48 --- /dev/null +++ b/include/infini/rt/detail/node_arena.h @@ -0,0 +1,156 @@ +#ifndef INFINI_RT_DETAIL_NODE_ARENA_H_ +#define INFINI_RT_DETAIL_NODE_ARENA_H_ + +#include +#include +#include +#include +#include + +namespace infini::rt::detail { + +/// ## Recycling node storage for a node-based container. +/// +/// A pool that keeps its free extents in an ordered container (`std::set`) hits +/// the same problem the pools themselves exist to solve: the container calls +/// `operator new` once per insert and `operator delete` once per erase. In a +/// steady-state alloc/free loop that is one host heap round trip per pooled +/// allocation, which is most of the overhead a pool is meant to remove. +/// +/// `NodeArena` hands out fixed-size nodes from bulk-allocated blocks and keeps +/// released nodes on an intrusive free list, so after warm-up a container +/// backed by it performs no host allocation at all. Blocks are never returned +/// individually; the whole arena is freed at destruction. +/// +/// The arena specializes itself to the first node size it sees, which is the +/// only size a given container ever asks for. Requests of any other size (or a +/// stricter alignment) fall through to the global allocation functions, so the +/// arena stays correct even if it is shared or reused. +/// +/// Not thread-safe: callers serialize access with their own lock. +class NodeArena { + public: + NodeArena() = default; + + NodeArena(const NodeArena&) = delete; + NodeArena& operator=(const NodeArena&) = delete; + + ~NodeArena() { + for (void* block : blocks_) { + ::operator delete(block, std::align_val_t{node_align_}); + } + } + + void* Allocate(std::size_t bytes, std::size_t alignment) { + if (node_size_ == 0) { + // First request fixes the pooled geometry. A node must be able to hold + // the free-list link while it is unused. + node_size_ = std::max(bytes, sizeof(void*)); + node_align_ = std::max(alignment, alignof(void*)); + } + + if (bytes > node_size_ || alignment > node_align_) { + return ::operator new(bytes, std::align_val_t{alignment}); + } + + if (free_ == nullptr) { + Grow(); + } + void* node = free_; + // The link lives in the node's own storage; `memcpy` reads it back without + // assuming anything about the object that used to be there. + std::memcpy(&free_, node, sizeof(void*)); + return node; + } + + void Deallocate(void* node, std::size_t bytes, std::size_t alignment) { + if (node == nullptr) { + return; + } + if (bytes > node_size_ || alignment > node_align_) { + ::operator delete(node, std::align_val_t{alignment}); + return; + } + std::memcpy(node, &free_, sizeof(void*)); + free_ = node; + } + + private: + // Blocks grow geometrically so a large live set costs a bounded number of + // host allocations, then capped so one huge burst does not reserve an + // unreasonable block. + static constexpr std::size_t kInitialNodes = 32; + static constexpr std::size_t kMaxNodesPerBlock = 4096; + + void Grow() { + const std::size_t count = next_count_; + next_count_ = std::min(next_count_ * 2, kMaxNodesPerBlock); + + const std::size_t stride = + (node_size_ + node_align_ - 1) / node_align_ * node_align_; + void* block = + ::operator new(stride * count, std::align_val_t{node_align_}); + blocks_.push_back(block); + + char* cursor = static_cast(block); + for (std::size_t i = 0; i < count; ++i) { + Deallocate(cursor + i * stride, node_size_, node_align_); + } + } + + std::vector blocks_; + void* free_ = nullptr; + std::size_t node_size_ = 0; + std::size_t node_align_ = alignof(std::max_align_t); + std::size_t next_count_ = kInitialNodes; +}; + +/// Standard-library allocator adaptor over a `NodeArena`. The arena is not +/// owned: it must outlive every container using it, which callers arrange by +/// declaring the arena before the container it backs. +template +class ArenaAllocator { + public: + using value_type = T; + + explicit ArenaAllocator(NodeArena* arena) : arena_(arena) {} + + template + ArenaAllocator(const ArenaAllocator& other) // NOLINT: allocator rebind + : arena_(other.arena()) {} + + T* allocate(std::size_t count) { + if (count != 1) { + return static_cast( + ::operator new(count * sizeof(T), std::align_val_t{alignof(T)})); + } + return static_cast(arena_->Allocate(sizeof(T), alignof(T))); + } + + void deallocate(T* ptr, std::size_t count) { + if (count != 1) { + ::operator delete(ptr, std::align_val_t{alignof(T)}); + return; + } + arena_->Deallocate(ptr, sizeof(T), alignof(T)); + } + + NodeArena* arena() const { return arena_; } + + template + bool operator==(const ArenaAllocator& other) const { + return arena_ == other.arena(); + } + + template + bool operator!=(const ArenaAllocator& other) const { + return arena_ != other.arena(); + } + + private: + NodeArena* arena_; +}; + +} // namespace infini::rt::detail + +#endif diff --git a/include/infini/rt/detail/pointer_table.h b/include/infini/rt/detail/pointer_table.h new file mode 100644 index 0000000..e6bd4ec --- /dev/null +++ b/include/infini/rt/detail/pointer_table.h @@ -0,0 +1,186 @@ +#ifndef INFINI_RT_DETAIL_POINTER_TABLE_H_ +#define INFINI_RT_DETAIL_POINTER_TABLE_H_ + +#include +#include +#include +#include + +namespace infini::rt::detail { + +/// ## Open-addressing table mapping a live pointer to a `Value`. +/// +/// Every pool in this directory needs the same structure: given the pointer a +/// caller hands back, find the bookkeeping record for it. `std::unordered_map` +/// is node-based, so it would call `operator new` on every insert and +/// `operator delete` on every erase -- meaning each pooled allocation performs +/// a host heap allocation of its own, which is most of what a pool is trying to +/// avoid. This table stores values inline in one vector and only allocates when +/// it grows, so a steady-state alloc/free loop performs no host allocation at +/// all. +/// +/// Linear probing with tombstones; the load factor is held at 1/2 so probe +/// sequences stay short and an empty slot always terminates a probe. +/// +/// Not thread-safe: callers serialize access with their own lock. +template +class PointerTable { + public: + void Insert(void* key, const Value& value) { + // Tombstones count toward the load factor: they still sit on probe paths, + // and a table saturated with them would break the empty-slot terminator. + if ((occupied_ + 1) * 2 > slots_.size()) { + Rehash(); + } + + const std::size_t mask = slots_.size() - 1; + std::size_t index = Hash(key) & mask; + std::size_t tombstone = kNoSlot; + + for (;; index = (index + 1) & mask) { + Slot& slot = slots_[index]; + if (slot.state == State::kOccupied) { + if (slot.key == key) { // Overwrite an existing entry. + slot.value = value; + return; + } + continue; + } + if (slot.state == State::kTombstone) { + if (tombstone == kNoSlot) { + tombstone = index; + } + continue; + } + break; // Empty: the key is absent. + } + + if (tombstone != kNoSlot) { + index = tombstone; // Reuse a tombstone ahead of the empty slot. + --tombstones_; + } else { + ++occupied_; + } + + slots_[index] = Slot{key, value, State::kOccupied}; + ++live_; + } + + /// Writes `key`'s value to `*out` without removing it. Returns false if `key` + /// is not present. Lets a caller inspect a record before deciding whether the + /// entry should come out, which `Take` alone cannot do -- it has already + /// tombstoned the slot by the time the value is available. + bool Find(void* key, Value* out) const { + if (live_ == 0) { + return false; + } + + const std::size_t mask = slots_.size() - 1; + for (std::size_t index = Hash(key) & mask;; index = (index + 1) & mask) { + const Slot& slot = slots_[index]; + if (slot.state == State::kEmpty) { + return false; + } + if (slot.state == State::kOccupied && slot.key == key) { + *out = slot.value; + return true; + } + } + } + + /// Removes `key` and writes its value to `*out`. Returns false if `key` is + /// not present, which is how a pool's `Deallocate` detects a foreign pointer. + bool Take(void* key, Value* out) { + if (live_ == 0) { + return false; + } + + const std::size_t mask = slots_.size() - 1; + for (std::size_t index = Hash(key) & mask;; index = (index + 1) & mask) { + Slot& slot = slots_[index]; + if (slot.state == State::kEmpty) { + return false; + } + if (slot.state == State::kOccupied && slot.key == key) { + *out = slot.value; + slot.state = State::kTombstone; + slot.key = nullptr; + ++tombstones_; + --live_; + return true; + } + } + } + + template + void ForEach(Visitor&& visitor) const { + for (const Slot& slot : slots_) { + if (slot.state == State::kOccupied) { + visitor(slot.value); + } + } + } + + /// Number of live entries. + std::size_t Size() const { return live_; } + + private: + enum class State : std::uint8_t { kEmpty, kOccupied, kTombstone }; + + struct Slot { + void* key = nullptr; + Value value{}; + State state = State::kEmpty; + }; + + static constexpr std::size_t kInitialSlots = 16; + static constexpr std::size_t kNoSlot = static_cast(-1); + + // Pointers from an allocator are aligned, so their low bits are mostly + // zero; a multiply-shift spreads the informative high bits down. + static std::size_t Hash(void* key) { + auto value = + static_cast(reinterpret_cast(key)); + value *= 0x9e3779b97f4a7c15ULL; + return static_cast(value >> 29); + } + + // Rebuilds the table, dropping tombstones. Capacity is sized for the live + // entries, not for `occupied_`: the tombstones counted there are discarded + // by this very rebuild, so sizing for them would buy room for what is about + // to be thrown away. Capacity stays a power of two. + // + // A churn-heavy workload reaches the load factor via tombstones rather than + // live entries, so this usually rebuilds at the same capacity instead of + // growing -- hence the name. + void Rehash() { + std::size_t capacity = kInitialSlots; + while (capacity <= (live_ + 1) * 2) { + capacity *= 2; + } + + std::vector old_slots(capacity); + old_slots.swap(slots_); + occupied_ = 0; + tombstones_ = 0; + live_ = 0; + + // Reusing `Insert` cannot recurse: `capacity` was chosen above the load + // factor for exactly this many entries, so the check in `Insert` stays + // false throughout. + for (const Slot& slot : old_slots) { + if (slot.state == State::kOccupied) { + Insert(slot.key, slot.value); + } + } + } + + std::vector slots_; + std::size_t occupied_ = 0; // live + tombstones, for the load factor + std::size_t tombstones_ = 0; + std::size_t live_ = 0; +}; + +} // namespace infini::rt::detail + +#endif diff --git a/scripts/compare_allocators.py b/scripts/compare_allocators.py new file mode 100644 index 0000000..eac235e --- /dev/null +++ b/scripts/compare_allocators.py @@ -0,0 +1,439 @@ +"""Compare every allocation strategy across the CPU and device backends. + +The question this answers is three questions, and they need different builds to +answer honestly: + + direct vs arena Is the arena worth having at all on this backend? + pool vs arena Which pool design wins, and on which shapes? + arena vs cuda_async Does a hand-written arena beat the vendor's own pool? + +The third only exists on a device, and the first has a completely different +answer on each backend -- a host `malloc` costs tens of nanoseconds, so nothing +the arena saves can pay for its bookkeeping, while a `cudaMalloc` costs hundreds +of microseconds and the same arena wins by orders of magnitude. Reporting one +number for "the allocator" would average those into something true of neither. +So this configures and runs both builds and prints each comparison per backend. + +Usage: + + python scripts/compare_allocators.py # build, run, compare + python scripts/compare_allocators.py --quick # shorter arms + python scripts/compare_allocators.py --no-build # reuse what is built + python scripts/compare_allocators.py --backend cpu # one backend only + +A note on why the builds are serialized rather than parallel: `generated/` is +written into the *source* tree at configure time and its contents depend on which +backends are enabled, so two configured build trees cannot be compiled +concurrently -- the second would compile against the first's headers. Each +backend is therefore configured, built, and only then is anything run. +""" + +import argparse +import json +import pathlib +import platform +import re +import subprocess +import sys + +# The executables that emit per-allocator rows. Only `perf_allocator_matrix` +# does: it is the one binary that runs every arm over one set of workloads, which +# is what makes its rows pivotable into the tables below. `perf_memory_pool` +# covers the general shapes, but only over `direct`/`pool`/`arena`, and it is +# driven by `scripts/run_performance_tests.py` instead. +_MATRIX_TESTS = ("perf_allocator_matrix",) + +# Each backend's build tree, the CMake options that select it, and whether a +# device is required. Ordered CPU-first so a machine without a GPU still gets a +# useful run before anything fails. +_BACKENDS = ( + { + "name": "cpu", + "build_dir": "build-perf-cpu", + "options": ["-DWITH_CPU=ON", "-DWITH_NVIDIA=OFF"], + }, + { + "name": "nvidia", + "build_dir": "build-perf-cuda", + "options": ["-DWITH_CPU=OFF", "-DWITH_NVIDIA=ON"], + }, +) + +# The three comparisons, as (baseline, candidate) arm pairs. `arena` is the +# candidate in every one -- including against `cuda_async`, where the natural +# phrasing would put the vendor's pool second. Keeping the arena in the candidate +# column means the win column always answers the same question, "should we adopt +# this thing", instead of flipping direction in the middle of the report. +_COMPARISONS = ( + ("direct", "arena", "the backend allocator vs the arena"), + ("pool", "arena", "the size-class pool vs the arena"), + ("cuda_async", "arena", "the vendor's stream-ordered pool vs the arena"), +) + +# Units where a larger number is the better outcome. Everything else here is a +# cost -- a duration, a byte count, a call count -- and lower wins. +_HIGHER_IS_BETTER = frozenset({"GiB/s", "seq_per_s"}) + +# Ratio rows are already a ratio; comparing two of them is meaningless. +# +# `x` covers both the retention-amplification rows and the fragmentation probe +# success rate. The latter is the one row here where a *ratio* is the primary +# result rather than a derived one, and it is deliberately not turned into a win +# column: 8/8 vs 8/8 is the expected outcome for both designs, and a 1.00x win +# column would read as "no difference measured" rather than "both succeeded". +_RATIO_UNITS = frozenset({"x"}) + +# Diagnostics rather than costs. `LedgerAccuracy` asks whether one allocator's +# self-reported retention matches what the driver says it took -- a question about +# that allocator's honesty, answered by the binary's own table. Racing two arms' +# answers would print a win column for a row where neither arm is competing. +_DIAGNOSTIC_BENCHMARKS = frozenset({"allocator_matrix.LedgerAccuracy"}) + + +def _repo_root(): + return pathlib.Path(__file__).resolve().parents[1] + + +def _run(command, **kwargs): + printable = " ".join(str(part) for part in command) + print(f"$ {printable}", flush=True) + return subprocess.run(command, check=True, **kwargs) + + +def _configure_and_build(backend, jobs): + """Configure and compile one backend's tree. + + Configuring rewrites `generated/` in the source tree, so this must complete + before another backend is configured. + """ + build_dir = _repo_root() / backend["build_dir"] + _run( + [ + "cmake", + "-S", + str(_repo_root()), + "-B", + str(build_dir), + "-DCMAKE_BUILD_TYPE=Release", + "-DINFINI_RT_BUILD_PERFORMANCE_TESTING=ON", + *backend["options"], + ], + stdout=subprocess.DEVNULL, + ) + for test in _MATRIX_TESTS: + _run( + [ + "cmake", + "--build", + str(build_dir), + "--target", + test, + "-j", + str(jobs), + ], + stdout=subprocess.DEVNULL, + ) + return build_dir + + +def _find_executable(build_dir, name): + for candidate in ( + build_dir / "tests" / "performance" / name, + build_dir / "tests" / "performance" / "Release" / name, + build_dir / name, + ): + if candidate.exists(): + return candidate + return None + + +def _run_matrix(build_dir, quick): + """Run one backend's benchmarks and return the parsed JSON result lines. + + stderr is forwarded rather than captured: it carries the binary's own + per-backend table and its skip messages, and a run that skipped an arm is + something the reader needs to see next to the numbers. + """ + results = [] + for test in _MATRIX_TESTS: + executable = _find_executable(build_dir, test) + if executable is None: + print(f" {test}: not built, skipping", file=sys.stderr) + continue + + command = [str(executable)] + if quick: + command.append("--quick") + completed = subprocess.run( + command, cwd=executable.parent, text=True, capture_output=True, + check=False, + ) + if completed.stderr: + sys.stderr.write(completed.stderr) + if completed.returncode != 0: + raise RuntimeError( + f"{executable} exited with status {completed.returncode}" + ) + for line in completed.stdout.splitlines(): + stripped = line.strip() + if stripped.startswith("{"): + results.append(json.loads(stripped)) + return results + + +def _git(args): + try: + return subprocess.check_output( + ["git", *args], cwd=_repo_root(), text=True, stderr=subprocess.DEVNULL + ).strip() + except (OSError, subprocess.CalledProcessError): + return "unknown" + + +def _row_key(result): + """Identity of a measurement with the allocator dimension removed. + + Two rows collapse to the same key exactly when they are the same workload at + the same parameters measured on different arms, which is what makes them + comparable. `arena_config` is deliberately part of the key: the reduced host + config and the production device config are not the same measurement, and + merging them would silently compare a 8 MB ramp against a 64 MB one. + """ + params = { + name: value + for name, value in result.get("params", {}).items() + if name != "allocator" + } + # `arena_config` stays in the key but out of the rendered label: it is + # constant within a backend, so printing it on every row is noise, while + # keeping it in the key is what stops a reduced-config host row from being + # merged with a production-config device row. + config = params.pop("arena_config", "") + rendered = ", ".join(f"{name}={value}" for name, value in sorted(params.items())) + return (result["benchmark"], rendered, config) + + +def _pivot(results): + """Group results into {(benchmark, params, config): {arm: result}}.""" + table = {} + for result in results: + arm = result.get("params", {}).get("allocator") + if arm is None: + continue + table.setdefault(_row_key(result), {})[arm] = result + return table + + +def _shorten(benchmark): + return re.sub(r"^allocator_matrix\.", "", benchmark) + + +def _format_value(value, unit): + if unit == "count": + return f"{value:,.0f}" + if unit == "bytes": + return f"{value / (1024 * 1024):,.1f} MiB" + # Sub-microsecond latencies and single-digit-percent probe rates both lose + # their meaning at two decimals, so small magnitudes get more. + if unit in ("us", "ms", "x") and 0.0 < abs(value) < 1.0: + return f"{value:,.4f}" + return f"{value:,.2f}" + + +def _speedup(baseline, candidate, unit): + """How many times better the candidate is. None where that has no meaning. + + A zero on either side is not reported as a ratio: for a call count zero is a + real and important outcome (the arena served a whole workload without going + upstream once), and dividing by it would turn the best result in the table + into a blank. + """ + if unit in _RATIO_UNITS: + return None + good, bad = (candidate, baseline) if unit in _HIGHER_IS_BETTER else ( + baseline, + candidate, + ) + if good <= 0.0 or bad <= 0.0: + return None + return good / bad + + +def _describe_zero(baseline, candidate, baseline_arm, candidate_arm): + if baseline == 0.0 and candidate == 0.0: + return "both zero" + if candidate == 0.0: + return f"{candidate_arm} zero" + if baseline == 0.0: + return f"{baseline_arm} zero" + return "-" + + +def _print_comparison(backend, table, baseline_arm, candidate_arm, caption): + rows = [] + configs = set() + for (benchmark, params, config), arms in table.items(): + if benchmark in _DIAGNOSTIC_BENCHMARKS: + continue + if baseline_arm not in arms or candidate_arm not in arms: + continue + configs.add(config) + rows.append( + (_shorten(benchmark), params, arms[baseline_arm], arms[candidate_arm]) + ) + + if not rows: + print( + f"\n[{backend}] {baseline_arm} vs {candidate_arm}: " + "no comparable rows (an arm did not run on this backend)." + ) + return + + rows.sort(key=lambda row: (row[0], row[1])) + + # Sized to the content rather than to a guess, so a long parameter list + # cannot push the numeric columns out of alignment. + name_width = max(len("workload"), *(len(row[0]) for row in rows)) + 2 + params_width = max(len("params"), *(len(row[1]) for row in rows)) + 2 + + config_note = f" (arena config: {', '.join(sorted(configs))})" if configs else "" + print(f"\n=== [{backend}] {caption}{config_note} ===") + win_label = f"{candidate_arm} win" + win_width = max(len(win_label), 12) + 2 + header = ( + f"{'workload':<{name_width}}{'params':<{params_width}}" + f"{baseline_arm:>16}{candidate_arm:>16}{'unit':>9}" + f"{win_label:>{win_width}}" + ) + print(header) + print("-" * len(header)) + + for workload, params, baseline, candidate in rows: + unit = baseline.get("unit", "") + base_value = baseline["median"] + cand_value = candidate["median"] + ratio = _speedup(base_value, cand_value, unit) + if ratio is None: + verdict = _describe_zero( + base_value, cand_value, baseline_arm, candidate_arm + ) + else: + verdict = f"{ratio:,.2f}x" + print( + f"{workload:<{name_width}}{params:<{params_width}}" + f"{_format_value(base_value, unit):>16}" + f"{_format_value(cand_value, unit):>16}" + f"{unit:>9}{verdict:>{win_width}}" + ) + + print( + f"\n`{candidate_arm} win` > 1 means {candidate_arm} is better on that row. " + "Rows in\nthe `x` unit are already ratios, so no win column is computed " + "for them." + ) + if candidate_arm == "cuda_async" or baseline_arm == "cuda_async": + print( + "cuda_async is stream-ordered: its release does not wait for pending\n" + "device work, so it offers a weaker guarantee than the other arms and\n" + "these ratios are not a drop-in speedup." + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Compare allocation strategies across backends." + ) + parser.add_argument( + "--quick", + action="store_true", + help="Fewer iterations and thread counts. Shapes hold, noise is higher.", + ) + parser.add_argument( + "--no-build", + action="store_true", + help="Run whatever is already built instead of configuring first.", + ) + parser.add_argument( + "--backend", + action="append", + dest="backends", + default=None, + choices=[backend["name"] for backend in _BACKENDS], + help="Limit to one backend. Can be passed more than once.", + ) + parser.add_argument("--jobs", type=int, default=16) + parser.add_argument( + "--output", + type=pathlib.Path, + default=None, + help="Write the merged raw results here as JSON.", + ) + args = parser.parse_args() + + selected = [ + backend + for backend in _BACKENDS + if args.backends is None or backend["name"] in args.backends + ] + + metadata = { + "commit": _git(["rev-parse", "HEAD"]), + "system": platform.platform(), + "quick": args.quick, + } + + # Build every selected backend before running any of them. Not an + # optimization -- configuring rewrites the shared `generated/` tree, so a + # build interleaved with another backend's configure would compile against + # the wrong headers. + build_dirs = {} + for backend in selected: + if args.no_build: + build_dirs[backend["name"]] = _repo_root() / backend["build_dir"] + continue + try: + build_dirs[backend["name"]] = _configure_and_build(backend, args.jobs) + except subprocess.CalledProcessError: + print( + f"{backend['name']}: configure or build failed, skipping it.", + file=sys.stderr, + ) + + merged = [] + for backend in selected: + build_dir = build_dirs.get(backend["name"]) + if build_dir is None or not build_dir.exists(): + continue + + print(f"\n### running {backend['name']} ###", flush=True) + try: + results = _run_matrix(build_dir, args.quick) + except (RuntimeError, OSError) as exc: + print(f"{backend['name']}: {exc}", file=sys.stderr) + continue + + for result in results: + result.update(metadata) + merged.extend(results) + + table = _pivot(results) + for baseline_arm, candidate_arm, caption in _COMPARISONS: + _print_comparison( + backend["name"], table, baseline_arm, candidate_arm, caption + ) + + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(merged, indent=2) + "\n", encoding="utf-8" + ) + print(f"\nwrote {len(merged)} raw results to {args.output}") + + if not merged: + print("\nno results: nothing ran successfully.", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/run_performance_tests.py b/scripts/run_performance_tests.py index d0b3776..2831ff0 100644 --- a/scripts/run_performance_tests.py +++ b/scripts/run_performance_tests.py @@ -145,6 +145,7 @@ def main(): tests = args.tests or [ "perf_runtime_dispatch", "perf_memory", + "perf_memory_pool", "perf_tensor_view", "perf_tensor_view_footprint", ] diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index bb05ed2..158fd1a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -39,11 +39,42 @@ function(add_infini_rt_backend_graph_test backend device_type supports_graph_cap "INFINI_RT_TEST_SUPPORTS_GRAPH_CAPTURE=${supports_graph_capture}") endfunction() +function(add_infini_rt_backend_memory_pool_test backend device_type + runtime_header) + string(TOLOWER "${backend}" backend_lower) + set(target "test_${backend_lower}_memory_pool") + add_infini_rt_test(${target} test_memory_pool_backend.cc) + target_compile_definitions(${target} + PRIVATE + "INFINI_RT_TEST_BACKEND_NAME=\"${backend}\"" + "INFINI_RT_TEST_DEVICE_TYPE=${device_type}" + "INFINI_RT_TEST_RUNTIME_HEADER=\"${runtime_header}\"") +endfunction() + +function(add_infini_rt_backend_arena_memory_pool_test backend device_type + runtime_header) + string(TOLOWER "${backend}" backend_lower) + set(target "test_${backend_lower}_arena_memory_pool") + add_infini_rt_test(${target} test_arena_memory_pool_backend.cc) + target_compile_definitions(${target} + PRIVATE + "INFINI_RT_TEST_BACKEND_NAME=\"${backend}\"" + "INFINI_RT_TEST_DEVICE_TYPE=${device_type}" + "INFINI_RT_TEST_RUNTIME_HEADER=\"${runtime_header}\"") +endfunction() + add_infini_rt_test(test_smoke test_smoke.cc) add_infini_rt_test(test_core test_core.cc) add_infini_rt_test(test_small_vector test_small_vector.cc) add_infini_rt_test(test_metadata_view test_metadata_view.cc) add_infini_rt_test(test_shape_strides_storage test_shape_strides_storage.cc) + +# The arena pool's concurrency test drives the allocator from several threads, +# which is the only way to observe that upstream calls happen with the pool's +# lock released without deadlocking or handing the same slice to two callers. +add_infini_rt_test(test_arena_memory_pool test_arena_memory_pool.cc) +find_package(Threads REQUIRED) +target_link_libraries(test_arena_memory_pool PRIVATE Threads::Threads) if(CMAKE_SYSTEM_NAME STREQUAL "Linux") add_infini_rt_test(test_tensor_view_allocations test_tensor_view_allocations.cc) @@ -56,6 +87,10 @@ if(WITH_CPU) add_infini_rt_backend_runtime_test( CPU infini::rt::Device::Type::kCpu infini/rt/cpu/runtime_.h 0 1 0 1 0 1 1 1) + add_infini_rt_backend_memory_pool_test( + CPU infini::rt::Device::Type::kCpu infini/rt/cpu/runtime_.h) + add_infini_rt_backend_arena_memory_pool_test( + CPU infini::rt::Device::Type::kCpu infini/rt/cpu/runtime_.h) endif() if(WITH_NVIDIA) @@ -66,6 +101,10 @@ if(WITH_NVIDIA) 1 1 1 1 1 1 1 1) add_infini_rt_backend_graph_test( NVIDIA infini::rt::Device::Type::kNvidia 1) + add_infini_rt_backend_memory_pool_test( + NVIDIA infini::rt::Device::Type::kNvidia infini/rt/nvidia/runtime_.h) + add_infini_rt_backend_arena_memory_pool_test( + NVIDIA infini::rt::Device::Type::kNvidia infini/rt/nvidia/runtime_.h) endif() if(WITH_ILUVATAR) diff --git a/tests/performance/CMakeLists.txt b/tests/performance/CMakeLists.txt index 7db86de..e5fbf3e 100644 --- a/tests/performance/CMakeLists.txt +++ b/tests/performance/CMakeLists.txt @@ -29,6 +29,59 @@ endfunction() add_infini_rt_performance_test(perf_runtime_dispatch perf_runtime_dispatch.cc) add_infini_rt_performance_test(perf_memory perf_memory.cc) +add_infini_rt_performance_test(perf_memory_pool perf_memory_pool.cc) add_infini_rt_performance_test(perf_tensor_view perf_tensor_view.cc) add_infini_rt_performance_test(perf_tensor_view_footprint perf_tensor_view_footprint.cc) + +find_package(Threads REQUIRED) +target_link_libraries(perf_memory_pool PRIVATE Threads::Threads) + +# Every allocation strategy the current backend offers, measured on one set of +# workloads. Unlike `arena_vs_pool` below, this one drives the real dispatch API, +# so the numbers are the backend's own -- which is the whole point on a device. +# +# `INFINI_RT_PERF_LARGE_ARENA_CONFIG` switches the arena to the production +# 64 MB -> 512 MB ramp and scales the footprints up to match. Only meaningful +# where memory is a device's: on a host build the same values would reserve +# gigabytes and measure the page allocator instead of the pool. +# +# Registered by hand rather than through the helper above, because the full sweep +# runs 64-thread and multi-stream arms across four allocators and the ctest +# target wants `--quick`: as a test the point is that every arm runs clean, not +# that the numbers are publication grade. Run the binary directly for those. +add_executable(perf_allocator_matrix perf_allocator_matrix.cc) +target_link_libraries(perf_allocator_matrix PRIVATE infinirt Threads::Threads) +target_include_directories(perf_allocator_matrix PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) +target_compile_definitions(perf_allocator_matrix + PRIVATE "INFINI_RT_PERF_BACKEND_NAME=\"${INFINI_RT_PERF_BACKEND_NAME}\"") +if(NOT INFINI_RT_PERF_BACKEND_NAME STREQUAL "cpu") + target_compile_definitions(perf_allocator_matrix + PRIVATE INFINI_RT_PERF_LARGE_ARENA_CONFIG) +endif() +add_test(NAME perf_allocator_matrix COMMAND perf_allocator_matrix --quick) +set_tests_properties(perf_allocator_matrix PROPERTIES + LABELS performance + TIMEOUT 900) + +# `MemoryPool` vs `ArenaMemoryPool` in one process. Unlike the targets above +# this one drives its own upstream stubs rather than the dispatch API -- one of +# them a calibrated busy-wait standing in for a synchronous `cudaMalloc`, which +# is the only condition under which the arena's amortization is visible on a +# machine without a device. It still links `infinirt` for the include paths and +# the backend-name definition the shared reporting header expects. +# +# Registered by hand rather than through the helper above, because the test +# needs `--quick`: the full run spends about half a minute in busy-waits, and as +# a ctest target the point is that the harness runs and its shared-contract +# checks pass, not that the numbers are publication grade. Run the binary +# directly for the longer arms. +add_executable(arena_vs_pool ab/arena_vs_pool.cc) +target_link_libraries(arena_vs_pool PRIVATE infinirt Threads::Threads) +target_include_directories(arena_vs_pool PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) +target_compile_definitions(arena_vs_pool + PRIVATE "INFINI_RT_PERF_BACKEND_NAME=\"${INFINI_RT_PERF_BACKEND_NAME}\"") +add_test(NAME arena_vs_pool COMMAND arena_vs_pool --quick) +set_tests_properties(arena_vs_pool PROPERTIES + LABELS performance + TIMEOUT 600) diff --git a/tests/performance/ab/arena_vs_pool.cc b/tests/performance/ab/arena_vs_pool.cc new file mode 100644 index 0000000..8ee9469 --- /dev/null +++ b/tests/performance/ab/arena_vs_pool.cc @@ -0,0 +1,1105 @@ +// A/B benchmark of `MemoryPool` against `ArenaMemoryPool` in one process. +// +// This is a different axis from `pool_ab.cc`, which compares two git revisions +// of one header. Here both headers are the current ones and the question is +// which *design* wins: a size-class cache that calls upstream once per miss, or +// an arena that reserves a large backing and slices it. +// +// The comparison is meaningless on a fast upstream. A host `malloc` costs tens +// of nanoseconds, so the arena's whole advantage -- turning N upstream calls into +// one -- is worth less than the bookkeeping it adds, and the arena loses on every +// timing. That is a real result for the CPU backend and it is reported as such, +// but it says nothing about a device. So three upstreams are used: +// +// `HostUpstream` - `std::malloc`, tens of ns. The CPU backend. The arena's +// floor: pure overhead, no amortization to earn back. +// `SlowUpstream` - a calibrated busy-wait, tens of us. A synchronous +// `cudaMalloc`. This is where the arena is supposed to win, +// and by how much is the number this harness exists to +// produce. +// Call counting - both upstreams tally calls, so every timing is reported +// next to the exact upstream call count that produced it. +// A count is immune to machine noise and is the honest +// summary of what the arena changes. +// +// Output format matches `tests/performance/perf_common.h`: one JSON object per +// line on stdout, a human-readable table on stderr. Every benchmark emits two +// rows differing only in the `allocator` param (`pool` vs `arena`). +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "perf_common.h" + +namespace { + +namespace perf = infini::rt::perf; + +// Counts calls into the global operator new so host heap traffic is an exact +// count rather than something inferred from a timing. Both pools claim to +// perform no host allocation in steady state -- the size-class pool via its +// inline `PointerTable`, the arena additionally via a `NodeArena` behind its +// ordered free set -- and this is what checks the claim. +std::atomic g_operator_new_calls{0}; + +} // namespace + +// Replacing the global allocation functions is the only portable way to observe +// per-insert container allocation. Defined at global scope because the standard +// requires these to be replaced, not overloaded. +void* operator new(std::size_t size) { + g_operator_new_calls.fetch_add(1, std::memory_order_relaxed); + void* ptr = std::malloc(size); + if (ptr == nullptr) { + throw std::bad_alloc(); + } + return ptr; +} + +void* operator new[](std::size_t size) { return operator new(size); } + +void operator delete(void* ptr) noexcept { std::free(ptr); } +void operator delete[](void* ptr) noexcept { std::free(ptr); } +void operator delete(void* ptr, std::size_t) noexcept { std::free(ptr); } +void operator delete[](void* ptr, std::size_t) noexcept { std::free(ptr); } +void operator delete(void* ptr, std::align_val_t) noexcept { std::free(ptr); } +void operator delete[](void* ptr, std::align_val_t) noexcept { std::free(ptr); } + +namespace { + +// -------------------------------------------------------------------------- +// Upstream allocators +// -------------------------------------------------------------------------- + +// Fast path: `std::malloc` directly, as the CPU backend's `Runtime::Malloc` +// does. Aligned to 256 B to match what a device allocator guarantees, so the +// pools' own alignment logic is what the timings observe. +struct HostUpstream { + using Error = int; + static constexpr Error kSuccess = 0; + + static std::atomic mallocs; + static std::atomic frees; + + static Error Malloc(void** ptr, std::size_t size) { + mallocs.fetch_add(1, std::memory_order_relaxed); + *ptr = std::aligned_alloc(256, (size + 255) / 256 * 256); + return (size != 0 && *ptr == nullptr) ? 2 : 0; + } + + static Error Free(void* ptr) { + frees.fetch_add(1, std::memory_order_relaxed); + std::free(ptr); + return 0; + } + + static void Reset() { + mallocs.store(0, std::memory_order_relaxed); + frees.store(0, std::memory_order_relaxed); + } +}; + +std::atomic HostUpstream::mallocs{0}; +std::atomic HostUpstream::frees{0}; + +// Slow path: models a synchronous device allocator. Busy-waits rather than +// sleeping, so the stall is CPU-bound like a real driver call and not a +// scheduler artifact that would let other threads run for free. +struct SlowUpstream { + using Error = int; + static constexpr Error kSuccess = 0; + + // Set from the command line; 50 us is the order of a `cudaMalloc`. + static double stall_us; + + static std::atomic mallocs; + static std::atomic frees; + + static void Stall() { + const auto start = std::chrono::steady_clock::now(); + while (std::chrono::duration( + std::chrono::steady_clock::now() - start) + .count() < stall_us) { + } + } + + static Error Malloc(void** ptr, std::size_t size) { + mallocs.fetch_add(1, std::memory_order_relaxed); + Stall(); + *ptr = std::aligned_alloc(256, (size + 255) / 256 * 256); + return (size != 0 && *ptr == nullptr) ? 2 : 0; + } + + static Error Free(void* ptr) { + frees.fetch_add(1, std::memory_order_relaxed); + Stall(); + std::free(ptr); + return 0; + } + + static void Reset() { + mallocs.store(0, std::memory_order_relaxed); + frees.store(0, std::memory_order_relaxed); + } +}; + +double SlowUpstream::stall_us = 50.0; +std::atomic SlowUpstream::mallocs{0}; +std::atomic SlowUpstream::frees{0}; + +// -------------------------------------------------------------------------- +// Configuration +// -------------------------------------------------------------------------- + +// A megabyte-scale arena config. The production default reserves 64 MB per +// backing and ramps to 512 MB; at that scale this harness would reserve +// gigabytes of host memory and measure the page allocator rather than the pool. +// The ratios that matter are preserved: an 8x doubling headroom to the cap, and +// a small/large threshold well below it. +struct AbArenaConfig { + static constexpr std::size_t kInitialCapacity = 8ull << 20; // 8 MB + static constexpr std::size_t kMaxCapacity = 64ull << 20; // 64 MB + static constexpr std::size_t kSmallThreshold = 1ull << 20; // 1 MB + static constexpr std::size_t kMinSliceAlignment = 512; + static constexpr std::size_t kMinSplitRemainder = 512; + static constexpr std::size_t kShrinkThreshold = 16; + static constexpr std::uint32_t kEmptyScansToDestroy = 2; +}; + +// Each arm is named by a tag carrying both the output label and the pool +// template, so every benchmark below is written once and instantiated twice. +struct SizeClassArm { + static constexpr const char* kName = "pool"; + template + using Pool = infini::rt::MemoryPool; +}; + +struct ArenaArm { + static constexpr const char* kName = "arena"; + template + using Pool = infini::rt::ArenaMemoryPool; +}; + +// -------------------------------------------------------------------------- +// Result collection +// -------------------------------------------------------------------------- + +// One comparison row. `higher_is_better` flips the ratio for metrics where a +// larger number is the good outcome; `pool_only` marks rows where the "no pool" +// notion does not apply. +struct Comparison { + std::string workload; + std::string params; + std::string unit; + double pool = 0.0; + double arena = 0.0; + bool higher_is_better = false; +}; + +std::vector g_comparisons; + +void Record(std::string workload, std::string params, std::string unit, + double pool, double arena, bool higher_is_better = false) { + g_comparisons.push_back(Comparison{std::move(workload), std::move(params), + std::move(unit), pool, arena, + higher_is_better}); +} + +std::vector WithArm(std::vector params, + const char* arm) { + params.push_back(perf::StringParam("allocator", arm)); + return params; +} + +std::string DescribeSize(std::size_t size) { + if (size >= 1024 * 1024) { + return std::to_string(size / (1024 * 1024)) + " MiB"; + } + if (size >= 1024) { + return std::to_string(size / 1024) + " KiB"; + } + return std::to_string(size) + " B"; +} + +// -------------------------------------------------------------------------- +// 1. Host heap traffic +// -------------------------------------------------------------------------- + +// Counts `operator new` calls across a fixed alloc/free loop, after a warmup so +// one-time table and node-block growth is excluded. This is the arena's biggest +// structural risk: it keeps free extents in a `std::set`, which without the +// `NodeArena` behind it would allocate once per insert and once per erase -- +// two host allocations per pooled allocation, worse than no pool at all. +template +double MeasureHostHeapTraffic(std::size_t iterations) { + typename Arm::template Pool pool; + + for (std::size_t i = 0; i < 2000; ++i) { + void* ptr = nullptr; + pool.Allocate(&ptr, 4096); + pool.Deallocate(ptr); + } + + const auto before = g_operator_new_calls.load(std::memory_order_relaxed); + for (std::size_t i = 0; i < iterations; ++i) { + void* ptr = nullptr; + pool.Allocate(&ptr, 4096); + pool.Deallocate(ptr); + } + const auto calls = static_cast( + g_operator_new_calls.load(std::memory_order_relaxed) - before); + + perf::PrintResult( + "arena_vs_pool.HostHeapTraffic", + WithArm({perf::NumberParam("size_bytes", 4096)}, Arm::kName), iterations, + "count", calls, calls); + return calls; +} + +// Same measurement under churn at mixed sizes, where the arena's free set is +// genuinely exercised: extents are split and coalesced on every operation, so +// the set sees inserts and erases rather than sitting on one entry. +template +double MeasureHostHeapTrafficUnderChurn(std::size_t iterations) { + constexpr std::size_t kLive = 64; + constexpr std::size_t kClasses = 32; + typename Arm::template Pool pool; + + std::mt19937 rng(4242); + std::vector live; + live.reserve(kLive); + + auto step = [&pool, &live, &rng](std::size_t index) { + if (live.size() < kLive && ((rng() & 3) != 0 || live.empty())) { + void* ptr = nullptr; + if (pool.Allocate(&ptr, (index % kClasses + 1) * 512) == 0) { + live.push_back(ptr); + } + } else { + const std::size_t at = rng() % live.size(); + pool.Deallocate(live[at]); + live.erase(live.begin() + static_cast(at)); + } + }; + + for (std::size_t i = 0; i < 20000; ++i) { // warm the tables and node blocks + step(i); + } + + const auto before = g_operator_new_calls.load(std::memory_order_relaxed); + for (std::size_t i = 0; i < iterations; ++i) { + step(i); + } + const auto calls = static_cast( + g_operator_new_calls.load(std::memory_order_relaxed) - before); + + for (void* ptr : live) { + pool.Deallocate(ptr); + } + + perf::PrintResult("arena_vs_pool.HostHeapTrafficUnderChurn", + WithArm({perf::NumberParam("live_blocks", kLive), + perf::NumberParam("size_classes", kClasses)}, + Arm::kName), + iterations, "count", calls, calls); + return calls; +} + +// -------------------------------------------------------------------------- +// 2. Steady state on a fast upstream -- the arena's overhead floor +// -------------------------------------------------------------------------- + +// Allocate one block, free it, repeat. Every iteration after the first is a hit +// in both designs, so this isolates per-operation bookkeeping with no upstream +// cost to amortize. The arena is expected to lose here: it does strictly more +// work per hit (an ordered-set lookup, a split, a coalesce) than a free-list +// pop, and this prices exactly that. +template +perf::Measurement BenchSteadyStateHit(std::size_t size, + std::size_t iterations) { + typename Arm::template Pool pool; + return perf::RunBenchmarkMeasured( + "arena_vs_pool.SteadyStateHit", + WithArm( + {perf::NumberParam("size_bytes", static_cast(size))}, + Arm::kName), + iterations, "ns", [&pool, size] { + void* ptr = nullptr; + auto status = pool.Allocate(&ptr, size); + perf::DoNotOptimize(status); + if (status == 0) { + status = pool.Deallocate(ptr); + perf::DoNotOptimize(status); + } + }); +} + +// A rolling window of live blocks: free the oldest, allocate a replacement -- +// the shape a layer-by-layer inference loop produces. Keeps both pools' live +// tables genuinely populated, and for the arena keeps its free set fragmented +// rather than collapsed to one extent. +template +perf::Measurement BenchLiveSetChurn(std::size_t live_blocks, + std::size_t iterations) { + constexpr std::size_t kSize = 4096; + typename Arm::template Pool pool; + + std::vector blocks(live_blocks, nullptr); + for (auto& block : blocks) { + if (pool.Allocate(&block, kSize) != 0) { + perf::SkipBenchmark("arena_vs_pool.LiveSetChurn", "prefill failed"); + return {}; + } + } + + std::size_t cursor = 0; + const auto measurement = perf::RunBenchmarkMeasured( + "arena_vs_pool.LiveSetChurn", + WithArm({perf::NumberParam("live_blocks", live_blocks), + perf::NumberParam("size_bytes", kSize)}, + Arm::kName), + iterations, "ns", [&pool, &blocks, &cursor, live_blocks] { + void*& slot = blocks[cursor]; + cursor = (cursor + 1) % live_blocks; + auto status = pool.Deallocate(slot); + slot = nullptr; + if (status == 0) { + status = pool.Allocate(&slot, kSize); + } + perf::DoNotOptimize(status); + }); + + for (void* block : blocks) { + if (block != nullptr) { + pool.Deallocate(block); + } + } + return measurement; +} + +// -------------------------------------------------------------------------- +// 3. The miss path on a slow upstream -- where the arena earns its keep +// -------------------------------------------------------------------------- + +// A growing set of live blocks at mixed sizes with nothing freed until the end. +// Neither pool can reuse anything, so every allocation is a miss. The size-class +// pool must call upstream once per block; the arena calls upstream once per +// backing and slices the rest. On a 50 us upstream that is the difference +// between N stalls and a handful. +// +// Reported per whole build-up rather than per block, since one sample is one +// cycle. The upstream call count is reported alongside: it is the cause, the +// timing is the effect. +template +perf::Measurement BenchFirstTouchGrowth(std::size_t blocks, + std::size_t samples, + std::size_t* upstream_calls) { + constexpr std::size_t kClasses = 64; + constexpr std::size_t kStride = 512; + + std::vector live; + live.reserve(blocks); + std::size_t last_calls = 0; + + const auto measurement = perf::RunBenchmarkMeasured( + "arena_vs_pool.FirstTouchGrowth", + WithArm({perf::NumberParam("blocks", blocks)}, Arm::kName), samples, "us", + [&live, &last_calls, blocks] { + // A fresh pool per sample: a warm one would serve the whole build-up + // from cache and measure the opposite of what this benchmark is for. + Upstream::Reset(); + typename Arm::template Pool pool; + for (std::size_t i = 0; i < blocks; ++i) { + void* ptr = nullptr; + if (pool.Allocate(&ptr, (i % kClasses + 1) * kStride) == 0) { + live.push_back(ptr); + } + } + last_calls = Upstream::mallocs.load(std::memory_order_relaxed); + for (void* ptr : live) { + pool.Deallocate(ptr); + } + live.clear(); + }); + + *upstream_calls = last_calls; + return measurement; +} + +// Trim the cache every iteration so nothing is ever reused: each allocation must +// go upstream. This is both pools' worst case, and it prices what a pool costs +// when its cache is useless -- the arena still amortizes, because one backing +// covers the whole iteration's slicing, but it also has a whole backing to +// release on every trim. +template +perf::Measurement BenchTrimmedMissPath(std::size_t size, + std::size_t iterations, + std::size_t* upstream_calls) { + Upstream::Reset(); + typename Arm::template Pool pool; + + const auto measurement = perf::RunBenchmarkMeasured( + "arena_vs_pool.TrimmedMissPath", + WithArm( + {perf::NumberParam("size_bytes", static_cast(size))}, + Arm::kName), + iterations, "us", [&pool, size] { + void* ptr = nullptr; + auto status = pool.Allocate(&ptr, size); + if (status == 0) { + status = pool.Deallocate(ptr); + } + perf::DoNotOptimize(status); + pool.ReleaseCached(); + }); + + *upstream_calls = Upstream::mallocs.load(std::memory_order_relaxed); + return measurement; +} + +// Warm a cold pool up across many size classes on a slow upstream. The +// size-class pool pays one upstream stall per class before it can start +// hitting; the arena pays once and then serves every class out of the same +// backing, because coalescing lets a freed block of one size feed a request of +// another. +// +// One sample is one whole cold-start rotation, not one allocation. It has to be: +// `RunBenchmarkMeasured` warms up before it times, so a long-lived pool would +// have paid every per-class stall outside the timed region and the timing would +// report a steady state where both designs only ever hit -- the opposite of what +// this benchmark is for. A fresh pool per sample is the only way the cold-start +// cost lands inside the measurement. +template +perf::Measurement BenchMixedClassesSlow(std::size_t classes, + std::size_t samples, + std::size_t* upstream_calls) { + constexpr std::size_t kStride = 512; + std::size_t last_calls = 0; + + const auto measurement = perf::RunBenchmarkMeasured( + "arena_vs_pool.MixedClassesSlow", + WithArm({perf::NumberParam("size_classes", classes), + perf::NumberParam("stride_bytes", kStride)}, + Arm::kName), + samples, "us", [&last_calls, classes] { + Upstream::Reset(); + typename Arm::template Pool pool; + // Two passes: the first is all misses, the second all hits. Both arms + // reach a warm state, so what separates them is the cost of getting + // there. + for (std::size_t pass = 0; pass < 2; ++pass) { + for (std::size_t i = 0; i < classes; ++i) { + void* ptr = nullptr; + auto status = pool.Allocate(&ptr, (i + 1) * kStride); + perf::DoNotOptimize(status); + if (status == 0) { + status = pool.Deallocate(ptr); + perf::DoNotOptimize(status); + } + } + } + last_calls = Upstream::mallocs.load(std::memory_order_relaxed); + }); + + *upstream_calls = last_calls; + return measurement; +} + +// -------------------------------------------------------------------------- +// 4. Memory amplification +// -------------------------------------------------------------------------- + +// How much upstream memory each design holds for the same live demand. Neither +// answer is strictly better -- the arena trades retention for upstream calls -- +// so this is reported rather than judged. It is a byte count, not a timing, so +// it is exact. +// +// The shape is deliberately adversarial to size classes: a long tail of distinct +// sizes, each seen once. The size-class pool retains a block per class forever; +// the arena's coalescing folds them back into reusable extents. +template +double MeasureAmplification(std::size_t classes) { + constexpr std::size_t kStride = 512; + constexpr std::size_t kLive = 32; + typename Arm::template Pool pool; + + std::mt19937 rng(1337); + std::vector live; + std::size_t peak_demand = 0; + std::size_t demand = 0; + + for (std::size_t i = 0; i < 20000; ++i) { + if (live.size() < kLive && ((rng() & 3) != 0 || live.empty())) { + const std::size_t size = (i % classes + 1) * kStride; + void* ptr = nullptr; + if (pool.Allocate(&ptr, size) == 0) { + live.push_back(ptr); + demand += size; + peak_demand = std::max(peak_demand, demand); + } + } else { + const std::size_t at = rng() % live.size(); + // Demand is tracked approximately: the exact size of the block being + // freed is not retained, so the mean class size stands in. Only the + // order of magnitude matters for an amplification ratio. + demand -= std::min(demand, (classes / 2 + 1) * kStride); + pool.Deallocate(live[at]); + live.erase(live.begin() + static_cast(at)); + } + } + + const auto reserved = static_cast(pool.GetStats().bytes_reserved); + for (void* ptr : live) { + pool.Deallocate(ptr); + } + + perf::PrintResult("arena_vs_pool.BytesReserved", + WithArm({perf::NumberParam("size_classes", classes), + perf::NumberParam("live_blocks", kLive)}, + Arm::kName), + 20000, "bytes", reserved, reserved); + return reserved; +} + +// -------------------------------------------------------------------------- +// 5. Concurrency +// -------------------------------------------------------------------------- + +// Runs `op(thread, index)` on `threads` threads and reports ns per operation. +// Threads park on `go` so thread creation stays out of the timed region. +template +perf::Measurement RunThreaded(const std::string& benchmark, + const std::vector& params, + std::size_t threads, std::size_t ops_per_thread, + Op&& op) { + constexpr std::size_t kSamples = 7; + const auto total_ops = threads * ops_per_thread; + std::vector samples; + samples.reserve(kSamples); + + for (std::size_t sample = 0; sample < kSamples + 1; ++sample) { + std::atomic go{false}; + std::vector workers; + workers.reserve(threads); + + for (std::size_t t = 0; t < threads; ++t) { + workers.emplace_back([&go, &op, ops_per_thread, t] { + while (!go.load(std::memory_order_acquire)) { + } + for (std::size_t i = 0; i < ops_per_thread; ++i) { + op(t, i); + } + }); + } + + const auto start = std::chrono::steady_clock::now(); + go.store(true, std::memory_order_release); + for (auto& worker : workers) { + worker.join(); + } + const auto end = std::chrono::steady_clock::now(); + + if (sample == 0) { // warmup + continue; + } + samples.push_back( + std::chrono::duration(end - start).count() / + static_cast(total_ops)); + } + + const perf::Measurement measurement{perf::Mean(samples), + perf::Median(samples)}; + perf::PrintResult(benchmark, params, total_ops, "ns", measurement.mean, + measurement.median); + return measurement; +} + +// Threads start at different size classes and rotate. Both pools serialize on +// one mutex, so this measures how long each holds it: the arena's critical +// section is longer (set lookup plus split plus coalesce), which is the cost it +// pays for needing upstream less often. +template +perf::Measurement BenchConcurrentMixedSizes(std::size_t threads, + std::size_t ops_per_thread) { + constexpr std::size_t kClasses = 16; + constexpr std::size_t kStride = 512; + typename Arm::template Pool pool; + + return RunThreaded( + "arena_vs_pool.ConcurrentMixedSizes", + WithArm({perf::NumberParam("threads", threads), + perf::NumberParam("size_classes", kClasses)}, + Arm::kName), + threads, ops_per_thread, [&pool](std::size_t thread, std::size_t op) { + const std::size_t size = ((thread + op) % kClasses + 1) * kStride; + void* ptr = nullptr; + auto status = pool.Allocate(&ptr, size); + if (status == 0) { + status = pool.Deallocate(ptr); + } + perf::DoNotOptimize(status); + }); +} + +// -------------------------------------------------------------------------- +// 6. Tail latency under a slow upstream +// -------------------------------------------------------------------------- + +struct Percentiles { + double p50 = 0.0; + double p99 = 0.0; + double max = 0.0; + double count = 0.0; +}; + +Percentiles Summarize(std::vector& samples) { + if (samples.empty()) { + return {}; + } + std::sort(samples.begin(), samples.end()); + const auto p99_index = std::min( + samples.size() - 1, static_cast(samples.size() * 0.99)); + return {samples[samples.size() / 2], samples[p99_index], samples.back(), + static_cast(samples.size())}; +} + +// One interfering thread repeatedly takes the miss path on a slow upstream while +// this thread only ever hits. Both pools call upstream with the lock released, so +// neither should let a hitter wait out a full stall -- this is the check that the +// arena did not regress that property while adding its shrink scan, which also +// runs on the allocation path. +template +Percentiles MeasureHitStallUnderMiss(double seconds) { + constexpr std::size_t kHitSize = 4096; + typename Arm::template Pool pool; + + // Prime the hit path so the measured allocation never misses. + void* warm = nullptr; + pool.Allocate(&warm, kHitSize); + pool.Deallocate(warm); + + std::atomic stop{false}; + std::thread misser([&pool, &stop] { + std::size_t i = 0; + while (!stop.load(std::memory_order_relaxed)) { + void* ptr = nullptr; + // A large fresh request every call, so neither design can serve it from + // what it already holds. + if (pool.Allocate(&ptr, (1u << 20) + (++i % 64) * 4096) == 0) { + pool.Deallocate(ptr); + } + pool.ReleaseCached(); + } + }); + + std::vector stalls; + stalls.reserve(1u << 20); + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::duration(seconds); + while (std::chrono::steady_clock::now() < deadline) { + void* ptr = nullptr; + const auto start = std::chrono::steady_clock::now(); + const auto status = pool.Allocate(&ptr, kHitSize); + const auto end = std::chrono::steady_clock::now(); + if (status == 0) { + pool.Deallocate(ptr); + } + stalls.push_back( + std::chrono::duration(end - start).count()); + } + + stop.store(true, std::memory_order_relaxed); + misser.join(); + return Summarize(stalls); +} + +void ReportPercentiles(const std::string& benchmark, + std::vector params, const char* arm, + const Percentiles& p) { + for (const auto& [suffix, value] : + {std::pair{"p50", p.p50}, std::pair{"p99", p.p99}, + std::pair{"max", p.max}}) { + auto row = params; + row.push_back(perf::StringParam("percentile", suffix)); + perf::PrintResult(benchmark, WithArm(std::move(row), arm), + static_cast(p.count), "us", value, value); + } +} + +// -------------------------------------------------------------------------- +// 7. Behavioral parity +// -------------------------------------------------------------------------- + +// A performance comparison between two allocators is only meaningful if both +// honor the same contract, so the shared parts of it are checked on each arm and +// any divergence is reported as a failure rather than left for the reader to +// infer from the timings. +// +// Only the *shared* contract is checked. The two designs deliberately differ on +// reuse geometry -- the size-class pool returns the same address for the same +// class, the arena returns whatever extent fits -- so pointer identity is not +// asserted here. +int g_parity_failures = 0; + +void Check(bool ok, const std::string& what, const char* arm) { + if (!ok) { + std::cerr << " PARITY FAIL [" << arm << "] " << what << "\n"; + ++g_parity_failures; + } +} + +template +void CheckSharedContract() { + const char* arm = Arm::kName; + typename Arm::template Pool pool; + + void* zero = reinterpret_cast(0x1234); + Check(pool.Allocate(&zero, 0) == 0 && zero == nullptr, + "zero size succeeds with a null pointer", arm); + Check(pool.Deallocate(nullptr) == 0, "freeing nullptr is a no-op", arm); + Check(pool.Allocate(nullptr, 64) != 0, "a null out-pointer is rejected", arm); + + int not_from_pool = 0; + Check(pool.Deallocate(¬_from_pool) != 0, "a foreign pointer is rejected", + arm); + + void* live = nullptr; + Check(pool.Allocate(&live, 8192) == 0, "an ordinary allocate succeeds", arm); + Check(pool.Deallocate(live) == 0, "the first free succeeds", arm); + Check(pool.Deallocate(live) != 0, "a double free is rejected", arm); + + void* aligned = nullptr; + Check(pool.Allocate(&aligned, 4096, 4096) == 0, "aligned allocate succeeds", + arm); + Check(reinterpret_cast(aligned) % 4096 == 0, + "requested alignment is honored", arm); + pool.Deallocate(aligned); + + // Distinct live blocks must not alias, which a slicing bug in the arena would + // violate in a way no counter would reveal. + std::vector> blocks; + for (std::size_t i = 0; i < 400; ++i) { + const std::size_t size = (i % 32 + 1) * 512; + void* ptr = nullptr; + if (pool.Allocate(&ptr, size) == 0) { + std::memset(ptr, static_cast(i & 0xff), size); + blocks.emplace_back(ptr, size); + } + } + bool distinct = true; + for (std::size_t i = 0; i < blocks.size(); ++i) { + const auto* bytes = static_cast(blocks[i].first); + for (std::size_t j = 0; j < blocks[i].second; ++j) { + if (bytes[j] != static_cast(i & 0xff)) { + distinct = false; + break; + } + } + } + Check(distinct, "concurrently live blocks never alias", arm); + Check(pool.GetStats().bytes_in_use > 0, "bytes_in_use is positive when live", + arm); + for (const auto& [ptr, size] : blocks) { + pool.Deallocate(ptr); + } + + const auto settled = pool.GetStats(); + Check(settled.bytes_in_use == 0, "bytes_in_use returns to zero", arm); + Check(settled.alloc_count == settled.free_count, + "alloc_count equals free_count", arm); + Check(settled.cache_hit_count + settled.cache_miss_count == + settled.alloc_count, + "hits plus misses equals allocs", arm); + + pool.ReleaseCached(); + Check(pool.GetStats().bytes_reserved == 0, + "bytes_reserved is zero after a trim", arm); +} + +// Every upstream `Malloc` must be matched by a `Free` once the pool dies, +// including blocks the caller never handed back. +template +void CheckNoUpstreamLeak() { + HostUpstream::Reset(); + { + typename Arm::template Pool pool; + std::vector live; + for (std::size_t i = 0; i < 600; ++i) { + void* ptr = nullptr; + if (pool.Allocate(&ptr, 512 * (i % 24 + 1)) == 0) { + live.push_back(ptr); + } + } + for (std::size_t i = 0; i < live.size() / 2; ++i) { + pool.Deallocate(live[i]); + } + } + Check(HostUpstream::mallocs.load() == HostUpstream::frees.load(), + "no upstream leak (mallocs == frees)", Arm::kName); +} + +// Concurrent smoke test: state must not corrupt and accounting must settle. +template +void CheckConcurrentIntegrity() { + typename Arm::template Pool pool; + std::atomic errors{0}; + std::vector workers; + + for (std::size_t t = 0; t < 8; ++t) { + workers.emplace_back([&pool, &errors, t] { + for (std::size_t i = 0; i < 20000; ++i) { + void* ptr = nullptr; + if (pool.Allocate(&ptr, 512 * ((t + i) % 24 + 1)) != 0 || + ptr == nullptr) { + errors.fetch_add(1, std::memory_order_relaxed); + continue; + } + if (pool.Deallocate(ptr) != 0) { + errors.fetch_add(1, std::memory_order_relaxed); + } + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + + Check(errors.load() == 0, "no errors under 8 concurrent threads", Arm::kName); + Check(pool.GetStats().bytes_in_use == 0, + "bytes_in_use settles to zero after concurrent churn", Arm::kName); +} + +// -------------------------------------------------------------------------- +// Summary +// -------------------------------------------------------------------------- + +void PrintSummary() { + if (g_comparisons.empty()) { + return; + } + + std::cerr << "\n=== MemoryPool vs ArenaMemoryPool ===\n"; + std::cerr << std::left << std::setw(30) << "workload" << std::setw(16) + << "params" << std::right << std::setw(16) << "pool" + << std::setw(16) << "arena" << std::setw(12) << "arena win" + << "\n"; + + for (const Comparison& row : g_comparisons) { + std::cerr << std::left << std::setw(30) << row.workload << std::setw(16) + << row.params << std::right << std::fixed << std::setprecision(2) + << std::setw(12) << row.pool << " " << std::setw(3) << row.unit + << std::setw(12) << row.arena << " " << std::setw(3) << row.unit; + + // >1 means the arena won. Inverted for higher-is-better metrics so the + // direction of "win" is the same in every row. A zero on either side is + // spelled out rather than reported as a ratio: for a count metric zero is a + // meaningful outcome, not missing data, and dividing by it would hide which + // arm reached it. + const double good = row.higher_is_better ? row.arena : row.pool; + const double bad = row.higher_is_better ? row.pool : row.arena; + if (good > 0.0 && bad > 0.0) { + std::cerr << std::setw(11) << std::setprecision(2) << (good / bad) << "x"; + } else if (row.pool == 0.0 && row.arena == 0.0) { + std::cerr << std::setw(12) << "both zero"; + } else if (row.arena == 0.0) { + std::cerr << std::setw(12) << "arena zero"; + } else { + std::cerr << std::setw(12) << "pool zero"; + } + std::cerr << "\n"; + } + + std::cerr << "\narena win > 1 means ArenaMemoryPool is better on that row.\n"; + std::cerr + << "Read the fast-upstream rows (SteadyStateHit, LiveSetChurn,\n" + "ConcurrentMixedSizes) as the arena's overhead floor: with upstream\n" + "at tens of ns there is nothing to amortize, so a ratio below 1 there\n" + "is expected and is the CPU backend's real answer. The slow-upstream\n" + "rows are the device case, and `calls` rows are the exact cause.\n"; + if (g_parity_failures == 0) { + std::cerr << "shared contract: all checks passed on both allocators.\n"; + } else { + std::cerr << "shared contract: " << g_parity_failures + << " check(s) FAILED -- treat the timings above as suspect.\n"; + } + std::cerr << std::endl; +} + +} // namespace + +int main(int argc, char** argv) { + bool quick = false; + double stall_seconds = 1.5; + + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "--quick") { + quick = true; + stall_seconds = 0.5; + } else if (arg.rfind("--stall-us=", 0) == 0) { + SlowUpstream::stall_us = std::atof(arg.c_str() + 11); + } else if (arg.rfind("--stall-seconds=", 0) == 0) { + stall_seconds = std::atof(arg.c_str() + 16); + } else if (arg == "--help" || arg == "-h") { + std::cerr << "usage: arena_vs_pool [--quick] [--stall-us=N] " + "[--stall-seconds=N]\n" + " --quick shorter slow-upstream arms\n" + " --stall-us=N simulated upstream latency " + "(default 50, a cudaMalloc)\n" + " --stall-seconds=N duration of each latency arm " + "(default 1.5)\n"; + return 0; + } + } + + std::cerr << "--- shared contract (must pass before timings mean " + "anything) ---\n"; + CheckSharedContract(); + CheckSharedContract(); + CheckNoUpstreamLeak(); + CheckNoUpstreamLeak(); + CheckConcurrentIntegrity(); + CheckConcurrentIntegrity(); + if (g_parity_failures == 0) { + std::cerr << " all contract checks passed.\n"; + } + + // 1. Host heap traffic -- exact counts, the least noisy signal here. + const std::size_t traffic_ops = quick ? 20000 : 100000; + { + const double p = MeasureHostHeapTraffic(traffic_ops); + const double a = MeasureHostHeapTraffic(traffic_ops); + Record("HostHeapTraffic", std::to_string(traffic_ops) + " ops", "calls", p, + a); + + const double pc = MeasureHostHeapTrafficUnderChurn( + traffic_ops); + const double ac = MeasureHostHeapTrafficUnderChurn(traffic_ops); + Record("HostHeapTraffic/churn", std::to_string(traffic_ops) + " ops", + "calls", pc, ac); + } + + // 2. Fast upstream: the arena's overhead floor. + const std::size_t iterations = quick ? 50000 : 200000; + for (const std::size_t size : {4096u, 65536u, 1u << 20}) { + const auto p = BenchSteadyStateHit(size, iterations); + const auto a = BenchSteadyStateHit(size, iterations); + Record("SteadyStateHit", DescribeSize(size), "ns", p.median, a.median); + } + + for (const std::size_t live : {1u, 8u, 64u, 512u}) { + const auto p = BenchLiveSetChurn(live, iterations); + const auto a = BenchLiveSetChurn(live, iterations); + Record("LiveSetChurn", std::to_string(live) + " live", "ns", p.median, + a.median); + } + + // 3. Memory amplification -- exact byte counts. + { + const double p = MeasureAmplification(64); + const double a = MeasureAmplification(64); + Record("BytesReserved", "64 classes", "B", p, a); + } + + // 4. Concurrency on a fast upstream. + const std::size_t ops_per_thread = quick ? 20000 : 60000; + for (const std::size_t threads : {1u, 2u, 4u, 8u}) { + const auto p = BenchConcurrentMixedSizes(threads, + ops_per_thread); + const auto a = BenchConcurrentMixedSizes(threads, ops_per_thread); + Record("ConcurrentMixedSizes", std::to_string(threads) + "T x16cls", "ns", + p.median, a.median); + } + + // 5. Slow upstream: the device case, and the reason the arena exists. + { + const std::size_t blocks = quick ? 200u : 800u; + const std::size_t samples = quick ? 3u : 5u; + std::size_t pool_calls = 0; + std::size_t arena_calls = 0; + + const auto p = BenchFirstTouchGrowth( + blocks, samples, &pool_calls); + const auto a = BenchFirstTouchGrowth( + blocks, samples, &arena_calls); + Record("FirstTouchGrowth", std::to_string(blocks) + " blocks", "us", + p.median, a.median); + + for (const auto& [arm, calls] : + {std::pair{"pool", pool_calls}, std::pair{"arena", arena_calls}}) { + perf::PrintResult( + "arena_vs_pool.FirstTouchGrowthUpstreamCalls", + WithArm({perf::NumberParam("blocks", blocks)}, arm), blocks, "count", + static_cast(calls), static_cast(calls)); + } + Record("FirstTouchGrowth/upstream", std::to_string(blocks) + " blocks", + "calls", static_cast(pool_calls), + static_cast(arena_calls)); + } + + { + const std::size_t classes = 64; + // One sample is a whole cold-start rotation over every class, so a handful + // of them is enough -- and on a 50 us upstream the size-class arm pays + // `classes` stalls per sample. + const std::size_t samples = quick ? 3u : 5u; + std::size_t pool_calls = 0; + std::size_t arena_calls = 0; + + const auto p = BenchMixedClassesSlow( + classes, samples, &pool_calls); + const auto a = BenchMixedClassesSlow( + classes, samples, &arena_calls); + Record("MixedClassesSlow", std::to_string(classes) + " classes", "us", + p.median, a.median); + Record("MixedClassesSlow/upstream", std::to_string(classes) + " classes", + "calls", static_cast(pool_calls), + static_cast(arena_calls)); + } + + { + const std::size_t iters = quick ? 100u : 300u; + std::size_t pool_calls = 0; + std::size_t arena_calls = 0; + + const auto p = BenchTrimmedMissPath( + 65536, iters, &pool_calls); + const auto a = + BenchTrimmedMissPath(65536, iters, &arena_calls); + Record("TrimmedMissPath", "64 KiB", "us", p.median, a.median); + } + + // 6. Tail latency: does a hitter wait out an unrelated thread's stall? + { + const std::vector params{perf::NumberParam( + "stall_us", static_cast(SlowUpstream::stall_us))}; + + const auto p = MeasureHitStallUnderMiss(stall_seconds); + const auto a = MeasureHitStallUnderMiss(stall_seconds); + ReportPercentiles("arena_vs_pool.HitStallUnderMiss", params, "pool", p); + ReportPercentiles("arena_vs_pool.HitStallUnderMiss", params, "arena", a); + Record("HitStall/Miss p50", "50us upstream", "us", p.p50, a.p50); + Record("HitStall/Miss p99", "50us upstream", "us", p.p99, a.p99); + Record("HitStall/Miss max", "50us upstream", "us", p.max, a.max); + } + + PrintSummary(); + // A contract failure fails the run: a faster allocator that behaves + // differently is not a win. + return g_parity_failures == 0 ? 0 : 1; +} diff --git a/tests/performance/perf_allocator_matrix.cc b/tests/performance/perf_allocator_matrix.cc new file mode 100644 index 0000000..4e96ffe --- /dev/null +++ b/tests/performance/perf_allocator_matrix.cc @@ -0,0 +1,1896 @@ +// Allocator matrix: the same workloads run across every allocation strategy +// available on the current backend, so the three comparisons the project cares +// about all come out of one binary. +// +// `direct` - the backend allocator itself (`malloc` / `cudaMalloc`). +// `pool` - `MemoryPool`, a size-class cache: one upstream call per miss. +// `arena` - `ArenaMemoryPool`, one upstream call per backing, then slices. +// `cuda_async` - `cudaMallocAsync`, CUDA's own stream-ordered pool. Present +// only where the backend supports it, and *not* semantically +// equivalent to the others: see `CudaAsyncArm`. +// +// Reading the output: +// direct vs arena - is the arena worth having at all on this backend? +// pool vs arena - which pool design wins, and on which shapes? +// arena vs cuda_async - does a hand-written arena beat the vendor's pool? +// +// `perf_memory_pool.cc` already covers the general shapes (single block, +// working-set churn, mixed size classes, first-touch growth, thread scaling). +// This file deliberately does *not* repeat them. What it adds is the set of +// measurements that file cannot make: +// +// - shapes that straddle `MemoryPool`'s 1 MB small/large boundary, where its +// rounding granularity jumps from 512 B to 2 MB; +// - large-block recycling at sizes that are *not* multiples of 2 MB, which is +// the only way the size-class rounding waste becomes visible; +// - growth to a gigabyte-scale high-water mark, far enough to walk the +// production 64 MB -> 512 MB ramp; +// - trim cost separated into its two components: bookkeeping traversal and +// the number of upstream frees, which on a device are wildly different +// costs because `cudaFree` implicitly synchronizes; +// - thread counts past 8, where lock contention actually bends; +// - device-only effects (implicit-sync cost, ledger accuracy against +// `MemGetInfo`, multi-stream traffic); +// - a synthetic layer-by-layer inference sequence, the only workload here +// that resembles what the library is for. +// +// stdout is one JSON object per line for `scripts/compare_allocators.py`; +// stderr carries the human-readable tables. +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "perf_common.h" + +namespace { + +namespace perf = infini::rt::perf; +namespace runtime = infini::rt::runtime; + +bool Success(runtime::Error status) { return status == runtime::kSuccess; } + +// -------------------------------------------------------------------------- +// Configuration +// -------------------------------------------------------------------------- + +// On a device the production values are the point: only at 64 MB -> 512 MB does +// the growth ramp, the oversize path, and the shrink heuristic behave the way +// they will in deployment. On the CPU backend the same values would reserve +// gigabytes of host memory and measure the page allocator, so the host build +// keeps the reduced scale and the difference is reported in the JSON so the two +// are never read as one series. +#if defined(INFINI_RT_PERF_LARGE_ARENA_CONFIG) +struct MatrixArenaConfig { + static constexpr std::size_t kInitialCapacity = 64ull << 20; // 64 MB + static constexpr std::size_t kMaxCapacity = 512ull << 20; // 512 MB + static constexpr std::size_t kSmallThreshold = 1ull << 20; // 1 MB + static constexpr std::size_t kMinSliceAlignment = 512; + static constexpr std::size_t kMinSplitRemainder = 512; + static constexpr std::size_t kShrinkThreshold = 16; + static constexpr std::uint32_t kEmptyScansToDestroy = 2; +}; +constexpr const char* kConfigName = "production"; +#else +struct MatrixArenaConfig { + static constexpr std::size_t kInitialCapacity = 8ull << 20; // 8 MB + static constexpr std::size_t kMaxCapacity = 64ull << 20; // 64 MB + static constexpr std::size_t kSmallThreshold = 1ull << 20; // 1 MB + static constexpr std::size_t kMinSliceAlignment = 512; + static constexpr std::size_t kMinSplitRemainder = 512; + static constexpr std::size_t kShrinkThreshold = 16; + static constexpr std::uint32_t kEmptyScansToDestroy = 2; +}; +constexpr const char* kConfigName = "reduced"; +#endif + +// Scales every footprint in this file. The high-water and inference workloads +// are sized in gigabytes on a device; on a host build that would measure the +// page allocator, so they shrink by this divisor. +#if defined(INFINI_RT_PERF_LARGE_ARENA_CONFIG) +constexpr std::size_t kFootprintDivisor = 1; +constexpr bool kDeviceBackend = true; +#else +constexpr std::size_t kFootprintDivisor = 16; +constexpr bool kDeviceBackend = false; +#endif + +// Forwards to the dispatch API, so one binary measures whichever backend the +// library was built with. +struct DispatchUpstream { + using Error = runtime::Error; + static constexpr Error kSuccess = runtime::kSuccess; + + static Error Malloc(void** ptr, std::size_t size) { + return runtime::Malloc(ptr, size); + } + static Error Free(void* ptr) { return runtime::Free(ptr); } +}; + +// -------------------------------------------------------------------------- +// Arms +// -------------------------------------------------------------------------- +// +// Every arm exposes the same surface so each benchmark below is written once +// and instantiated per arm. `UpstreamAllocs`/`UpstreamFrees` are the exact +// counters that explain the timings; `Sync` is a no-op for the synchronous arms +// and the stream synchronization point for the asynchronous one. + +class DirectArm { + public: + static constexpr const char* kName = "direct"; + static constexpr bool kStreamOrdered = false; + // Nothing is retained past a `Deallocate`, so `ReleaseCached` is a no-op and + // the benchmarks that exist to price a trim have nothing to price. + static constexpr bool kHasCache = false; + + static bool Available() { return true; } + + runtime::Error Allocate(void** ptr, std::size_t size) { + ++upstream_allocs_; + return runtime::Malloc(ptr, size); + } + + runtime::Error Deallocate(void* ptr) { + if (ptr == nullptr) { + return runtime::kSuccess; + } + ++upstream_frees_; + return runtime::Free(ptr); + } + + // No cache to trim, and nothing retained beyond what is live. + void ReleaseCached() {} + void Sync() {} + + std::size_t UpstreamAllocs() const { return upstream_allocs_; } + std::size_t UpstreamFrees() const { return upstream_frees_; } + // Direct allocation reserves exactly what is live, so retention is not a + // meaningful axis for this arm and every byte metric is reported as absent + // rather than as zero. + bool TracksBytes() const { return false; } + std::size_t BytesReserved() const { return 0; } + + private: + std::size_t upstream_allocs_ = 0; + std::size_t upstream_frees_ = 0; +}; + +class PoolArm { + public: + static constexpr const char* kName = "pool"; + static constexpr bool kStreamOrdered = false; + static constexpr bool kHasCache = true; + + static bool Available() { return true; } + + runtime::Error Allocate(void** ptr, std::size_t size) { + return pool_.Allocate(ptr, size); + } + runtime::Error Deallocate(void* ptr) { return pool_.Deallocate(ptr); } + void ReleaseCached() { pool_.ReleaseCached(); } + void Sync() {} + + std::size_t UpstreamAllocs() const { + return pool_.GetStats().upstream_alloc_count; + } + std::size_t UpstreamFrees() const { + return pool_.GetStats().upstream_free_count; + } + bool TracksBytes() const { return true; } + std::size_t BytesReserved() const { return pool_.GetStats().bytes_reserved; } + + private: + infini::rt::MemoryPool pool_; +}; + +class ArenaArm { + public: + static constexpr const char* kName = "arena"; + static constexpr bool kStreamOrdered = false; + static constexpr bool kHasCache = true; + + static bool Available() { return true; } + + runtime::Error Allocate(void** ptr, std::size_t size) { + return pool_.Allocate(ptr, size); + } + runtime::Error Deallocate(void* ptr) { return pool_.Deallocate(ptr); } + void ReleaseCached() { pool_.ReleaseCached(); } + void Sync() {} + + std::size_t UpstreamAllocs() const { + return pool_.GetStats().upstream_alloc_count; + } + std::size_t UpstreamFrees() const { + return pool_.GetStats().upstream_free_count; + } + bool TracksBytes() const { return true; } + std::size_t BytesReserved() const { return pool_.GetStats().bytes_reserved; } + + private: + infini::rt::ArenaMemoryPool pool_; +}; + +// CUDA's own stream-ordered pool, present as the fairest available reference: +// the honest question is not whether an arena beats `cudaMalloc` -- of course it +// does -- but whether it beats the pool the vendor already ships. +// +// It is NOT semantically equivalent to the other three arms and its numbers must +// not be read as a drop-in speedup. `FreeAsync` only *orders* the release behind +// the stream's current work; it does not wait for it, and reuse is likewise +// stream-ordered. The synchronous arms return memory that is immediately safe +// for any consumer. So this arm gets to overlap release with compute in a way +// the others cannot, and every benchmark synchronizes it once at the sample +// boundary rather than per operation -- measuring what it actually offers, at +// the cost of a weaker guarantee. +class CudaAsyncArm { + public: + static constexpr const char* kName = "cuda_async"; + static constexpr bool kStreamOrdered = true; + // The driver's pool has a cache, but exposes no way to release it on demand + // short of destroying the pool, so `ReleaseCached` is a no-op here too and the + // trim benchmarks have nothing to measure. + static constexpr bool kHasCache = false; + + // Probed rather than assumed: the CPU backend's `MallocAsync` returns an + // error, and not every device backend implements the stream-ordered API. + static bool Available() { + runtime::Stream stream{}; + if (!Success(runtime::StreamCreate(&stream))) { + return false; + } + void* ptr = nullptr; + const bool ok = Success(runtime::MallocAsync(&ptr, 4096, stream)) && + Success(runtime::StreamSynchronize(stream)); + if (ok) { + runtime::FreeAsync(ptr, stream); + runtime::StreamSynchronize(stream); + } + runtime::StreamDestroy(stream); + return ok; + } + + CudaAsyncArm() { runtime::StreamCreate(&stream_); } + ~CudaAsyncArm() { + if (stream_ != runtime::Stream{}) { + runtime::StreamSynchronize(stream_); + runtime::StreamDestroy(stream_); + } + } + + CudaAsyncArm(const CudaAsyncArm&) = delete; + CudaAsyncArm& operator=(const CudaAsyncArm&) = delete; + + runtime::Error Allocate(void** ptr, std::size_t size) { + ++upstream_allocs_; + return runtime::MallocAsync(ptr, size, stream_); + } + + runtime::Error Deallocate(void* ptr) { + if (ptr == nullptr) { + return runtime::kSuccess; + } + ++upstream_frees_; + return runtime::FreeAsync(ptr, stream_); + } + + void ReleaseCached() {} + void Sync() { runtime::StreamSynchronize(stream_); } + + // These count API calls, not driver allocations: the whole point of the + // stream-ordered pool is that it services most of them from its own cache, + // and it exposes no counter for how often it went to the driver. Reported + // anyway so the column is not blank, but it is not comparable to the pools' + // upstream counts. + std::size_t UpstreamAllocs() const { return upstream_allocs_; } + std::size_t UpstreamFrees() const { return upstream_frees_; } + bool TracksBytes() const { return false; } + std::size_t BytesReserved() const { return 0; } + + private: + runtime::Stream stream_{}; + std::size_t upstream_allocs_ = 0; + std::size_t upstream_frees_ = 0; +}; + +// -------------------------------------------------------------------------- +// Result collection +// -------------------------------------------------------------------------- + +// One measured value for one arm. Kept as a flat list and pivoted at print +// time, so adding an arm needs no change to the table code. +struct Cell { + std::string workload; + std::string params; + std::string arm; + std::string unit; + double value = 0.0; + bool present = true; +}; + +std::vector g_cells; +std::vector g_arm_order; + +void RecordCell(const std::string& workload, const std::string& params, + const std::string& arm, const std::string& unit, double value, + bool present = true) { + g_cells.push_back(Cell{workload, params, arm, unit, value, present}); + if (std::find(g_arm_order.begin(), g_arm_order.end(), arm) == + g_arm_order.end()) { + g_arm_order.push_back(arm); + } +} + +std::vector WithArm(std::vector params, + const char* arm) { + params.push_back(perf::StringParam("allocator", arm)); + params.push_back(perf::StringParam("arena_config", kConfigName)); + return params; +} + +std::string DescribeSize(std::size_t size) { + if (size >= 1024 * 1024) { + return std::to_string(size / (1024 * 1024)) + " MiB"; + } + if (size >= 1024) { + return std::to_string(size / 1024) + " KiB"; + } + return std::to_string(size) + " B"; +} + +// -------------------------------------------------------------------------- +// 1. Cross-threshold rotation +// -------------------------------------------------------------------------- + +// Rotates 1 KB / 2 MB / 4 KB. The sizes are chosen to straddle `MemoryPool`'s +// 1 MB small/large boundary, where its rounding granularity jumps from 512 B to +// 2 MB: the two small sizes land in distinct 512 B classes and the 2 MB one in +// its own large class, so no request can ever reuse another's block and the pool +// must hold one live block per class forever. The arena serves all three out of +// one backing, and coalescing means the 2 MB hole can be re-split into small +// requests. +// +// `perf_memory_pool.cc`'s `MixedSizeClasses` rotates 512 B-strided sizes that +// all stay on the small side of the boundary, so it never exercises the +// granularity jump this benchmark exists for. +template +void BenchCrossThresholdRotation(std::size_t iterations) { + const std::size_t sizes[] = {1024, 2ull << 20, 4096}; + constexpr std::size_t kCount = 3; + + Arm arm; + std::size_t index = 0; + const auto measurement = perf::RunBenchmarkMeasured( + "allocator_matrix.CrossThresholdRotation", + WithArm({perf::NumberParam("size_count", kCount)}, Arm::kName), iterations, + "us", [&arm, &index, &sizes] { + const std::size_t size = sizes[index]; + index = (index + 1) % kCount; + void* ptr = nullptr; + auto status = arm.Allocate(&ptr, size); + perf::DoNotOptimize(status); + if (Success(status)) { + status = arm.Deallocate(ptr); + perf::DoNotOptimize(status); + } + }); + arm.Sync(); + + RecordCell("CrossThreshold", "1K/2M/4K", Arm::kName, "us", + measurement.median); + RecordCell("CrossThreshold/upstream", "1K/2M/4K", Arm::kName, "calls", + static_cast(arm.UpstreamAllocs())); + RecordCell("CrossThreshold/reserved", "1K/2M/4K", Arm::kName, "B", + static_cast(arm.BytesReserved()), arm.TracksBytes()); + + perf::PrintResult("allocator_matrix.CrossThresholdRotationUpstreamCalls", + WithArm({perf::NumberParam("size_count", kCount)}, + Arm::kName), + iterations, "count", + static_cast(arm.UpstreamAllocs()), + static_cast(arm.UpstreamAllocs())); +} + +// -------------------------------------------------------------------------- +// 2. Large-block recycling +// -------------------------------------------------------------------------- + +// Rotates 9 / 11 / 13 MiB: deliberately *not* multiples of `MemoryPool`'s 2 MB +// large-size granularity. Each request is rounded up to 10 / 12 / 14 MiB, so the +// pool wastes 1 MB inside every block (about 8%) and, because the rounded sizes +// are distinct classes, retains one block of each forever. A rotation at 10 MiB +// would show none of this -- 10 is already a multiple of 2 -- which is why the +// sizes are odd. +// +// The arena rounds to its 512 B slice granularity instead, and coalescing lets +// one freed 13 MiB extent serve the next 9 MiB request. +template +void BenchLargeBlockRecycle(std::size_t iterations) { + const std::size_t sizes[] = {9ull << 20, 11ull << 20, 13ull << 20}; + constexpr std::size_t kCount = 3; + + // Probe first: three live blocks at once must fit, or the numbers would be an + // OOM rather than a measurement. + { + std::vector probe; + bool ok = true; + for (std::size_t i = 0; i < kCount && ok; ++i) { + void* ptr = nullptr; + ok = Success(runtime::Malloc(&ptr, sizes[i])) && ptr != nullptr; + if (ok) { + probe.push_back(ptr); + } + } + for (void* ptr : probe) { + runtime::Free(ptr); + } + if (!ok) { + perf::SkipBenchmark("allocator_matrix.LargeBlockRecycle", + "device cannot hold the working set"); + return; + } + } + + Arm arm; + std::size_t index = 0; + const auto measurement = perf::RunBenchmarkMeasured( + "allocator_matrix.LargeBlockRecycle", + WithArm({perf::NumberParam("size_count", kCount)}, Arm::kName), iterations, + "us", [&arm, &index, &sizes] { + const std::size_t size = sizes[index]; + index = (index + 1) % kCount; + void* ptr = nullptr; + auto status = arm.Allocate(&ptr, size); + perf::DoNotOptimize(status); + if (Success(status)) { + status = arm.Deallocate(ptr); + perf::DoNotOptimize(status); + } + }); + arm.Sync(); + + // Demand is the largest single request, since only one block is live at a + // time. Anything reserved beyond that is the design's retention. + const std::size_t demand = sizes[kCount - 1]; + RecordCell("LargeRecycle", "9/11/13 MiB", Arm::kName, "us", + measurement.median); + RecordCell("LargeRecycle/reserved", "9/11/13 MiB", Arm::kName, "B", + static_cast(arm.BytesReserved()), arm.TracksBytes()); + RecordCell("LargeRecycle/amplif", "9/11/13 MiB", Arm::kName, "x", + static_cast(arm.BytesReserved()) / + static_cast(demand), + arm.TracksBytes()); + RecordCell("LargeRecycle/upstream", "9/11/13 MiB", Arm::kName, "calls", + static_cast(arm.UpstreamAllocs())); +} + +// -------------------------------------------------------------------------- +// 3. High-water growth +// -------------------------------------------------------------------------- + +// Allocates without ever freeing until a gigabyte-scale high-water mark, which +// on the production config is far enough to walk the whole 64 MB -> 512 MB ramp +// and then keep adding capped backings. Nothing can be reused, so every request +// is a miss: `MemoryPool` needs one upstream call per block, the arena one per +// backing. +// +// The call count is the honest summary. It is exact, immune to machine noise, +// and on a device it is also the timing: at hundreds of microseconds per +// `cudaMalloc`, thousands of calls versus a handful is the entire result. +template +void BenchHighWaterGrowth() { + const std::size_t target = (1ull << 30) / kFootprintDivisor; + constexpr std::size_t kBlock = 1ull << 20; // 1 MiB per block + const std::size_t blocks = target / kBlock; + + Arm arm; + std::vector live; + live.reserve(blocks); + + const auto start = std::chrono::steady_clock::now(); + for (std::size_t i = 0; i < blocks; ++i) { + void* ptr = nullptr; + if (!Success(arm.Allocate(&ptr, kBlock))) { + break; + } + live.push_back(ptr); + } + arm.Sync(); + const auto end = std::chrono::steady_clock::now(); + + const auto elapsed_us = + std::chrono::duration(end - start).count(); + const auto reserved = arm.BytesReserved(); + const auto upstream = arm.UpstreamAllocs(); + const bool tracks = arm.TracksBytes(); + + for (void* ptr : live) { + arm.Deallocate(ptr); + } + arm.Sync(); + + const std::string params = DescribeSize(live.size() * kBlock) + " live"; + RecordCell("HighWater", params, Arm::kName, "us", elapsed_us); + RecordCell("HighWater/upstream", params, Arm::kName, "calls", + static_cast(upstream)); + RecordCell("HighWater/reserved", params, Arm::kName, "B", + static_cast(reserved), tracks); + + perf::PrintResult("allocator_matrix.HighWaterGrowth", + WithArm({perf::NumberParam( + "live_bytes", static_cast( + live.size() * kBlock))}, + Arm::kName), + live.size(), "us", elapsed_us, elapsed_us); + perf::PrintResult("allocator_matrix.HighWaterGrowthUpstreamCalls", + WithArm({perf::NumberParam( + "live_bytes", static_cast( + live.size() * kBlock))}, + Arm::kName), + live.size(), "count", static_cast(upstream), + static_cast(upstream)); +} + +// -------------------------------------------------------------------------- +// 4. Trim cost, split into its two components +// -------------------------------------------------------------------------- + +// Fills a cache, then trims it. The interesting part is not the wall time but +// its decomposition: `MemoryPool` walks its free lists and issues one upstream +// free *per cached block*, while the arena walks an ordered set plus its backing +// vector and issues one *per backing*. On a host those are similar; on a device +// they are not remotely, because every `cudaFree` implicitly synchronizes the +// whole device. Both numbers are reported so the cause is visible next to the +// effect. +template +void BenchTrimCost(std::size_t cached_blocks) { + constexpr std::size_t kSize = 64 * 1024; + constexpr std::size_t kIterations = 100; + + // An arm with no releasable cache would report the cost of the alloc/free + // loop with the trim removed, which is a different measurement wearing this + // one's name. Skipped rather than printed, so the row cannot be read as + // "trimming is free here". + if (!Arm::kHasCache) { + perf::SkipBenchmark("allocator_matrix.TrimCost", + std::string(Arm::kName) + " has no releasable cache"); + return; + } + + Arm arm; + std::vector blocks(cached_blocks, nullptr); + + auto fill_and_trim = [&arm, &blocks] { + for (void*& block : blocks) { + arm.Allocate(&block, kSize); + } + for (void*& block : blocks) { + arm.Deallocate(block); + block = nullptr; + } + arm.ReleaseCached(); + }; + + const auto measurement = perf::RunBenchmarkMeasured( + "allocator_matrix.TrimCost", + WithArm({perf::NumberParam("cached_blocks", cached_blocks), + perf::NumberParam("size_bytes", kSize)}, + Arm::kName), + kIterations, "us", fill_and_trim); + arm.Sync(); + + // Counted in its own cycle rather than divided out of the timed run: the + // runner's warmup and sample counts are its business, and dividing by an + // assumed total would silently go wrong the moment either changes. + const auto before_frees = arm.UpstreamFrees(); + fill_and_trim(); + arm.Sync(); + const auto frees_per_trim = + static_cast(arm.UpstreamFrees() - before_frees); + + const std::string params = std::to_string(cached_blocks) + " cached"; + RecordCell("TrimCost", params, Arm::kName, "us", measurement.median); + RecordCell("TrimCost/frees", params, Arm::kName, "calls", frees_per_trim); +} + +// -------------------------------------------------------------------------- +// 5. Thread scaling past 8 +// -------------------------------------------------------------------------- + +// Runs `op(thread, index)` on `threads` threads and reports ns per operation. +// Threads park on `go` so thread creation stays out of the timed region. +template +perf::Measurement RunThreaded(const std::string& benchmark, + const std::vector& params, + std::size_t threads, std::size_t ops_per_thread, + Op&& op) { + constexpr std::size_t kSamples = 7; + const auto total_ops = threads * ops_per_thread; + std::vector samples; + samples.reserve(kSamples); + + for (std::size_t sample = 0; sample < kSamples + 1; ++sample) { + std::atomic go{false}; + std::vector workers; + workers.reserve(threads); + + for (std::size_t t = 0; t < threads; ++t) { + workers.emplace_back([&go, &op, ops_per_thread, t] { + while (!go.load(std::memory_order_acquire)) { + } + for (std::size_t i = 0; i < ops_per_thread; ++i) { + op(t, i); + } + }); + } + + const auto start = std::chrono::steady_clock::now(); + go.store(true, std::memory_order_release); + for (auto& worker : workers) { + worker.join(); + } + const auto end = std::chrono::steady_clock::now(); + + if (sample == 0) { // warmup + continue; + } + samples.push_back( + std::chrono::duration(end - start).count() / + static_cast(total_ops)); + } + + const perf::Measurement measurement{perf::Mean(samples), + perf::Median(samples)}; + perf::PrintResult(benchmark, params, total_ops, "ns", measurement.mean, + measurement.median); + return measurement; +} + +// Both pools serialize on one mutex, so past a handful of threads the curve is +// about how long each design holds it. `perf_memory_pool.cc` stops at 8, which +// on a 160-core host is well before the knee; this goes to 64 to find where the +// arena's longer critical section (ordered-set lookup, split, coalesce) starts +// to dominate its fewer upstream calls. +template +void BenchThreadScaling(std::size_t threads, std::size_t ops_per_thread) { + constexpr std::size_t kClasses = 16; + constexpr std::size_t kStride = 512; + + Arm arm; + const auto measurement = RunThreaded( + "allocator_matrix.ThreadScaling", + WithArm({perf::NumberParam("threads", threads), + perf::NumberParam("size_classes", kClasses)}, + Arm::kName), + threads, ops_per_thread, [&arm](std::size_t thread, std::size_t op) { + const std::size_t size = ((thread + op) % kClasses + 1) * kStride; + void* ptr = nullptr; + auto status = arm.Allocate(&ptr, size); + if (Success(status)) { + status = arm.Deallocate(ptr); + } + perf::DoNotOptimize(status); + }); + arm.Sync(); + + RecordCell("ThreadScaling", std::to_string(threads) + "T", Arm::kName, "ns", + measurement.median); +} + +// -------------------------------------------------------------------------- +// 6. Device-only: implicit synchronization cost +// -------------------------------------------------------------------------- + +// `cudaFree` implicitly synchronizes the whole device: it waits for every +// previously enqueued operation on every stream. That is the entire reason the +// arena has shrink hysteresis, and until now nothing measured it. +// +// The shape: keep a stream loaded with asynchronous work, then trim the +// allocator and time it. Compared against the same trim against an idle stream, +// the delta is what the trim cost the pending work. An allocator that trims by +// issuing one upstream free per cached block pays this once per block; one that +// frees whole backings pays it once per backing. +// +// On a backend whose `MemsetAsync` is synchronous (the CPU one) there is no +// pending work to stall, and the delta is reported as approximately zero -- a +// real answer for that backend, not a missing measurement. +template +void BenchImplicitSyncCost() { + constexpr std::size_t kSize = 64 * 1024; + constexpr std::size_t kCachedBlocks = 64; + // Enough queued work that the stream is still busy when the trim lands. + constexpr std::size_t kBusyOps = 200; + const std::size_t busy_bytes = (32ull << 20) / kFootprintDivisor; + + // Same reasoning as `BenchTrimCost`: with no cache to release there is no + // trim, so there is no stall to attribute to one. + if (!Arm::kHasCache) { + perf::SkipBenchmark("allocator_matrix.ImplicitSyncCost", + std::string(Arm::kName) + " has no releasable cache"); + return; + } + + runtime::Stream stream{}; + if (!Success(runtime::StreamCreate(&stream))) { + perf::SkipBenchmark("allocator_matrix.ImplicitSyncCost", + "stream creation failed"); + return; + } + + void* busy_buffer = nullptr; + if (!Success(runtime::Malloc(&busy_buffer, busy_bytes))) { + runtime::StreamDestroy(stream); + perf::SkipBenchmark("allocator_matrix.ImplicitSyncCost", + "could not reserve the interference buffer"); + return; + } + + Arm arm; + std::vector blocks(kCachedBlocks, nullptr); + + // Fills the cache and returns everything, leaving the allocator holding + // memory that the next trim will release. + auto fill_cache = [&arm, &blocks] { + for (void*& block : blocks) { + arm.Allocate(&block, kSize); + } + for (void*& block : blocks) { + arm.Deallocate(block); + block = nullptr; + } + }; + + auto time_trim = [&](bool with_interference) { + constexpr std::size_t kSamples = 5; + std::vector samples; + samples.reserve(kSamples); + + for (std::size_t sample = 0; sample < kSamples + 1; ++sample) { + fill_cache(); + if (with_interference) { + for (std::size_t i = 0; i < kBusyOps; ++i) { + runtime::MemsetAsync(busy_buffer, static_cast(i & 0xff), + busy_bytes, stream); + } + } + + const auto start = std::chrono::steady_clock::now(); + arm.ReleaseCached(); + arm.Sync(); + const auto end = std::chrono::steady_clock::now(); + + // Drain before the next sample so leftovers cannot bleed across. + runtime::StreamSynchronize(stream); + if (sample == 0) { // warmup + continue; + } + samples.push_back( + std::chrono::duration(end - start).count()); + } + return perf::Median(samples); + }; + + const double idle = time_trim(false); + const double loaded = time_trim(true); + + runtime::Free(busy_buffer); + runtime::StreamDestroy(stream); + + const auto params = WithArm( + {perf::NumberParam("cached_blocks", kCachedBlocks), + perf::NumberParam("queued_ops", kBusyOps)}, + Arm::kName); + perf::PrintResult("allocator_matrix.TrimIdleStream", params, kCachedBlocks, + "us", idle, idle); + perf::PrintResult("allocator_matrix.TrimBusyStream", params, kCachedBlocks, + "us", loaded, loaded); + + RecordCell("Trim/idle stream", "64 cached", Arm::kName, "us", idle); + RecordCell("Trim/busy stream", "64 cached", Arm::kName, "us", loaded); + // The stall the trim imposed on pending work. Negative values are noise on a + // backend with no asynchrony to stall, and are clamped so the table reads as + // "no measurable stall" rather than as a nonsensical negative cost. + RecordCell("Trim/sync stall", "64 cached", Arm::kName, "us", + std::max(0.0, loaded - idle)); +} + +// -------------------------------------------------------------------------- +// 7. Device-only: ledger accuracy +// -------------------------------------------------------------------------- + +// Checks the pool's self-reported `bytes_reserved` against what the device says +// it lost. A pool that under-reports retention would look good on the +// amplification rows for the wrong reason, and no unit test can catch that +// because only the driver knows the truth. +// +// Every block is written to before the second reading. On a device that is +// merely belt-and-braces, but on a host it is required: `malloc` returns +// untouched pages that consume nothing until first touch, so without the write +// the ledger would look like a gross over-report of memory that genuinely had +// not been committed yet. +// +// The device figure also includes the driver's own per-allocation overhead (page +// rounding, internal metadata), so the two are not expected to match exactly. +// The ratio is what matters: near 1.0 means the ledger is honest. +// +// Device builds only. The CPU backend's `MemGetInfo` reports system-wide free +// memory from `/proc/meminfo`, so every other process on the machine moves the +// reading and the ratio would describe the machine rather than the allocator. +// Skipped there rather than printed with a caveat, because a number nobody +// should act on is worse than no number. +template +void MeasureLedgerAccuracy() { + constexpr std::size_t kBlock = 1ull << 20; + const std::size_t blocks = (256ull << 20) / kFootprintDivisor / kBlock; + + if (!kDeviceBackend) { + perf::SkipBenchmark("allocator_matrix.LedgerAccuracy", + "MemGetInfo is system-wide on the host backend"); + return; + } + + std::size_t free_before = 0; + std::size_t total = 0; + if (!Success(runtime::MemGetInfo(&free_before, &total)) || free_before == 0) { + perf::SkipBenchmark("allocator_matrix.LedgerAccuracy", + "the backend does not report device memory"); + return; + } + + Arm arm; + if (!arm.TracksBytes()) { + return; // Nothing to check against for the direct and async arms. + } + + std::vector live; + live.reserve(blocks); + for (std::size_t i = 0; i < blocks; ++i) { + void* ptr = nullptr; + if (!Success(arm.Allocate(&ptr, kBlock))) { + break; + } + // Commits the pages, so the reading below reflects what was reserved rather + // than what happens to have been faulted in. + runtime::Memset(ptr, 0, kBlock); + live.push_back(ptr); + } + arm.Sync(); + runtime::DeviceSynchronize(); + + std::size_t free_after = 0; + runtime::MemGetInfo(&free_after, &total); + const double device_consumed = + free_before > free_after ? static_cast(free_before - free_after) + : 0.0; + const double reported = static_cast(arm.BytesReserved()); + + for (void* ptr : live) { + arm.Deallocate(ptr); + } + arm.Sync(); + + const std::string params = DescribeSize(live.size() * kBlock) + " live"; + RecordCell("Ledger/reported", params, Arm::kName, "B", reported); + RecordCell("Ledger/device", params, Arm::kName, "B", device_consumed); + RecordCell("Ledger/ratio", params, Arm::kName, "x", + reported == 0.0 ? 0.0 : device_consumed / reported); + + perf::PrintResult( + "allocator_matrix.LedgerAccuracy", + WithArm({perf::NumberParam( + "live_bytes", + static_cast(live.size() * kBlock))}, + Arm::kName), + live.size(), "bytes", device_consumed, device_consumed); +} + +// -------------------------------------------------------------------------- +// 8. Device-only: multi-stream traffic +// -------------------------------------------------------------------------- + +// The other concurrency benchmark shares one allocator across threads that do +// nothing but allocate. A GPU server's real shape is different: each worker owns +// a stream, and allocation is interleaved with enqueued device work. That work +// is what an allocator can stall -- so this measures per-operation cost when +// every thread also has a stream to keep fed. +template +void BenchMultiStream(std::size_t streams, std::size_t ops_per_stream) { + constexpr std::size_t kSize = 256 * 1024; + + std::vector handles(streams, runtime::Stream{}); + for (auto& stream : handles) { + if (!Success(runtime::StreamCreate(&stream))) { + for (auto& created : handles) { + if (created != runtime::Stream{}) { + runtime::StreamDestroy(created); + } + } + perf::SkipBenchmark("allocator_matrix.MultiStream", + "stream creation failed"); + return; + } + } + + Arm arm; + const auto measurement = RunThreaded( + "allocator_matrix.MultiStream", + WithArm({perf::NumberParam("streams", streams), + perf::NumberParam("size_bytes", kSize)}, + Arm::kName), + streams, ops_per_stream, + [&arm, &handles](std::size_t index, std::size_t) { + void* ptr = nullptr; + if (!Success(arm.Allocate(&ptr, kSize))) { + return; + } + // Enqueue work against the block, then wait for it before releasing: + // the pools make no stream guarantees, so a caller must synchronize + // before handing memory back. This is the cost of using them correctly. + runtime::MemsetAsync(ptr, 0, kSize, handles[index]); + runtime::StreamSynchronize(handles[index]); + const auto status = arm.Deallocate(ptr); + perf::DoNotOptimize(status); + }); + arm.Sync(); + + for (auto& stream : handles) { + runtime::StreamDestroy(stream); + } + + RecordCell("MultiStream", std::to_string(streams) + " streams", Arm::kName, + "ns", measurement.median); +} + +// -------------------------------------------------------------------------- +// 9. Layer-by-layer inference +// -------------------------------------------------------------------------- + +// Everything above is a microbenchmark. This is the only workload here shaped +// like what the library is actually for, and the only one whose number answers +// "how much faster does inference get". +// +// The sequence mirrors a transformer forward pass: +// - weights allocated once and held for the whole run, so the arena's +// resident backing carries a permanent live block; +// - prefill walks the layers, allocating this layer's activations before +// releasing the previous layer's, which is what keeps two layers live at +// once and prevents a trivially reusable single-block pattern; +// - the KV cache grows monotonically, one allocation per layer per step, never +// freed until the end -- the shape that defeats a size-class cache, since +// each step's cache slab is a different size; +// - then many decode steps, each a small activation per layer. +// +// Reported per whole sequence, since one sample is one inference run. +template +void BenchLayerwiseInference(std::size_t decode_steps) { + constexpr std::size_t kLayers = 32; + const std::size_t weight_bytes = (4ull << 20) / kFootprintDivisor; + const std::size_t activation_bytes = (2ull << 20) / kFootprintDivisor; + const std::size_t kv_step_bytes = (128ull << 10) / kFootprintDivisor; + const std::size_t decode_bytes = (64ull << 10) / kFootprintDivisor; + + Arm arm; + + // Weights: allocated up front, released only at the very end. + std::vector weights; + weights.reserve(kLayers); + for (std::size_t layer = 0; layer < kLayers; ++layer) { + void* ptr = nullptr; + if (!Success(arm.Allocate(&ptr, weight_bytes))) { + for (void* held : weights) { + arm.Deallocate(held); + } + perf::SkipBenchmark("allocator_matrix.LayerwiseInference", + "device cannot hold the model weights"); + return; + } + weights.push_back(ptr); + } + arm.Sync(); + + std::vector kv_cache; + kv_cache.reserve(kLayers * (decode_steps + 1)); + + const auto before_upstream = arm.UpstreamAllocs(); + const auto start = std::chrono::steady_clock::now(); + + // Prefill: two layers' activations live at once. + void* previous = nullptr; + for (std::size_t layer = 0; layer < kLayers; ++layer) { + void* activation = nullptr; + if (Success(arm.Allocate(&activation, activation_bytes))) { + if (previous != nullptr) { + arm.Deallocate(previous); + } + previous = activation; + } + void* kv = nullptr; + // Each layer's slab differs slightly in size, as a real cache's does with + // sequence length -- and as a size-class cache cannot reuse. + if (Success(arm.Allocate(&kv, kv_step_bytes + layer * 512))) { + kv_cache.push_back(kv); + } + } + if (previous != nullptr) { + arm.Deallocate(previous); + previous = nullptr; + } + + // Decode: one small activation per layer per step, plus a growing cache. + for (std::size_t step = 0; step < decode_steps; ++step) { + for (std::size_t layer = 0; layer < kLayers; ++layer) { + void* activation = nullptr; + if (Success(arm.Allocate(&activation, decode_bytes))) { + arm.Deallocate(activation); + } + void* kv = nullptr; + if (Success(arm.Allocate(&kv, kv_step_bytes + step * 256))) { + kv_cache.push_back(kv); + } + } + } + arm.Sync(); + + const auto end = std::chrono::steady_clock::now(); + const auto elapsed_ms = + std::chrono::duration(end - start).count(); + const auto upstream = arm.UpstreamAllocs() - before_upstream; + const auto reserved = arm.BytesReserved(); + const bool tracks = arm.TracksBytes(); + + for (void* ptr : kv_cache) { + arm.Deallocate(ptr); + } + for (void* ptr : weights) { + arm.Deallocate(ptr); + } + arm.Sync(); + + const std::string params = std::to_string(decode_steps) + " steps"; + RecordCell("Inference", params, Arm::kName, "ms", elapsed_ms); + RecordCell("Inference/upstream", params, Arm::kName, "calls", + static_cast(upstream)); + RecordCell("Inference/reserved", params, Arm::kName, "B", + static_cast(reserved), tracks); + + const auto json_params = + WithArm({perf::NumberParam("layers", kLayers), + perf::NumberParam("decode_steps", decode_steps)}, + Arm::kName); + perf::PrintResult("allocator_matrix.LayerwiseInference", json_params, 1, "ms", + elapsed_ms, elapsed_ms); + perf::PrintResult("allocator_matrix.LayerwiseInferenceUpstreamCalls", + json_params, 1, "count", static_cast(upstream), + static_cast(upstream)); +} + +// -------------------------------------------------------------------------- +// 10. Concurrent inference: the shape of a real server +// -------------------------------------------------------------------------- + +// The gap the rest of this file leaves. `ThreadScaling` is multi-threaded but +// its sizes (512 B - 8 KiB) all sit inside the arena's fast-bin range, so every +// thread serves itself out of its own front cache and the measurement is close +// to a best case for that cache. `LayerwiseInference` has the right shapes but +// runs on one thread. Neither answers "does the arena still win when several +// threads each drive a real forward pass", which is the question a serving +// deployment actually asks. +// +// The arena's fast bins top out at `kFastBinCount * kMinSliceAlignment` = 64 KiB. +// Of the shapes below only the decode activation is at or under that, so the +// front cache covers roughly one allocation in four and everything else +// serializes on the pool mutex with a best-fit lookup and a split. That is the +// point: the win here has to come from amortizing upstream calls, not from the +// cache, and this is where we find out whether it does. +// +// Each thread owns a stream and runs an independent sequence, so the threads +// contend for one allocator exactly as concurrent requests would. Latency +// percentiles rather than a mean: a serving system is bought on its tail, and a +// pool that occasionally stalls a thread behind a `cudaMalloc` shows up in p99 +// while a mean hides it. +struct ConcurrentInferenceResult { + double p50 = 0.0; + double p99 = 0.0; + double max = 0.0; + double throughput = 0.0; // sequences per second + std::size_t upstream_allocs = 0; + std::size_t bytes_reserved = 0; + bool tracks_bytes = false; +}; + +template +ConcurrentInferenceResult BenchConcurrentInference(std::size_t threads, + std::size_t sequences, + std::size_t decode_steps) { + // Per-thread footprints, so `threads` of them are live at once. Divided by the + // thread count rather than fixed: a 32-thread run at the single-threaded + // sizes would need 32 models resident, which is an OOM rather than a + // measurement. What stays constant is total pressure on the allocator. + constexpr std::size_t kLayers = 8; + const std::size_t weight_bytes = (4ull << 20) / kFootprintDivisor; + const std::size_t activation_bytes = (2ull << 20) / kFootprintDivisor; + const std::size_t kv_step_bytes = (128ull << 10) / kFootprintDivisor; + const std::size_t decode_bytes = (64ull << 10) / kFootprintDivisor; + + std::vector handles(threads, runtime::Stream{}); + for (auto& stream : handles) { + if (!Success(runtime::StreamCreate(&stream))) { + for (auto& created : handles) { + if (created != runtime::Stream{}) { + runtime::StreamDestroy(created); + } + } + perf::SkipBenchmark("allocator_matrix.ConcurrentInference", + "stream creation failed"); + return {}; + } + } + + Arm arm; + + // Weights are per-thread and held for the whole run, the way a replica's + // parameters are. Allocated before the timed region so the ramp-up cost of + // creating the first backings is not charged to a sequence's latency. + std::vector> weights(threads); + bool weights_ok = true; + for (std::size_t t = 0; t < threads && weights_ok; ++t) { + for (std::size_t layer = 0; layer < kLayers; ++layer) { + void* ptr = nullptr; + if (!Success(arm.Allocate(&ptr, weight_bytes))) { + weights_ok = false; + break; + } + weights[t].push_back(ptr); + } + } + if (!weights_ok) { + for (const auto& held : weights) { + for (void* ptr : held) { + arm.Deallocate(ptr); + } + } + for (auto& stream : handles) { + runtime::StreamDestroy(stream); + } + perf::SkipBenchmark("allocator_matrix.ConcurrentInference", + "device cannot hold one model per thread"); + return {}; + } + arm.Sync(); + + // One vector per thread, so recording a sample takes no lock and the + // measurement does not add contention of its own on top of the allocator's. + std::vector> samples(threads); + for (auto& per_thread : samples) { + per_thread.reserve(sequences); + } + + const auto before_upstream = arm.UpstreamAllocs(); + std::atomic go{false}; + std::vector workers; + workers.reserve(threads); + + for (std::size_t t = 0; t < threads; ++t) { + workers.emplace_back([&, t] { + const runtime::Stream stream = handles[t]; + std::vector kv_cache; + kv_cache.reserve(kLayers * (decode_steps + 1)); + + while (!go.load(std::memory_order_acquire)) { + } + + for (std::size_t sequence = 0; sequence < sequences; ++sequence) { + const auto start = std::chrono::steady_clock::now(); + + // Prefill: two layers' activations live at once, and a KV slab per + // layer whose size varies with the layer -- the shape a size-class + // cache cannot reuse. + void* previous = nullptr; + for (std::size_t layer = 0; layer < kLayers; ++layer) { + void* activation = nullptr; + if (Success(arm.Allocate(&activation, activation_bytes))) { + // Enqueued work against the block, then waited on before release: + // the pools make no stream guarantees, so this is what using them + // correctly costs. + runtime::MemsetAsync(activation, 0, activation_bytes, stream); + if (previous != nullptr) { + runtime::StreamSynchronize(stream); + arm.Deallocate(previous); + } + previous = activation; + } + void* kv = nullptr; + if (Success(arm.Allocate(&kv, kv_step_bytes + layer * 512))) { + kv_cache.push_back(kv); + } + } + if (previous != nullptr) { + runtime::StreamSynchronize(stream); + arm.Deallocate(previous); + } + + // Decode: a small activation per layer per step, plus a growing cache. + for (std::size_t step = 0; step < decode_steps; ++step) { + for (std::size_t layer = 0; layer < kLayers; ++layer) { + void* activation = nullptr; + if (Success(arm.Allocate(&activation, decode_bytes))) { + runtime::MemsetAsync(activation, 0, decode_bytes, stream); + runtime::StreamSynchronize(stream); + arm.Deallocate(activation); + } + void* kv = nullptr; + if (Success(arm.Allocate(&kv, kv_step_bytes + step * 256))) { + kv_cache.push_back(kv); + } + } + } + + // The sequence ends: its whole KV cache goes back at once. This is what + // makes the benchmark more than a longer `ThreadScaling` -- it is the + // moment a backing can drain, which is what arms the arena's shrink + // scan and its cache reclamation, and those run inside the lock while + // every other thread is still allocating. + runtime::StreamSynchronize(stream); + for (void* ptr : kv_cache) { + arm.Deallocate(ptr); + } + kv_cache.clear(); + + const auto end = std::chrono::steady_clock::now(); + samples[t].push_back( + std::chrono::duration(end - start).count()); + } + }); + } + + const auto wall_start = std::chrono::steady_clock::now(); + go.store(true, std::memory_order_release); + for (auto& worker : workers) { + worker.join(); + } + const auto wall_end = std::chrono::steady_clock::now(); + arm.Sync(); + + ConcurrentInferenceResult result; + result.upstream_allocs = arm.UpstreamAllocs() - before_upstream; + result.bytes_reserved = arm.BytesReserved(); + result.tracks_bytes = arm.TracksBytes(); + + const double wall_seconds = + std::chrono::duration(wall_end - wall_start).count(); + result.throughput = wall_seconds > 0.0 + ? static_cast(threads * sequences) / + wall_seconds + : 0.0; + + std::vector all; + all.reserve(threads * sequences); + for (const auto& per_thread : samples) { + all.insert(all.end(), per_thread.begin(), per_thread.end()); + } + if (!all.empty()) { + std::sort(all.begin(), all.end()); + result.p50 = all[all.size() / 2]; + const auto p99_index = + std::min(all.size() - 1, static_cast( + static_cast(all.size()) * 0.99)); + result.p99 = all[p99_index]; + result.max = all.back(); + } + + for (const auto& held : weights) { + for (void* ptr : held) { + arm.Deallocate(ptr); + } + } + arm.Sync(); + for (auto& stream : handles) { + runtime::StreamDestroy(stream); + } + + const auto params = + WithArm({perf::NumberParam("threads", threads), + perf::NumberParam("layers", kLayers), + perf::NumberParam("decode_steps", decode_steps)}, + Arm::kName); + const std::size_t total = threads * sequences; + perf::PrintResult("allocator_matrix.ConcurrentInferenceP50", params, total, + "ms", result.p50, result.p50); + perf::PrintResult("allocator_matrix.ConcurrentInferenceP99", params, total, + "ms", result.p99, result.p99); + perf::PrintResult("allocator_matrix.ConcurrentInferenceThroughput", params, + total, "seq_per_s", result.throughput, result.throughput); + perf::PrintResult("allocator_matrix.ConcurrentInferenceUpstreamCalls", params, + total, "count", + static_cast(result.upstream_allocs), + static_cast(result.upstream_allocs)); + + const std::string label = std::to_string(threads) + "T"; + RecordCell("ConcInfer/p50", label, Arm::kName, "ms", result.p50); + RecordCell("ConcInfer/p99", label, Arm::kName, "ms", result.p99); + RecordCell("ConcInfer/max", label, Arm::kName, "ms", result.max); + RecordCell("ConcInfer/upstream", label, Arm::kName, "calls", + static_cast(result.upstream_allocs)); + RecordCell("ConcInfer/reserved", label, Arm::kName, "B", + static_cast(result.bytes_reserved), result.tracks_bytes); + return result; +} + +// -------------------------------------------------------------------------- +// 11. Fragmentation under a long random-lifetime run +// -------------------------------------------------------------------------- + +// Everything else here allocates in a pattern. This one does not: random sizes +// spanning four orders of magnitude, random lifetimes, sustained long enough +// that the allocator's internal state is whatever the run made it rather than +// whatever it was designed for. Then it asks the question that matters at the +// end of such a run -- can you still get a large contiguous block? +// +// The two designs fail differently, which is why both the success rate and the +// retention are reported. `MemoryPool` never splits, so it cannot fragment +// internally at all: a large request either finds a matching size class or goes +// upstream, and it succeeds as long as the *device* has room. What it does +// instead is retain a block of every size class it ever saw, so its reserved +// bytes climb toward the sum of the whole size distribution. The arena splits +// and coalesces, so it reuses far more, but a large request needs a contiguous +// run inside one backing -- and if live blocks are scattered across every +// backing, that run may not exist even though the free bytes are there. +// +// The deterministic LCG is deliberate: two arms must see the identical sequence, +// or the comparison is between two workloads rather than two allocators. +template +void BenchFragmentation(std::size_t operations) { + // Held live at any moment; each slot is replaced when its lifetime expires. + constexpr std::size_t kSlots = 256; + // The probe: can a large contiguous block still be had at the end? + const std::size_t probe_bytes = (32ull << 20) / kFootprintDivisor; + constexpr std::size_t kProbes = 8; + + struct Slot { + void* ptr = nullptr; + std::size_t expires_at = 0; + }; + + Arm arm; + std::vector slots(kSlots); + + // Same constants as `std::minstd_rand`, inlined so the sequence cannot change + // with the standard library. + std::uint64_t state = 0x2545f4914f6cdd1dull; + auto next = [&state] { + state = state * 6364136223846793005ull + 1442695040888963407ull; + return static_cast(state >> 33); + }; + + // Sizes spanning 1 KiB to about 4 MiB, log-distributed so small allocations + // dominate by count and large ones by bytes -- the shape a real mix has, and + // the one that scatters small live blocks through the backings that a large + // request needs whole. + auto random_size = [&next] { + const std::uint32_t decade = next() % 4; // 1 KiB, 16 KiB, 256 KiB, 4 MiB + const std::size_t base = 1024ull << (4 * decade); + return base + (next() % base); + }; + + std::size_t failures = 0; + for (std::size_t op = 0; op < operations; ++op) { + Slot& slot = slots[next() % kSlots]; + if (slot.ptr != nullptr) { + if (slot.expires_at > op) { + continue; // Not due yet: leave it live and let the hole persist. + } + arm.Deallocate(slot.ptr); + slot.ptr = nullptr; + } + void* ptr = nullptr; + if (Success(arm.Allocate(&ptr, random_size()))) { + slot.ptr = ptr; + // Lifetimes from a few operations to a few thousand, so short-lived + // blocks churn through the holes long-lived ones leave behind. + slot.expires_at = op + 1 + (next() % 4096); + } else { + ++failures; + } + } + arm.Sync(); + + const auto churn_reserved = arm.BytesReserved(); + const auto churn_upstream = arm.UpstreamAllocs(); + + // Probe the fragmented state: several large blocks at once, so the question is + // whether the pool can assemble contiguous runs and not merely find one. + std::vector probes; + probes.reserve(kProbes); + for (std::size_t i = 0; i < kProbes; ++i) { + void* ptr = nullptr; + if (!Success(arm.Allocate(&ptr, probe_bytes)) || ptr == nullptr) { + break; + } + probes.push_back(ptr); + } + arm.Sync(); + const double probe_rate = + static_cast(probes.size()) / static_cast(kProbes); + const auto probe_upstream = arm.UpstreamAllocs() - churn_upstream; + + for (void* ptr : probes) { + arm.Deallocate(ptr); + } + for (Slot& slot : slots) { + if (slot.ptr != nullptr) { + arm.Deallocate(slot.ptr); + slot.ptr = nullptr; + } + } + arm.Sync(); + + const std::string label = std::to_string(operations / 1000) + "k ops"; + RecordCell("Fragment/reserved", label, Arm::kName, "B", + static_cast(churn_reserved), arm.TracksBytes()); + RecordCell("Fragment/upstream", label, Arm::kName, "calls", + static_cast(churn_upstream)); + RecordCell("Fragment/fail", label, Arm::kName, "calls", + static_cast(failures)); + RecordCell("Fragment/probe ok", label, Arm::kName, "x", probe_rate); + RecordCell("Fragment/probe up", label, Arm::kName, "calls", + static_cast(probe_upstream)); + + const auto params = + WithArm({perf::NumberParam("operations", operations), + perf::NumberParam("probe_bytes", probe_bytes)}, + Arm::kName); + perf::PrintResult("allocator_matrix.FragmentationReserved", params, + operations, "bytes", static_cast(churn_reserved), + static_cast(churn_reserved)); + perf::PrintResult("allocator_matrix.FragmentationProbeSuccess", params, + kProbes, "x", probe_rate, probe_rate); + perf::PrintResult("allocator_matrix.FragmentationProbeUpstreamCalls", params, + kProbes, "count", static_cast(probe_upstream), + static_cast(probe_upstream)); +} + +// -------------------------------------------------------------------------- +// 12. Allocation latency tail +// -------------------------------------------------------------------------- + +// Every timing above is a mean or a median over a loop, which is the right +// summary for throughput and the wrong one for a serving deployment: what a +// request feels is its own allocation, and the allocation that goes upstream +// costs three orders of magnitude more than the one that hits. A design with a +// better median and a worse tail is worse for serving, and no row here could +// currently tell you that. +// +// Timed per operation with a monotonic clock, which on a device backend is +// sound: the pools are synchronous, so a `Malloc` that misses blocks until the +// driver returns and the interval is the real cost. The clock's own overhead +// (tens of nanoseconds) is a visible fraction of a cache hit, so the p50 here +// reads slightly high compared with the loop-averaged rows -- consistently +// across arms, which is what keeps the comparison fair. +template +void BenchLatencyTail(std::size_t size, std::size_t operations) { + Arm arm; + std::vector samples; + samples.reserve(operations); + + // Warm the arena's ramp and the pool's size class, so the measurement is of + // the steady state rather than of first-touch growth. The cold path is what + // `HighWaterGrowth` measures. + for (std::size_t i = 0; i < 64; ++i) { + void* ptr = nullptr; + if (Success(arm.Allocate(&ptr, size))) { + arm.Deallocate(ptr); + } + } + arm.Sync(); + + for (std::size_t op = 0; op < operations; ++op) { + void* ptr = nullptr; + const auto start = std::chrono::steady_clock::now(); + const auto status = arm.Allocate(&ptr, size); + const auto end = std::chrono::steady_clock::now(); + if (Success(status)) { + arm.Deallocate(ptr); + } + samples.push_back( + std::chrono::duration(end - start).count()); + } + arm.Sync(); + + if (samples.empty()) { + return; + } + std::sort(samples.begin(), samples.end()); + auto quantile = [&samples](double q) { + const auto index = std::min( + samples.size() - 1, + static_cast(static_cast(samples.size()) * q)); + return samples[index]; + }; + + const std::string label = DescribeSize(size); + RecordCell("Latency/p50", label, Arm::kName, "us", quantile(0.50)); + RecordCell("Latency/p99", label, Arm::kName, "us", quantile(0.99)); + RecordCell("Latency/p999", label, Arm::kName, "us", quantile(0.999)); + RecordCell("Latency/max", label, Arm::kName, "us", samples.back()); + + const auto params = WithArm( + {perf::NumberParam("size_bytes", static_cast(size))}, + Arm::kName); + perf::PrintResult("allocator_matrix.LatencyP50", params, operations, "us", + quantile(0.50), quantile(0.50)); + perf::PrintResult("allocator_matrix.LatencyP99", params, operations, "us", + quantile(0.99), quantile(0.99)); + perf::PrintResult("allocator_matrix.LatencyP999", params, operations, "us", + quantile(0.999), quantile(0.999)); + perf::PrintResult("allocator_matrix.LatencyMax", params, operations, "us", + samples.back(), samples.back()); +} + +// -------------------------------------------------------------------------- +// 13. Memory bandwidth through pooled memory +// -------------------------------------------------------------------------- + +// Every other row prices the allocator. This one asks whether using it costs +// anything *afterwards* -- whether a kernel reading a sliced block runs as fast +// as one reading a dedicated upstream allocation. +// +// There is a real mechanism to check for, not just due diligence. A block from +// `cudaMalloc` starts at a 256 B (in practice much coarser) boundary; a slice +// out of an arena backing is only guaranteed `kMinSliceAlignment` = 512 B, and +// after a split it can start at an arbitrary multiple of that. If that landed +// mid-page or misaligned against the memory transaction size, sustained +// bandwidth would drop. So the arena is measured on a *split* slice rather than +// on a fresh backing's first chunk, which is the case that could actually differ. +// +// Reported as GiB/s of `MemsetAsync` traffic, which is bandwidth-bound on a +// device and the closest thing to a STREAM kernel available through the +// dispatch API without a kernel-launch surface. +template +void BenchBandwidth() { + const std::size_t size = (64ull << 20) / kFootprintDivisor; + constexpr std::size_t kIterations = 32; + + runtime::Stream stream{}; + if (!Success(runtime::StreamCreate(&stream))) { + perf::SkipBenchmark("allocator_matrix.Bandwidth", "stream creation failed"); + return; + } + + Arm arm; + + // Force the block under test to be a split remainder rather than a whole + // backing: allocate a small block first so the large one starts at an offset, + // which is the alignment case a fresh allocation would never exercise. + void* leading = nullptr; + arm.Allocate(&leading, 4096); + + void* buffer = nullptr; + if (!Success(arm.Allocate(&buffer, size)) || buffer == nullptr) { + if (leading != nullptr) { + arm.Deallocate(leading); + } + runtime::StreamDestroy(stream); + perf::SkipBenchmark("allocator_matrix.Bandwidth", + "device cannot hold the bandwidth buffer"); + return; + } + + // Warm: first touch on a host backing faults pages in, and on a device the + // first launch pays context setup. Neither is bandwidth. + runtime::MemsetAsync(buffer, 0, size, stream); + runtime::StreamSynchronize(stream); + + std::vector samples; + samples.reserve(kIterations); + for (std::size_t i = 0; i < kIterations; ++i) { + const auto start = std::chrono::steady_clock::now(); + runtime::MemsetAsync(buffer, static_cast(i & 0xff), size, stream); + runtime::StreamSynchronize(stream); + const auto end = std::chrono::steady_clock::now(); + const double seconds = std::chrono::duration(end - start).count(); + if (seconds > 0.0) { + samples.push_back(static_cast(size) / seconds / + (1024.0 * 1024.0 * 1024.0)); + } + } + + arm.Deallocate(buffer); + if (leading != nullptr) { + arm.Deallocate(leading); + } + arm.Sync(); + runtime::StreamDestroy(stream); + + if (samples.empty()) { + return; + } + const double median = perf::Median(samples); + RecordCell("Bandwidth", DescribeSize(size), Arm::kName, "GiB/s", median); + perf::PrintResult( + "allocator_matrix.Bandwidth", + WithArm({perf::NumberParam("size_bytes", + static_cast(size))}, + Arm::kName), + kIterations, "GiB/s", perf::Mean(samples), median); +} + +// -------------------------------------------------------------------------- +// Reporting +// -------------------------------------------------------------------------- + +void PrintMatrix() { + if (g_cells.empty()) { + return; + } + + // Rows in first-seen order, so the table follows the benchmark order rather + // than an alphabetical one that would separate a timing from its cause. + std::vector> rows; + for (const Cell& cell : g_cells) { + const auto key = std::make_pair(cell.workload, cell.params); + if (std::find(rows.begin(), rows.end(), key) == rows.end()) { + rows.push_back(key); + } + } + + std::cerr << "\n=== " << INFINI_RT_PERF_BACKEND_NAME + << " allocator matrix (median; arena config: " << kConfigName + << ") ===\n"; + std::cerr << std::left << std::setw(24) << "workload" << std::setw(14) + << "params"; + for (const std::string& arm : g_arm_order) { + std::cerr << std::right << std::setw(16) << arm; + } + std::cerr << std::right << std::setw(12) << "unit" << "\n"; + + for (const auto& [workload, params] : rows) { + std::cerr << std::left << std::setw(24) << workload << std::setw(14) + << params; + std::string unit; + for (const std::string& arm : g_arm_order) { + const auto found = std::find_if( + g_cells.begin(), g_cells.end(), [&](const Cell& cell) { + return cell.workload == workload && cell.params == params && + cell.arm == arm; + }); + if (found == g_cells.end()) { + std::cerr << std::right << std::setw(16) << "-"; + } else if (!found->present) { + // Absent by construction, not missing: `direct` reserves exactly what + // is live, so retention is not an axis it has. + std::cerr << std::right << std::setw(16) << "n/a"; + unit = found->unit; + } else { + std::cerr << std::right << std::setw(16) << std::fixed + << std::setprecision(found->unit == "calls" ? 0 : 2) + << found->value; + unit = found->unit; + } + } + std::cerr << std::right << std::setw(12) << unit << "\n"; + } + + std::cerr + << "\nLower is better except `x` ratio rows. `calls` rows are exact\n" + "counts, not timings, and are the cause behind the timing above them.\n" + "`n/a` means the metric does not apply to that arm; `-` means the arm\n" + "did not run that workload.\n" + "cuda_async is stream-ordered: its `Deallocate` does not wait for\n" + "pending device work, so it offers a weaker guarantee than the other\n" + "three arms and its timings are not a drop-in speedup.\n"; + std::cerr << std::endl; +} + +// -------------------------------------------------------------------------- +// Driver +// -------------------------------------------------------------------------- + +bool PrepareRuntime() { + int device_count = 0; + if (!Success(runtime::GetDeviceCount(&device_count)) || device_count <= 0) { + std::cerr << "perf_allocator_matrix skipped: no available device." + << std::endl; + return false; + } + if (!Success(runtime::SetDevice(0))) { + std::cerr << "perf_allocator_matrix skipped: device 0 is not available." + << std::endl; + return false; + } + return true; +} + +// Runs `body` for each arm in turn, skipping the stream-ordered one where the +// backend does not support it. +template