首页 文章

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

提问于
浏览 次
1

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

这可能吗？如果是这样，是否可以使用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函数中 .

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

c++ templates cuda malloc thrust

1 回答

  • 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 迭代器的类型来知道使用自定义执行策略 .

    回复于

相关问题