c++ - 如何在cuda中实现一个简化的argwhere

Question

在 cuda 中实现简化的 argwhere 的最佳方式是什么。基本上，我想编写一个内核，它接收两个相同大小的图像并返回一个图像位置数组，两个图像上的值相等。比较是微不足道的，但是，我一直在尝试生成结果数组。我不知道如何跨所有线程同步附加到数组。

Answer 1

一种可能的方法是使用缩减/流压缩方法，在这种情况下我认为 thrust将是一个显而易见的选择。

另一种可能的方法是使用基于原子的方法，如果匹配密度低，它应该比并行减少更快，其中每个找到匹配的线程自动请求全局数组中的一个槽，以存储匹配索引。

这是两种方法的一个有效示例:

$ cat t909.cu
#include <thrust/device_vector.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/copy.h>

#include <iostream>

#define MAX_SIZE 1048576
#define nTPB 256

__device__ int match_indices[MAX_SIZE];
__device__ int next_idx = 0;

__device__ void add_idx(int idx){

  int my_idx = atomicAdd(&next_idx, 1);
  if (my_idx < MAX_SIZE) match_indices[my_idx] = idx;
}

template <typename T>
__device__ bool match_func(T &d1, T &d2){

  return (d1 == d2);
}

template <typename T>
__global__ void k1(const T * __restrict__ d1, const T * __restrict__ d2, const int dsize){

  int idx = threadIdx.x+blockDim.x*blockIdx.x;
  if (idx < dsize){
    if (match_func(d1[idx], d2[idx])) add_idx(idx);
    }
}

typedef thrust::tuple<int, int> mytuple;
struct my_comp : public thrust::unary_function<mytuple, int>
{
  __host__ __device__
  int operator()(mytuple &t1){
    if (thrust::get<0>(t1) == thrust::get<1>(t1)) return 0;
    else return 1;
  }
};

using namespace thrust::placeholders;

int main(){

  thrust::device_vector<int> d1(MAX_SIZE, 1);
  thrust::device_vector<int> d2(MAX_SIZE, 2);
  d1[12] = 2; d1[16] = 2; d1[MAX_SIZE-1] = 2;

  //method 2
  k1<<<(MAX_SIZE+nTPB-1)/nTPB,nTPB>>>(thrust::raw_pointer_cast(d1.data()), thrust::raw_pointer_cast(d2.data()), MAX_SIZE);
  int total_matches;
  cudaMemcpyFromSymbol(&total_matches, next_idx, sizeof(int));
  int *matches = new int[total_matches];
  cudaMemcpyFromSymbol(matches, match_indices, total_matches*sizeof(int));
  std::cout << "Kernel results: " << std::endl;
  for (int i = 0; i < total_matches; i++)
    std::cout << matches[i] << ",";
  std::cout << std::endl;

  //method 1
  thrust::device_vector<int> result(MAX_SIZE);
  int result_size = thrust::copy_if(thrust::counting_iterator<int>(0), thrust::counting_iterator<int>(MAX_SIZE), thrust::make_transform_iterator(thrust::make_zip_iterator(thrust::make_tuple(d1.begin(), d2.begin())), my_comp()), result.begin(), _1 == 0) - result.begin();
  std::cout << "Thrust results: " << std::endl;
  thrust::copy_n(result.begin(), result_size, std::ostream_iterator<int>(std::cout, ","));
  std::cout << std::endl;
}
$ nvcc -o t909 t909.cu
$ cuda-memcheck ./t909
========= CUDA-MEMCHECK
Kernel results:
12,16,1048575,
Thrust results:
12,16,1048575,
========= ERROR SUMMARY: 0 errors
$

作为最后的说明，如果您对这两种方法进行时序比较，您会发现当匹配密度较低(比如 ~1% 或更少)时，原子/内核方法明显更快，但推力方法当匹配密度高(比如 ~50% 或更高)时速度更快。确切的比较将取决于您运行的 GPU 以及可能还有其他因素，例如整体数据集大小。