Mini-Infer (30): CUDA 后端支持 (中) — CUDADeviceContext 与异构执行环境

1. DeviceContext 抽象回顾

在 Mini-Infer 的架构中，DeviceContext 是执行环境的抽象基类：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// mini_infer/backends/device_context.h

class DeviceContext {
public:
    virtual ~DeviceContext() = default;

    /**
     * @brief 获取设备类型
     */
    virtual core::DeviceType device_type() const = 0;

    /**
     * @brief 同步设备（等待所有操作完成）
     */
    virtual void synchronize() = 0;
};

对于 CPU，CPUDeviceContext 的实现非常简单（几乎是空的）。但对于 CUDA，我们需要管理更多资源。

2. CUDADeviceContext 核心设计

CUDADeviceContext 是 GPU 推理的执行环境，管理以下资源：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
// mini_infer/backends/cuda/cuda_device_context.h

class CUDADeviceContext : public DeviceContext {
public:
    explicit CUDADeviceContext(int device_id = 0);
    ~CUDADeviceContext() override;

    // DeviceContext 接口
    core::DeviceType device_type() const override { return core::DeviceType::CUDA; }
    void synchronize() override;

    // CUDA 特有接口
    int device_id() const { return device_id_; }
    cudaStream_t stream() const { return stream_; }
    cudnnHandle_t cudnn_handle();      // 懒初始化
    cublasHandle_t cublas_handle();    // 懒初始化
    void* get_workspace(size_t size);  // TensorRT 风格
    const cudaDeviceProp& device_properties() const;
    size_t available_memory() const;

private:
    int device_id_;
    cudaStream_t stream_{nullptr};
    cudnnHandle_t cudnn_handle_{nullptr};
    cublasHandle_t cublas_handle_{nullptr};
    cudaDeviceProp device_prop_;
    void* workspace_{nullptr};
    size_t workspace_size_{0};
};

3. CUDA Stream 管理

A. 什么是 CUDA Stream？

CUDA Stream 是一个命令队列，GPU 按顺序执行队列中的操作。不同 Stream 之间的操作可以并行执行。

1
2
3
4
Stream 0: [Kernel A] → [Kernel B] → [Kernel C]
Stream 1: [Kernel D] → [Kernel E]
                ↓
          可以并行执行

B. Stream 创建

1
2
3
4
5
// mini_infer/backends/cuda/cuda_device_context.cpp

void CUDADeviceContext::init_stream() {
    CUDA_CHECK_THROW(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking));
}

为什么使用 cudaStreamNonBlocking？

  • 默认 Stream（Stream 0）会与其他 Stream 同步。
  • cudaStreamNonBlocking 创建的 Stream 不会与默认 Stream 同步，允许更好的并行性。

C. 同步操作

1
2
3
4
5
6
7
void CUDADeviceContext::synchronize() {
    cudaError_t status = cudaStreamSynchronize(stream_);
    if (status != cudaSuccess) {
        MI_LOG_ERROR("[CUDADeviceContext] Stream synchronization failed: " +
                     std::string(cudaGetErrorString(status)));
    }
}

使用场景：

  • 在读取 GPU 计算结果之前。
  • 在测量执行时间时。
  • 在跨 Stream 依赖时。

4. cuDNN Handle 管理

A. 为什么需要 Handle？

cuDNN 是 NVIDIA 的深度学习原语库，提供高度优化的卷积、池化等操作。每个 cuDNN 操作都需要一个 cudnnHandle_t。

Handle 的作用：

  1. 状态管理：保存 cuDNN 的内部状态。
  2. Stream 绑定：指定操作在哪个 Stream 上执行。
  3. 资源复用：避免重复创建/销毁的开销。

B. 懒初始化 (Lazy Initialization)

1
2
3
4
5
6
7
8
9
10
11
12
cudnnHandle_t CUDADeviceContext::cudnn_handle() {
    if (!cudnn_handle_) {
        init_cudnn();
    }
    return cudnn_handle_;
}

void CUDADeviceContext::init_cudnn() {
    CUDNN_CHECK_THROW(cudnnCreate(&cudnn_handle_));
    CUDNN_CHECK_THROW(cudnnSetStream(cudnn_handle_, stream_));
    MI_LOG_INFO("[CUDADeviceContext] cuDNN handle initialized");
}

为什么懒初始化？

  1. 节省资源：如果不使用 cuDNN 操作，就不创建 Handle。
  2. 加快启动：Context 创建时不需要等待 cuDNN 初始化。
  3. 按需加载：某些模型可能只用 cuBLAS，不用 cuDNN。

C. Stream 绑定

1
CUDNN_CHECK_THROW(cudnnSetStream(cudnn_handle_, stream_));

这确保所有 cuDNN 操作都在我们的 Stream 上执行，与其他操作正确同步。

5. cuBLAS Handle 管理

A. cuBLAS 的作用

cuBLAS 是 NVIDIA 的线性代数库，提供矩阵乘法（GEMM）、向量操作等。在深度学习中，全连接层和注意力机制大量使用 GEMM。

B. 实现

1
2
3
4
5
6
7
8
9
10
11
12
cublasHandle_t CUDADeviceContext::cublas_handle() {
    if (!cublas_handle_) {
        init_cublas();
    }
    return cublas_handle_;
}

void CUDADeviceContext::init_cublas() {
    CUBLAS_CHECK_THROW(cublasCreate(&cublas_handle_));
    CUBLAS_CHECK_THROW(cublasSetStream(cublas_handle_, stream_));
    MI_LOG_INFO("[CUDADeviceContext] cuBLAS handle initialized");
}

C. 与 cuDNN 的协同

cuDNN 和 cuBLAS 可以共享同一个 Stream，确保操作按正确顺序执行：

1
2
3
Stream: [cuDNN Conv] → [cuBLAS GEMM] → [cuDNN Pooling]
                    ↓
              顺序执行，无需显式同步

6. Workspace 管理 (TensorRT 风格)

A. 什么是 Workspace？

许多 cuDNN 操作需要临时内存（Workspace）来存储中间结果。例如，卷积的 im2col 变换需要额外空间。

B. 动态扩展策略

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
void* CUDADeviceContext::get_workspace(size_t size) {
    if (size > workspace_size_) {
        // 需要重新分配
        if (workspace_) {
            cudaFree(workspace_);
            workspace_ = nullptr;
        }

        cudaError_t status = cudaMalloc(&workspace_, size);
        if (status != cudaSuccess) {
            MI_LOG_ERROR("[CUDADeviceContext] Failed to allocate workspace of size " +
                         std::to_string(size) + " bytes");
            workspace_size_ = 0;
            return nullptr;
        }

        workspace_size_ = size;
        MI_LOG_INFO("[CUDADeviceContext] Allocated workspace: " +
                    std::to_string(size / 1024.0 / 1024.0) + " MB");
    }

    return workspace_;
}

设计思想：

  1. 按需分配：只在需要时分配。
  2. 只增不减：Workspace 只会增大，不会缩小。
  3. 复用内存：多个操作共享同一块 Workspace。

C. 为什么不预分配固定大小？

  1. 模型差异大：不同模型需要的 Workspace 大小差异很大。
  2. 节省显存：小模型不需要大 Workspace。
  3. 灵活性：动态形状场景下，Workspace 需求可能变化。

TensorRT 的 IExecutionContext::setDeviceMemory 也采用类似策略。

7. 设备属性查询

A. cudaDeviceProp 缓存

1
2
3
4
5
6
7
8
9
10
11
12
13
14
void CUDADeviceContext::query_device_properties() {
    CUDA_CHECK_THROW(cudaGetDeviceProperties(&device_prop_, device_id_));

    std::ostringstream oss;
    oss << "[CUDADeviceContext] Device " << device_id_ << " properties:\n"
        << "  Name: " << device_prop_.name << "\n"
        << "  Compute Capability: " << device_prop_.major << "." << device_prop_.minor << "\n"
        << "  Total Global Memory: " << (device_prop_.totalGlobalMem / 1024.0 / 1024.0) << " MB\n"
        << "  Multiprocessors: " << device_prop_.multiProcessorCount << "\n"
        << "  Max Threads Per Block: " << device_prop_.maxThreadsPerBlock << "\n"
        << "  Warp Size: " << device_prop_.warpSize;

    MI_LOG_INFO(oss.str());
}

常用属性：

属性 含义 用途
name GPU 名称 日志和调试
major/minor 计算能力 选择 Kernel 实现
totalGlobalMem 总显存 内存规划
multiProcessorCount SM 数量 并行度优化
maxThreadsPerBlock 最大线程数 Kernel 配置
warpSize Warp 大小 内存对齐

B. available_memory() 实现

1
2
3
4
5
6
7
8
9
10
11
12
size_t CUDADeviceContext::available_memory() const {
    size_t free_bytes = 0;
    size_t total_bytes = 0;

    cudaError_t status = cudaMemGetInfo(&free_bytes, &total_bytes);
    if (status != cudaSuccess) {
        MI_LOG_WARNING("[CUDADeviceContext] Failed to query memory info");
        return 0;
    }

    return free_bytes;
}

使用场景：

  • 在分配大块内存前检查是否有足够空间。
  • 内存不足时的降级策略。
  • 性能分析和监控。

8. 资源生命周期管理

A. 构造函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
CUDADeviceContext::CUDADeviceContext(int device_id)
    : device_id_(device_id) {

    // 1. 设置当前设备
    CUDA_CHECK_THROW(cudaSetDevice(device_id_));

    // 2. 查询设备属性
    query_device_properties();

    // 3. 初始化 Stream
    init_stream();

    MI_LOG_INFO("[CUDADeviceContext] Initialized on device " +
                std::to_string(device_id_) + ": " +
                std::string(device_prop_.name));
}

注意：cuDNN 和 cuBLAS Handle 不在构造函数中初始化（懒初始化）。

B. 析构函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
CUDADeviceContext::~CUDADeviceContext() {
    // 按相反顺序释放资源

    // 1. 释放 Workspace
    if (workspace_) {
        cudaFree(workspace_);
        workspace_ = nullptr;
    }

    // 2. 销毁 cuBLAS Handle
    if (cublas_handle_) {
        cublasDestroy(cublas_handle_);
        cublas_handle_ = nullptr;
    }

    // 3. 销毁 cuDNN Handle
    if (cudnn_handle_) {
        cudnnDestroy(cudnn_handle_);
        cudnn_handle_ = nullptr;
    }

    // 4. 销毁 Stream
    if (stream_) {
        cudaStreamDestroy(stream_);
        stream_ = nullptr;
    }

    MI_LOG_INFO("[CUDADeviceContext] Destroyed context for device " +
                std::to_string(device_id_));
}

RAII 原则：资源在构造时获取，在析构时释放，确保不会泄漏。

9. 错误处理宏

A. CUDNN_CHECK / CUDNN_CHECK_THROW

1
2
3
4
5
6
7
8
9
#define CUDNN_CHECK(call) \
    do { \
        cudnnStatus_t status = call; \
        if (status != CUDNN_STATUS_SUCCESS) { \
            MI_LOG_ERROR("[cuDNN] " + std::string(cudnnGetErrorString(status)) + \
                         " at " + std::string(__FILE__) + ":" + std::to_string(__LINE__)); \
            return core::Status::ERROR_RUNTIME; \
        } \
    } while(0)

B. CUBLAS_CHECK / CUBLAS_CHECK_THROW

1
2
3
4
5
6
7
8
9
#define CUBLAS_CHECK(call) \
    do { \
        cublasStatus_t status = call; \
        if (status != CUBLAS_STATUS_SUCCESS) { \
            MI_LOG_ERROR("[cuBLAS] Error code " + std::to_string(status) + \
                         " at " + std::string(__FILE__) + ":" + std::to_string(__LINE__)); \
            return core::Status::ERROR_RUNTIME; \
        } \
    } while(0)

注意：cuBLAS 没有 cublasGetErrorString 函数，只能输出错误码。

10. 总结

本篇我们实现了 CUDA 后端的核心组件——CUDADeviceContext：

  • CUDA Stream 管理：异步执行的基础。
  • cuDNN Handle 懒初始化：按需创建，节省资源。
  • cuBLAS Handle 懒初始化：与 cuDNN 共享 Stream。
  • Workspace 动态扩展：TensorRT 风格的临时内存管理。
  • 设备属性查询：为优化决策提供信息。
  • RAII 资源管理：确保资源正确释放。

下一篇，我们将看看如何在 Build-Time 预加载权重到 GPU，实现 TensorRT 风格的零开销推理。