From 7508938e0eded0564696ad5360a977ac6f4db00b Mon Sep 17 00:00:00 2001 From: atchen <1920510674@qq.com> Date: Fri, 24 Jul 2026 10:07:51 +0000 Subject: [PATCH 1/8] feat: add backend-agnostic caching memory pool --- include/infini/rt/memory_pool.h | 278 ++++++++++++++++++++++++++++++ tests/CMakeLists.txt | 16 ++ tests/test_memory_pool_backend.cc | 224 ++++++++++++++++++++++++ 3 files changed, 518 insertions(+) create mode 100644 include/infini/rt/memory_pool.h create mode 100644 tests/test_memory_pool_backend.cc diff --git a/include/infini/rt/memory_pool.h b/include/infini/rt/memory_pool.h new file mode 100644 index 0000000..103251c --- /dev/null +++ b/include/infini/rt/memory_pool.h @@ -0,0 +1,278 @@ +#ifndef INFINI_RT_MEMORY_POOL_H_ +#define INFINI_RT_MEMORY_POOL_H_ + +#include +#include +#include +#include +#include +#include +#include + +namespace infini::rt { + +/// ## Backend-agnostic caching allocator. +/// +/// `cudaMalloc`/`cudaFree` (and other device allocators) are synchronous and +/// expensive, so runtimes typically layer a caching allocator on top. A +/// `MemoryPool` keeps freed blocks in per-size-class free lists and hands them +/// back on the next matching request, so hot allocation loops pay the upstream +/// allocator only on a cache miss. +/// +/// The pool is a pure composition over an `Upstream` allocator: any type that +/// provides `Malloc(void**, size_t)`, `Free(void*)`, an `Error` type alias, and +/// a `static constexpr Error kSuccess` satisfies the contract. Every +/// `runtime::Runtime<...>` device specialization (CPU, NVIDIA, ...) qualifies, +/// and the same interface serves CPU aligned allocations. Tests can inject a +/// mock upstream to exercise the pool without any device. +/// +/// Blocks are reused only when both the rounded size and the requested +/// alignment match, so a reused block is always geometrically identical to the +/// request; the pool never splits or coalesces, which keeps reuse free of +/// fragmentation hazards at the cost of some retained-but-unused memory (call +/// `ReleaseCached` to hand that back to the upstream allocator). +/// +/// The pool is thread-safe: every public method takes an internal mutex. It is +/// neither copyable nor movable. +template +class MemoryPool { + public: + using Error = typename Upstream::Error; + + /// Runtime statistics. Byte counters are cumulative live totals; `peak_*` + /// track high-water marks. The remaining counters are monotonic tallies. + struct Stats { + /// Bytes currently handed out to callers (sum of rounded block sizes). + std::size_t bytes_in_use = 0; + + /// Bytes currently held from the upstream allocator (in use + cached). + 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 a cached free block. + std::size_t cache_hit_count = 0; + + /// Allocations that required a fresh upstream allocation. + std::size_t cache_miss_count = 0; + + /// Calls into `Upstream::Malloc`. + std::size_t upstream_alloc_count = 0; + + /// Calls into `Upstream::Free`. + std::size_t upstream_free_count = 0; + }; + + MemoryPool() = default; + + MemoryPool(const MemoryPool&) = delete; + MemoryPool& operator=(const MemoryPool&) = delete; + + /// Frees every block still held from the upstream allocator, including + /// blocks that were never handed back via `Deallocate`. Any outstanding + /// pointer from `Allocate` dangles after destruction. + ~MemoryPool() { + for (auto& [key, blocks] : free_lists_) { + for (const Block& block : blocks) { + Upstream::Free(block.base); + } + } + for (auto& [ptr, block] : allocated_) { + Upstream::Free(block.base); + } + } + + /// Allocates at least `size` bytes, reusing a cached block when one with a + /// matching size class and alignment is available. `alignment` of `0` uses + /// the upstream allocator's natural alignment; 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; + } + + std::lock_guard lock(mutex_); + + const std::size_t rounded = RoundSize(size); + const BucketKey key{rounded, alignment}; + + Block block{}; + auto list_it = free_lists_.find(key); + if (list_it != free_lists_.end() && !list_it->second.empty()) { + block = list_it->second.back(); + list_it->second.pop_back(); + ++stats_.cache_hit_count; + } else { + const std::size_t upstream_size = + alignment == 0 ? rounded : rounded + alignment; + void* base = nullptr; + const Error status = Upstream::Malloc(&base, upstream_size); + ++stats_.upstream_alloc_count; + if (status != Upstream::kSuccess) { + return status; + } + + block.base = base; + block.aligned = alignment == 0 ? base : AlignUp(base, alignment); + block.rounded_size = rounded; + block.upstream_size = upstream_size; + block.alignment = alignment; + + stats_.bytes_reserved += upstream_size; + if (stats_.bytes_reserved > stats_.peak_bytes_reserved) { + stats_.peak_bytes_reserved = stats_.bytes_reserved; + } + ++stats_.cache_miss_count; + } + + allocated_.emplace(block.aligned, block); + stats_.bytes_in_use += block.rounded_size; + if (stats_.bytes_in_use > stats_.peak_bytes_in_use) { + stats_.peak_bytes_in_use = stats_.bytes_in_use; + } + ++stats_.alloc_count; + + *ptr = block.aligned; + return Upstream::kSuccess; + } + + /// Returns a block from `Allocate` to the pool's free list for reuse. The + /// block is not handed back to the upstream allocator until `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; + } + + std::lock_guard lock(mutex_); + + auto it = allocated_.find(ptr); + if (it == allocated_.end()) { + return InvalidValue(); + } + + const Block block = it->second; + allocated_.erase(it); + + stats_.bytes_in_use -= block.rounded_size; + ++stats_.free_count; + + free_lists_[BucketKey{block.rounded_size, block.alignment}].push_back( + block); + return Upstream::kSuccess; + } + + /// Hands every cached (freed but not-yet-returned) block back to the upstream + /// allocator. Blocks currently in use are untouched. This is the pool's + /// defragmentation / trim knob: call it to release retained memory back to + /// the device. + void ReleaseCached() { + std::lock_guard lock(mutex_); + + for (auto& [key, blocks] : free_lists_) { + for (const Block& block : blocks) { + Upstream::Free(block.base); + ++stats_.upstream_free_count; + stats_.bytes_reserved -= block.upstream_size; + } + } + free_lists_.clear(); + } + + /// Returns a snapshot of the pool's statistics. + Stats GetStats() const { + std::lock_guard lock(mutex_); + 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`."); + + // A single upstream allocation tracked by the pool. `base` is the pointer + // owned by the upstream allocator; `aligned` is what the caller sees. + struct Block { + void* base = nullptr; + void* aligned = nullptr; + std::size_t rounded_size = 0; + std::size_t upstream_size = 0; + std::size_t alignment = 0; + }; + + // Free lists are keyed by rounded size and alignment so a reused block is + // always geometrically identical to the request. + struct BucketKey { + std::size_t size = 0; + std::size_t alignment = 0; + + bool operator==(const BucketKey& other) const { + return size == other.size && alignment == other.alignment; + } + }; + + struct BucketKeyHash { + std::size_t operator()(const BucketKey& key) const { + // Mix the two fields; alignment is small so a shift keeps it out of the + // low bits that size dominates. + return key.size ^ (key.alignment << 1); + } + }; + + // Small allocations round to 512 B; large ones to 2 MB. This keeps the number + // of distinct size classes bounded so freed blocks are likely to be reused. + static constexpr std::size_t kSmallThreshold = 1u << 20; // 1 MB + static constexpr std::size_t kSmallGranularity = 512; + static constexpr std::size_t kLargeGranularity = 1u << 21; // 2 MB + + static std::size_t RoundUp(std::size_t size, std::size_t granularity) { + return (size + granularity - 1) / granularity * granularity; + } + + static std::size_t RoundSize(std::size_t size) { + return size <= kSmallThreshold ? RoundUp(size, kSmallGranularity) + : RoundUp(size, kLargeGranularity); + } + + static void* AlignUp(void* ptr, std::size_t alignment) { + const auto address = reinterpret_cast(ptr); + const auto aligned = (address + alignment - 1) & ~(alignment - 1); + return reinterpret_cast(aligned); + } + + static Error InvalidValue() { return static_cast(1); } + + mutable std::mutex mutex_; + std::unordered_map allocated_; + std::unordered_map, BucketKeyHash> free_lists_; + Stats stats_; +}; + +} // namespace infini::rt + +#endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index bb05ed2..4e0cc23 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -39,6 +39,18 @@ 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() + 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) @@ -56,6 +68,8 @@ 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) endif() if(WITH_NVIDIA) @@ -66,6 +80,8 @@ 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) endif() if(WITH_ILUVATAR) diff --git a/tests/test_memory_pool_backend.cc b/tests/test_memory_pool_backend.cc new file mode 100644 index 0000000..4df7370 --- /dev/null +++ b/tests/test_memory_pool_backend.cc @@ -0,0 +1,224 @@ +// Exercises `MemoryPool` over a *real* runtime backend (CPU, NVIDIA, ...). +// +// `test_memory_pool.cc` already covers the pool's bookkeeping against a mock +// upstream. This test instead instantiates the pool over the backend's actual +// `runtime::Runtime` specialization, so it proves the two compose correctly and +// that pool-handed pointers are genuine device memory. Device pointers cannot +// be dereferenced from the host, so usability is checked through `Memcpy` +// round-trips. The whole suite is skipped when no device is present. +#include +#include +#include INFINI_RT_TEST_RUNTIME_HEADER + +#include +#include +#include +#include + +#include "test_helper.h" + +namespace { + +using Runtime = infini::rt::runtime::Runtime; +using Pool = infini::rt::MemoryPool; + +constexpr const char* kBackend = INFINI_RT_TEST_BACKEND_NAME; + +bool SelectDevice() { + int device_count = 0; + if (Runtime::GetDeviceCount(&device_count) != Runtime::kSuccess || + device_count <= 0) { + std::cout << kBackend << " memory pool skipped: no available device." + << std::endl; + return false; + } + if (Runtime::SetDevice(0) != Runtime::kSuccess) { + std::cout << kBackend << " memory pool skipped: device 0 unavailable." + << std::endl; + return false; + } + return true; +} + +// Writes `input` into device memory `ptr` and reads it back, asserting the +// bytes survive the round trip. This is the only host-safe way to confirm a +// device pointer is real and usable. +template +void ExpectUsable(infini::rt::test::TestContext* context, void* ptr, + const std::array& input, + const char* message) { + if (!context->Expect(ptr != nullptr, message)) { + return; + } + std::array output{}; + context->ExpectEqual( + Runtime::Memcpy(ptr, input.data(), N, Runtime::kMemcpyHostToDevice), + Runtime::kSuccess, "memcpy host-to-device should succeed"); + context->ExpectEqual( + Runtime::Memcpy(output.data(), ptr, N, Runtime::kMemcpyDeviceToHost), + Runtime::kSuccess, "memcpy device-to-host should succeed"); + context->ExpectEqual(output, input, + "pool-allocated memory should round-trip bytes"); +} + +// A block returned by the pool must be real, usable device memory. +void TestAllocationIsUsable(infini::rt::test::TestContext* context) { + Pool pool; + void* ptr = nullptr; + context->ExpectEqual(pool.Allocate(&ptr, 256), Runtime::kSuccess, + "allocate should succeed on a real backend"); + const std::array input{0, 1, 2, 3, 4, 5, 6, 7}; + ExpectUsable(context, ptr, input, "allocation should produce a pointer"); + context->ExpectEqual(pool.Deallocate(ptr), Runtime::kSuccess, + "deallocate should succeed"); +} + +// Freeing then re-requesting the same size class reuses the cached block +// without touching the upstream device allocator. +void TestCacheReuse(infini::rt::test::TestContext* context) { + Pool pool; + void* first = nullptr; + context->ExpectEqual(pool.Allocate(&first, 4096), Runtime::kSuccess, + "first allocate should succeed"); + context->ExpectEqual(pool.Deallocate(first), Runtime::kSuccess, + "deallocate should cache the block"); + + void* second = nullptr; + context->ExpectEqual(pool.Allocate(&second, 4096), Runtime::kSuccess, + "second allocate should succeed"); + context->ExpectEqual(second, first, "same size class should reuse the block"); + + const Pool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.cache_hit_count, std::size_t{1}, + "one cache hit expected"); + context->ExpectEqual(stats.cache_miss_count, std::size_t{1}, + "only the first allocation misses"); + context->ExpectEqual(stats.upstream_alloc_count, std::size_t{1}, + "reuse must not call the device allocator again"); + pool.Deallocate(second); +} + +// Two sizes that round to the same class share a block; a distinct class does +// not, and each remains independently usable. +void TestSizeClasses(infini::rt::test::TestContext* context) { + Pool pool; + void* a = nullptr; + pool.Allocate(&a, 100); // rounds to the 512 B class + pool.Deallocate(a); + void* b = nullptr; + pool.Allocate(&b, 500); // same 512 B class + context->ExpectEqual(b, a, "100 and 500 share a size class"); + + void* c = nullptr; + pool.Allocate(&c, 8192); // a different class + context->Expect(c != b, "a distinct size class must not reuse the block"); + + const std::array input{9, 8, 7, 6}; + ExpectUsable(context, b, input, "reused block should be usable"); + ExpectUsable(context, c, input, "fresh block should be usable"); + pool.Deallocate(b); + pool.Deallocate(c); +} + +// A requested power-of-two alignment must be honored by the returned pointer, +// which must still be usable device memory. +void TestAlignment(infini::rt::test::TestContext* context) { + Pool pool; + constexpr std::size_t kAlignment = 256; + void* ptr = nullptr; + context->ExpectEqual(pool.Allocate(&ptr, 100, kAlignment), Runtime::kSuccess, + "aligned allocate should succeed"); + context->ExpectEqual(reinterpret_cast(ptr) % kAlignment, + std::uintptr_t{0}, + "returned pointer should honor the alignment"); + const std::array input{1, 1, 2, 3, 5, 8, 13, 21}; + ExpectUsable(context, ptr, input, "aligned allocation should be usable"); + pool.Deallocate(ptr); +} + +// Statistics track live/reserved bytes, peaks, and call counts across the +// allocate/deallocate cycle. +void TestStats(infini::rt::test::TestContext* context) { + Pool pool; + void* a = nullptr; + void* b = nullptr; + pool.Allocate(&a, 1024); // rounds to 1024 + pool.Allocate(&b, 2048); // rounds to 2048 + + Pool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.bytes_in_use, std::size_t{1024 + 2048}, + "bytes_in_use tracks rounded sizes"); + context->ExpectEqual(stats.peak_bytes_in_use, std::size_t{1024 + 2048}, + "peak matches the high-water mark"); + context->ExpectEqual(stats.alloc_count, std::size_t{2}, + "two allocations counted"); + + pool.Deallocate(a); + stats = pool.GetStats(); + context->ExpectEqual(stats.bytes_in_use, std::size_t{2048}, + "bytes_in_use drops on free"); + context->ExpectEqual(stats.peak_bytes_in_use, std::size_t{1024 + 2048}, + "peak stays at the high-water mark"); + context->Expect(stats.bytes_reserved >= 1024 + 2048, + "reserved memory retained while cached"); + pool.Deallocate(b); +} + +// ReleaseCached hands cached blocks back to the device; live blocks are +// untouched. Verified through stats and the upstream free counter. +void TestReleaseCached(infini::rt::test::TestContext* context) { + Pool pool; + void* a = nullptr; + void* b = nullptr; + pool.Allocate(&a, 1024); + pool.Allocate(&b, 4096); + pool.Deallocate(a); + pool.Deallocate(b); + + Pool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.upstream_free_count, std::size_t{0}, + "cached blocks are not yet freed upstream"); + + pool.ReleaseCached(); + stats = pool.GetStats(); + context->ExpectEqual(stats.upstream_free_count, std::size_t{2}, + "release should free both cached blocks upstream"); + context->ExpectEqual(stats.bytes_reserved, std::size_t{0}, + "reserved bytes drop to zero after release"); +} + +// Concurrently live blocks must be distinct and independently usable. +void TestDistinctLiveBlocks(infini::rt::test::TestContext* context) { + Pool pool; + void* a = nullptr; + void* b = nullptr; + pool.Allocate(&a, 512); + pool.Allocate(&b, 512); + context->Expect(a != b, "two live blocks must not alias"); + const std::array first{1, 2, 3, 4}; + const std::array second{5, 6, 7, 8}; + ExpectUsable(context, a, first, "first live block should be usable"); + ExpectUsable(context, b, second, "second live block should be usable"); + pool.Deallocate(a); + pool.Deallocate(b); +} + +} // namespace + +int main() { + infini::rt::test::TestContext context; + + if (!SelectDevice()) { + return context.ExitCode(); + } + + TestAllocationIsUsable(&context); + TestCacheReuse(&context); + TestSizeClasses(&context); + TestAlignment(&context); + TestStats(&context); + TestReleaseCached(&context); + TestDistinctLiveBlocks(&context); + + return context.ExitCode(); +} From b24fd73bbd079022093193efd7990662ee6ad28c Mon Sep 17 00:00:00 2001 From: atchen <1920510674@qq.com> Date: Fri, 7 Aug 2026 10:08:20 +0000 Subject: [PATCH 2/8] perf: replace the pool's block map with an open-addressing table `std::unordered_map` is node-based, so tracking a live block called `operator new` on every insert and `operator delete` on every erase -- a host heap round trip on the exact path a device allocator exists to keep short. `BlockTable` is an open-addressing table with linear probing and tombstones, held at a load factor of 1/2 in one flat array. Insert and erase touch no allocator at all once the array has grown. `ReleaseCached` also stops holding the lock across upstream frees: it now settles the stats and detaches the blocks under the lock, then frees outside it. `cudaFree` implicitly synchronizes the whole device, so holding a mutex across one stalls every other thread for the duration. --- include/infini/rt/memory_pool.h | 272 ++++++++++++++++++++++++++------ 1 file changed, 227 insertions(+), 45 deletions(-) diff --git a/include/infini/rt/memory_pool.h b/include/infini/rt/memory_pool.h index 103251c..4799ed1 100644 --- a/include/infini/rt/memory_pool.h +++ b/include/infini/rt/memory_pool.h @@ -87,9 +87,7 @@ class MemoryPool { Upstream::Free(block.base); } } - for (auto& [ptr, block] : allocated_) { - Upstream::Free(block.base); - } + allocated_.ForEach([](const Block& block) { Upstream::Free(block.base); }); } /// Allocates at least `size` bytes, reusing a cached block when one with a @@ -110,48 +108,57 @@ class MemoryPool { return Upstream::kSuccess; } - std::lock_guard lock(mutex_); - const std::size_t rounded = RoundSize(size); const BucketKey key{rounded, alignment}; - Block block{}; - auto list_it = free_lists_.find(key); - if (list_it != free_lists_.end() && !list_it->second.empty()) { - block = list_it->second.back(); - list_it->second.pop_back(); - ++stats_.cache_hit_count; - } else { - const std::size_t upstream_size = - alignment == 0 ? rounded : rounded + alignment; - void* base = nullptr; - const Error status = Upstream::Malloc(&base, upstream_size); - ++stats_.upstream_alloc_count; - if (status != Upstream::kSuccess) { - return status; + { + std::lock_guard lock(mutex_); + Block cached{}; + if (TakeCached(key, &cached)) { + ++stats_.cache_hit_count; + *ptr = Register(cached); + return Upstream::kSuccess; } + } - block.base = base; - block.aligned = alignment == 0 ? base : AlignUp(base, alignment); - block.rounded_size = rounded; - block.upstream_size = upstream_size; - block.alignment = alignment; + // Cache miss. 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 a + // cache hit -- for the duration of one `cudaMalloc`. + const std::size_t upstream_size = + alignment == 0 ? rounded : rounded + alignment; + void* base = nullptr; + const Error status = Upstream::Malloc(&base, upstream_size); - stats_.bytes_reserved += upstream_size; - if (stats_.bytes_reserved > stats_.peak_bytes_reserved) { - stats_.peak_bytes_reserved = stats_.bytes_reserved; + std::lock_guard lock(mutex_); + ++stats_.upstream_alloc_count; + + if (status != Upstream::kSuccess) { + // Another thread may have returned a matching block while the lock was + // released, which turns an upstream failure into a hit. + Block cached{}; + if (TakeCached(key, &cached)) { + ++stats_.cache_hit_count; + *ptr = Register(cached); + return Upstream::kSuccess; } - ++stats_.cache_miss_count; + return status; } - allocated_.emplace(block.aligned, block); - stats_.bytes_in_use += block.rounded_size; - if (stats_.bytes_in_use > stats_.peak_bytes_in_use) { - stats_.peak_bytes_in_use = stats_.bytes_in_use; + Block block{}; + block.base = base; + block.aligned = alignment == 0 ? base : AlignUp(base, alignment); + block.rounded_size = rounded; + block.upstream_size = upstream_size; + block.alignment = alignment; + + stats_.bytes_reserved += upstream_size; + if (stats_.bytes_reserved > stats_.peak_bytes_reserved) { + stats_.peak_bytes_reserved = stats_.bytes_reserved; } - ++stats_.alloc_count; + ++stats_.cache_miss_count; - *ptr = block.aligned; + *ptr = Register(block); return Upstream::kSuccess; } @@ -166,14 +173,11 @@ class MemoryPool { std::lock_guard lock(mutex_); - auto it = allocated_.find(ptr); - if (it == allocated_.end()) { + Block block{}; + if (!allocated_.Take(ptr, &block)) { return InvalidValue(); } - const Block block = it->second; - allocated_.erase(it); - stats_.bytes_in_use -= block.rounded_size; ++stats_.free_count; @@ -187,16 +191,28 @@ class MemoryPool { /// defragmentation / trim knob: call it to release retained memory back to /// the device. void ReleaseCached() { - std::lock_guard lock(mutex_); + // Detach the free lists under the lock, then call upstream without it. As + // in `Allocate`, an upstream call is far more expensive than the + // bookkeeping, so it must not block other threads. Stats are settled while + // the lock is held: once detached the blocks are no longer the pool's, so a + // concurrent `GetStats` sees them gone even though the frees are in flight. + std::unordered_map, BucketKeyHash> detached; + { + std::lock_guard lock(mutex_); + detached.swap(free_lists_); + for (const auto& [key, blocks] : detached) { + for (const Block& block : blocks) { + ++stats_.upstream_free_count; + stats_.bytes_reserved -= block.upstream_size; + } + } + } - for (auto& [key, blocks] : free_lists_) { + for (const auto& [key, blocks] : detached) { for (const Block& block : blocks) { Upstream::Free(block.base); - ++stats_.upstream_free_count; - stats_.bytes_reserved -= block.upstream_size; } } - free_lists_.clear(); } /// Returns a snapshot of the pool's statistics. @@ -225,6 +241,149 @@ class MemoryPool { std::size_t alignment = 0; }; + // Open-addressing table mapping a live pointer to its `Block`. + // + // `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 the pool is trying to avoid. This table stores blocks 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. + class BlockTable { + public: + void Insert(void* key, const Block& block) { + // 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.block = block; + 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, block, State::kOccupied}; + ++live_; + } + + // Removes `key` and writes its block to `*out`. Returns false if `key` is + // not present, which is how `Deallocate` detects a foreign pointer. + bool Take(void* key, Block* 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.block; + 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.block); + } + } + } + + private: + enum class State : std::uint8_t { kEmpty, kOccupied, kTombstone }; + + struct Slot { + void* key = nullptr; + Block block{}; + 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.block); + } + } + } + + std::vector slots_; + std::size_t occupied_ = 0; // live + tombstones, for the load factor + std::size_t tombstones_ = 0; + std::size_t live_ = 0; + }; + // Free lists are keyed by rounded size and alignment so a reused block is // always geometrically identical to the request. struct BucketKey { @@ -267,8 +426,31 @@ class MemoryPool { static Error InvalidValue() { return static_cast(1); } + // Pops a cached block for `key`. Caller must hold `mutex_`. + bool TakeCached(const BucketKey& key, Block* out) { + auto it = free_lists_.find(key); + if (it == free_lists_.end() || it->second.empty()) { + return false; + } + *out = it->second.back(); + it->second.pop_back(); + return true; + } + + // Marks `block` live and returns the pointer the caller sees. Caller must + // hold `mutex_`. + void* Register(const Block& block) { + allocated_.Insert(block.aligned, block); + stats_.bytes_in_use += block.rounded_size; + if (stats_.bytes_in_use > stats_.peak_bytes_in_use) { + stats_.peak_bytes_in_use = stats_.bytes_in_use; + } + ++stats_.alloc_count; + return block.aligned; + } + mutable std::mutex mutex_; - std::unordered_map allocated_; + BlockTable allocated_; std::unordered_map, BucketKeyHash> free_lists_; Stats stats_; }; From 76b8f09b24916698dffbb42c33c09b523bd45c99 Mon Sep 17 00:00:00 2001 From: atchen <1920510674@qq.com> Date: Fri, 7 Aug 2026 10:09:01 +0000 Subject: [PATCH 3/8] refactor: add NodeArena and PointerTable allocation helpers Two pieces the arena pool in the next commit needs, separated out because neither is specific to it. `NodeArena` is a bump allocator with a free list, for the fixed-size nodes a pool's own bookkeeping needs. An allocator that calls the host heap to describe its own state has a dependency it cannot control the latency of, and under a lock that latency is every thread's. `PointerTable` is the previous commit's open-addressing table generalized over the mapped type. `MemoryPool` keeps its own `BlockTable`, which maps to a fixed record type and stays private to that header; this one is the reusable form the arena needs, mapping a pointer to an arbitrary `Value`. --- include/infini/rt/detail/node_arena.h | 156 +++++++++++++++++++ include/infini/rt/detail/pointer_table.h | 186 +++++++++++++++++++++++ 2 files changed, 342 insertions(+) create mode 100644 include/infini/rt/detail/node_arena.h create mode 100644 include/infini/rt/detail/pointer_table.h 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 From 264926c854789e9e37e8542ba1053950ebec4792 Mon Sep 17 00:00:00 2001 From: atchen <1920510674@qq.com> Date: Fri, 7 Aug 2026 10:10:50 +0000 Subject: [PATCH 4/8] feat: add a sub-allocating arena memory pool `MemoryPool` calls upstream once per cache miss. On a device that is hundreds of microseconds, so a workload whose sizes keep missing pays it over and over -- and a size-class cache misses structurally on the shapes inference produces, because a KV cache slab that grows with sequence length lands in a different class every step. `ArenaMemoryPool` requests large backing stores and slices them, so upstream calls scale with the high-water footprint rather than with the allocation count. Each backing is an address-ordered chunk list, so splitting and coalescing make a released block reusable at any size, not only at the size class it was allocated with. Three things make it fast enough for the win to survive contention: - Coalescing is deferred. `Deallocate` files a chunk into its exact-size bin and returns; merging runs in one pass only when a request cannot be served from what is already indexed. A loop that cycles through a handful of sizes never merges at all. - The free index is in two parts: 128 exact-size fast bins with an occupancy bitmap, plus a tree ordered by (size, address). An occupied exact-fit bin is already the best fit, so the common case never touches the tree. - In front of both sits a per-thread cache of exact-size blocks. A hit takes no lock. `Deallocate` parks into the releasing thread's cache, so the next allocation of that size on that thread skips both the acquisition and the best-fit lookup. Growth follows a doubling ramp to a cap, past which further backings are added at the cap size; a single request larger than the cap gets its own oversize backing. The first non-oversize backing is resident so a steady small workload keeps a warm arena. Everything else is subject to automatic shrink, with two scans of hysteresis -- `cudaFree` implicitly synchronizes the whole device, so thrashing it from the allocation path costs more than holding the memory one more round. The pool does not track streams: a block is reusable the instant `Deallocate` returns, and coalescing means it may come back at a different offset and size. Callers must ensure device-side access has completed first. This is the same contract `MemoryPool` has, documented here because coalescing makes violating it corrupt an unrelated allocation rather than merely reuse a block early. Neither pool replaces the other. The arena wins where upstream is expensive or the shapes defeat size classes; the size-class pool wins on large-block recycling and keeps a shorter critical section at low thread counts. Both stay. --- include/infini/rt/arena_memory_pool.h | 1618 +++++++++++++++++++++++ tests/CMakeLists.txt | 23 + tests/test_arena_memory_pool.cc | 937 +++++++++++++ tests/test_arena_memory_pool_backend.cc | 393 ++++++ 4 files changed, 2971 insertions(+) create mode 100644 include/infini/rt/arena_memory_pool.h create mode 100644 tests/test_arena_memory_pool.cc create mode 100644 tests/test_arena_memory_pool_backend.cc 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/tests/CMakeLists.txt b/tests/CMakeLists.txt index 4e0cc23..158fd1a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -51,11 +51,30 @@ function(add_infini_rt_backend_memory_pool_test backend 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) @@ -70,6 +89,8 @@ if(WITH_CPU) 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) @@ -82,6 +103,8 @@ if(WITH_NVIDIA) 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/test_arena_memory_pool.cc b/tests/test_arena_memory_pool.cc new file mode 100644 index 0000000..659d243 --- /dev/null +++ b/tests/test_arena_memory_pool.cc @@ -0,0 +1,937 @@ +// Exercises `ArenaMemoryPool` against a mock upstream allocator. +// +// The pool's interesting behaviors -- the growth ramp, oversize backings, +// automatic shrink -- only trigger at multi-hundred-megabyte scale with the +// production config, which no CI machine can allocate (and a GPU-less one +// certainly cannot). That is what `Config` is a template parameter for: these +// tests instantiate the pool at kilobyte scale and drive exactly the same code +// paths, with a mock upstream that counts calls and hands out host memory so +// slices can be written through to prove they do not overlap. +// +// `test_arena_memory_pool_backend.cc` covers the same pool over a real device +// runtime. +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "test_helper.h" + +namespace { + +using infini::rt::ArenaMemoryPool; + +// -------------------------------------------------------------------------- +// Mock upstream +// -------------------------------------------------------------------------- + +// Counting upstream over `std::aligned_alloc`. Real host memory rather than +// fake pointers, so tests can write through the slices the pool hands out and +// catch a slicing bug that overlapping ranges would otherwise hide. +// +// State is static because the pool takes its upstream as a type, not an +// instance -- the same shape every `runtime::Runtime<...>` specialization has. +struct MockUpstream { + using Error = int; + static constexpr Error kSuccess = 0; + static constexpr Error kFailure = 7; + + // Atomic because the pool calls upstream with its own lock released -- by + // design, since a device allocator is a slow synchronous call. A real + // `cudaMalloc` is thread-safe, so the mock has to be too. + static std::atomic malloc_calls; + static std::atomic free_calls; + // When non-zero, requests strictly larger than this fail. Drives the OOM + // fallback chain. Only set while no other thread is running. + static std::size_t capacity_limit; + + static void Reset() { + malloc_calls.store(0); + free_calls.store(0); + capacity_limit = 0; + } + + static Error Malloc(void** ptr, std::size_t size) { + malloc_calls.fetch_add(1, std::memory_order_relaxed); + if (capacity_limit != 0 && size > capacity_limit) { + *ptr = nullptr; + return kFailure; + } + // 256 B matches what `cudaMalloc` guarantees, so the pool's own slice + // alignment is what the tests below are actually observing. + void* base = std::aligned_alloc(256, RoundUp(size, 256)); + if (base == nullptr) { + *ptr = nullptr; + return kFailure; + } + *ptr = base; + return kSuccess; + } + + static Error Free(void* ptr) { + free_calls.fetch_add(1, std::memory_order_relaxed); + std::free(ptr); + return kSuccess; + } + + static std::size_t RoundUp(std::size_t size, std::size_t granularity) { + return (size + granularity - 1) / granularity * granularity; + } +}; + +std::atomic MockUpstream::malloc_calls{0}; +std::atomic MockUpstream::free_calls{0}; +std::size_t MockUpstream::capacity_limit = 0; + +// Kilobyte-scale mirror of `DefaultArenaConfig`, preserving every ratio that +// matters: initial capacity, an 8x doubling headroom to the cap, a small/large +// threshold below the cap, and the same slice granularity. +struct TinyConfig { + static constexpr std::size_t kInitialCapacity = 8 * 1024; + static constexpr std::size_t kMaxCapacity = 64 * 1024; + static constexpr std::size_t kSmallThreshold = 1024; + static constexpr std::size_t kMinSliceAlignment = 64; + static constexpr std::size_t kMinSplitRemainder = 64; + static constexpr std::size_t kShrinkThreshold = 4; + static constexpr std::uint32_t kEmptyScansToDestroy = 2; +}; + +using Pool = ArenaMemoryPool; + +// A config that never shrinks automatically, for tests that want to observe +// fragmentation and reuse without backings disappearing underneath them. +struct NoShrinkConfig : TinyConfig { + static constexpr std::size_t kShrinkThreshold = + static_cast(-1) / 2; +}; + +using StablePool = ArenaMemoryPool; + +// -------------------------------------------------------------------------- +// Helpers +// -------------------------------------------------------------------------- + +template +void* Alloc(infini::rt::test::TestContext* context, P* pool, std::size_t size, + std::size_t alignment = 0) { + void* ptr = nullptr; + context->ExpectEqual(pool->Allocate(&ptr, size, alignment), + MockUpstream::kSuccess, "allocate should succeed"); + context->Expect(ptr != nullptr, "allocate should produce a pointer"); + return ptr; +} + +// Writes a byte pattern over `[ptr, ptr + size)`. Combined with `Verify`, this +// is how the tests prove two live slices do not overlap: an overlap shows up as +// one block reading back another's pattern. +void Fill(void* ptr, std::size_t size, std::uint8_t seed) { + auto* bytes = static_cast(ptr); + for (std::size_t i = 0; i < size; ++i) { + bytes[i] = static_cast(seed + (i & 0x3f)); + } +} + +bool Verify(const void* ptr, std::size_t size, std::uint8_t seed) { + const auto* bytes = static_cast(ptr); + for (std::size_t i = 0; i < size; ++i) { + if (bytes[i] != static_cast(seed + (i & 0x3f))) { + return false; + } + } + return true; +} + +// -------------------------------------------------------------------------- +// Basics +// -------------------------------------------------------------------------- + +// One allocation carves a backing store out of the upstream allocator; the +// pointer is usable and the reported capacity follows the configured ramp. +void TestFirstAllocation(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + Pool pool; + + void* ptr = Alloc(context, &pool, 256); + Fill(ptr, 256, 0x11); + context->Expect(Verify(ptr, 256, 0x11), "slice should be writable memory"); + + const Pool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.upstream_alloc_count, std::size_t{1}, + "first allocation creates one backing store"); + context->ExpectEqual(stats.backing_count, std::size_t{1}, + "one backing store live"); + context->Expect(stats.bytes_reserved >= TinyConfig::kInitialCapacity, + "reserved bytes cover the initial capacity"); + context->ExpectEqual(stats.bytes_in_use, std::size_t{256}, + "one 256 B slice in use"); + context->ExpectEqual(stats.cache_miss_count, std::size_t{1}, + "the first allocation misses"); + pool.Deallocate(ptr); +} + +// A zero-byte request succeeds with a null pointer and touches nothing, and a +// null `Deallocate` is a no-op -- the same contract `MemoryPool` offers. +void TestDegenerateRequests(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + Pool pool; + + void* ptr = reinterpret_cast(std::uintptr_t{0xdeadbeef}); + context->ExpectEqual(pool.Allocate(&ptr, 0), MockUpstream::kSuccess, + "zero-byte allocate should succeed"); + context->ExpectEqual(ptr, static_cast(nullptr), + "zero-byte allocate yields nullptr"); + context->ExpectEqual(pool.Deallocate(nullptr), MockUpstream::kSuccess, + "deallocating nullptr is a no-op"); + context->Expect(pool.Allocate(nullptr, 64) != MockUpstream::kSuccess, + "a null out-pointer is an invalid value"); + context->ExpectEqual(MockUpstream::malloc_calls.load(), std::size_t{0}, + "degenerate requests touch no upstream memory"); +} + +// A pointer the pool never handed out, and a double free, are both rejected +// rather than corrupting the chunk lists. +void TestForeignPointer(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + Pool pool; + + int stack_object = 0; + context->Expect(pool.Deallocate(&stack_object) != MockUpstream::kSuccess, + "a foreign pointer must be rejected"); + + void* ptr = Alloc(context, &pool, 128); + context->ExpectEqual(pool.Deallocate(ptr), MockUpstream::kSuccess, + "first free should succeed"); + context->Expect(pool.Deallocate(ptr) != MockUpstream::kSuccess, + "a double free must be rejected"); +} + +// Many allocations come out of one backing store: this is the whole point of +// the arena, so the upstream call count must stay at one. +void TestSlicingAvoidsUpstream(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + std::vector blocks; + for (std::size_t i = 0; i < 16; ++i) { + void* ptr = Alloc(context, &pool, 256); + Fill(ptr, 256, static_cast(i + 1)); + blocks.push_back(ptr); + } + + // Every block still reads back its own pattern, so no two slices overlap. + bool distinct = true; + for (std::size_t i = 0; i < blocks.size(); ++i) { + distinct = distinct && Verify(blocks[i], 256, + static_cast(i + 1)); + } + context->Expect(distinct, "concurrent slices must not overlap"); + + const StablePool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.upstream_alloc_count, std::size_t{1}, + "16 slices of 256 B fit in one 8 KB backing store"); + context->ExpectEqual(stats.cache_hit_count, std::size_t{15}, + "only the first allocation creates a backing"); + + for (void* ptr : blocks) { + pool.Deallocate(ptr); + } +} + +// `alignment` is honored, and the gap skipped to reach the aligned offset comes +// back as reusable free space rather than leaking. +void TestAlignment(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + // Occupy an offset that leaves the next chunk misaligned for a 1 KB request. + void* filler = Alloc(context, &pool, 64); + void* aligned = Alloc(context, &pool, 256, 1024); + context->ExpectEqual(reinterpret_cast(aligned) % 1024, + std::uintptr_t{0}, + "the returned pointer must honor the alignment"); + + const StablePool::Stats before = pool.GetStats(); + pool.Deallocate(aligned); + pool.Deallocate(filler); + const StablePool::Stats after = pool.GetStats(); + + context->ExpectEqual(after.bytes_in_use, std::size_t{0}, + "everything has been returned"); + // Coalescing must reunite the alignment gap with its neighbors, leaving the + // backing as one extent again. + context->ExpectEqual(after.largest_free_chunk, + before.bytes_reserved - before.bytes_unusable, + "the whole backing coalesces back into one chunk"); +} + +// With no explicit alignment the pool still guarantees its slice granularity -- +// callers and vectorized kernels rely on the natural alignment a device +// allocator would have given them. +void TestNaturalAlignment(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + bool all_aligned = true; + std::vector blocks; + for (std::size_t size : {std::size_t{1}, std::size_t{7}, std::size_t{65}, + std::size_t{130}, std::size_t{999}}) { + void* ptr = Alloc(context, &pool, size); + all_aligned = all_aligned && + reinterpret_cast(ptr) % + TinyConfig::kMinSliceAlignment == + 0; + blocks.push_back(ptr); + } + context->Expect(all_aligned, + "every slice honors the minimum slice alignment"); + for (void* ptr : blocks) { + pool.Deallocate(ptr); + } +} + +// -------------------------------------------------------------------------- +// Coalescing +// -------------------------------------------------------------------------- + +// The property a bump-pointer arena cannot provide: after N adjacent blocks are +// released, the backing can serve one request spanning all of them again. +void TestCoalescingRestoresContiguity(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + std::vector blocks; + for (std::size_t i = 0; i < 8; ++i) { + blocks.push_back(Alloc(context, &pool, 512)); + } + const std::size_t reserved = pool.GetStats().bytes_reserved; + const std::size_t unusable = pool.GetStats().bytes_unusable; + + // Free in an interleaved order so the merge happens from both sides, not just + // as a tidy right-to-left unwind. + for (std::size_t i : {std::size_t{3}, std::size_t{0}, std::size_t{7}, + std::size_t{1}, std::size_t{5}, std::size_t{2}, + std::size_t{6}, std::size_t{4}}) { + pool.Deallocate(blocks[i]); + } + + const StablePool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.bytes_in_use, std::size_t{0}, + "all blocks returned"); + context->ExpectEqual(stats.largest_free_chunk, reserved - unusable, + "the backing coalesces back into a single extent"); + + // And it can actually be handed out as one block again. + void* whole = nullptr; + context->ExpectEqual(pool.Allocate(&whole, reserved - unusable), + MockUpstream::kSuccess, + "the coalesced extent serves a full-span request"); + context->ExpectEqual(pool.GetStats().upstream_alloc_count, std::size_t{1}, + "serving it needed no new backing store"); + pool.Deallocate(whole); +} + +// A released block is reusable at a completely different size, which exact +// size-class matching could not do. +void TestReuseAcrossSizes(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + void* big = Alloc(context, &pool, 2048); + pool.Deallocate(big); + + // Four 512 B requests should come out of the 2 KB hole, not a new backing. + std::vector blocks; + for (std::size_t i = 0; i < 4; ++i) { + blocks.push_back(Alloc(context, &pool, 512)); + } + context->ExpectEqual(pool.GetStats().upstream_alloc_count, std::size_t{1}, + "a freed 2 KB block serves four 512 B requests"); + for (void* ptr : blocks) { + pool.Deallocate(ptr); + } +} + +// -------------------------------------------------------------------------- +// Growth +// -------------------------------------------------------------------------- + +// Exhausting a backing adds another, and capacity follows the doubling ramp up +// to the configured cap. +void TestGrowthRamp(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + std::vector blocks; + std::vector reserved_after; + // Keep allocating 2 KB blocks; each backing holds a few, so this walks the + // ramp 8 KB -> 16 KB -> 32 KB -> 64 KB. + for (std::size_t i = 0; i < 64; ++i) { + blocks.push_back(Alloc(context, &pool, 2048)); + reserved_after.push_back(pool.GetStats().bytes_reserved); + } + + const StablePool::Stats stats = pool.GetStats(); + context->Expect(stats.backing_count > 1, "growth added backing stores"); + context->ExpectEqual(stats.upstream_alloc_count, stats.backing_count, + "one upstream call per backing store"); + // 64 x 2 KB = 128 KB of demand served by far fewer than 64 upstream calls -- + // the point of the ramp. + context->Expect(stats.upstream_alloc_count <= 8, + "the doubling ramp keeps upstream calls sublinear"); + + // No individual step may exceed the cap. + bool capped = true; + for (std::size_t i = 1; i < reserved_after.size(); ++i) { + const std::size_t step = reserved_after[i] - reserved_after[i - 1]; + capped = capped && step <= NoShrinkConfig::kMaxCapacity; + } + context->Expect(capped, "no backing store exceeds the capacity cap"); + + for (void* ptr : blocks) { + pool.Deallocate(ptr); + } +} + +// A single request larger than the cap gets its own exactly sized backing: +// several capped backings cannot serve it, because a slice never spans two +// upstream allocations. +void TestOversizeRequest(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + constexpr std::size_t kHuge = NoShrinkConfig::kMaxCapacity * 3; + void* ptr = Alloc(context, &pool, kHuge); + Fill(ptr, kHuge, 0x5a); + context->Expect(Verify(ptr, kHuge, 0x5a), + "an oversize slice is fully writable"); + + const StablePool::Stats stats = pool.GetStats(); + context->Expect(stats.bytes_reserved >= kHuge, + "the oversize backing covers the request"); + context->ExpectEqual(stats.bytes_in_use, kHuge, "the whole request is live"); + + // The ramp must not have been advanced by the exception: the next ordinary + // backing stays at cap size rather than jumping to 6x the cap. + pool.Deallocate(ptr); + const std::size_t before = pool.GetStats().bytes_reserved; + std::vector blocks; + for (std::size_t i = 0; i < 200; ++i) { + blocks.push_back(Alloc(context, &pool, 512)); + } + const std::size_t grew = pool.GetStats().bytes_reserved - before; + context->Expect(grew <= NoShrinkConfig::kMaxCapacity, + "an oversize backing must not advance the growth ramp"); + for (void* ptr2 : blocks) { + pool.Deallocate(ptr2); + } +} + +// -------------------------------------------------------------------------- +// Shrink +// -------------------------------------------------------------------------- + +// The headline scenario: a burst allocates extra backings, and once the +// workload goes back to small requests the extras are returned upstream while +// the resident backing stays warm. +void TestShrinkAfterBurst(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + Pool pool; + + // Warm up the resident backing with a small allocation. + void* resident_block = Alloc(context, &pool, 128); + + // Burst: large requests force several more backings. + std::vector burst; + for (std::size_t i = 0; i < 24; ++i) { + burst.push_back(Alloc(context, &pool, 4096)); + } + const Pool::Stats peak = pool.GetStats(); + context->Expect(peak.backing_count > 1, "the burst added backing stores"); + + // The burst finishes. + for (void* ptr : burst) { + pool.Deallocate(ptr); + } + context->ExpectEqual(pool.GetStats().backing_count, peak.backing_count, + "freeing alone does not release backings upstream"); + + // Now a run of small requests. Each scan needs `kShrinkThreshold` + // allocations, and a non-oversize backing needs `kEmptyScansToDestroy` scans, + // so drive enough small work for the hysteresis to play out. + std::vector small; + for (std::size_t i = 0; + i < TinyConfig::kShrinkThreshold * (TinyConfig::kEmptyScansToDestroy + 2); + ++i) { + small.push_back(Alloc(context, &pool, 64)); + } + + const Pool::Stats after = pool.GetStats(); + context->ExpectEqual(after.backing_count, std::size_t{1}, + "idle backings are destroyed automatically"); + context->Expect(after.shrink_count > 0, "shrink was recorded"); + context->ExpectEqual(after.upstream_free_count, after.shrink_count, + "every shrink is one upstream free"); + context->Expect(after.bytes_reserved <= TinyConfig::kInitialCapacity, + "reserved memory falls back to the resident backing"); + + // The resident block was never disturbed. + context->ExpectEqual(pool.Deallocate(resident_block), MockUpstream::kSuccess, + "the resident backing survived the shrink"); + for (void* ptr : small) { + pool.Deallocate(ptr); + } +} + +// A backing with live blocks is never destroyed, however long the small-request +// run gets. +void TestShrinkSpareLiveBackings(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + Pool pool; + + void* pinned = nullptr; + std::vector burst; + for (std::size_t i = 0; i < 24; ++i) { + void* ptr = Alloc(context, &pool, 4096); + if (pool.GetStats().backing_count > 1 && pinned == nullptr) { + pinned = ptr; // Lives in a non-resident backing. + Fill(pinned, 4096, 0x33); + } else { + burst.push_back(ptr); + } + } + context->Expect(pinned != nullptr, "the burst reached a second backing"); + for (void* ptr : burst) { + pool.Deallocate(ptr); + } + + std::vector small; + for (std::size_t i = 0; i < TinyConfig::kShrinkThreshold * 8; ++i) { + small.push_back(Alloc(context, &pool, 64)); + } + + context->Expect(pool.GetStats().backing_count >= 2, + "a backing holding a live block is never destroyed"); + context->Expect(Verify(pinned, 4096, 0x33), + "the live block's contents are intact"); + + pool.Deallocate(pinned); + for (void* ptr : small) { + pool.Deallocate(ptr); + } +} + +// A large request resets the run, so an ongoing burst is never shrunk out from +// under itself. +void TestLargeRequestResetsShrinkRun(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + Pool pool; + + std::vector live; + live.push_back(Alloc(context, &pool, 64)); + for (std::size_t i = 0; i < 24; ++i) { + live.push_back(Alloc(context, &pool, 4096)); + } + const std::size_t peak_backings = pool.GetStats().backing_count; + + // Interleave: never `kShrinkThreshold` small requests in a row. + for (std::size_t round = 0; round < 12; ++round) { + for (std::size_t i = 0; i < TinyConfig::kShrinkThreshold - 1; ++i) { + void* ptr = Alloc(context, &pool, 64); + pool.Deallocate(ptr); + } + void* big = Alloc(context, &pool, 4096); + pool.Deallocate(big); + } + + context->ExpectEqual(pool.GetStats().shrink_count, std::size_t{0}, + "an interleaved large request keeps the arena warm"); + context->ExpectEqual(pool.GetStats().backing_count, peak_backings, + "no backing was destroyed mid-burst"); + for (void* ptr : live) { + pool.Deallocate(ptr); + } +} + +// Hysteresis: an alternating big/small workload must not destroy and re-create +// a backing every iteration. `cudaFree` implicitly synchronizes the device, so +// thrashing would cost more than the memory it reclaims. +void TestShrinkHysteresis(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + Pool pool; + + // Force a second backing to exist, then drive many alternating rounds. + void* anchor = Alloc(context, &pool, 64); + std::vector burst; + for (std::size_t i = 0; i < 24; ++i) { + burst.push_back(Alloc(context, &pool, 4096)); + } + for (void* ptr : burst) { + pool.Deallocate(ptr); + } + + constexpr std::size_t kRounds = 40; + for (std::size_t round = 0; round < kRounds; ++round) { + void* big = Alloc(context, &pool, 4096); + for (std::size_t i = 0; i < 6; ++i) { + void* ptr = Alloc(context, &pool, 64); + pool.Deallocate(ptr); + } + pool.Deallocate(big); + } + + const Pool::Stats stats = pool.GetStats(); + // Without hysteresis this would be ~one destroy plus one create per round. + context->Expect(stats.shrink_count < kRounds / 2, + "hysteresis prevents per-iteration backing thrash"); + context->Expect(stats.upstream_alloc_count < kRounds, + "upstream calls do not grow linearly with iterations"); + pool.Deallocate(anchor); +} + +// An oversize backing skips the hysteresis: holding that much memory idle costs +// more than the extra upstream call. +void TestOversizeShrinksPromptly(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + Pool pool; + + void* anchor = Alloc(context, &pool, 64); + void* huge = Alloc(context, &pool, TinyConfig::kMaxCapacity * 3); + const std::size_t peak_reserved = pool.GetStats().bytes_reserved; + pool.Deallocate(huge); + + // Exactly one scan's worth of small requests. + std::vector small; + for (std::size_t i = 0; i < TinyConfig::kShrinkThreshold; ++i) { + small.push_back(Alloc(context, &pool, 64)); + } + + const Pool::Stats stats = pool.GetStats(); + context->Expect(stats.bytes_reserved < peak_reserved, + "the oversize backing goes back on the first idle scan"); + context->Expect(stats.shrink_count >= 1, "shrink was recorded"); + + pool.Deallocate(anchor); + for (void* ptr : small) { + pool.Deallocate(ptr); + } +} + +// -------------------------------------------------------------------------- +// ReleaseCached, stats, OOM +// -------------------------------------------------------------------------- + +// `ReleaseCached` returns every drained backing, resident one included, and +// leaves backings with live blocks alone. +void TestReleaseCached(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + std::vector blocks; + for (std::size_t i = 0; i < 24; ++i) { + blocks.push_back(Alloc(context, &pool, 4096)); + } + void* keep = blocks.front(); + for (std::size_t i = 1; i < blocks.size(); ++i) { + pool.Deallocate(blocks[i]); + } + + pool.ReleaseCached(); + StablePool::Stats stats = pool.GetStats(); + context->Expect(stats.backing_count >= 1, + "the backing holding a live block is retained"); + context->Expect(stats.upstream_free_count > 0, + "drained backings were released"); + + pool.Deallocate(keep); + pool.ReleaseCached(); + stats = pool.GetStats(); + context->ExpectEqual(stats.backing_count, std::size_t{0}, + "an empty pool releases everything"); + context->ExpectEqual(stats.bytes_reserved, std::size_t{0}, + "reserved bytes drop to zero"); + context->ExpectEqual(stats.upstream_free_count, stats.upstream_alloc_count, + "every backing store was freed exactly once"); + + // And the pool is still usable, restarting the growth ramp. + void* fresh = Alloc(context, &pool, 128); + context->ExpectEqual(pool.GetStats().bytes_reserved, + std::size_t{NoShrinkConfig::kInitialCapacity}, + "the growth ramp restarts after a full release"); + pool.Deallocate(fresh); +} + +// The byte counters must always close: reserved memory is either in use, free +// inside a backing, or trimmed off a backing's head. +void TestStatsBalance(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + std::mt19937 rng(1234); + std::uniform_int_distribution size_dist(1, 4096); + std::vector live; + bool balanced = true; + bool waste_bounded = true; + + for (std::size_t step = 0; step < 2000; ++step) { + if (live.empty() || (rng() & 1) != 0) { + live.push_back(Alloc(context, &pool, size_dist(rng))); + } else { + const std::size_t index = rng() % live.size(); + pool.Deallocate(live[index]); + live.erase(live.begin() + static_cast(index)); + } + + const StablePool::Stats stats = pool.GetStats(); + balanced = balanced && + stats.bytes_reserved == stats.bytes_in_use + + stats.bytes_free_in_backings + + stats.bytes_unusable; + waste_bounded = waste_bounded && stats.bytes_internal_waste <= + stats.bytes_in_use; + } + + context->Expect(balanced, + "reserved == in_use + free_in_backings + unusable, always"); + context->Expect(waste_bounded, + "internal waste is a subset of the bytes in use"); + + for (void* ptr : live) { + pool.Deallocate(ptr); + } + const StablePool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.bytes_in_use, std::size_t{0}, + "everything returned"); + context->ExpectEqual(stats.bytes_internal_waste, std::size_t{0}, + "internal waste clears with the last live chunk"); + context->ExpectEqual(stats.alloc_count, stats.free_count, + "every allocation was matched by a free"); +} + +// Randomized churn at mixed sizes must not degrade into unusable fragments: +// the arena has to keep serving requests without an upstream call per +// allocation, and free space has to stay in usable extents. +void TestFragmentationUnderChurn(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + std::mt19937 rng(98765); + std::uniform_int_distribution size_dist(64, 8192); + std::vector live; + constexpr std::size_t kSteps = 20000; + + for (std::size_t step = 0; step < kSteps; ++step) { + if (live.size() < 48 && ((rng() & 3) != 0 || live.empty())) { + void* ptr = nullptr; + if (pool.Allocate(&ptr, size_dist(rng)) != MockUpstream::kSuccess) { + context->Expect(false, "churn must not fail to allocate"); + break; + } + live.push_back(ptr); + } else { + const std::size_t index = rng() % live.size(); + pool.Deallocate(live[index]); + live.erase(live.begin() + static_cast(index)); + } + } + + const StablePool::Stats stats = pool.GetStats(); + // Upstream calls must scale with the footprint, not the operation count. + context->Expect(stats.upstream_alloc_count < kSteps / 100, + "churn does not drive an upstream call per allocation"); + // Live demand peaked around 48 x 8 KB = 384 KB; a healthy allocator holds a + // small multiple of that. + context->Expect(stats.bytes_reserved < 4 * 1024 * 1024, + "memory amplification stays bounded under churn"); + + for (void* ptr : live) { + pool.Deallocate(ptr); + } + // Fully drained, every backing must have collapsed to one extent, which + // `ReleaseCached` can then hand back in full. + pool.ReleaseCached(); + const StablePool::Stats drained = pool.GetStats(); + context->ExpectEqual(drained.bytes_reserved, std::size_t{0}, + "a fully drained pool releases every backing"); + context->ExpectEqual(drained.upstream_free_count, + drained.upstream_alloc_count, + "no backing store leaked across the churn"); +} + +// When the upstream allocator cannot satisfy the ramp capacity, the pool falls +// back to a smaller backing rather than failing the request. +void TestOomFallbackShrinksRequest(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + // Below the 8 KB initial capacity, so the first ramp attempt must fail. + MockUpstream::capacity_limit = 4096; + + StablePool pool; + void* ptr = nullptr; + context->ExpectEqual(pool.Allocate(&ptr, 1024), MockUpstream::kSuccess, + "the pool falls back to a capacity upstream can serve"); + context->Expect(ptr != nullptr, "fallback still yields a pointer"); + context->Expect(pool.GetStats().bytes_reserved <= 4096, + "the fallback backing fits within the upstream limit"); + pool.Deallocate(ptr); + MockUpstream::capacity_limit = 0; +} + +// A request the upstream allocator cannot serve at any size fails cleanly, with +// a null pointer and the pool still usable. +void TestOomFailurePropagates(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + MockUpstream::capacity_limit = 1024; + + StablePool pool; + void* ptr = reinterpret_cast(std::uintptr_t{0xabcd}); + context->Expect(pool.Allocate(&ptr, 64 * 1024) != MockUpstream::kSuccess, + "an unservable request returns the upstream error"); + context->ExpectEqual(ptr, static_cast(nullptr), + "a failed allocation yields nullptr"); + + MockUpstream::capacity_limit = 0; + void* recovered = Alloc(context, &pool, 512); + context->Expect(recovered != nullptr, + "the pool still works after an upstream failure"); + pool.Deallocate(recovered); +} + +// Concurrent allocate/free traffic across threads. Each thread writes a pattern +// unique to itself into every block it holds and checks it before releasing, so +// a slice handed to two threads at once shows up as corrupted data rather than +// as a merely suspicious counter. Run under a thread sanitizer this also covers +// the locking discipline. +void TestConcurrentTraffic(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + StablePool pool; + + constexpr std::size_t kThreads = 8; + constexpr std::size_t kOpsPerThread = 4000; + std::atomic corruption{0}; + std::atomic failures{0}; + + std::vector threads; + for (std::size_t t = 0; t < kThreads; ++t) { + threads.emplace_back([&pool, &corruption, &failures, t] { + std::mt19937 rng(static_cast(t * 7919 + 13)); + std::uniform_int_distribution size_dist(64, 4096); + const auto seed = static_cast(t * 31 + 1); + std::vector> live; + + for (std::size_t op = 0; op < kOpsPerThread; ++op) { + if (live.size() < 12 && ((rng() & 3) != 0 || live.empty())) { + const std::size_t size = size_dist(rng); + void* ptr = nullptr; + if (pool.Allocate(&ptr, size) != MockUpstream::kSuccess || + ptr == nullptr) { + failures.fetch_add(1, std::memory_order_relaxed); + continue; + } + Fill(ptr, size, seed); + live.emplace_back(ptr, size); + } else { + const std::size_t index = rng() % live.size(); + const auto [ptr, size] = live[index]; + if (!Verify(ptr, size, seed)) { + corruption.fetch_add(1, std::memory_order_relaxed); + } + pool.Deallocate(ptr); + live.erase(live.begin() + static_cast(index)); + } + } + + for (const auto& [ptr, size] : live) { + if (!Verify(ptr, size, seed)) { + corruption.fetch_add(1, std::memory_order_relaxed); + } + pool.Deallocate(ptr); + } + }); + } + for (std::thread& thread : threads) { + thread.join(); + } + + context->ExpectEqual(corruption.load(), std::size_t{0}, + "concurrent slices must never alias"); + context->ExpectEqual(failures.load(), std::size_t{0}, + "concurrent allocation must not fail"); + + const StablePool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.bytes_in_use, std::size_t{0}, + "all threads returned their blocks"); + context->ExpectEqual(stats.alloc_count, stats.free_count, + "allocation and free counts agree across threads"); + pool.ReleaseCached(); + context->ExpectEqual(pool.GetStats().upstream_free_count, + pool.GetStats().upstream_alloc_count, + "no backing store leaked under concurrency"); +} + +// Destruction frees each backing store exactly once. Blocks left outstanding +// are slices, not upstream pointers, so they must not produce their own frees. +void TestDestructorFreesBackingsOnce(infini::rt::test::TestContext* context) { + MockUpstream::Reset(); + std::size_t allocs = 0; + { + StablePool pool; + std::vector blocks; + for (std::size_t i = 0; i < 32; ++i) { + blocks.push_back(Alloc(context, &pool, 2048)); + } + // Deliberately leave half outstanding. + for (std::size_t i = 0; i < blocks.size(); i += 2) { + pool.Deallocate(blocks[i]); + } + allocs = pool.GetStats().upstream_alloc_count; + context->Expect(allocs > 1, "the test spans several backing stores"); + } + + context->ExpectEqual(MockUpstream::free_calls.load(), allocs, + "destruction frees each backing store exactly once"); + context->ExpectEqual(MockUpstream::malloc_calls.load(), + MockUpstream::free_calls.load(), + "no upstream allocation leaked"); +} + +} // namespace + +int main() { + infini::rt::test::TestContext context; + + TestFirstAllocation(&context); + TestDegenerateRequests(&context); + TestForeignPointer(&context); + TestSlicingAvoidsUpstream(&context); + TestAlignment(&context); + TestNaturalAlignment(&context); + + TestCoalescingRestoresContiguity(&context); + TestReuseAcrossSizes(&context); + + TestGrowthRamp(&context); + TestOversizeRequest(&context); + + TestShrinkAfterBurst(&context); + TestShrinkSpareLiveBackings(&context); + TestLargeRequestResetsShrinkRun(&context); + TestShrinkHysteresis(&context); + TestOversizeShrinksPromptly(&context); + + TestReleaseCached(&context); + TestStatsBalance(&context); + TestFragmentationUnderChurn(&context); + TestOomFallbackShrinksRequest(&context); + TestOomFailurePropagates(&context); + TestConcurrentTraffic(&context); + TestDestructorFreesBackingsOnce(&context); + + return context.ExitCode(); +} diff --git a/tests/test_arena_memory_pool_backend.cc b/tests/test_arena_memory_pool_backend.cc new file mode 100644 index 0000000..539440c --- /dev/null +++ b/tests/test_arena_memory_pool_backend.cc @@ -0,0 +1,393 @@ +// Exercises `ArenaMemoryPool` over a *real* runtime backend (CPU, NVIDIA, ...). +// +// `test_arena_memory_pool.cc` already covers the pool's bookkeeping against a +// mock upstream at kilobyte scale. This test instead instantiates the pool over +// the backend's actual `runtime::Runtime` specialization, which is where the +// arena's central claim has to hold: a slice is an *interior offset* into one +// upstream allocation, so nothing but a real device round trip can prove that +// the pointer arithmetic lands inside genuine device memory and that two +// adjacent slices out of the same backing do not overwrite each other. Device +// pointers cannot be dereferenced from the host, so every check goes through +// `Memcpy`. The whole suite is skipped when no device is present. +#include +#include +#include INFINI_RT_TEST_RUNTIME_HEADER + +#include +#include +#include +#include +#include + +#include "test_helper.h" + +namespace { + +using Runtime = infini::rt::runtime::Runtime; + +// A megabyte-scale config. The production default reserves 64 MB per backing +// and ramps to 512 MB, which a shared CI device may not have to spare -- and at +// that scale a test would never reach the growth, oversize, or shrink paths +// within a reasonable number of allocations. The ratios that matter are +// preserved: an 8x doubling headroom to the cap, and a small/large threshold +// well below it. +struct BackendConfig { + static constexpr std::size_t kInitialCapacity = 4ull << 20; // 4 MB + static constexpr std::size_t kMaxCapacity = 32ull << 20; // 32 MB + static constexpr std::size_t kSmallThreshold = 64ull << 10; // 64 KB + static constexpr std::size_t kMinSliceAlignment = 512; + static constexpr std::size_t kMinSplitRemainder = 512; + static constexpr std::size_t kShrinkThreshold = 8; + static constexpr std::uint32_t kEmptyScansToDestroy = 2; +}; + +using Pool = infini::rt::ArenaMemoryPool; + +constexpr const char* kBackend = INFINI_RT_TEST_BACKEND_NAME; + +bool SelectDevice() { + int device_count = 0; + if (Runtime::GetDeviceCount(&device_count) != Runtime::kSuccess || + device_count <= 0) { + std::cout << kBackend << " arena memory pool skipped: no available device." + << std::endl; + return false; + } + if (Runtime::SetDevice(0) != Runtime::kSuccess) { + std::cout << kBackend << " arena memory pool skipped: device 0 unavailable." + << std::endl; + return false; + } + return true; +} + +// Writes `input` into device memory `ptr` and reads it back, asserting the +// bytes survive the round trip. This is the only host-safe way to confirm a +// device pointer is real and usable. +template +void ExpectUsable(infini::rt::test::TestContext* context, void* ptr, + const std::array& input, + const char* message) { + if (!context->Expect(ptr != nullptr, message)) { + return; + } + std::array output{}; + context->ExpectEqual( + Runtime::Memcpy(ptr, input.data(), N, Runtime::kMemcpyHostToDevice), + Runtime::kSuccess, "memcpy host-to-device should succeed"); + context->ExpectEqual( + Runtime::Memcpy(output.data(), ptr, N, Runtime::kMemcpyDeviceToHost), + Runtime::kSuccess, "memcpy device-to-host should succeed"); + context->ExpectEqual(output, input, + "pool-allocated memory should round-trip bytes"); +} + +// Fills device memory `[ptr, ptr + size)` with a pattern derived from `seed`. +bool FillDevice(void* ptr, std::size_t size, std::uint8_t seed) { + std::vector host(size); + for (std::size_t i = 0; i < size; ++i) { + host[i] = static_cast(seed + (i & 0x3f)); + } + return Runtime::Memcpy(ptr, host.data(), size, + Runtime::kMemcpyHostToDevice) == Runtime::kSuccess; +} + +bool VerifyDevice(void* ptr, std::size_t size, std::uint8_t seed) { + std::vector host(size, 0); + if (Runtime::Memcpy(host.data(), ptr, size, Runtime::kMemcpyDeviceToHost) != + Runtime::kSuccess) { + return false; + } + for (std::size_t i = 0; i < size; ++i) { + if (host[i] != static_cast(seed + (i & 0x3f))) { + return false; + } + } + return true; +} + +// A slice returned by the pool must be real, usable device memory -- not just a +// plausible-looking address computed off a backing base pointer. +void TestSliceIsUsableDeviceMemory(infini::rt::test::TestContext* context) { + Pool pool; + void* ptr = nullptr; + context->ExpectEqual(pool.Allocate(&ptr, 256), Runtime::kSuccess, + "allocate should succeed on a real backend"); + const std::array input{0, 1, 2, 3, 4, 5, 6, 7}; + ExpectUsable(context, ptr, input, "allocation should produce a pointer"); + + const Pool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.upstream_alloc_count, std::size_t{1}, + "one device allocation backs the slice"); + context->Expect(stats.bytes_reserved >= BackendConfig::kInitialCapacity, + "the backing covers the configured initial capacity"); + context->ExpectEqual(pool.Deallocate(ptr), Runtime::kSuccess, + "deallocate should succeed"); +} + +// The arena's reason for existing: many slices come out of one device +// allocation, and every one of them addresses its own disjoint bytes. An +// off-by-one in the split arithmetic shows up here as one slice reading back a +// neighbor's pattern. +void TestSlicesAreDisjointOnDevice(infini::rt::test::TestContext* context) { + Pool pool; + constexpr std::size_t kSlices = 24; + constexpr std::size_t kSize = 8192; + + std::vector blocks; + for (std::size_t i = 0; i < kSlices; ++i) { + void* ptr = nullptr; + if (pool.Allocate(&ptr, kSize) != Runtime::kSuccess || ptr == nullptr) { + context->Expect(false, "slicing allocation should succeed"); + break; + } + if (!FillDevice(ptr, kSize, static_cast(i * 7 + 1))) { + context->Expect(false, "writing a slice should succeed"); + } + blocks.push_back(ptr); + } + + bool disjoint = true; + for (std::size_t i = 0; i < blocks.size(); ++i) { + disjoint = disjoint && + VerifyDevice(blocks[i], kSize, + static_cast(i * 7 + 1)); + } + context->Expect(disjoint, "slices out of one backing must not overlap"); + + const Pool::Stats stats = pool.GetStats(); + // 24 x 8 KB = 192 KB, comfortably inside one 4 MB backing. + context->ExpectEqual(stats.upstream_alloc_count, std::size_t{1}, + "24 slices need only one device allocation"); + context->ExpectEqual(stats.cache_hit_count, kSlices - 1, + "only the first allocation creates a backing"); + + for (void* ptr : blocks) { + pool.Deallocate(ptr); + } +} + +// Coalescing has to work on device memory too: after the slices are released +// the backing must serve one request spanning all of them, and that whole span +// must still round-trip bytes. +void TestCoalescedSpanIsUsable(infini::rt::test::TestContext* context) { + Pool pool; + std::vector blocks; + for (std::size_t i = 0; i < 8; ++i) { + void* ptr = nullptr; + pool.Allocate(&ptr, 64 * 1024); + blocks.push_back(ptr); + } + for (void* ptr : blocks) { + pool.Deallocate(ptr); + } + + // Round down to the slice granularity: a request is rounded *up* before it is + // fitted, and the extent's own size need not be a multiple of it -- the + // backing's head is trimmed by however much the device base pointer was + // misaligned. Asking for the raw extent size would round past it and + // legitimately need a second backing. + const std::size_t span = pool.GetStats().largest_free_chunk / + BackendConfig::kMinSliceAlignment * + BackendConfig::kMinSliceAlignment; + context->Expect(span >= 8 * 64 * 1024, + "the released slices coalesce into one extent"); + + void* whole = nullptr; + context->ExpectEqual(pool.Allocate(&whole, span), Runtime::kSuccess, + "the coalesced extent serves a full-span request"); + context->ExpectEqual(pool.GetStats().upstream_alloc_count, std::size_t{1}, + "serving it needed no new device allocation"); + // Probe both ends of the span: the arithmetic that produced it has to be + // right at the tail, not just at the base. + const std::array input{2, 4, 6, 8, 10, 12, 14, 16}; + ExpectUsable(context, whole, input, "the span's head is usable"); + ExpectUsable(context, static_cast(whole) + span - 8, input, + "the span's tail is usable"); + pool.Deallocate(whole); +} + +// A requested power-of-two alignment must be honored by the slice, which must +// still be usable device memory. +void TestAlignment(infini::rt::test::TestContext* context) { + Pool pool; + constexpr std::size_t kAlignment = 4096; + // Offset the arena first so the natural next slice is misaligned. + void* filler = nullptr; + pool.Allocate(&filler, 512); + + void* ptr = nullptr; + context->ExpectEqual(pool.Allocate(&ptr, 1024, kAlignment), Runtime::kSuccess, + "aligned allocate should succeed"); + context->ExpectEqual(reinterpret_cast(ptr) % kAlignment, + std::uintptr_t{0}, + "returned pointer should honor the alignment"); + const std::array input{1, 1, 2, 3, 5, 8, 13, 21}; + ExpectUsable(context, ptr, input, "aligned slice should be usable"); + pool.Deallocate(ptr); + pool.Deallocate(filler); +} + +// A request larger than the capacity cap gets its own exactly sized device +// allocation, and every byte of it is addressable. +void TestOversizeRequest(infini::rt::test::TestContext* context) { + Pool pool; + constexpr std::size_t kHuge = BackendConfig::kMaxCapacity + (4ull << 20); + void* ptr = nullptr; + if (pool.Allocate(&ptr, kHuge) != Runtime::kSuccess || ptr == nullptr) { + std::cout << kBackend + << " oversize case skipped: device cannot serve the request." + << std::endl; + return; + } + + const Pool::Stats stats = pool.GetStats(); + context->Expect(stats.bytes_reserved >= kHuge, + "the oversize backing covers the request"); + // At least the request, possibly a little more: the tail left over in the + // oversize backing is below `kMinSplitRemainder`, so it stays with the served + // chunk as internal waste rather than becoming an unusable sliver. + context->Expect(stats.bytes_in_use >= kHuge, + "the whole request is live"); + context->Expect( + stats.bytes_in_use - kHuge < BackendConfig::kMinSplitRemainder + + BackendConfig::kMinSliceAlignment, + "an oversize backing wastes at most a sliver on the served chunk"); + const std::array input{9, 9, 8, 8, 7, 7, 6, 6}; + ExpectUsable(context, ptr, input, "the oversize slice's head is usable"); + ExpectUsable(context, static_cast(ptr) + kHuge - 8, input, + "the oversize slice's tail is usable"); + pool.Deallocate(ptr); +} + +// The headline behavior on a real device: a burst reserves extra backings, and +// once the workload returns to small requests the extras go back to the device +// while the resident backing stays warm. +void TestShrinkAfterBurst(infini::rt::test::TestContext* context) { + Pool pool; + void* resident = nullptr; + pool.Allocate(&resident, 1024); + + std::vector burst; + for (std::size_t i = 0; i < 24; ++i) { + void* ptr = nullptr; + if (pool.Allocate(&ptr, 1ull << 20) != Runtime::kSuccess) { + break; + } + burst.push_back(ptr); + } + const std::size_t peak_backings = pool.GetStats().backing_count; + if (peak_backings <= 1) { + std::cout << kBackend + << " shrink case skipped: the burst stayed in one backing." + << std::endl; + pool.Deallocate(resident); + for (void* ptr : burst) { + pool.Deallocate(ptr); + } + return; + } + + for (void* ptr : burst) { + pool.Deallocate(ptr); + } + context->ExpectEqual(pool.GetStats().backing_count, peak_backings, + "freeing alone does not release device memory"); + + std::vector small; + for (std::size_t i = 0; + i < BackendConfig::kShrinkThreshold * + (BackendConfig::kEmptyScansToDestroy + 2); + ++i) { + void* ptr = nullptr; + pool.Allocate(&ptr, 1024); + small.push_back(ptr); + } + + const Pool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.backing_count, std::size_t{1}, + "idle backings are returned to the device"); + context->Expect(stats.shrink_count > 0, "shrink was recorded"); + context->ExpectEqual(stats.upstream_free_count, stats.shrink_count, + "every shrink is one device free"); + + // The resident backing survived, and the block living in it is intact. + const std::array input{3, 1, 4, 1, 5, 9, 2, 6}; + ExpectUsable(context, resident, input, + "the resident slice outlived the shrink"); + pool.Deallocate(resident); + for (void* ptr : small) { + pool.Deallocate(ptr); + } +} + +// `ReleaseCached` hands every drained backing back to the device, and the pool +// keeps working afterwards. +void TestReleaseCached(infini::rt::test::TestContext* context) { + Pool pool; + std::vector blocks; + for (std::size_t i = 0; i < 8; ++i) { + void* ptr = nullptr; + pool.Allocate(&ptr, 128 * 1024); + blocks.push_back(ptr); + } + for (void* ptr : blocks) { + pool.Deallocate(ptr); + } + + pool.ReleaseCached(); + Pool::Stats stats = pool.GetStats(); + context->ExpectEqual(stats.backing_count, std::size_t{0}, + "an empty pool releases every backing"); + context->ExpectEqual(stats.bytes_reserved, std::size_t{0}, + "reserved bytes drop to zero"); + context->ExpectEqual(stats.upstream_free_count, stats.upstream_alloc_count, + "every device allocation was freed exactly once"); + + void* fresh = nullptr; + context->ExpectEqual(pool.Allocate(&fresh, 4096), Runtime::kSuccess, + "the pool still works after a full release"); + const std::array input{7, 7, 7, 7}; + ExpectUsable(context, fresh, input, "a fresh slice is usable"); + pool.Deallocate(fresh); +} + +// The production configuration has to instantiate and work over a real backend, +// not just the reduced one the rest of this file uses. One small allocation is +// enough to prove it: it reserves a default-sized backing and slices it. +void TestDefaultConfigInstantiates(infini::rt::test::TestContext* context) { + infini::rt::ArenaMemoryPool pool; + void* ptr = nullptr; + if (pool.Allocate(&ptr, 4096) != Runtime::kSuccess || ptr == nullptr) { + std::cout << kBackend + << " default-config case skipped: device cannot reserve " + << (infini::rt::DefaultArenaConfig::kInitialCapacity >> 20) + << " MB." << std::endl; + return; + } + const std::array input{1, 2, 3, 4, 5, 6, 7, 8}; + ExpectUsable(context, ptr, input, "a default-config slice is usable"); + context->ExpectEqual(pool.Deallocate(ptr), Runtime::kSuccess, + "deallocate should succeed"); +} + +} // namespace + +int main() { + infini::rt::test::TestContext context; + + if (!SelectDevice()) { + return context.ExitCode(); + } + + TestSliceIsUsableDeviceMemory(&context); + TestSlicesAreDisjointOnDevice(&context); + TestCoalescedSpanIsUsable(&context); + TestAlignment(&context); + TestOversizeRequest(&context); + TestShrinkAfterBurst(&context); + TestReleaseCached(&context); + TestDefaultConfigInstantiates(&context); + + return context.ExitCode(); +} From e616bf52010f20483e49a69a701c8b425b9ed1a5 Mon Sep 17 00:00:00 2001 From: atchen <1920510674@qq.com> Date: Fri, 7 Aug 2026 10:11:50 +0000 Subject: [PATCH 5/8] test: add the allocator comparison harness Neither pool is universally better, so the claim in the previous commit is only worth as much as the measurement behind it. This is that measurement. `perf_allocator_matrix` runs one set of workloads across every allocation strategy the current backend offers -- the backend allocator itself, `MemoryPool`, `ArenaMemoryPool`, and `cudaMallocAsync` where it exists -- through the real dispatch API. Arms share a surface, so each benchmark is written once and instantiated per arm. The workloads are chosen to be the ones a general pool benchmark cannot make: shapes that straddle the size-class boundary, large-block recycling at sizes that are not multiples of the large granularity, growth to a gigabyte-scale high-water mark, trim cost split into traversal and upstream frees, thread counts past 8, a layer-by-layer inference sequence, a multi-threaded one, latency quantiles, and a long random-lifetime fragmentation run with a large-block probe at the end. Two design notes that matter for reading the output. Timings are reported next to the exact upstream call counts that cause them, because on a device the call count *is* the timing and it carries no machine noise. And `cuda_async` is marked stream-ordered everywhere it appears: its release does not wait for pending device work, so it offers a weaker guarantee than the other three arms and its ratios are not a drop-in speedup. `scripts/compare_allocators.py` configures and runs both the host and device builds and pivots the rows into the three comparisons that have distinct answers -- direct vs arena, pool vs arena, vendor pool vs arena. Serialized rather than parallel because `generated/` is written into the source tree at configure time and its contents depend on which backends are enabled. The arena's config scale is part of each row's identity, so the reduced host config and the production device config are never averaged into one series. `arena_vs_pool` measures the same two designs in one process against its own upstream stubs, one of them a calibrated busy-wait standing in for a synchronous `cudaMalloc`. That is the only way to see the arena's amortization on a machine with no device, and it keeps both arms on one CPU at one thermal state. --- .gitignore | 1 + docs/build.md | 39 + scripts/compare_allocators.py | 439 +++++ scripts/run_performance_tests.py | 1 + tests/performance/CMakeLists.txt | 53 + tests/performance/ab/arena_vs_pool.cc | 1105 ++++++++++++ tests/performance/perf_allocator_matrix.cc | 1896 ++++++++++++++++++++ tests/performance/perf_common.h | 32 +- tests/performance/perf_memory_pool.cc | 1252 +++++++++++++ 9 files changed, 4813 insertions(+), 5 deletions(-) create mode 100644 scripts/compare_allocators.py create mode 100644 tests/performance/ab/arena_vs_pool.cc create mode 100644 tests/performance/perf_allocator_matrix.cc create mode 100644 tests/performance/perf_memory_pool.cc 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/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/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