cuda - 扫描并生成 map (与 Thrust::scatter 一起使用)

Question

我想使用推力创建一个数组，其中包含通过测试的所有索引。我想我在这里遗漏了一些东西，因为我似乎无法弄清楚如何做到这一点。有没有一种简单的方法可以利用推力来做到这一点？

struct is_odd
{
  __device__
  bool operator()(int &x)
  {
    return (x % 2) == 1;
  }
};

int A[] = {1, 2, 1, 1, 4, 1};
int result[] = {-1, -1, -1, -1, -1};

thrust::find_map_if(A, A+5, result, is_odd()); // this function doesn't exist!

// result = {0, 2, 3, 5, -1} 

我需要这个 map 来分散任意数据数组(不是 A)。

Answer 1

可能有很多方法可以解决这个问题。这是一种可能的方法:

1. 创建索引数组(或使用计数迭代器)

   int  A[] = {1, 2, 1, 1, 4, 1};
   int iA[] = {0, 1, 2, 3, 4, 5};  
   

   例如，您可以使用 thrust::sequence 来执行此操作。或者，您可以跳过 iA 的显式生成，并在下一步中使用counting_iterator。

2. 使用thrust::remove_copy_if获取索引数组并将其缩减为与测试结果相对应的元素。

这是一个完整的示例。请注意，remove_copy_if 复制仿函数测试为 false 的元素:

$ cat t596.cu
#include <iostream>
#include <thrust/device_vector.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/remove.h>
#include <thrust/copy.h>

#define DSIZE 6

struct is_odd
{
  __device__
  bool operator()(int x) const
  {
    return !(x&1);
  }
};

int main(){

  int A[] = {1, 2, 1, 1, 4, 1};
  thrust::device_vector<int> d_A(A, A+DSIZE);
  thrust::device_vector<int> result(DSIZE, -1);
  thrust::counting_iterator<int> first(0);
  thrust::counting_iterator<int> last = first+DSIZE;
  thrust::device_vector<int>::iterator r_end = thrust::remove_copy_if(first, last, d_A.begin(), result.begin(), is_odd());
  thrust::copy(result.begin(), r_end, std::ostream_iterator<int>(std::cout, " "));
  std::cout << std::endl;
  thrust::copy(result.begin(), result.end(), std::ostream_iterator<int>(std::cout, " "));
  std::cout << std::endl;
}

$ nvcc -arch=sm_20 -o t596 t596.cu
$ ./t596
0 2 3 5
0 2 3 5 -1 -1
$