在 CUDA 中混合自定义内存管理和推力

Question

在我的项目中，我实现了一个自定义内存分配器，以避免在应用程序“预热”后对cudaMalloc 进行不必要的调用。此外，我使用自定义内核进行基本的数组填充、数组之间的算术运算等，并希望通过使用Thrust 并摆脱这些内核来简化我的代码。设备上的每个数组都是通过原始指针创建和访问的（目前），我想在这些对象上使用device_vector 和Thrusts 方法，但我发现自己在原始指针和device_ptr<> 之间进行转换时间，让我的代码有些混乱。

我相当模糊的问题：您将/如何以最易读的方式组织自定义内存管理、Thrusts 数组方法和对自定义内核的调用的使用？

Answer 1

与所有标准 c++ 容器一样，您可以通过为 thrust::device_vector 提供自己的 "allocator" 来自定义分配存储的方式。默认情况下，thrust::device_vector 的分配器为thrust::device_malloc_allocator，当 Thrust 的后端系统为 CUDA 时，它使用cudaMalloc (cudaFree) 分配（释放）存储。

有时，需要自定义 device_vector 分配内存的方式，例如在 OP 的情况下，他们希望在程序初始化时执行的单个大分配中子分配存储。这可以避免许多单独调用底层分配方案可能产生的开销，在本例中为cudaMalloc。

提供device_vector 自定义分配器的一种简单方法是从device_malloc_allocator 继承。原则上可以从头开始编写整个分配器，但使用继承方法，只需要提供allocate 和deallocate 成员函数。一旦定义了自定义分配器，就可以将其作为第二个模板参数提供给device_vector。

此示例代码演示了如何提供一个自定义分配器，该分配器在分配和解除分配时打印一条消息：

#include <thrust/device_malloc_allocator.h>
#include <thrust/device_vector.h>
#include <iostream>

template<typename T>
  struct my_allocator : thrust::device_malloc_allocator<T>
{
  // shorthand for the name of the base class
  typedef thrust::device_malloc_allocator<T> super_t;

  // get access to some of the base class's typedefs

  // note that because we inherited from device_malloc_allocator,
  // pointer is actually thrust::device_ptr<T>
  typedef typename super_t::pointer   pointer;

  typedef typename super_t::size_type size_type;

  // customize allocate
  pointer allocate(size_type n)
  {
    std::cout << "my_allocator::allocate(): Hello, world!" << std::endl;

    // defer to the base class to allocate storage for n elements of type T
    // in practice, you'd do something more interesting here
    return super_t::allocate(n);
  }

  // customize deallocate
  void deallocate(pointer p, size_type n)
  {
    std::cout << "my_allocator::deallocate(): Hello, world!" << std::endl;

    // defer to the base class to deallocate n elements of type T at address p
    // in practice, you'd do something more interesting here
    super_t::deallocate(p,n);
  }
};

int main()
{
  // create a device_vector which uses my_allocator
  thrust::device_vector<int, my_allocator<int> > vec;

  // create 10 ints
  vec.resize(10, 13);

  return 0;
}

这是输出：

$ nvcc my_allocator_test.cu -arch=sm_20 -run
my_allocator::allocate(): Hello, world!
my_allocator::deallocate(): Hello, world!

在此示例中，请注意，我们曾在 vec.resize(10,13) 上收到过来自 my_allocator::allocate() 的消息。 my_allocator::deallocate() 会在 vec 因销毁其元素而超出范围时被调用一次。