1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
|
#include <thrust/device_vector.h>
#include <thrust/functional.h>
#include <thrust/generate.h>
#include <thrust/host_vector.h>
#include <thrust/random.h>
#include <thrust/reduce.h>
#include <iostream>
#include "include/host_device.h"
// convert a linear index to a row index
template <typename T>
struct linear_index_to_row_index : public thrust::unary_function<T, T>
{
T C; // number of columns
__host__ __device__ linear_index_to_row_index(T C)
: C(C)
{}
__host__ __device__ T operator()(T i)
{
return i / C;
}
};
int main()
{
int R = 5; // number of rows
int C = 8; // number of columns
thrust::default_random_engine rng;
thrust::uniform_int_distribution<int> dist(10, 99);
// initialize data
thrust::device_vector<int> array(R * C);
for (size_t i = 0; i < array.size(); i++)
{
array[i] = dist(rng);
}
// allocate storage for row sums and indices
thrust::device_vector<int> row_sums(R);
thrust::device_vector<int> row_indices(R);
// compute row sums by summing values with equal row indices
thrust::reduce_by_key(
thrust::make_transform_iterator(thrust::counting_iterator<int>(0), linear_index_to_row_index<int>(C)),
thrust::make_transform_iterator(thrust::counting_iterator<int>(0), linear_index_to_row_index<int>(C)) + (R * C),
array.begin(),
row_indices.begin(),
row_sums.begin(),
thrust::equal_to<int>(),
thrust::plus<int>());
// print data
for (int i = 0; i < R; i++)
{
std::cout << "[ ";
for (int j = 0; j < C; j++)
{
std::cout << array[i * C + j] << " ";
}
std::cout << "] = " << row_sums[i] << "\n";
}
return 0;
}
|
