Implement CUDA event pool to minimize runtime resource allocation overhead (libcurl multi poll-based backend 2/n)

cpp/CMakeLists.txt

@@ -169,6 +169,7 @@ set(SOURCES
     "src/logger.cpp"
     "src/mmap.cpp"
     "src/detail/env.cpp"
+    "src/detail/event.cpp"
     "src/detail/nvtx.cpp"
     "src/detail/posix_io.cpp"
     "src/detail/stream.cpp"

cpp/include/kvikio/detail/event.hpp

@@ -0,0 +1,195 @@
+/*
+ * 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 `CudaEventPool::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.
+ * `CudaEventPool::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 CudaEventPool {
+ 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 CudaEvent {
+    friend class CudaEventPool;
+
+   private:
+    CudaEventPool* _pool{};
+    CUevent _event{};
+    CUcontext _cuda_context{};
+
+    /**
+     * @brief Construct a CudaEvent wrapping a CUDA event handle
+     *
+     * @param pool The owning CudaEventPool 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 CudaEvent(CudaEventPool* pool, CUevent event, CUcontext context) noexcept;
+
+   public:
+    ~CudaEvent() noexcept;
+
+    // Move-only
+    CudaEvent(CudaEvent const&)            = delete;
+    CudaEvent& operator=(CudaEvent const&) = delete;
+    CudaEvent(CudaEvent&& o) noexcept;
+    CudaEvent& operator=(CudaEvent&& 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 CudaEvent.
+     */
+    [[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 check whether all work captured by the event has completed.
+     *
+     * 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()`.
+     *
+     * @note An event that has never been recorded by `cudaAPI::instance().EventRecord()` reports
+     * `is_done() == true`, since CUDA's `cuEventQuery` returns `CUDA_SUCCESS` when there is no
+     * captured work. Users that rely on `is_done() == true` as a clean baseline signal should
+     * `synchronize()` before dropping the event so that the next pool acquirer
+     * `CudaEventPool::instance().get()` sees an idle state.
+     *
+     * @return true if the event has completed (or has never been recorded), 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 is_done() const;
+  };
+
+ private:
+  std::mutex mutable _mutex;
+  // Per-context pools of free events
+  std::unordered_map<CUcontext, std::vector<CUevent>> _pools;
+
+  CudaEventPool() = 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.
+  ~CudaEventPool() noexcept = default;
+
+  /**
+   * @brief Return an event to the pool for reuse
+   *
+   * Called by CudaEvent'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
+  CudaEventPool(CudaEventPool const&)            = delete;
+  CudaEventPool& operator=(CudaEventPool const&) = delete;
+  CudaEventPool(CudaEventPool&&)                 = delete;
+  CudaEventPool& operator=(CudaEventPool&&)      = 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 CudaEvent object will automatically return the event to the pool when it goes out
+   * of scope.
+   *
+   * @return RAII CudaEvent object wrapping the acquired CUDA event
+   * @exception kvikio::CUfileException if no CUDA context is current or event creation fails
+   */
+  [[nodiscard]] CudaEvent 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 CudaEventPool instance
+   */
+  static CudaEventPool& instance();
+};
+}  // namespace kvikio::detail

cpp/include/kvikio/shim/cuda.hpp

@@ -92,6 +92,8 @@ class cudaAPI {
 
   decltype(cuPointerGetAttribute)* PointerGetAttribute{nullptr};
   decltype(cuPointerGetAttributes)* PointerGetAttributes{nullptr};
+  decltype(cuCtxCreate)* CtxCreate{nullptr};
+  decltype(cuCtxDestroy)* CtxDestroy{nullptr};
   decltype(cuCtxPushCurrent)* CtxPushCurrent{nullptr};
   decltype(cuCtxPopCurrent)* CtxPopCurrent{nullptr};
   decltype(cuCtxGetCurrent)* CtxGetCurrent{nullptr};
@@ -108,6 +110,11 @@ class cudaAPI {
   decltype(cuStreamCreate)* StreamCreate{nullptr};
   decltype(cuStreamDestroy)* StreamDestroy{nullptr};
   decltype(cuDriverGetVersion)* DriverGetVersion{nullptr};
+  decltype(cuEventSynchronize)* EventSynchronize{nullptr};
+  decltype(cuEventCreate)* EventCreate{nullptr};
+  decltype(cuEventDestroy)* EventDestroy{nullptr};
+  decltype(cuEventRecord)* EventRecord{nullptr};
+  decltype(cuEventQuery)* EventQuery{nullptr};
 
  private:
   cudaAPI();

cpp/src/detail/event.cpp

@@ -0,0 +1,146 @@
+/*
+ * 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 {
+
+CudaEventPool::CudaEvent::CudaEvent(CudaEventPool* pool,
+                                    CUevent event,
+                                    CUcontext cuda_context) noexcept
+  : _pool(pool), _event(event), _cuda_context(cuda_context)
+{
+}
+
+CudaEventPool::CudaEvent::~CudaEvent() noexcept
+{
+  if (_event != nullptr) { _pool->put(_event, _cuda_context); }
+}
+
+CudaEventPool::CudaEvent::CudaEvent(CudaEvent&& o) noexcept
+  : _pool(std::exchange(o._pool, nullptr)),
+    _event(std::exchange(o._event, nullptr)),
+    _cuda_context(std::exchange(o._cuda_context, nullptr))
+{
+}
+
+CudaEventPool::CudaEvent& CudaEventPool::CudaEvent::operator=(CudaEvent&& 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 CudaEventPool::CudaEvent::get() const noexcept { return _event; }
+
+CUcontext CudaEventPool::CudaEvent::cuda_context() const noexcept { return _cuda_context; }
+
+void CudaEventPool::CudaEvent::record(CUstream stream)
+{
+  KVIKIO_CUDA_DRIVER_TRY(cudaAPI::instance().EventRecord(_event, stream));
+}
+
+void CudaEventPool::CudaEvent::synchronize()
+{
+  KVIKIO_NVTX_FUNC_RANGE();
+  KVIKIO_CUDA_DRIVER_TRY(cudaAPI::instance().EventSynchronize(_event));
+}
+
+bool CudaEventPool::CudaEvent::is_done() 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;
+}
+
+CudaEventPool::CudaEvent CudaEventPool::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 CudaEvent 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));
+  }
+
+  return CudaEvent(this, event, ctx);
+}
+
+void CudaEventPool::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);
+  } 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 CudaEventPool::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 CudaEventPool::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;
+}
+
+CudaEventPool& CudaEventPool::instance()
+{
+  static CudaEventPool pool;
+  return pool;
+}
+
+}  // namespace kvikio::detail

cpp/src/shim/cuda.cpp

@@ -26,6 +26,8 @@ cudaAPI::cudaAPI()
   get_symbol(MemcpyAsync, lib, KVIKIO_STRINGIFY(cuMemcpyAsync));
   get_symbol(PointerGetAttribute, lib, KVIKIO_STRINGIFY(cuPointerGetAttribute));
   get_symbol(PointerGetAttributes, lib, KVIKIO_STRINGIFY(cuPointerGetAttributes));
+  get_symbol(CtxCreate, lib, KVIKIO_STRINGIFY(cuCtxCreate));
+  get_symbol(CtxDestroy, lib, KVIKIO_STRINGIFY(cuCtxDestroy));
   get_symbol(CtxPushCurrent, lib, KVIKIO_STRINGIFY(cuCtxPushCurrent));
   get_symbol(CtxPopCurrent, lib, KVIKIO_STRINGIFY(cuCtxPopCurrent));
   get_symbol(CtxGetCurrent, lib, KVIKIO_STRINGIFY(cuCtxGetCurrent));
@@ -42,6 +44,11 @@ cudaAPI::cudaAPI()
   get_symbol(StreamCreate, lib, KVIKIO_STRINGIFY(cuStreamCreate));
   get_symbol(StreamDestroy, lib, KVIKIO_STRINGIFY(cuStreamDestroy));
   get_symbol(DriverGetVersion, lib, KVIKIO_STRINGIFY(cuDriverGetVersion));
+  get_symbol(EventSynchronize, lib, KVIKIO_STRINGIFY(cuEventSynchronize));
+  get_symbol(EventCreate, lib, KVIKIO_STRINGIFY(cuEventCreate));
+  get_symbol(EventDestroy, lib, KVIKIO_STRINGIFY(cuEventDestroy));
+  get_symbol(EventRecord, lib, KVIKIO_STRINGIFY(cuEventRecord));
+  get_symbol(EventQuery, lib, KVIKIO_STRINGIFY(cuEventQuery));
 
   KVIKIO_CUDA_DRIVER_TRY(DriverGetVersion(&driver_version));