diff --git a/include/infini/rt/device_guard.h b/include/infini/rt/device_guard.h new file mode 100644 index 0000000..dd0d247 --- /dev/null +++ b/include/infini/rt/device_guard.h @@ -0,0 +1,80 @@ +#ifndef INFINI_RT_DEVICE_GUARD_H_ +#define INFINI_RT_DEVICE_GUARD_H_ + +#include + +#include +#include + +namespace infini::rt { + +// Select the owning resource's device and restore this thread's prior +// selection. +class DeviceGuard { + public: + DeviceGuard(Device::Type target_type, int target_index) noexcept + : previous_type_(runtime_device_type()) { + if (!enabled(target_type, AllDeviceTypes{})) { + std::fputs("InfiniRT device scope unsupported target device type\n", + stderr); + restored_ = true; + return; + } + const auto previous_status = runtime::GetDevice(&previous_index_); + restore_index_ = previous_status == runtime::kSuccess; + warn(previous_status, "read previous device"); + if (restore_index_ && previous_type_ == target_type && + previous_index_ == target_index) { + active_ = true; + return; + } + set_runtime_device_type(target_type); + const auto status = runtime::SetDevice(target_index); + active_ = status == runtime::kSuccess; + warn(status, "select target device"); + if (!active_) restore(); + } + + DeviceGuard(Device::Type target_type, int target_index, bool required) + : DeviceGuard(target_type, target_index) { + if (required && !active_) + throw std::runtime_error("cannot select runtime device"); + } + + DeviceGuard(const DeviceGuard&) = delete; + DeviceGuard& operator=(const DeviceGuard&) = delete; + DeviceGuard(DeviceGuard&&) = delete; + DeviceGuard& operator=(DeviceGuard&&) = delete; + ~DeviceGuard() noexcept { restore(); } + + bool active() const noexcept { return active_; } + + private: + template + static bool enabled(Device::Type type, List) noexcept { + return ((type == types && DeviceEnabled::value) || ...); + } + + static void warn(runtime::Error status, const char* operation) noexcept { + if (status != runtime::kSuccess) + std::fprintf(stderr, "InfiniRT device scope %s error=%lld\n", operation, + static_cast(status)); + } + void restore() noexcept { + if (restored_) return; + restored_ = true; + set_runtime_device_type(previous_type_); + if (restore_index_) + warn(runtime::SetDevice(previous_index_), "restore previous device"); + } + + Device::Type previous_type_; + int previous_index_ = 0; + bool restore_index_ = false; + bool active_ = false; + bool restored_ = false; +}; + +} // namespace infini::rt + +#endif diff --git a/include/infini/rt/reusable_memory_pool.h b/include/infini/rt/reusable_memory_pool.h new file mode 100644 index 0000000..d5c0d07 --- /dev/null +++ b/include/infini/rt/reusable_memory_pool.h @@ -0,0 +1,61 @@ +#ifndef INFINI_RT_REUSABLE_MEMORY_POOL_H_ +#define INFINI_RT_REUSABLE_MEMORY_POOL_H_ + +#include + +#include +#include + +namespace infini::rt { + +// Block reuse and capture retention, independent of models/requests. The pool +// owns one ordinary execution stream; consumers borrow stream(), never destroy +// it. Raw allocation pointers require a live pool or a committed pin lease. +// Reuse/deallocation is ordered on this stream by the caller. A foreign-thread +// release must first prove completion; uses on other streams must be joined. +class ReusableMemoryPool { + struct State; + struct PinState; + + public: + class PinLease { + public: + ~PinLease() noexcept; + PinLease(const PinLease&) = delete; + PinLease& operator=(const PinLease&) = delete; + + private: + friend class ReusableMemoryPool; + explicit PinLease(std::shared_ptr state) noexcept; + std::shared_ptr state_; + }; + + explicit ReusableMemoryPool(Device device); + ~ReusableMemoryPool() noexcept; + ReusableMemoryPool(const ReusableMemoryPool&) = delete; + ReusableMemoryPool& operator=(const ReusableMemoryPool&) = delete; + + Device device() const; + runtime::Stream stream() const; + // During native CUDA stream capture, only cached blocks may be reused. + // Preallocate before capture; a cache miss throws without invalidating it. + std::byte* allocate(std::size_t size); + std::byte* allocate_exact(std::size_t size); + void deallocate(std::byte* ptr); + // Retain an additional logical owner and return the physical block size. + std::size_t retain_allocation(void* ptr); + + void begin_pin_mode(); + std::shared_ptr commit_pin_mode(); + void cancel_pin_mode(); + // True if the logical allocation predates this pin transaction. + bool retain_for_capture(void* ptr); + std::size_t pinned_bytes(); + void trim(); + + private: + std::shared_ptr state_; +}; + +} // namespace infini::rt +#endif diff --git a/scripts/generate_public_headers.py b/scripts/generate_public_headers.py index 516d4aa..22c49d8 100644 --- a/scripts/generate_public_headers.py +++ b/scripts/generate_public_headers.py @@ -221,6 +221,12 @@ def _write_generated_header(include_root, source_root, devices): runtime_declarations = "\n\n".join( f"{function.signature()};" for function in public_runtime_functions ) + if any(function.name == "GetDeviceCount" for function in public_runtime_functions): + runtime_declarations += ( + "\n\nError GetDeviceCount(Device::Type type, int* count);" + ) + if any(function.name == "EventQuery" for function in public_runtime_functions): + runtime_declarations += "\n\nbool IsEventNotReady(Error status);" path = include_root / "infini" / "rt" / "generated.h" path.parent.mkdir(parents=True, exist_ok=True) @@ -521,6 +527,38 @@ def _write_runtime_dispatch_function(function, devices): """ +def _write_explicit_device_count_query(devices): + cases = "\n".join( + f" case {_DEVICE_TYPES[device]}:\n" + f" return QueryDeviceCount<{_DEVICE_TYPES[device]}>(count);" + for device in devices + ) + return f"""namespace {{ + +template +Error QueryDeviceCount(int* count) {{ + using DeviceRuntime = Runtime; + int queried_count = -1; + const auto status = DeviceRuntime::GetDeviceCount(&queried_count); + if (status != DeviceRuntime::kSuccess) return static_cast(status); + if (queried_count < 0) return InvalidValueError(); + *count = queried_count; + return kSuccess; +}} + +}} // namespace + +Error GetDeviceCount(Device::Type type, int* count) {{ + if (count == nullptr) return InvalidValueError(); + switch (type) {{ +{cases} + default: + return InvalidValueError(); + }} +}} +""" + + def _runtime_header_for_device(source_root, device): for _, header_name, target in _DEVICE_HEADERS[device]: if header_name == "runtime_.h": @@ -549,6 +587,30 @@ def _public_runtime_functions_for_devices(devices, source_root): def _write_runtime_dispatch(source_path, source_root, devices): functions = _public_runtime_functions_for_devices(devices, source_root) + event_not_ready_dispatch = "" + event_query = next( + (function for function in functions if function.name == "EventQuery"), None + ) + if event_query is not None: + event_not_ready_cases = "\n".join( + f" case {_DEVICE_TYPES[device]}:\n" + f" return status == static_cast(Runtime<{_DEVICE_TYPES[device]}>::kEventNotReady);" + for device in _devices_for_function(event_query, devices, source_root) + if re.search( + r"\bkEventNotReady\b", + _runtime_header_for_device(source_root, device).read_text(), + ) + ) + event_not_ready_dispatch = f""" +bool IsEventNotReady(Error status) {{ + if (status == kSuccess) return false; + switch (infini::rt::runtime_device_type()) {{ +{event_not_ready_cases} + default: + return false; + }} +}} +""" stream_capture_mode_helper = ( """ template @@ -582,6 +644,16 @@ def _write_runtime_dispatch(source_path, source_root, devices): ) for function in functions ) + device_count_query = next( + (function for function in functions if function.name == "GetDeviceCount"), None + ) + explicit_device_count_query = ( + _write_explicit_device_count_query( + _devices_for_function(device_count_query, devices, source_root) + ) + if device_count_query is not None + else "" + ) set_device_type_cases = "\n".join( f""" case {_DEVICE_TYPES[device]}: runtime_device_type_ = device_type; @@ -660,6 +732,8 @@ def _write_runtime_dispatch(source_path, source_root, devices): }} // namespace {dispatch_functions} +{explicit_device_count_query} +{event_not_ready_dispatch} }} // namespace infini::rt::runtime """ ) diff --git a/src/data_type.h b/src/data_type.h index 9762508..c653535 100644 --- a/src/data_type.h +++ b/src/data_type.h @@ -93,12 +93,14 @@ struct Float16 { // Inf return {static_cast(sign | 0x7C00)}; } else if (exponent >= -14) { - return {static_cast(sign | ((exponent + 15) << 10) | - (mantissa >> 13))}; - } else if (exponent >= -24) { + mantissa += 0x0FFF + ((mantissa >> 13) & 1); + return {static_cast( + sign | (((exponent + 15) << 10) + (mantissa >> 13)))}; + } else if (exponent >= -25) { mantissa |= 0x800000; - mantissa >>= (-14 - exponent); - return {static_cast(sign | (mantissa >> 13))}; + const unsigned shift = static_cast(-exponent - 1); + const auto bias = ((1u << (shift - 1)) - 1) + ((mantissa >> shift) & 1); + return {static_cast(sign | ((mantissa + bias) >> shift))}; } // Too small for subnormal: return signed zero. return {sign}; diff --git a/src/native/ascend/runtime_.h b/src/native/ascend/runtime_.h index 6503279..4129e53 100644 --- a/src/native/ascend/runtime_.h +++ b/src/native/ascend/runtime_.h @@ -35,12 +35,17 @@ struct Runtime static constexpr Error kSuccess = ACL_SUCCESS; - static constexpr auto SetDevice = aclrtSetDevice; + static Error SetDevice(int device) { + const auto status = EnsureInitialized(); + return status == kSuccess ? aclrtSetDevice(device) : status; + } static constexpr auto GetDevice = aclrtGetDevice; static auto GetDeviceCount(int* count) { assert(count != nullptr); + const auto initialized = EnsureInitialized(); + if (initialized != kSuccess) return initialized; std::uint32_t device_count = 0; auto status = aclrtGetDeviceCount(&device_count); *count = static_cast(device_count); @@ -182,6 +187,16 @@ struct Runtime } private: + static Error EnsureInitialized() { + // Cache one live epoch; external owners must not finalize while RT is live. + // No static destructor finalization or post-finalize reinitialization. + static const Error status = [] { + const auto result = aclInit(nullptr); + return result == ACL_ERROR_REPEAT_INITIALIZE ? kSuccess : result; + }(); + return status; + } + struct RiApi { using CaptureBeginFn = Error (*)(Stream, int); using CaptureGetInfoFn = Error (*)(Stream, int*, void**); diff --git a/src/native/cambricon/data_type_.h b/src/native/cambricon/data_type_.h index b35a936..e6b28eb 100644 --- a/src/native/cambricon/data_type_.h +++ b/src/native/cambricon/data_type_.h @@ -1,8 +1,6 @@ #ifndef INFINI_RT_CAMBRICON_DATA_TYPE__H_ #define INFINI_RT_CAMBRICON_DATA_TYPE__H_ -#include "bang_bf16.h" -#include "bang_fp16.h" #include "data_type.h" #include "native/cambricon/device_.h" @@ -10,12 +8,12 @@ namespace infini::rt { template <> struct TypeMap { - using type = __half; + using type = Float16; }; template <> struct TypeMap { - using type = __bang_bfloat16; + using type = BFloat16; }; } // namespace infini::rt diff --git a/src/native/cuda/hygon/runtime_.h b/src/native/cuda/hygon/runtime_.h index ce4e879..df88c44 100644 --- a/src/native/cuda/hygon/runtime_.h +++ b/src/native/cuda/hygon/runtime_.h @@ -32,6 +32,8 @@ struct Runtime static constexpr Error kSuccess = cudaSuccess; + static constexpr Error kEventNotReady = cudaErrorNotReady; + static constexpr auto SetDevice = cudaSetDevice; static constexpr auto GetDevice = cudaGetDevice; diff --git a/src/native/cuda/iluvatar/runtime_.h b/src/native/cuda/iluvatar/runtime_.h index f864482..db94c3d 100644 --- a/src/native/cuda/iluvatar/runtime_.h +++ b/src/native/cuda/iluvatar/runtime_.h @@ -32,6 +32,8 @@ struct Runtime static constexpr Error kSuccess = cudaSuccess; + static constexpr Error kEventNotReady = cudaErrorNotReady; + static constexpr auto SetDevice = cudaSetDevice; static constexpr auto GetDevice = cudaGetDevice; diff --git a/src/native/cuda/mars/runtime_.h b/src/native/cuda/mars/runtime_.h index 1bb9249..067707f 100644 --- a/src/native/cuda/mars/runtime_.h +++ b/src/native/cuda/mars/runtime_.h @@ -30,6 +30,8 @@ struct Runtime static constexpr Error kSuccess = hcSuccess; + static constexpr Error kEventNotReady = hcErrorNotReady; + static constexpr auto SetDevice = hcSetDevice; static constexpr auto GetDevice = hcGetDevice; diff --git a/src/native/cuda/metax/runtime_.h b/src/native/cuda/metax/runtime_.h index 7796157..994c517 100644 --- a/src/native/cuda/metax/runtime_.h +++ b/src/native/cuda/metax/runtime_.h @@ -30,6 +30,8 @@ struct Runtime static constexpr Error kSuccess = mcSuccess; + static constexpr Error kEventNotReady = mcErrorNotReady; + static constexpr auto SetDevice = mcSetDevice; static constexpr auto GetDevice = mcGetDevice; diff --git a/src/native/cuda/moore/runtime_.h b/src/native/cuda/moore/runtime_.h index 3100285..2c7340d 100644 --- a/src/native/cuda/moore/runtime_.h +++ b/src/native/cuda/moore/runtime_.h @@ -31,6 +31,8 @@ struct Runtime static constexpr Error kSuccess = musaSuccess; + static constexpr Error kEventNotReady = musaErrorNotReady; + static constexpr auto SetDevice = musaSetDevice; static constexpr auto GetDevice = [](auto&&... args) { diff --git a/src/native/cuda/nvidia/runtime_.h b/src/native/cuda/nvidia/runtime_.h index 99e87f6..d369d46 100644 --- a/src/native/cuda/nvidia/runtime_.h +++ b/src/native/cuda/nvidia/runtime_.h @@ -32,6 +32,8 @@ struct Runtime static constexpr Error kSuccess = cudaSuccess; + static constexpr Error kEventNotReady = cudaErrorNotReady; + static constexpr auto SetDevice = cudaSetDevice; static constexpr auto GetDevice = cudaGetDevice; diff --git a/src/native/cuda/thead/runtime_.h b/src/native/cuda/thead/runtime_.h index d96be70..34decb0 100644 --- a/src/native/cuda/thead/runtime_.h +++ b/src/native/cuda/thead/runtime_.h @@ -33,6 +33,8 @@ struct Runtime static constexpr Error kSuccess = cudaSuccess; + static constexpr Error kEventNotReady = cudaErrorNotReady; + static constexpr auto SetDevice = cudaSetDevice; static constexpr auto GetDevice = cudaGetDevice; diff --git a/src/reusable_memory_pool.cc b/src/reusable_memory_pool.cc new file mode 100644 index 0000000..a67a121 --- /dev/null +++ b/src/reusable_memory_pool.cc @@ -0,0 +1,414 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace infini::rt { +namespace { +void check(runtime::Error status, const char* operation) { + if (status != runtime::kSuccess) + throw std::runtime_error(std::string(operation) + + " failed with error code " + + std::to_string(static_cast(status))); +} +void warn(runtime::Error status, const char* operation) noexcept { + if (status != runtime::kSuccess) + std::fprintf(stderr, "ReusableMemoryPool cleanup %s error=%lld\n", + operation, static_cast(status)); +} +std::size_t align_up(std::size_t size, std::size_t alignment) { + if (size > std::numeric_limits::max() - (alignment - 1)) + throw std::length_error("allocator request exceeds addressable size"); + return (size + alignment - 1) / alignment * alignment; +} +} // namespace + +struct ReusableMemoryPool::State { + struct Block { + enum class Kind { SizeClass, Large, Exact }; + void* ptr = nullptr; + std::size_t size = 0; + Kind kind = Kind::Large; + std::size_t class_index = 0; + std::size_t pin_count = 0; + bool in_use = false; + std::size_t use_count = 0; + uint64_t allocation_capture_generation = 0; + }; + struct SizeClass { + std::size_t block_size; + std::vector> free_blocks; + }; + Device device_; + runtime::Stream stream_ = nullptr; + runtime::Stream allocation_stream_ = nullptr; + bool pinned_mode_ = false; + std::thread::id pin_owner_; + uint64_t capture_generation_ = 0; + std::vector size_classes_; + std::vector> large_blocks_; + std::unordered_map> all_blocks_; + std::vector> capture_frozen_blocks_; + std::unordered_set capture_frozen_block_set_; + std::mutex mutex_; + + explicit State(Device device) : device_(device) { + size_classes_ = { + {32 * 1024, {}}, {256 * 1024, {}}, {1 * 1024 * 1024, {}}, + {2 * 1024 * 1024, {}}, {4 * 1024 * 1024, {}}, {8 * 1024 * 1024, {}}, + {16 * 1024 * 1024, {}}, {32 * 1024 * 1024, {}}, {64 * 1024 * 1024, {}}, + {128 * 1024 * 1024, {}}, {256 * 1024 * 1024, {}}}; + if (device_.type() == Device::Type::kCount) + throw std::invalid_argument("invalid pool device type"); + DeviceGuard guard{device_.type(), device_.index(), true}; + const auto status = runtime::StreamCreate(&stream_); + if (status != runtime::kSuccess && stream_ != nullptr) { + warn(runtime::StreamDestroy(stream_), "partial stream destroy"); + stream_ = nullptr; + } + check(status, "create pool stream"); + allocation_stream_ = + device_.type() == Device::Type::kNvidia ? stream_ : nullptr; + } + + runtime::Error allocate_block_(void** ptr, std::size_t size) { +#if defined(WITH_NVIDIA) + if (allocation_stream_ != nullptr) { + cudaStreamCaptureStatus status; + check(static_cast(cudaStreamIsCapturing( + reinterpret_cast(allocation_stream_), &status)), + "query pool stream capture"); + // A captured cudaMallocAsync belongs to the graph, not this pool. + if (status != cudaStreamCaptureStatusNone) + throw std::runtime_error( + "pool cannot grow during CUDA stream capture; preallocate blocks " + "before capture"); + } +#endif + return allocation_stream_ != nullptr + ? runtime::MallocAsync(ptr, size, allocation_stream_) + : runtime::Malloc(ptr, size); + } + runtime::Error free_block_(void* ptr) { + return allocation_stream_ != nullptr + ? runtime::FreeAsync(ptr, allocation_stream_) + : runtime::Free(ptr); + } + void freeze_for_capture_(const std::shared_ptr& block) { + if (pinned_mode_ && pin_owner_ == std::this_thread::get_id() && + !capture_frozen_block_set_.count(block.get())) { + capture_frozen_blocks_.push_back(block); + try { + capture_frozen_block_set_.insert(block.get()); + } catch (...) { + capture_frozen_blocks_.pop_back(); + throw; + } + ++block->pin_count; + } + } + void release_frozen_blocks_( + const std::vector>& blocks) noexcept { + try { + std::lock_guard lock(mutex_); + for (const auto& block : blocks) + if (block->pin_count > 0) --block->pin_count; + } catch (...) { + } + } + std::byte* allocate_impl_(std::size_t size, bool exact) { + if (size == 0) { + if (exact) + throw std::invalid_argument("exact allocator request must be nonzero"); + return nullptr; + } + std::lock_guard lock(mutex_); + size = align_up(size, 256); + const uint64_t generation = + pinned_mode_ && pin_owner_ == std::this_thread::get_id() + ? capture_generation_ + : 0; + std::shared_ptr block; + for (std::size_t class_index = 0; + !exact && class_index < size_classes_.size(); ++class_index) { + auto& cls = size_classes_[class_index]; + if (size <= cls.block_size) { + const auto free_block = std::find_if( + cls.free_blocks.begin(), cls.free_blocks.end(), + [](const auto& candidate) { return !candidate->in_use; }); + if (free_block != cls.free_blocks.end()) { + block = *free_block; + freeze_for_capture_(block); + block->allocation_capture_generation = generation; + cls.free_blocks.erase(free_block); + block->in_use = true; + block->use_count = 1; + return reinterpret_cast(block->ptr); + } + block = std::make_shared(); + block->size = cls.block_size; + block->kind = Block::Kind::SizeClass; + block->class_index = class_index; + block->in_use = true; + block->use_count = 1; + try { + check(allocate_block_(&block->ptr, block->size), + "allocate size-class block"); + all_blocks_.emplace(block->ptr, block); + freeze_for_capture_(block); + block->allocation_capture_generation = generation; + } catch (...) { + if (block->ptr) { + all_blocks_.erase(block->ptr); + (void)free_block_(block->ptr); + } + throw; + } + return reinterpret_cast(block->ptr); + } + } + // Pinned free blocks remain reusable after their logical owners release + // them: pinning preserves the address, not exclusive graph ownership. + auto it = std::find_if( + large_blocks_.begin(), large_blocks_.end(), + [size, exact](const std::shared_ptr& candidate) { + return !candidate->in_use && + candidate->kind == + (exact ? Block::Kind::Exact : Block::Kind::Large) && + (exact ? candidate->size == size : candidate->size >= size); + }); + if (it != large_blocks_.end()) { + block = *it; + freeze_for_capture_(block); + block->allocation_capture_generation = generation; + block->in_use = true; + block->use_count = 1; + return reinterpret_cast(block->ptr); + } + block = std::make_shared(); + block->size = size; + block->kind = exact ? Block::Kind::Exact : Block::Kind::Large; + block->in_use = true; + block->use_count = 1; + try { + check(allocate_block_(&block->ptr, block->size), + "allocate large/exact block"); + all_blocks_.emplace(block->ptr, block); + large_blocks_.push_back(block); + freeze_for_capture_(block); + block->allocation_capture_generation = generation; + } catch (...) { + if (block->ptr) { + large_blocks_.erase( + std::remove(large_blocks_.begin(), large_blocks_.end(), block), + large_blocks_.end()); + all_blocks_.erase(block->ptr); + (void)free_block_(block->ptr); + } + throw; + } + return reinterpret_cast(block->ptr); + } + void deallocate(std::byte* ptr) { + if (!ptr) return; + std::lock_guard lock(mutex_); + auto it = all_blocks_.find(reinterpret_cast(ptr)); + if (it == all_blocks_.end()) + throw std::runtime_error("Pointer not allocated by this allocator"); + auto block = it->second; + if (!block->in_use || block->use_count == 0) + throw std::runtime_error("Double free detected in ReusableMemoryPool"); + if (block->use_count > 1) { + --block->use_count; + return; + } + if (block->kind == Block::Kind::SizeClass) + size_classes_[block->class_index].free_blocks.push_back(block); + block->use_count = 0; + block->in_use = false; + } + std::size_t retain_allocation(void* ptr) { + std::lock_guard lock(mutex_); + auto it = all_blocks_.find(ptr); + if (it == all_blocks_.end()) + throw std::runtime_error("Pointer not allocated by this allocator"); + auto block = it->second; + if (block->kind == Block::Kind::SizeClass) { + auto& free_blocks = size_classes_[block->class_index].free_blocks; + free_blocks.erase( + std::remove(free_blocks.begin(), free_blocks.end(), block), + free_blocks.end()); + } + block->in_use = true; + ++block->use_count; + return block->size; + } + ~State() noexcept { + // The final shared owner, not the user-facing pool object, closes the + // queue. Never free backing storage if completion cannot be established. + try { + DeviceGuard guard{device_.type(), device_.index(), true}; + if (stream_) { + const auto status = runtime::StreamSynchronize(stream_); + if (status != runtime::kSuccess) { + warn(status, "drain before release; resources retained"); + return; + } + } + std::lock_guard lock(mutex_); + for (auto& entry : all_blocks_) + if (entry.second->ptr) + warn(free_block_(entry.second->ptr), "free block"); + if (stream_) { + const auto status = runtime::StreamSynchronize(stream_); + if (status != runtime::kSuccess) { + warn(status, "complete frees; stream retained"); + return; + } + warn(runtime::StreamDestroy(stream_), "destroy stream"); + stream_ = nullptr; + } + } catch (const std::exception& error) { + std::fprintf(stderr, + "ReusableMemoryPool cleanup resources retained: %s\n", + error.what()); + } catch (...) { + std::fprintf( + stderr, + "ReusableMemoryPool cleanup resources retained: unknown error\n"); + } + } +}; + +struct ReusableMemoryPool::PinState { + std::shared_ptr owner; + std::vector> blocks; + PinState(std::shared_ptr state, + std::vector> retained) noexcept + : owner(std::move(state)), blocks(std::move(retained)) {} + ~PinState() noexcept { owner->release_frozen_blocks_(blocks); } +}; +ReusableMemoryPool::PinLease::PinLease(std::shared_ptr state) noexcept + : state_(std::move(state)) {} +ReusableMemoryPool::PinLease::~PinLease() noexcept = default; +ReusableMemoryPool::ReusableMemoryPool(Device device) + : state_(std::make_shared(device)) {} +ReusableMemoryPool::~ReusableMemoryPool() noexcept = default; +Device ReusableMemoryPool::device() const { return state_->device_; } +runtime::Stream ReusableMemoryPool::stream() const { return state_->stream_; } +std::byte* ReusableMemoryPool::allocate(std::size_t size) { + DeviceGuard guard{state_->device_.type(), state_->device_.index(), true}; + return state_->allocate_impl_(size, false); +} +std::byte* ReusableMemoryPool::allocate_exact(std::size_t size) { + DeviceGuard guard{state_->device_.type(), state_->device_.index(), true}; + return state_->allocate_impl_(size, true); +} +void ReusableMemoryPool::deallocate(std::byte* ptr) { state_->deallocate(ptr); } +std::size_t ReusableMemoryPool::retain_allocation(void* ptr) { + return state_->retain_allocation(ptr); +} +void ReusableMemoryPool::begin_pin_mode() { + std::lock_guard lock(state_->mutex_); + if (state_->pinned_mode_) + throw std::runtime_error("allocator graph capture is already active"); + if (!state_->capture_frozen_blocks_.empty() || + !state_->capture_frozen_block_set_.empty()) + throw std::runtime_error("allocator pin transaction contains stale blocks"); + if (state_->capture_generation_ == std::numeric_limits::max()) + throw std::overflow_error("allocator capture generation exhausted"); + ++state_->capture_generation_; + state_->pinned_mode_ = true; + state_->pin_owner_ = std::this_thread::get_id(); +} +std::shared_ptr +ReusableMemoryPool::commit_pin_mode() { + std::vector> blocks; + { + std::lock_guard lock(state_->mutex_); + if (!state_->pinned_mode_) return nullptr; + if (state_->pin_owner_ != std::this_thread::get_id()) + throw std::runtime_error( + "cannot commit allocator pin mode: another thread owns the graph " + "capture"); + state_->pinned_mode_ = false; + state_->pin_owner_ = {}; + blocks = std::move(state_->capture_frozen_blocks_); + state_->capture_frozen_block_set_.clear(); + } + try { + auto retained = std::make_shared(state_, std::move(blocks)); + return std::shared_ptr(new PinLease(std::move(retained))); + } catch (...) { + state_->release_frozen_blocks_(blocks); + throw; + } +} +void ReusableMemoryPool::cancel_pin_mode() { + std::lock_guard lock(state_->mutex_); + if (state_->pinned_mode_ && state_->pin_owner_ != std::this_thread::get_id()) + throw std::runtime_error( + "cannot cancel allocator pin mode: another thread owns the graph " + "capture"); + for (const auto& block : state_->capture_frozen_blocks_) + if (block->pin_count > 0) --block->pin_count; + state_->capture_frozen_blocks_.clear(); + state_->capture_frozen_block_set_.clear(); + state_->pinned_mode_ = false; + state_->pin_owner_ = {}; +} +bool ReusableMemoryPool::retain_for_capture(void* ptr) { + std::lock_guard lock(state_->mutex_); + if (!state_->pinned_mode_ || state_->pin_owner_ != std::this_thread::get_id()) + throw std::runtime_error( + "allocator graph capture is not active on this thread"); + auto it = state_->all_blocks_.find(ptr); + if (it == state_->all_blocks_.end()) + throw std::runtime_error("capture requires pool-owned memory"); + state_->freeze_for_capture_(it->second); + return it->second->allocation_capture_generation != + state_->capture_generation_; +} +std::size_t ReusableMemoryPool::pinned_bytes() { + std::lock_guard lock(state_->mutex_); + std::size_t bytes = 0; + for (const auto& entry : state_->all_blocks_) + if (entry.second->pin_count > 0) bytes += entry.second->size; + return bytes; +} +void ReusableMemoryPool::trim() { + DeviceGuard guard{state_->device_.type(), state_->device_.index(), true}; + std::lock_guard lock(state_->mutex_); + for (auto& cls : state_->size_classes_) { + for (auto it = cls.free_blocks.begin(); it != cls.free_blocks.end();) { + if ((*it)->pin_count == 0) { + check(state_->free_block_((*it)->ptr), "trim size-class block"); + state_->all_blocks_.erase((*it)->ptr); + it = cls.free_blocks.erase(it); + } else + ++it; + } + } + for (auto it = state_->large_blocks_.begin(); + it != state_->large_blocks_.end();) { + if ((*it)->pin_count == 0 && !(*it)->in_use) { + check(state_->free_block_((*it)->ptr), "trim large/exact block"); + state_->all_blocks_.erase((*it)->ptr); + it = state_->large_blocks_.erase(it); + } else + ++it; + } +} + +} // namespace infini::rt diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index bf759d2..6bb2dba 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -41,6 +41,7 @@ endfunction() add_infini_rt_test(test_smoke test_smoke.cc) add_infini_rt_test(test_core test_core.cc) +add_infini_rt_test(test_device_guard test_device_guard.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) @@ -59,6 +60,8 @@ if(WITH_CPU) endif() if(WITH_NVIDIA) + add_infini_rt_test(test_reusable_memory_pool_capture + test_reusable_memory_pool_capture.cc) set(INFINI_RT_TEST_HAS_RUNTIME_BACKEND ON) add_infini_rt_backend_runtime_test( NVIDIA infini::rt::Device::Type::kNvidia diff --git a/tests/test_device_guard.cc b/tests/test_device_guard.cc new file mode 100644 index 0000000..76cdb13 --- /dev/null +++ b/tests/test_device_guard.cc @@ -0,0 +1,112 @@ +#include +#include + +#include + +#include "test_helper.h" + +namespace { +using namespace infini::rt; +struct TestUnwind {}; + +void CheckRejected(test::TestContext& context, Device::Type type) { + const auto previous = runtime_device_type(); + { + DeviceGuard guard(type, 0); + context.Expect(!guard.active(), "Unavailable device must not be active."); + context.Expect(runtime_device_type() == previous, + "Failed selection must preserve the current backend."); + } + bool threw = false; + try { + DeviceGuard guard(type, 0, true); + } catch (const std::runtime_error&) { + threw = true; + } + context.Expect(threw, "Required unavailable device must throw."); + threw = false; + try { + ReusableMemoryPool pool(Device{type, 0}); + } catch (const std::exception&) { + threw = true; + } + context.Expect(threw, "Pool must reject an unavailable device."); + context.Expect(runtime_device_type() == previous, + "Rejected pool construction must preserve the backend."); +} + +template +void CheckBackend(test::TestContext& context) { + if constexpr (!DeviceEnabled::value) { + CheckRejected(context, type); + } else { + int count = 0; + if (runtime::GetDeviceCount(type, &count) != runtime::kSuccess || + count == 0) + return; + // Exercise both the no-switch path and an index switch, where available. + DeviceGuard outer(type, 0, true); + { + DeviceGuard inner(type, count - 1, true); + int selected = -1; + context.Expect(runtime::GetDevice(&selected) == runtime::kSuccess && + selected == count - 1, + "Guard must select the requested device."); + } + int restored = -1; + context.Expect(runtime_device_type() == type && + runtime::GetDevice(&restored) == runtime::kSuccess && + restored == 0, + "Nested guard must restore the prior selection."); + + for (bool unwind : {false, true}) { + if (count > 1) { + set_runtime_device_type(type); + context.Expect(runtime::SetDevice(0) == runtime::kSuccess, + "Failed to select the initial device."); + try { + DeviceGuard same_device(type, 0, true); + context.Expect(runtime::SetDevice(count - 1) == runtime::kSuccess, + "Failed to change device inside the guard."); + if (unwind) throw TestUnwind{}; + } catch (const TestUnwind&) { + } + context.Expect( + runtime::GetDevice(&restored) == runtime::kSuccess && restored == 0, + "Same-device guard must restore an internal index change."); + } + if constexpr (type != Device::Type::kCpu && + DeviceEnabled::value) { + set_runtime_device_type(type); + context.Expect(runtime::SetDevice(0) == runtime::kSuccess, + "Failed to select the initial backend."); + try { + DeviceGuard same_device(type, 0, true); + set_runtime_device_type(Device::Type::kCpu); + if (unwind) throw TestUnwind{}; + } catch (const TestUnwind&) { + } + context.Expect( + runtime_device_type() == type && + runtime::GetDevice(&restored) == runtime::kSuccess && + restored == 0, + "Same-device guard must restore an internal backend change."); + } + } + } +} + +template +void CheckBackends(test::TestContext& context, List) { + (CheckBackend(context), ...); +} +} // namespace + +int main() { + test::TestContext context; + CheckRejected(context, Device::Type::kCount); + CheckRejected(context, static_cast(-1)); + CheckRejected(context, static_cast(999)); + CheckBackends(context, AllDeviceTypes{}); + return context.ExitCode(); +} diff --git a/tests/test_native_runtime.cc b/tests/test_native_runtime.cc index 3ae7ea9..ecc5a87 100644 --- a/tests/test_native_runtime.cc +++ b/tests/test_native_runtime.cc @@ -375,6 +375,21 @@ void TestEvent(infini::rt::test::TestContext* context) { } } +template +void TestEventStatus(infini::rt::test::TestContext* context) { + infini::rt::set_runtime_device_type(INFINI_RT_TEST_DEVICE_TYPE); + context->Expect( + !infini::rt::runtime::IsEventNotReady(infini::rt::runtime::kSuccess), + "Success must not be classified as an incomplete event."); + if constexpr (kSupportsEvent && + INFINI_RT_TEST_DEVICE_TYPE != infini::rt::Device::Type::kCpu) { + context->Expect(infini::rt::runtime::IsEventNotReady( + static_cast( + DeviceRuntime::kEventNotReady)), + "Dispatch must recognize this backend's not-ready status."); + } +} + void TestStreamWaitEvent(infini::rt::test::TestContext* context) { typename Runtime::Stream stream{}; if (!CreateStream(context, &stream)) { @@ -478,6 +493,7 @@ int main() { TestMemsetAsync(&context); TestStream(&context); TestEvent(&context); + TestEventStatus(&context); TestStreamWaitEvent(&context); TestEventElapsedTime(&context); diff --git a/tests/test_reusable_memory_pool_capture.cc b/tests/test_reusable_memory_pool_capture.cc new file mode 100644 index 0000000..fee74be --- /dev/null +++ b/tests/test_reusable_memory_pool_capture.cc @@ -0,0 +1,99 @@ +#include +#include + +#include +#include +#include + +#include "test_helper.h" + +namespace { +using namespace infini::rt; + +void Check(runtime::Error status) { + if (status != runtime::kSuccess) + throw std::runtime_error("runtime error " + std::to_string(status)); +} + +void TestCapture(test::TestContext& context) { + DeviceGuard guard(Device::Type::kNvidia, 0, true); + ReusableMemoryPool pool(Device{Device::Type::kNvidia, 0}); + auto* cached = pool.allocate(32); + auto* exact = pool.allocate_exact(32); + pool.deallocate(cached); + pool.deallocate(exact); + Check(runtime::StreamSynchronize(pool.stream())); + + pool.begin_pin_mode(); + Check(runtime::StreamBeginCapture( + pool.stream(), runtime::StreamCaptureMode::kStreamCaptureModeRelaxed)); + auto* reused = pool.allocate(32); + auto* reused_exact = pool.allocate_exact(32); + context.Expect(reused == cached && reused_exact == exact, + "Capture must still reuse cached blocks."); + for (bool exact_request : {false, true}) { + bool rejected = false; + try { + auto* unexpected = + exact_request ? pool.allocate_exact(1024) : pool.allocate(64 * 1024); + pool.deallocate(unexpected); + } catch (const std::runtime_error&) { + rejected = true; + } + context.Expect(rejected, "Capture must reject a pool cache miss."); + } + // A rejected allocation must leave capture usable, not invalidated. + Check(runtime::MemsetAsync(reused, 42, 32, pool.stream())); + Check(runtime::MemsetAsync(reused_exact, 17, 32, pool.stream())); + pool.deallocate(reused); + pool.deallocate(reused_exact); + runtime::Graph graph = nullptr; + Check(runtime::StreamEndCapture(pool.stream(), &graph)); + auto lease = pool.commit_pin_mode(); + runtime::GraphExec exec = nullptr; + Check(runtime::GraphInstantiate(&exec, graph)); + for (int replay = 0; replay < 3; ++replay) { + Check(runtime::MemsetAsync(cached, 0, 32, pool.stream())); + Check(runtime::MemsetAsync(exact, 0, 32, pool.stream())); + Check(runtime::GraphLaunch(exec, pool.stream())); + Check(runtime::StreamSynchronize(pool.stream())); + std::array result{}; + Check(runtime::Memcpy(result.data(), cached, result.size(), + runtime::kMemcpyDeviceToHost)); + for (auto value : result) + context.Expect(value == 42, "Replayed size-class buffer must be valid."); + Check(runtime::Memcpy(result.data(), exact, result.size(), + runtime::kMemcpyDeviceToHost)); + for (auto value : result) + context.Expect(value == 17, "Replayed exact buffer must be valid."); + } + Check(runtime::GraphExecDestroy(exec)); + Check(runtime::GraphDestroy(graph)); + lease.reset(); + pool.trim(); + + // Pin mode alone is LM recording, not native CUDA stream capture. + pool.begin_pin_mode(); + auto* outside_capture = pool.allocate(64 * 1024); + pool.deallocate(outside_capture); + pool.cancel_pin_mode(); + pool.trim(); +} +} // namespace + +int main() { + test::TestContext context; + int count = 0; + if (runtime::GetDeviceCount(Device::Type::kNvidia, &count) != + runtime::kSuccess || + count == 0) { + std::cout << "Skipping pool capture test: no NVIDIA device.\n"; + return 0; + } + try { + TestCapture(context); + } catch (const std::exception& error) { + context.Expect(false, error.what()); + } + return context.ExitCode(); +}