我最近在使用该Thrust
库时遇到了性能问题。这些来自于在大型嵌套循环结构的基础上分配内存的推力。这显然是不希望的,理想的执行是使用预先分配的全局内存板。我想通过以下三种方式之一删除或改进有问题的代码:
- 实现自定义推力内存分配器
- 用 CUB 代码替换推力代码(带有预先分配的临时存储)
- 编写一个自定义内核来做我想做的事
虽然第三个选项是我通常的首选,但我想要执行的操作是copy_if
/select_if
类型操作,其中返回数据和索引。编写自定义内核可能会重新发明轮子,因此我更愿意选择其他两个选项之一。
我一直听到关于 CUB 的好消息,所以我认为这是在愤怒中使用它的理想机会。我想知道的是:
如何实现select_if
带有返回索引的 CUB?
这可以用一个ArgIndexInputIterator
和这样的仿函数来完成吗?
struct GreaterThan
{
int compare;
__host__ __device__ __forceinline__
GreaterThan(int compare) : compare(compare) {}
__host__ __device__ __forceinline__
bool operator()(const cub::ArgIndexInputIterator<int> &a) const {
return (a.value > compare);
}
};
在代码的主体中包含以下内容:
//d_in = device int array
//d_temp_storage = some preallocated block
int threshold_value;
GreaterThan select_op(threshold_value);
cub::ArgIndexInputIterator<int> input_itr(d_in);
cub::ArgIndexInputIterator<int> output_itr(d_out); //????
CubDebugExit(DeviceSelect::If(d_temp_storage, temp_storage_bytes, input_itr, output_itr, d_num_selected, num_items, select_op));
这会尝试在后台进行任何内存分配吗?
编辑:
因此,根据 Robert Crovella 的评论,仿函数应该采用取消引用 a 的产品cub::ArgIndexInputIterator<int>
,这应该是cub::ItemOffsetPair<int>
现在的仿函数:
struct GreaterThan
{
int compare;
__host__ __device__ __forceinline__
GreaterThan(int compare) : compare(compare) {}
__host__ __device__ __forceinline__
bool operator()(const cub::ItemOffsetPair<int,int> &a) const {
return (a.value > compare);
}
};
在代码中,d_out
应该是一个设备数组cub::ItemOffsetPair<int,int>
:
//d_in = device int array
//d_temp_storage = some preallocated block
cub::ItemOffsetPair<int,int> * d_out;
//allocate d_out
int threshold_value;
GreaterThan select_op(threshold_value);
cub::ArgIndexInputIterator<int,int> input_itr(d_in);
CubDebugExit(DeviceSelect::If(d_temp_storage, temp_storage_bytes, input_itr, d_out, d_num_selected, num_items, select_op));