使用shared_ptr管理cudaStream_t实现CUDA流自动销毁的方法

核心问题解答 #

关于cudaStreamCreateWithFlags的内存分配逻辑 #

调用cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)创建流时：

• CUDA运行时会在内部托管堆上分配流对应的状态、命令队列等核心资源，这部分内存由CUDA驱动统一管理
• 代码中声明的cudaStream_t stream是栈上分配的句柄变量，本质是指向CUDA内部流对象的不透明指针，本身不存储流的实际数据

最初编写的智能指针管理版本存在悬空指针bug：循环内声明的stream是单次迭代作用域的栈变量，离开当前循环块后该栈内存就会被系统回收，将这个栈变量的地址传给shared_ptr托管，后续智能指针解引用时访问的是已经失效的内存地址，会触发段错误等未定义行为，根本无法正常执行流销毁逻辑。

推荐实现方案 #

通过RAII机制封装CUDA流，将流的创建、销毁逻辑和C++对象生命周期绑定，配合shared_ptr的引用计数能力，即可实现所有流引用被移除后自动调用cudaStreamDestroy释放资源的效果，参考实现如下：

#include <iostream>
#include <vector>
#include <memory>
#include <cuda_runtime.h>

// CUDA错误检查宏，可根据自身错误处理需求替换
#define gpuErrchk(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)
{
   if (code != cudaSuccess) 
   {
      fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
      if (abort) exit(code);
   }
}

class StreamWrapper {
    public:
        StreamWrapper() {
            gpuErrchk(cudaStreamCreateWithFlags(&this->stream, cudaStreamNonBlocking));
        }

        ~StreamWrapper() { gpuErrchk(cudaStreamDestroy(stream)); }

        // 获取流句柄，用于常规CUDA Runtime API调用
        cudaStream_t& get() { return stream; }

        // 获取流句柄指针，兼容需要传入二级指针的CUDA接口
        cudaStream_t* ref() { return &this->stream; }

    // 禁止拷贝，避免重复销毁流
    StreamWrapper(const StreamWrapper&) = delete;
    StreamWrapper& operator=(const StreamWrapper&) = delete;

    private:
        cudaStream_t stream;
};

int main(){
    int numberOfStreams = 10;
    std::vector<std::shared_ptr<StreamWrapper>> streamsTemp(numberOfStreams);

    for (size_t i = 0; i < numberOfStreams; ++i)
    {
        // 优先用make_shared创建，避免异常场景下的内存泄漏
        streamsTemp[i] = std::make_shared<StreamWrapper>();
    }

    // 流功能测试
    for (size_t i = 0; i < numberOfStreams; ++i)
    {
        int * d_i;
        gpuErrchk(cudaMallocAsync(&d_i, sizeof(int), streamsTemp[i]->get()));
        gpuErrchk(cudaMemcpyAsync(d_i, &i, sizeof(int), cudaMemcpyHostToDevice, streamsTemp[i]->get()));

        int out;
        gpuErrchk(cudaMemcpyAsync(&out, d_i, sizeof(int), cudaMemcpyDeviceToHost, streamsTemp[i]->get()));
        gpuErrchk(cudaFreeAsync(d_i, streamsTemp[i]->get()));
        gpuErrchk(cudaStreamSynchronize(streamsTemp[i]->get()));

        std::cout << "Out: " << std::to_string(out) << " In: " << std::to_string(i) << std::endl;
    }

    // 离开作用域时，所有shared_ptr引用计数归0，自动触发StreamWrapper析构，销毁对应CUDA流
    return 0;
}

实现提示：

• 如果不需要跨模块共享流的所有权，直接用std::vector<StreamWrapper>存储流包装对象即可，不需要额外套shared_ptr，性能开销更低
• 包装类需要禁用拷贝构造、拷贝赋值逻辑，避免多个对象持有同一个流句柄导致重复销毁

内容的提问来源于stack exchange，提问作者DysonSphere