How to find indices of a vector given another vector in Thrust

I have two thrust device vectors, let's say a and b. I would like to find the indices of vector a where it is smaller/greater than the pair-wise component of vector b ($a_{ij}>b_{ij}$).

I found the following links on how to do it with a scalar:

My main problem is that transform only takes one vector as input while as for the other approach, I am not sure how to do a pair-wise comparison.

Example: a = {1,3,5,6,9} and b = {2,1,4,7,8}. I am looking for indices where a_ij >b_ij.

Therefore, the output should be {1,2,4} with 0-indexed as a is larger at these places than the pair-wise component of b.

Solution

This is at least partially a stream-compaction problem. Thrust provides algorithms for stream compaction such as copy_if.

Thrust does have a transform variant that accepts two vectors as input. You can also do something similar with zip_iterator as input to any thrust algorithm.

I think the most compact method will be to use copy_if with both a zip iterator and a counting iterator. There is a copy_if variant that accepts a stencil array, and that works well for our purpose here.

Here is an example:

# cat t61.cu
#include <thrust/copy.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>


using mt = int;  // data type
using it = size_t; // index type

struct my_functor
{
  template <typename T>
  __host__ __device__
  bool operator()(T t) { return (thrust::get<0>(t) > thrust::get<1>(t));}
};

int main(){

  mt a[] = {1,3,5,6,9};
  mt b[] = {2,1,4,7,8};
  it len = sizeof(a)/sizeof(a[0]);
  thrust::device_vector<mt> da(a, a+len);
  thrust::device_vector<mt> db(b, b+len);
  thrust::device_vector<it> dr(len);
  auto my_idx = thrust::counting_iterator<it>(0);
  auto my_zip = thrust::make_zip_iterator(thrust::make_tuple(da.begin(), db.begin()));
  it lr = thrust::copy_if(my_idx, my_idx+len, my_zip, dr.begin(), my_functor()) - dr.begin();
  thrust::host_vector<it> hr = dr;
  thrust::copy_n(hr.begin(), lr, std::ostream_iterator<it>(std::cout, ","));
  std::cout << std::endl;
}
# nvcc -o t61 t61.cu
# compute-sanitizer ./t61
========= COMPUTE-SANITIZER
1,2,4,
========= ERROR SUMMARY: 0 errors