Search code examples
c++cudathrust

Count reduction using Thrust


Given some input keys and values, I am trying to count how many consecutive values with the same key exist. I will give an example to make this more clear.

Input keys: { 1, 4, 4, 4, 2, 2, 1 }

Input values: { 9, 8, 7, 6, 5, 4, 3 }

Expected output keys: { 1, 4, 2, 1 }

Expected output values: { 1, 3, 2, 1 }

I am trying to solve this problem on a GPU using CUDA. The reduction capabilities of the Thrust library seemed like a good solution for this and I got to the following:

#include <thrust/reduce.h>
#include <thrust/functional.h>

struct count_functor : public thrust::binary_function<int, int, int>
{
    __host__ __device__
        int operator()(int input, int counter)
    {
        return counter + 1;
    }
};

const int N = 7;
int A[N] = { 1, 4, 4, 4, 2, 2, 1 }; // input keys
int B[N] = { 9, 8, 7, 6, 5, 4, 3 }; // input values
int C[N];                         // output keys
int D[N];                         // output values

thrust::pair<int*, int*> new_end;
thrust::equal_to<int> binary_pred;
count_functor binary_op;
new_end = thrust::reduce_by_key(A, A + N, B, C, D, binary_pred, binary_op);
for (int i = 0; i < new_end.first - C; i++) {
    std::cout << C[i] << " - " << D[i] << "\n";
}

This code is pretty similar to an example from the Thrust documentation. However, instead of the plus operation, I am trying to count. The output from this code is the following:

1 - 9
4 - 7
2 - 5
1 - 3

However, I would expected the second column to contain the values 1, 3, 2, 1. I think the counts are off because the reduction starts with the first value it finds and does not apply the operator until it has a second value, but I am not sure this is the case.

Am I overlooking something about the reduce_by_key function that could solve this problem or should I use a completely different function to achieve what I want?


Solution

  • For your use case you don't need the values of B, the values of D are only dependent on the values of A.

    In order to count how many consecutive values are in A you can supply a thrust::constant_iterator as the input values and apply thrust::reduce_by_key:

    #include <thrust/reduce.h>
    #include <thrust/functional.h>
    #include <iostream>
    #include <thrust/iterator/constant_iterator.h>
    
    int main()
    {
    const int N = 7;
    int A[N] = { 1, 4, 4, 4, 2, 2, 1 }; 
    int C[N];
    int D[N];
    
    thrust::pair<int*, int*> new_end;
    thrust::equal_to<int> binary_pred;
    thrust::plus<int> binary_op;
    new_end = thrust::reduce_by_key(A, A + N, thrust::make_constant_iterator(1), C, D, binary_pred, binary_op);
    
    for (int i = 0; i < new_end.first - C; i++) {
        std::cout << C[i] << " - " << D[i] << "\n";
    }
    return 0;
    }
    

    output

    1 - 1
    4 - 3
    2 - 2
    1 - 1