-
Notifications
You must be signed in to change notification settings - Fork 91
Implement CUDA event pool to minimize runtime resource allocation overhead (libcurl multi poll-based backend 2/n) #919
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 23 commits
08dd7f9
3948400
43a38d9
bf20744
7e0c71f
74e6740
812c8a9
c9d7acd
d878f79
df8a4a8
9d73727
b71c857
dab118e
3103863
4208e5c
badc794
075dbd8
64ee425
87e8fe2
e741a51
c323302
a97196b
4f820a3
d1609a0
ad87ea0
5d22b95
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,188 @@ | ||
| /* | ||
| * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
| #pragma once | ||
|
|
||
| #include <mutex> | ||
| #include <unordered_map> | ||
| #include <vector> | ||
|
|
||
| #include <kvikio/shim/cuda.hpp> | ||
|
|
||
| namespace kvikio::detail { | ||
| /** | ||
| * @brief Thread-safe singleton pool for reusable CUDA events | ||
| * | ||
| * Manages a pool of CUDA events organized by CUDA context. Events are retained and reused across | ||
| * calls to minimize allocation overhead. Each context maintains its own separate pool of events | ||
| * since CUDA events are context-specific resources. | ||
| * | ||
| * All events are created with `CU_EVENT_DISABLE_TIMING` for minimal overhead. | ||
| * | ||
| * Call `EventPool::instance().get()` to acquire an event that is bound to the CUDA context | ||
| * currently set on the calling thread. The event will be automatically returned to the pool when it | ||
| * goes out of scope (RAII). | ||
| * | ||
| * @note The destructor intentionally does NOT call `cuEventDestroy` on cached events. | ||
| * `EventPool::instance()` is a function-local static destructed after `main` returns, and | ||
| * making CUDA driver API calls from a static object's destructor at that point is undefined | ||
| * behavior per the CUDA programming guide: | ||
| * https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/intro-to-cuda-cpp.html#runtime-initialization | ||
| * The OS reclaims process memory at exit regardless. | ||
| */ | ||
| class EventPool { | ||
| public: | ||
| /** | ||
| * @brief RAII wrapper for a pooled CUDA event | ||
| * | ||
| * Automatically returns the event to the pool when destroyed. Provides access to the underlying | ||
| * CUevent handle and common event operations (record, synchronize). | ||
| * | ||
| * @note Non-copyable but movable to allow transfer of ownership while maintaining RAII | ||
| */ | ||
| class Event { | ||
| friend class EventPool; | ||
|
|
||
| private: | ||
| EventPool* _pool{}; | ||
| CUevent _event{}; | ||
| CUcontext _cuda_context{}; | ||
|
|
||
| /** | ||
| * @brief Construct an Event wrapping a CUDA event handle | ||
| * | ||
| * @param pool The owning EventPool to return this event to on destruction | ||
| * @param event The CUDA event handle to wrap | ||
| * @param context The CUDA context associated with this event | ||
| */ | ||
| explicit Event(EventPool* pool, CUevent event, CUcontext context) noexcept; | ||
|
|
||
| public: | ||
| ~Event() noexcept; | ||
|
|
||
| // Move-only | ||
| Event(Event const&) = delete; | ||
| Event& operator=(Event const&) = delete; | ||
| Event(Event&& o) noexcept; | ||
| Event& operator=(Event&& o) noexcept; | ||
|
|
||
| /** | ||
| * @brief Get the underlying CUDA event handle | ||
| * | ||
| * @return The CUevent handle wrapped by this object | ||
| */ | ||
| [[nodiscard]] CUevent get() const noexcept; | ||
|
|
||
| /** | ||
| * @brief Get the CUDA context associated with this event | ||
| * | ||
| * @return The CUcontext this event belongs to. Returns nullptr for a moved-from Event. | ||
| */ | ||
| [[nodiscard]] CUcontext cuda_context() const noexcept; | ||
|
|
||
| /** | ||
| * @brief Record the event on a CUDA stream | ||
| * | ||
| * Records the event to capture the current state of the stream. The event will be signaled when | ||
| * all preceding operations on the stream have completed. | ||
| * | ||
| * @param stream The CUDA stream to record the event on. Must belong to the same context as this | ||
| * event. Otherwise CUDA returns an error. | ||
| * | ||
| * @exception kvikio::CUfileException if the record operation fails | ||
| */ | ||
| void record(CUstream stream); | ||
|
|
||
| /** | ||
| * @brief Block the calling thread until the event has been signaled | ||
| * | ||
| * Waits for all work captured by a preceding record() call to complete. | ||
| * | ||
| * @exception kvikio::CUfileException if the synchronize operation fails | ||
| */ | ||
| void synchronize(); | ||
|
|
||
| /** | ||
| * @brief Non-blocking query whether the event has been signaled | ||
| * | ||
| * Returns true if all work captured by a preceding record() call has completed, false if work | ||
| * is still pending. This is the non-blocking counterpart to synchronize(). | ||
| * | ||
| * @return true if the event has been signaled, false if work is still in progress | ||
| * | ||
| * @exception kvikio::CUfileException if the underlying cuEventQuery returns an error other than | ||
| * CUDA_SUCCESS or CUDA_ERROR_NOT_READY | ||
| */ | ||
| [[nodiscard]] bool query() const; | ||
| }; | ||
|
|
||
| private: | ||
| std::mutex mutable _mutex; | ||
| // Per-context pools of free events | ||
| std::unordered_map<CUcontext, std::vector<CUevent>> _pools; | ||
|
|
||
| EventPool() = default; | ||
|
|
||
| // Intentionally `noexcept = default`. See the class-level @note above: issuing CUDA driver | ||
| // API calls (e.g., cuEventDestroy) from this destructor would be UB because the singleton is | ||
| // destructed after main returns. The defaulted destructor runs ~_pools, which tears down the | ||
| // std::vector<CUevent> entries without touching the handles. | ||
| ~EventPool() noexcept = default; | ||
|
|
||
| /** | ||
| * @brief Return an event to the pool for reuse | ||
| * | ||
| * Called by Event's destructor (and move-assignment operator) via the friend declaration. Adds | ||
| * the event to the pool associated with its context for future reuse. | ||
| * | ||
| * @param event The CUDA event handle to return | ||
| * @param context The CUDA context associated with the event | ||
| * | ||
| * @note noexcept: any failure inside push_back (e.g., allocator failure) is caught and logged. | ||
| * The event is then destroyed instead of being cached. | ||
| */ | ||
| void put(CUevent event, CUcontext context) noexcept; | ||
|
|
||
| public: | ||
| // Non-copyable, non-movable singleton | ||
| EventPool(EventPool const&) = delete; | ||
| EventPool& operator=(EventPool const&) = delete; | ||
| EventPool(EventPool&&) = delete; | ||
| EventPool& operator=(EventPool&&) = delete; | ||
|
|
||
| /** | ||
| * @brief Acquire a CUDA event for the CUDA context currently set on the calling thread. | ||
| * | ||
| * Returns a cached event for the current CUDA context if available, otherwise creates a new one. | ||
| * The returned Event object will automatically return the event to the pool when it goes out of | ||
| * scope. | ||
| * | ||
| * @return RAII Event object wrapping the acquired CUDA event | ||
| * @exception kvikio::CUfileException if no CUDA context is current or event creation fails | ||
| */ | ||
| [[nodiscard]] Event get(); | ||
|
|
||
| /** | ||
| * @brief Get the number of free events for a specific context | ||
| * | ||
| * @param context The CUDA context to query | ||
| * @return The number of events available for reuse in that context's pool | ||
| */ | ||
| [[nodiscard]] std::size_t num_free_events(CUcontext context) const; | ||
|
|
||
| /** | ||
| * @brief Get the total number of free events across all contexts | ||
| * | ||
| * @return The total count of events available for reuse | ||
| */ | ||
| [[nodiscard]] std::size_t total_free_events() const; | ||
|
|
||
| /** | ||
| * @brief Get the singleton instance of the event pool | ||
| * | ||
| * @return Reference to the singleton EventPool instance | ||
| */ | ||
| static EventPool& instance(); | ||
| }; | ||
| } // namespace kvikio::detail | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,144 @@ | ||
| /* | ||
| * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| #include <exception> | ||
| #include <utility> | ||
|
|
||
| #include <kvikio/detail/event.hpp> | ||
| #include <kvikio/detail/nvtx.hpp> | ||
| #include <kvikio/error.hpp> | ||
| #include <kvikio/logger.hpp> | ||
| #include <kvikio/shim/cuda.hpp> | ||
|
|
||
| namespace kvikio::detail { | ||
|
|
||
| EventPool::Event::Event(EventPool* pool, CUevent event, CUcontext cuda_context) noexcept | ||
| : _pool(pool), _event(event), _cuda_context(cuda_context) | ||
| { | ||
| } | ||
|
|
||
| EventPool::Event::~Event() noexcept | ||
| { | ||
| if (_event != nullptr) { _pool->put(_event, _cuda_context); } | ||
| } | ||
|
|
||
| EventPool::Event::Event(Event&& o) noexcept | ||
| : _pool(std::exchange(o._pool, nullptr)), | ||
| _event(std::exchange(o._event, nullptr)), | ||
| _cuda_context(std::exchange(o._cuda_context, nullptr)) | ||
| { | ||
| } | ||
|
|
||
| EventPool::Event& EventPool::Event::operator=(Event&& o) noexcept | ||
| { | ||
| if (this != &o) { | ||
| if (_event != nullptr) { | ||
| // Return this event to the pool | ||
| _pool->put(_event, _cuda_context); | ||
| } | ||
| _pool = std::exchange(o._pool, nullptr); | ||
| _event = std::exchange(o._event, nullptr); | ||
| _cuda_context = std::exchange(o._cuda_context, nullptr); | ||
| } | ||
| return *this; | ||
| } | ||
|
|
||
| CUevent EventPool::Event::get() const noexcept { return _event; } | ||
|
|
||
| CUcontext EventPool::Event::cuda_context() const noexcept { return _cuda_context; } | ||
|
|
||
| void EventPool::Event::record(CUstream stream) | ||
| { | ||
| KVIKIO_CUDA_DRIVER_TRY(cudaAPI::instance().EventRecord(_event, stream)); | ||
| } | ||
|
|
||
| void EventPool::Event::synchronize() | ||
| { | ||
| KVIKIO_NVTX_FUNC_RANGE(); | ||
| KVIKIO_CUDA_DRIVER_TRY(cudaAPI::instance().EventSynchronize(_event)); | ||
| } | ||
|
|
||
| bool EventPool::Event::query() const | ||
| { | ||
| auto const status = cudaAPI::instance().EventQuery(_event); | ||
| if (status == CUDA_SUCCESS) { return true; } | ||
| if (status == CUDA_ERROR_NOT_READY) { return false; } | ||
| // Any other return code is an error. | ||
| KVIKIO_CUDA_DRIVER_TRY(status); | ||
| // Unreachable. Macro throws on non-success codes. | ||
| return false; | ||
| } | ||
|
|
||
| EventPool::Event EventPool::get() | ||
| { | ||
| KVIKIO_NVTX_FUNC_RANGE(); | ||
| CUcontext ctx{}; | ||
| KVIKIO_CUDA_DRIVER_TRY(cudaAPI::instance().CtxGetCurrent(&ctx)); | ||
| KVIKIO_EXPECT(ctx != nullptr, "No CUDA context is current"); | ||
|
|
||
| CUevent event{}; | ||
| { | ||
| std::lock_guard const lock(_mutex); | ||
| // If the key (`ctx`) is found from the pool, assign the search result to `event` | ||
| if (auto it = _pools.find(ctx); it != _pools.end() && !it->second.empty()) { | ||
| event = it->second.back(); | ||
| it->second.pop_back(); | ||
| } | ||
| } | ||
|
|
||
| if (event == nullptr) { | ||
| // Create an event outside the lock to improve performance. The pool is not updated here. The | ||
| // returned Event object will automatically return the event to the pool when it goes out of | ||
| // scope | ||
| KVIKIO_CUDA_DRIVER_TRY(cudaAPI::instance().EventCreate(&event, CU_EVENT_DISABLE_TIMING)); | ||
| } | ||
|
Comment on lines
+93
to
+98
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. question: Should we have a max size on this event pool? Since we never destroy events, could the pool unboundedly grow in a long-running application? I suppose it will depend on the maximum concurrency of threads issuing reads?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think it is probably fine to let it grow unbounded. The idea I'm trying to implement in #921 is that each
This event is reused for all the chunks on the same thread originating from the same But I do think that in the future a limitation on the resource pools (currently we have bounce buffer pool, this event pool, and libcurl easy handle pool) is a good feature. |
||
|
|
||
| return Event(this, event, ctx); | ||
| } | ||
|
|
||
| void EventPool::put(CUevent event, CUcontext cuda_context) noexcept | ||
| { | ||
| KVIKIO_NVTX_FUNC_RANGE(); | ||
| if (event == nullptr) { return; } | ||
|
|
||
| try { | ||
| std::lock_guard const lock(_mutex); | ||
| _pools[cuda_context].push_back(event); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This note may not hold for pooled events: /**
* @note An event that has never been recorded reports `is_done() == true`,
* since CUDA's `cuEventQuery` returns `CUDA_SUCCESS` when there is no
* captured work.
*/If an Either tighten the note to acknowledge that pooled events may retain pending work from a previous user, or make
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. That makes sense. I've updated the comment. |
||
| } catch (std::exception const& e) { | ||
| // push_back can throw on allocator failure (e.g., out-of-memory). The event cannot stay | ||
| // cached, so destroy it to release its CUDA resources. | ||
| KVIKIO_LOG_ERROR(e.what()); | ||
| try { | ||
| KVIKIO_CUDA_DRIVER_TRY(cudaAPI::instance().EventDestroy(event)); | ||
| } catch (std::exception const& e) { | ||
| KVIKIO_LOG_ERROR(e.what()); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| std::size_t EventPool::num_free_events(CUcontext cuda_context) const | ||
| { | ||
| std::lock_guard const lock(_mutex); | ||
| auto it = _pools.find(cuda_context); | ||
| return (it != _pools.end()) ? it->second.size() : 0; | ||
| } | ||
|
|
||
| std::size_t EventPool::total_free_events() const | ||
| { | ||
| std::lock_guard const lock(_mutex); | ||
| std::size_t total{0}; | ||
| for (auto const& [_, events] : _pools) { | ||
| total += events.size(); | ||
| } | ||
| return total; | ||
| } | ||
|
|
||
| EventPool& EventPool::instance() | ||
| { | ||
| static EventPool pool; | ||
| return pool; | ||
| } | ||
|
|
||
| } // namespace kvikio::detail | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Consider naming them
CudaEventPoolandCudaEvent.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure! Done!