4、内存管理基础：GPU显存分配与释放，cudaMalloc/cudaFree，TensorRT中的IGpuAllocator接口

聊到TensorRT推理，绕不开一个话题——显存管理。

说实话，很多同学刚开始接触GPU编程时，觉得显存分配不就是调个cudaMalloc嘛，有啥好讲的？但我在实际项目中踩过不少坑，才意识到：显存管理的好坏，直接决定了你的模型能不能跑起来、跑得快不快。

4.1 从cudaMalloc说起

先看最基础的操作。在CUDA里，分配显存用cudaMalloc，释放用cudaFree。这俩函数，说白了就是GPU版的malloc和free。

float* d_data = nullptr;
size_t size = 1024 * 1024 * sizeof(float); // 4MB

cudaError_t err = cudaMalloc(&d_data, size);
if (err != cudaSuccess) {
    // 处理错误
}

// ... 使用显存 ...

cudaFree(d_data);
d_data = nullptr; // 我习惯释放后立即置空，防止野指针

嗯，这里有个细节要注意：cudaMalloc返回的是设备指针，你不能在主机端直接读写它。我见过有人试图用memcpy去拷贝设备指针指向的内容，结果程序直接崩了。

警告：cudaMalloc分配失败时，返回的不是nullptr，而是错误码。一定要检查返回值！我曾经在一个生产环境里忽略了检查，结果模型加载到一半静默失败，排查了整整两天。

4.2 显存分配的坑与经验

你以为显存分配就这么简单？来，我分享几个实际项目中遇到的坑。

4.2.1 显存碎片化

频繁地cudaMalloc和cudaFree，会导致显存碎片化。就像你反复往一个箱子里塞不同大小的积木，再拿出来，最后箱子里的空隙越来越多。

我建议的做法是：一次性分配大块显存，然后自己管理。比如用内存池技术，或者直接用TensorRT提供的接口。

4.2.2 异步操作与同步

还有一个经典问题：cudaFree是同步操作吗？答案是：它会等待所有未完成的CUDA操作完成后再释放。这意味着如果你在推理循环里频繁分配释放，性能会大打折扣。

技巧：我个人习惯在模型初始化阶段就把所有需要的显存分配好，推理过程中只复用，不释放。这样既避免了碎片化，也减少了同步开销。

4.3 TensorRT中的IGpuAllocator接口

好了，现在进入正题。TensorRT作为一个推理优化引擎，它自己也要管理显存。但它没有硬编码使用cudaMalloc，而是提供了一个接口——IGpuAllocator。

这个接口长这样：

class IGpuAllocator {
public:
    virtual void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) = 0;
    virtual void free(void* memory) = 0;
};

说白了，TensorRT把显存分配的权利交给你了。你可以实现自己的分配策略，比如用内存池、用预分配的大块内存，甚至用自定义的显存管理库。

4.3.1 为什么要用IGpuAllocator？

你想想看，如果TensorRT内部直接调cudaMalloc，那你就没法控制它的显存行为。但通过这个接口，你可以：

  • 控制显存上限：比如限制TensorRT最多用4GB显存
  • 复用显存：多个模型实例共享同一块显存池
  • 调试和监控：在allocate/free里加日志，跟踪显存使用情况

我记得有一次，客户要求在一个只有6GB显存的卡上同时跑两个模型。如果不实现自定义分配器，TensorRT各自独立分配，很容易就超了。后来我写了一个共享内存池的分配器，两个模型共用4GB，剩下的留给其他任务，完美解决。

3.3.2 实现一个简单的分配器

来看一个最简实现：

class SimpleAllocator : public nvinfer1::IGpuAllocator {
private:
    size_t totalAllocated = 0;
    size_t maxMemory = 4ULL * 1024 * 1024 * 1024; // 4GB上限

public:
    void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override {
        // 检查是否超过上限
        if (totalAllocated + size > maxMemory) {
            return nullptr; // 返回nullptr表示分配失败
        }

        void* ptr = nullptr;
        cudaError_t err = cudaMalloc(&ptr, size);
        if (err != cudaSuccess) {
            return nullptr;
        }

        totalAllocated += size;
        // 我习惯在这里打日志，方便调试
        // printf("Allocated %zu bytes, total: %zu\n", size, totalAllocated);
        return ptr;
    }

    void free(void* memory) override {
        if (memory) {
            cudaFree(memory);
            // 注意：这里没法精确减去size，因为free接口不传大小
            // 所以实际项目中需要自己维护一个映射表
        }
    }
};
核心要点：实现IGpuAllocator时，allocate返回nullptr表示分配失败，TensorRT会处理这个错误。不要在里面抛异常，CUDA环境里异常处理很麻烦。

4.4 如何让TensorRT使用自定义分配器

实现好了分配器，怎么让TensorRT用上它？很简单，在创建builder或runtime时传入：

SimpleAllocator myAllocator;
nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
builder->setGpuAllocator(&myAllocator);

// 或者对runtime设置
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
runtime->setGpuAllocator(&myAllocator);

嗯，这里要注意：分配器的生命周期必须长于builder/runtime。我见过有人把分配器定义在局部作用域里，builder还在用的时候分配器就被析构了，结果程序崩溃得一塌糊涂。

4.5 实战建议

总结一下我的经验：

  1. 能不自己分配就别自己分配：TensorRT内部管理已经做得不错了，除非你有特殊需求（比如显存限制、共享池），否则用默认的就好。
  2. 如果自己实现分配器，一定要做压力测试：我曾经实现过一个分配器，在小模型上跑得好好的，换到大模型就崩了。原因是分配器里有个整数溢出bug。
  3. 善用日志：在allocate和free里加日志，可以帮你定位显存泄漏问题。我一般会在分配时记录调用栈，这样泄漏了也能知道是谁分配的。
避坑指南：我曾经在实现分配器时，忘了处理alignment参数。结果TensorRT内部要求128字节对齐，我直接返回了cudaMalloc的地址（默认256字节对齐），侥幸没出问题。但如果你用cudaMalloc分配小内存，对齐可能不满足要求。建议用cudaMalloc时，自己再做一次对齐处理。

好了，这一章就到这里。显存管理看似基础，但往往是性能瓶颈和bug的源头。下一章我们会讲如何用这些知识，实现模型的分块加载——说白了，就是怎么把一个大模型拆开，一块一块地塞进显存里。

← 上一章 📖 返回目录 下一章 →