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 ab0211e..2c53232 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) if(CMAKE_SYSTEM_NAME STREQUAL "Linux") @@ -53,6 +65,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) @@ -63,6 +77,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(); +}