Implementation:NVIDIA DALI CUDAStreamPool

Knowledge Sources	NVIDIA_DALI
Domains	Core, CUDA
Last Updated	2026-02-08 16:00 GMT

Overview

This file implements the CUDAStreamPool class, a thread-safe pool for reusing CUDA stream objects across devices, with lazy initialization and busy-stream tracking.

Description

CUDA stream creation and destruction carry significant overhead. The CUDAStreamPool manages per-device pools of reusable CUDA streams to amortize this cost. Unlike the simpler CUDAEventPool, the stream pool differentiates between "ready" streams (idle, immediately reusable) and "busy" streams (still executing previously submitted work). When Get is called, the pool first checks the ready list; if empty, it queries busy streams via cudaStreamQuery to find any that have completed, promoting them to ready status. If no stream is available, a new non-blocking stream is created.

When a stream is returned via Put, its completion status is checked: completed streams go directly to the ready list, while still-active streams are placed on the busy list. This design avoids blocking on stream completion during pool return operations, keeping the critical path fast.

The pool uses lazy initialization via Init(), deferring the cudaGetDeviceCount call until the first Get or Put operation. The constructor pre-reserves 128 entries in the device stream vector to avoid allocation on the first call. Thread safety is provided by a spinlock, and an unused_ free list recycles StreamEntry linked-list nodes to minimize heap allocations. A singleton accessor (instance()) provides process-wide access.

The Get method returns a CUDAStreamLease (defined in the header) rather than a raw CUDAStream. The lease automatically returns the stream to the pool when destroyed or reset, implementing RAII-based stream management.

Usage

Use CUDAStreamPool::instance().Get(device_id) to obtain a CUDAStreamLease. The lease converts implicitly to cudaStream_t and AccessOrder, so it can be used directly with CUDA runtime APIs and DALI ordering primitives. When the lease goes out of scope, the stream is automatically returned to the pool.

Code Reference

Source Location

• Repository: NVIDIA_DALI
• File: dali/core/cuda_stream_pool.cc
• Lines: 1-141

Signature

class DLL_PUBLIC CUDAStreamPool {
 public:
  CUDAStreamPool();
  ~CUDAStreamPool();

  CUDAStreamLease Get(int device_id = -1);
  void Put(CUDAStream &&stream, int device_id = -1);
  void Purge();

  static CUDAStreamPool &instance();

 private:
  void Init();
  CUDAStream GetFromPool(int device_id);

  struct StreamEntry {
    StreamEntry() = default;
    explicit StreamEntry(CUDAStream stream, StreamEntry *next = nullptr);
    CUDAStream stream;
    StreamEntry *next = nullptr;
  };

  StreamEntry *unused_ = nullptr;
  std::atomic_int lease_count_{0};
  std::vector<StreamEntry *> dev_streams_;
  std::vector<StreamEntry *> busy_streams_;
  spinlock lock_;

  static StreamEntry *Pop(StreamEntry *&head);
  static void Push(StreamEntry *&head, StreamEntry *new_entry);
  void DeleteList(StreamEntry *&head);
};

class CUDAStreamLease {
 public:
  constexpr CUDAStreamLease() = default;
  CUDAStreamLease(CUDAStreamLease &&other);
  CUDAStreamLease &operator=(CUDAStreamLease &&other);
  ~CUDAStreamLease();

  cudaStream_t get() const & noexcept;
  operator cudaStream_t() const & noexcept;
  operator AccessOrder() const & noexcept;
  explicit operator bool() const noexcept;

  CUDAStream release() noexcept;
  int device_id() const noexcept;
  CUDAStreamPool *pool() const noexcept;
  void reset();
};

Import

#include "dali/core/cuda_stream_pool.h"

I/O Contract

Inputs

Name	Type	Required	Description
device_id	int	No	CUDA device ordinal; defaults to -1 (current device)
stream	CUDAStream &&	Yes (Put)	CUDA stream wrapper to return to the pool; ownership is transferred

Outputs

Name	Type	Description
CUDAStreamLease	CUDAStreamLease	RAII wrapper that holds a leased stream and returns it to the pool on destruction
instance	CUDAStreamPool &	Reference to the process-wide singleton pool

Usage Examples

Leasing a stream with RAII

{
  auto lease = CUDAStreamPool::instance().Get(0);
  // Use with CUDA APIs directly
  cudaMemcpyAsync(dst, src, size, cudaMemcpyDeviceToDevice, lease.get());
  // Stream is automatically returned to the pool when lease goes out of scope
}

Using a lease as AccessOrder

auto lease = CUDAStreamPool::instance().Get();
AccessOrder order = lease;  // implicit conversion
some_dali_operation(order);

Detaching a stream from the pool

auto lease = CUDAStreamPool::instance().Get(device_id);
CUDAStream owned_stream = lease.release();
// Stream is now owned by the caller and will NOT be returned to the pool

Related Pages

• Environment:NVIDIA_DALI_CUDA_GPU_Environment