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

cpp/CMakeLists.txt

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
@@ -153,8 +153,10 @@ set(SOURCES
     "src/file_utils.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"
     "src/shim/cuda.cpp"
     "src/shim/cufile.cpp"
     "src/shim/utils.cpp"

cpp/include/kvikio/detail/event.hpp

@@ -0,0 +1,152 @@
+/*
+ * 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 will be automatically returned to the
+ * pool when it goes out of scope (RAII).
+ *
+ * @note The destructor intentionally leaks events to avoid CUDA cleanup issues when static
+ * destructors run after CUDA context destruction. @sa BounceBufferPool
+ */
+class EventPool {
+ public:
+  class Event {
+    friend class EventPool;
+
+   private:
+    CUevent _event{nullptr};
+    CUcontext _context{nullptr};
+
+    /**
+     * @brief Construct an Event wrapping a CUDA event handle
+     *
+     * @param event The CUDA event handle to wrap
+     * @param context The CUDA context associated with this event
+     */
+    explicit Event(CUevent event, CUcontext context) noexcept;
+
+   public:
+    ~Event() noexcept;
+
+    // Move-only
+    Event(Event const&)            = delete;
+    Event& operator=(Event const&) = delete;
+    Event(Event&& other) noexcept;
+    Event& operator=(Event&& other) 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
+     */
+    [[nodiscard]] CUcontext 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)
+     */
+    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 CudaError if the synchronize operation fails
+     */
+    void synchronize();
+  };
+
+ private:
+  std::mutex mutable _mutex;
+  // Per-context pools of free events
+  std::unordered_map<CUcontext, std::vector<CUevent>> _pools;
+
+  EventPool() = default;
+
+  // Intentionally leak events during static destruction. @sa BounceBufferPool
+  ~EventPool() noexcept = default;
+
+ 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 from the pool
+   *
+   * 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 CudaError if no CUDA context is current or event creation fails
+   */
+  [[nodiscard]] Event get();
+
+  /**
+   * @brief Return an event to the pool for reuse
+   *
+   * Typically called automatically by Event's destructor. 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
+   */
+  void put(CUevent event, CUcontext context) noexcept;
+
+  /**
+   * @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

cpp/include/kvikio/detail/posix_io.hpp

@@ -1,18 +1,18 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
 #include <unistd.h>
+
 #include <cstddef>
 #include <cstdlib>
-#include <map>
-#include <thread>
 #include <type_traits>
 
 #include <kvikio/bounce_buffer.hpp>
 #include <kvikio/detail/nvtx.hpp>
+#include <kvikio/detail/stream.hpp>
 #include <kvikio/detail/utils.hpp>
 #include <kvikio/error.hpp>
 #include <kvikio/shim/cuda.hpp>
@@ -36,37 +36,6 @@ enum class PartialIO : uint8_t {
   NO,   ///< POSIX read/write is called repeatedly until all requested bytes are processed.
 };
 
-/**
- * @brief Singleton class to retrieve a CUDA stream for device-host copying
- *
- * Call `StreamsByThread::get` to get the CUDA stream assigned to the current
- * CUDA context and thread.
- */
-class StreamsByThread {
- private:
-  std::map<std::pair<CUcontext, std::thread::id>, CUstream> _streams;
-
- public:
-  StreamsByThread() = default;
-
-  // Here we intentionally do not destroy in the destructor the CUDA resources
-  // (e.g. CUstream) with static storage duration, but instead let them leak
-  // on program termination. This is to prevent undefined behavior in CUDA. See
-  // <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#initialization>
-  // This also prevents crash (segmentation fault) if clients call
-  // cuDevicePrimaryCtxReset() or cudaDeviceReset() before program termination.
-  ~StreamsByThread() = default;
-
-  KVIKIO_EXPORT static CUstream get(CUcontext ctx, std::thread::id thd_id);
-
-  static CUstream get();
-
-  StreamsByThread(StreamsByThread const&)            = delete;
-  StreamsByThread& operator=(StreamsByThread const&) = delete;
-  StreamsByThread(StreamsByThread&& o)               = delete;
-  StreamsByThread& operator=(StreamsByThread&& o)    = delete;
-};
-
 /**
  * @brief Read or write host memory to or from disk using POSIX with opportunistic Direct I/O
  *
@@ -238,7 +207,7 @@ std::size_t posix_device_io(int fd_direct_off,
   off_t const chunk_size2 = convert_size2off(bounce_buffer.size());
 
   // Get a stream for the current CUDA context and thread
-  CUstream stream = StreamsByThread::get();
+  CUstream stream = StreamCachePerThreadAndContext::get();
 
   while (bytes_remaining > 0) {
     off_t const nbytes_requested = std::min(chunk_size2, bytes_remaining);

cpp/include/kvikio/detail/stream.hpp

@@ -0,0 +1,67 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <map>
+#include <mutex>
+#include <thread>
+
+#include <kvikio/shim/cuda.hpp>
+
+namespace kvikio::detail {
+/**
+ * @brief Singleton cache that provides one CUDA stream per (context, thread) pair.
+ *
+ * This class manages CUDA streams used for host-device memory transfers. Each unique combination of
+ * CUDA context and calling thread is assigned a dedicated stream, which is created lazily on first
+ * access and reused for subsequent calls.
+ *
+ * The cache is thread-safe and handles concurrent access from multiple threads.
+ *
+ * @note CUDA streams are intentionally leaked on program termination rather than destroyed in the
+ * destructor. This avoids undefined behavior that can occur when destroying CUDA resources during
+ * static destruction, and prevents crashes (segmentation faults) if clients call
+ * cuDevicePrimaryCtxReset() or cudaDeviceReset() before program termination. See:
+ * https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#initialization
+ */
+class StreamCachePerThreadAndContext {
+ private:
+  std::map<std::pair<CUcontext, std::thread::id>, CUstream> _streams;
+  std::mutex mutable _mutex;
+
+ private:
+  StreamCachePerThreadAndContext()  = default;
+  ~StreamCachePerThreadAndContext() = default;
+
+ public:
+  /**
+   * @brief Get or create a CUDA stream for the specified context and thread.
+   *
+   * If a stream already exists for the given (context, thread) pair, it is returned. Otherwise, a
+   * new stream is created, cached, and returned.
+   *
+   * @param ctx The CUDA context. If null, the null stream is returned.
+   * @param thd_id The thread identifier.
+   * @return The CUDA stream associated with this (context, thread) pair, or nullptr if @p ctx is
+   * null.
+   */
+  KVIKIO_EXPORT static CUstream get(CUcontext ctx, std::thread::id thd_id);
+
+  /**
+   * @brief Get or create a CUDA stream for the current context and thread.
+   *
+   * Convenience overload that uses the current CUDA context and calling thread's ID.
+   *
+   * @return The CUDA stream associated with the current (context, thread) pair, or nullptr if no
+   * CUDA context is current.
+   */
+  static CUstream get();
+
+  StreamCachePerThreadAndContext(StreamCachePerThreadAndContext const&)            = delete;
+  StreamCachePerThreadAndContext& operator=(StreamCachePerThreadAndContext const&) = delete;
+  StreamCachePerThreadAndContext(StreamCachePerThreadAndContext&& o)               = delete;
+  StreamCachePerThreadAndContext& operator=(StreamCachePerThreadAndContext&& o)    = delete;
+};
+}  // namespace kvikio::detail

cpp/include/kvikio/shim/cuda.hpp

@@ -107,6 +107,10 @@ 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};
 
  private:
   cudaAPI();