Summing the entries of a matrix using a stencil with Thrust

stencil_tiled_range.cu

#undef _GLIBCXX_ATOMIC_BUILTINS
#undef _GLIBCXX_USE_INT128 

#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/iterator/permutation_iterator.h>
#include <thrust/functional.h>
#include <thrust/inner_product.h>

#include <thrust/fill.h>
#include <thrust/sequence.h>
#include <thrust/host_vector.h>

// for printing
#include <thrust/copy.h>
#include <ostream>
#include <iomanip>
#include <iostream>

template <typename Iterator>
class tiled_range
{
    public:

    typedef typename thrust::iterator_difference<Iterator>::type data_type;

    struct tile_functor : public thrust::unary_function<data_type,data_type>
    {
        data_type tile_size_x;
        data_type leading_dimension;

        tile_functor(data_type tile_size_x, data_type leading_dimension)
            : tile_size_x(tile_size_x), leading_dimension(leading_dimension) {}

        __host__ __device__
        data_type operator()(const data_type& i) const
        { 
            int x=i % tile_size_x;
            int y=i / tile_size_x;
            return leading_dimension*y + x;
        }
    };

    typedef typename thrust::counting_iterator<data_type>                   CountingIterator;
    typedef typename thrust::transform_iterator<tile_functor, CountingIterator>   TransformIterator;
    typedef typename thrust::permutation_iterator<Iterator,TransformIterator>     PermutationIterator;

    // construct repeated_range for the range [first,last)
    tiled_range(Iterator first, Iterator last, data_type tile_size_x, data_type tile_size_y, data_type leading_dimension)
        : first(first), last(last), tile_size_x(tile_size_x), tile_size_y(tile_size_y), leading_dimension(leading_dimension) {}
   
    PermutationIterator begin(void) const
    {
        return PermutationIterator(first, TransformIterator(CountingIterator(0), tile_functor(tile_size_x, leading_dimension)));
    }

    PermutationIterator end(void) const
    {
        return begin() + tile_size_y * tile_size_x;
    }
    
    protected:
    Iterator first;
    Iterator last;
    data_type tile_size_x, tile_size_y, leading_dimension;
  };

// print an M-by-N array
template <typename T>
void print(size_t m, size_t n, thrust::host_vector<T>& h_data)
{
  for(size_t i = 0; i < m; i++)
  {
    for(size_t j = 0; j < n; j++)
      std::cout << std::setw(8) << h_data[i * n + j] << " ";
    std::cout << "\n";
  }
}


int main(void)
{
  int n_columns = 10;
  int n_rows = 2;
  int n_stencil_columns = 5;
  int n_stencil_rows = 2;
  int offset=2;
  
  thrust::host_vector<int> data(n_columns*n_rows);
  thrust::host_vector<int> stencil(n_stencil_columns*n_stencil_rows);

  thrust::sequence(data.begin(), data.end());
  thrust::fill(stencil.begin(), stencil.end(), 0);
  for (int stencil_row=0; stencil_row < n_stencil_rows; stencil_row++)
  {
    stencil[stencil_row*n_stencil_columns+1]=1;
    stencil[stencil_row*n_stencil_columns+2]=1;
  }

  std::cout << "The matrix:\n";
  print(n_rows, n_columns, data);
  std::cout << "\nThe stenctil:\n";
  print(n_stencil_rows, n_stencil_columns, stencil);

  typedef thrust::host_vector<int>::iterator Iterator;
  tiled_range<Iterator> tile(data.begin()+offset, data.end(), n_stencil_columns, n_stencil_rows, n_columns);
  int sum=thrust::inner_product(tile.begin(), tile.end(), stencil.begin(), 0);
  std::cout << "\nThe sum according to the stencil at offset " 
            << offset << " is " << sum << ".\n";

  return 0;
}