CODEWITHSUNDEEP
c++cudathrust
TK Omble
TK Omble

Reputation: 409

Get index of vector inside CUDA thrust::transform operator() function

In CUDA Thrust transform, is it possible to get the index of a vector, passed into the operator() function, inside the function?

say, we have, 

struct op{
    float operator()(const float& f){
        //do something like return the index
    }
};
vector<float> v(100);
thrust::transform(v.begin(),v.end(),v.begin(),op());

how do i get the index of the vector inside the operator()? basically i want a easy way to make a identity matrix in CUDA.

Upvotes: 3

Views: 3426

Answers (1)

Robert Crovella
Robert Crovella

Reputation: 152053

There are probably many ways to do this. One approach would be:

  1. use thrust::sequence to create a vector of indices of the same length as your data vector (or instead just use a counting_iterator)
  2. use a zip_iterator to return a thrust::tuple, combining the data vector and the index vector, returning a tuple of a data item plus its index
  3. Define the operator op() to take the particular tuple as one of it's arguments
  4. In the operator, use thrust::get<> to retrieve either the data element, or the index as needed, from the tuple

You can read more about most of these concepts in the thrust quickstart guide.

EDIT: In response to question below, here's a worked example. Although this doesn't actually use any device_vector, if we were doing this on the GPU (using device_vector) the only activity that would generate any significant GPU activity would be the call to thrust::transform, ie. there would be only 1 "pass" on the GPU.

(Yes, the thrust::sequence call would also generate a GPU kernel, but I'm just using that to create some data for this example).

#include <thrust/host_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/sequence.h>
#include <thrust/copy.h>

#define N 30
#define SELECT 3

typedef thrust::tuple<int, int>            tpl2int;
typedef thrust::host_vector<int>::iterator intiter;
typedef thrust::counting_iterator<int>     countiter;
typedef thrust::tuple<intiter, countiter>  tpl2intiter;
typedef thrust::zip_iterator<tpl2intiter>  idxzip;



struct select_unary_op : public thrust::unary_function<tpl2int, int>
{
  __host__ __device__
  int operator()(const tpl2int& x) const
  {
    if ((x.get<1>() %SELECT) == 0)
      return x.get<0>();
    else return -1;
   }
};

int main() {

  thrust::host_vector<int> A(N);
  thrust::host_vector<int> result(N);
  thrust::sequence(A.begin(), A.end());
  thrust::counting_iterator<int> idxfirst(0);
  thrust::counting_iterator<int> idxlast = idxfirst +N;

  idxzip first = thrust::make_zip_iterator(thrust::make_tuple(A.begin(), idxfirst));
  idxzip  last = thrust::make_zip_iterator(thrust::make_tuple(A.end(), idxlast));
  select_unary_op my_unary_op;

  thrust::transform(first, last, result.begin(), my_unary_op);
  std::cout << "Results :" << std::endl;
  thrust::copy(result.begin(), result.end(), std::ostream_iterator<int>( std::cout, " "));
  std::cout << std::endl;


  return 0;

}

Upvotes: 8

Related Questions