Mini-Infer (29): CUDA 后端支持 (上) — CUDAAllocator 与显存管理

Mini-Infer (29): CUDA 后端支持 (上) — CUDAAllocator 与显存管理

1. 问题背景：CPU 与 GPU 内存管理的差异

在之前的实现中，Mini-Infer 只支持 CPU 推理。现在我们要添加 CUDA 后端支持，首先需要解决的就是 GPU 显存管理。

CPU 和 GPU 内存管理有几个关键差异：

特性	CPU 内存	GPU 显存
分配函数	malloc / aligned_alloc	cudaMalloc
释放函数	free	cudaFree
分配开销	较低	较高（需要驱动调用）
默认对齐	通常 16 字节	256 字节
访问方式	直接访问	需要 Kernel 或 cudaMemcpy
错误处理	返回 nullptr	返回 cudaError_t

为了统一管理，我们需要一个抽象的 Allocator 接口，以及针对 CUDA 的具体实现。

---

2. Allocator 抽象接口回顾

// mini_infer/core/allocator.h

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

    /**
     * @brief 分配内存
     * @param size 字节数
     * @param alignment 对齐字节数
     * @return 分配的内存指针，失败返回 nullptr
     */
    virtual void* allocate(size_t size, size_t alignment) = 0;

    /**
     * @brief 释放内存
     * @param ptr 要释放的指针
     */
    virtual void deallocate(void* ptr) = 0;

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

这个接口足够简单，但足以支持不同的内存分配策略。

---

3. CUDAAllocator 实现

A. 类定义

// mini_infer/backends/cuda/cuda_allocator.h

#ifdef MINI_INFER_USE_CUDA

#include <cuda_runtime.h>

class CUDAAllocator : public core::Allocator {
public:
    /**
     * @brief 构造 CUDA 分配器
     * @param device_id CUDA 设备 ID (默认: 0)
     */
    explicit CUDAAllocator(int device_id = 0);
    ~CUDAAllocator() override = default;

    // 禁用拷贝
    CUDAAllocator(const CUDAAllocator&) = delete;
    CUDAAllocator& operator=(const CUDAAllocator&) = delete;

    void* allocate(size_t size, size_t alignment) override;
    void deallocate(void* ptr) override;

    core::DeviceType device_type() const override {
        return core::DeviceType::CUDA;
    }

    int device_id() const { return device_id_; }

private:
    int device_id_;
};

#endif  // MINI_INFER_USE_CUDA

B. allocate 实现

// mini_infer/backends/cuda/cuda_allocator.cpp

CUDAAllocator::CUDAAllocator(int device_id) : device_id_(device_id) {
    // 设置当前设备
    cudaSetDevice(device_id_);
}

void* CUDAAllocator::allocate(size_t size, size_t alignment) {
    if (size == 0) {
        return nullptr;
    }

    // 设置当前设备（多 GPU 场景）
    cudaSetDevice(device_id_);

    void* ptr = nullptr;
    cudaError_t status = cudaMalloc(&ptr, size);

    if (status != cudaSuccess) {
        MI_LOG_ERROR("[CUDAAllocator] cudaMalloc failed: " +
                     std::string(cudaGetErrorString(status)) +
                     " (requested " + std::to_string(size) + " bytes)");
        return nullptr;
    }

    return ptr;
}

注意：alignment 参数在当前实现中被忽略，因为 cudaMalloc 默认返回 256 字节对齐的内存。

C. deallocate 实现

void CUDAAllocator::deallocate(void* ptr) {
    if (ptr == nullptr) {
        return;
    }

    // 设置当前设备
    cudaSetDevice(device_id_);

    cudaError_t status = cudaFree(ptr);
    if (status != cudaSuccess) {
        MI_LOG_ERROR("[CUDAAllocator] cudaFree failed: " +
                     std::string(cudaGetErrorString(status)));
    }
}

D. device_id 管理

在多 GPU 系统中，每个 CUDAAllocator 实例绑定到一个特定的 GPU：

// 为 GPU 0 创建分配器
auto allocator0 = std::make_shared<CUDAAllocator>(0);

// 为 GPU 1 创建分配器
auto allocator1 = std::make_shared<CUDAAllocator>(1);

每次分配/释放前都调用 cudaSetDevice，确保操作在正确的 GPU 上执行。

---

4. 错误处理宏设计

CUDA API 调用可能失败，我们需要统一的错误处理机制。

A. CUDA_CHECK (返回 Status)

// mini_infer/backends/cuda/cuda_device_context.h

#define CUDA_CHECK(call) \
    do { \
        cudaError_t status = call; \
        if (status != cudaSuccess) { \
            MI_LOG_ERROR("[CUDA] " + std::string(cudaGetErrorString(status)) + \
                         " at " + std::string(__FILE__) + ":" + std::to_string(__LINE__)); \
            return core::Status::ERROR_RUNTIME; \
        } \
    } while(0)

使用场景：在返回 core::Status 的函数中。

core::Status some_function() {
    CUDA_CHECK(cudaMalloc(&ptr, size));  // 失败时自动返回 ERROR_RUNTIME
    CUDA_CHECK(cudaMemcpy(dst, src, size, cudaMemcpyHostToDevice));
    return core::Status::SUCCESS;
}

B. CUDA_CHECK_THROW (抛出异常)

#define CUDA_CHECK_THROW(call) \
    do { \
        cudaError_t status = call; \
        if (status != cudaSuccess) { \
            std::string error_msg = "[CUDA] " + std::string(cudaGetErrorString(status)) + \
                                   " at " + std::string(__FILE__) + ":" + std::to_string(__LINE__); \
            MI_LOG_ERROR(error_msg); \
            throw std::runtime_error(error_msg); \
        } \
    } while(0)

使用场景：在构造函数或无法返回 Status 的函数中。

CUDADeviceContext::CUDADeviceContext(int device_id) {
    CUDA_CHECK_THROW(cudaSetDevice(device_id));  // 失败时抛出异常
    CUDA_CHECK_THROW(cudaStreamCreate(&stream_));
}

C. 返回 Status vs 抛出异常的选择

场景	推荐方式	原因
构造函数	抛出异常	构造函数无法返回值
普通成员函数	返回 Status	调用者可以优雅处理
析构函数	静默处理	析构函数不应抛出异常
性能关键路径	返回 Status	异常有性能开销

---

5. 与 CPUAllocator 的对比

// CPU 分配器 (单例模式)
class CPUAllocator : public Allocator {
public:
    static CPUAllocator* instance() {
        static CPUAllocator allocator;
        return &allocator;
    }

    void* allocate(size_t size, size_t alignment) override {
#ifdef _WIN32
        return _aligned_malloc(size, alignment);
#else
        void* ptr = nullptr;
        posix_memalign(&ptr, alignment, size);
        return ptr;
#endif
    }

    void deallocate(void* ptr) override {
#ifdef _WIN32
        _aligned_free(ptr);
#else
        free(ptr);
#endif
    }
};

对比：

特性	CPUAllocator	CUDAAllocator
模式	单例	每设备一个实例
对齐	显式指定	固定 256 字节
错误处理	返回 nullptr	返回 nullptr + 日志
跨平台	需要条件编译	CUDA API 统一

---

6. 与 Storage 的集成

在 Storage::reset 中，根据设备类型选择分配器：

void Storage::reset(size_t capacity_bytes, DeviceType device, size_t alignment) {
    // 根据设备类型选择分配器
    auto allocator_type = device == DeviceType::CPU
                              ? AllocatorFactory::AllocatorType::CPU
                              : AllocatorFactory::AllocatorType::CUDA;
    auto allocator = AllocatorFactory::get_allocator(allocator_type);

    void* ptr = allocator->allocate(aligned_bytes, alignment);

    // 使用自定义删除器
    buffer_.reset(ptr, [allocator](void* p) {
        allocator->deallocate(p);
    });

    // 零初始化
#ifdef MINI_INFER_USE_CUDA
    if (device == DeviceType::CUDA) {
        cudaMemset(ptr, 0, aligned_bytes);
    } else
#endif
    {
        std::memset(ptr, 0, aligned_bytes);
    }
}

---

7. 未来优化方向

当前的 CUDAAllocator 是最简单的实现，每次分配都调用 cudaMalloc。在生产环境中，这可能成为性能瓶颈。

A. 内存池 (Memory Pooling)

预分配一大块显存，然后从中分配小块：

class PooledCUDAAllocator : public Allocator {
    void* pool_;           // 预分配的大块显存
    size_t pool_size_;
    FreeList free_list_;   // 空闲块链表

    void* allocate(size_t size, size_t alignment) override {
        // 从 free_list_ 中找到合适的块
        // 如果没有，从 pool_ 中切分
    }
};

B. 缓存分配器 (Caching Allocator)

类似 PyTorch 的 CUDACachingAllocator：

class CachingCUDAAllocator : public Allocator {
    std::map<size_t, std::vector<void*>> cached_blocks_;

    void* allocate(size_t size, size_t alignment) override {
        // 1. 查找缓存中是否有合适大小的块
        // 2. 如果有，直接返回
        // 3. 如果没有，调用 cudaMalloc
    }

    void deallocate(void* ptr) override {
        // 不真正释放，而是放入缓存
        cached_blocks_[size].push_back(ptr);
    }
};

C. 内存统计和追踪

class TrackedCUDAAllocator : public Allocator {
    std::atomic<size_t> allocated_bytes_{0};
    std::atomic<size_t> peak_bytes_{0};

    void* allocate(size_t size, size_t alignment) override {
        void* ptr = cudaMalloc(...);
        allocated_bytes_ += size;
        peak_bytes_ = std::max(peak_bytes_.load(), allocated_bytes_.load());
        return ptr;
    }
};

---

8. 总结

本篇我们实现了 CUDA 后端的第一个组件——CUDAAllocator：

• 简单封装：cudaMalloc / cudaFree 的 RAII 封装。
• 多设备支持：每个分配器绑定到特定的 GPU。
• 错误处理宏：CUDA_CHECK 和 CUDA_CHECK_THROW 统一错误处理。
• 与 Storage 集成：通过 AllocatorFactory 自动选择分配器。

下一篇，我们将实现 CUDADeviceContext，管理 CUDA Stream、cuDNN Handle 和 cuBLAS Handle，为 GPU 推理提供完整的执行环境。