混合自定义内存管理和CUDA中的Thrust

Question

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

我的相当模糊的问题:如何Thrust以最可读的方式组织自定义内存管理,数组方法和调用自定义内核的用法？

Answer 1

与所有标准c ++容器一样,您可以thrust::device_vector通过为存储提供自己的"分配器"来自定义分配存储的方式.默认情况下,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!

在这个例子中,请注意我们从my_allocator::allocate()一次听到vec.resize(10,13).my_allocator::deallocate()当vec超出范围时会调用一次,因为它会破坏其元素.