问如何在cuda中实现一个简化的argwhere

EN

$ 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
$