如何使用 CUDA 推力执行策略覆盖推力的低级设备内存分配器

Question

我想覆盖低级CUDA设备内存分配器（实现为thrust:：system:：CUDA:：detail:：malloc（）），以便它在主机（CPU）线程上调用时使用自定义分配器，而不是直接调用cudaMalloc（。

这可能吗？如果是这样的话，是否可以使用Thrust"执行策略"机制来做到这一点？我试过这样一个模型：

struct eptCGA : thrust::system::cuda::detail::execution_policy<eptCGA>
{
};
/// overload the Thrust malloc() template function implementation
template<typename eptCGA> __host__ __device__ void* malloc( eptCGA, size_t n )
{
#ifndef __CUDA_ARCH__
    return MyMalloc( n );   /* (called from a host thread) */
#else
    return NULL;            /* (called from a device GPU thread) */
#endif
}

/* called as follows, for example */
eptCGA epCGA;
thrust::remove_if( epCGA, ... );

这是有效的。但是Thrust的其他组件调用底层malloc实现，似乎没有使用"执行策略"机制。例如，

    thrust::device_vector<UINT64> MyDeviceVector( ... );

不使用"执行策略"参数公开重载。相反，malloc（）在15个嵌套函数调用的底部被调用，使用的执行策略似乎是硬连接到调用堆栈中间某个Thrust函数中的执行策略。

有人能澄清我所采取的方法是不正确的吗？并解释一个可行的实施应该做什么？

Answer 1

这对我很有用。你可以一次性创建一个自定义执行策略和分配器，使用你的自定义malloc：

#include <thrust/system/cuda/execution_policy.h>
#include <thrust/system/cuda/memory.h>
#include <thrust/system/cuda/vector.h>
#include <thrust/remove.h>
// create a custom execution policy by deriving from the existing cuda::execution_policy
struct my_policy : thrust::cuda::execution_policy<my_policy> {};
// provide an overload of malloc() for my_policy
__host__ __device__ void* malloc(my_policy, size_t n )
{
  printf("hello, world from my special malloc!n");
  return thrust::raw_pointer_cast(thrust::cuda::malloc(n));
}
// create a custom allocator which will use our malloc
// we can inherit from cuda::allocator to reuse its existing functionality
template<class T>
struct my_allocator : thrust::cuda::allocator<T>
{
  using super_t = thrust::cuda::allocator<T>;
  using pointer = typename super_t::pointer;
  pointer allocate(size_t n)
  {
    T* raw_ptr = reinterpret_cast<T*>(malloc(my_policy{}, sizeof(T) * n));
    // wrap the raw pointer in the special pointer wrapper for cuda pointers
    return pointer(raw_ptr);
  }
};
template<class T>
using my_vector = thrust::cuda::vector<T, my_allocator<T>>;
int main()
{
  my_vector<int> vec(10, 13);
  vec.push_back(7);
  assert(thrust::count(vec.begin(), vec.end(), 13) == 10);
  // because we're superstitious
  my_policy policy;
  auto new_end = thrust::remove(policy, vec.begin(), vec.end(), 13);
  vec.erase(new_end, vec.end());
  assert(vec.size() == 1);
  return 0;
}

这是我系统的输出：

$ nvcc -std=c++11 -I. test.cu -run
hello, world from my special malloc!
hello, world from my special malloc!
hello, world from my special malloc!
hello, world from my special malloc!

您可以更喜欢使用thrust::pointer<T,Tag>包装器将my_policy合并到自定义的pointer类型中。这将产生用my_policy而不是CUDA执行策略标记my_vector的迭代器的效果。这样，您就不必为每个算法调用提供显式的执行策略（就像本例对thrust::remove的调用所做的那样）。相反，Thrust只需查看my_vector迭代器的类型就可以知道如何使用自定义执行策略。