如何使用 CUDA Thrust 执行策略覆盖 Thrust 的低级设备内存分配器
How to use CUDA Thrust execution policy to override Thrust's low-level device memory allocator
我想覆盖低级 CUDA 设备内存分配器(实现为 thrust::system::cuda::detail::malloc()),以便它在主机上调用时使用自定义分配器而不是直接调用 cudaMalloc() (CPU) 线程.
这可能吗?如果可以,是不是可以用Thrust"execution policy"机制来做呢?我试过这样的模型:
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 实现,似乎没有使用 "execution policy" 机制。例如,
thrust::device_vector<UINT64> MyDeviceVector( ... );
不会公开带有 "execution policy" 参数的重载。相反,malloc() 在 15 个嵌套函数调用的底部被调用,使用的执行策略似乎硬连线到该调用堆栈中间某处的 Thrust 函数之一。
有人可以澄清我所采用的方法是如何不正确的,并解释一个可行的实现应该做什么吗?
这是对我有用的东西。您可以一次性创建自定义执行策略和使用自定义 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
的迭代器的类型就知道使用您的自定义执行策略。
我想覆盖低级 CUDA 设备内存分配器(实现为 thrust::system::cuda::detail::malloc()),以便它在主机上调用时使用自定义分配器而不是直接调用 cudaMalloc() (CPU) 线程.
这可能吗?如果可以,是不是可以用Thrust"execution policy"机制来做呢?我试过这样的模型:
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 实现,似乎没有使用 "execution policy" 机制。例如,
thrust::device_vector<UINT64> MyDeviceVector( ... );
不会公开带有 "execution policy" 参数的重载。相反,malloc() 在 15 个嵌套函数调用的底部被调用,使用的执行策略似乎硬连线到该调用堆栈中间某处的 Thrust 函数之一。
有人可以澄清我所采用的方法是如何不正确的,并解释一个可行的实现应该做什么吗?
这是对我有用的东西。您可以一次性创建自定义执行策略和使用自定义 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
的迭代器的类型就知道使用您的自定义执行策略。