如何在Thrust执行策略中使用内存池减少cudaMalloc调用？

问题背景 #

你需要在Thrust中通过内存池减少cudaMalloc调用次数，目前遇到的核心障碍是：

• thrust::device_vector支持自定义分配器，但thrust::sort等算法内部分配的临时缓冲区无法通过容器分配器控制
• 早期Stack Overflow回答提到的执行策略挂钩方法缺乏最新文档支持，而官方示例custom_temporary_allocation.cu采用全局变量挂钩的方式存在全局状态污染问题

最新解决方案（Thrust 1.15+/CUDA 11.7+） #

从CUDA 11.7附带的Thrust 1.15版本开始，官方正式支持执行策略级别的自定义临时内存分配器注入，无需依赖全局状态，步骤如下：

1. 实现符合Thrust规范的内存池分配器 #

自定义分配器需要继承thrust::device_malloc_allocator并重载核心分配/释放方法：

#include <thrust/device_malloc_allocator.h>
#include <thrust/system/cuda/execution_policy.h>
#include <cuda_runtime_api.h>

template <typename T>
struct pool_allocator : thrust::device_malloc_allocator<T> {
    using super_t = thrust::device_malloc_allocator<T>;
    using pointer = typename super_t::pointer;
    using size_type = typename super_t::size_type;

    pointer allocate(size_type n) override {
        // 替换为你的内存池分配逻辑，例如从预分配的大块内存中切分
        void* ptr = nullptr;
        cudaError_t err = cudaMalloc(&ptr, n * sizeof(T));
        if (err != cudaSuccess) {
            throw thrust::bad_alloc();
        }
        return pointer(static_cast<T*>(ptr));
    }

    void deallocate(pointer p, size_type n) override {
        // 替换为内存池的内存回收逻辑
        cudaFree(static_cast<void*>(p.get()));
    }
};

2. 将分配器绑定到执行策略 #

使用thrust::cuda::par的with_allocator方法，创建绑定了自定义内存池的执行策略，所有基于该策略的算法都会使用内存池分配临时内存：

// 初始化绑定了内存池分配器的执行策略
auto pool_exec = thrust::cuda::par(pool_allocator<int>());

// 使用该策略执行sort，临时缓冲区将从内存池获取
thrust::device_vector<int> d_vec = {3, 1, 4, 1, 5};
thrust::sort(pool_exec, d_vec.begin(), d_vec.end());

3. 关键注意事项 #

• 线程安全性：如果在多线程环境下使用，内存池的分配/回收逻辑必须实现线程安全（如使用CUDA流本地内存池或加锁机制）
• 内存池预优化：建议在程序启动阶段预先分配足够的大块设备内存，避免内存池内部频繁调用cudaMalloc
• 类型通用性：通过模板参数让分配器支持任意数据类型，无需为每种类型单独实现

旧版本Thrust兼容方案 #

如果使用CUDA 11.7之前的版本，无法使用执行策略绑定分配器，可采用以下折中方案：

• 手动传入临时缓冲区：部分Thrust算法（如thrust::sort的重载版本）支持用户手动传入预分配的临时存储，可提前从内存池分配好缓冲区传入
• 改进全局钩子：将官方示例中的全局分配器替换为线程局部存储（TLS）的分配器实例，减少全局状态冲突

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