// matrix transpose, using scatter()

// similar to (though less efficient than) code included in the examples
// in the Thrust package

// matrices assumed stored in one dimension, row-major order

#include <stdio.h>
#include <thrust/device_vector.h>
#include <thrust/scatter.h>
#include <thrust/sequence.h>

struct transidx {
   const int nr;  // number of rows in input
   const int nc;  // number of columns in input
   // set nr, nc
   __host__ __device__
   transidx(int _nr, int _nc): nr(_nr), nc(_nc) {};
   // element i in input should map to which element in output?
   __host__ __device__
   int operator()(const int i)
   {  int r = i / nc; int c = i % nc;  // row r, col c in input
      // that will be row c and col r in output, which has nr cols
      return c * nr + r;
   }
};

int main()
{  int mat[6] = {  // test data
      5, 12, 13,
      3, 4, 5};
   int nrow=2,ncol=3,n=nrow*ncol;
   thrust::device_vector<int> dmat(mat,mat+n);
   // allocate map vector
   thrust::device_vector<int> dmap(n);
   // allocate vector for output of gather
   thrust::device_vector<int> ddst(n);
   // construct map; element r of input matrix goes to s of output
   thrust::device_vector<int> seq(n);
   thrust::sequence (seq.begin(),seq.end());
   thrust::transform(seq.begin(),seq.end(),dmap.begin(),transidx(nrow,ncol));
   thrust::scatter(dmat.begin(),dmat.end(),dmap.begin(),ddst.begin());
   // ddst should now hold the transposed matrix, 5,3,12,4,13,5; check it:
   thrust::copy(ddst.begin(), ddst.end(), std::ostream_iterator<int>(std::cout, " "));
   std::cout << "\n";
}
