Search code examples
cudathrust

Finding the first index of every distinct value in CUDA array


Assume we have an array like this:

0, 0, 0, 1, 2, 2, 2, 3, 3, 4, ...

I would like to have the index of every first occurrence of every value, so in this example [0, 3, 4, 7, 9]. The array is sorted and all possible values are known and consecutive.

Possible solutions I have is using a kernel for every element in this array and use an atomicmin to save the lowest index. But I assume a better approach is possible.


Solution

  • You can do this with a single call to thrust::unique_by_key() if you provide a vector of indices e.g. via thrust::sequence(). Here's a worked example:

    $ cat t3.cu
    #include <thrust/device_vector.h>
    #include <thrust/copy.h>
    #include <thrust/unique.h>
    #include <thrust/sequence.h>
    #include <iostream>
    
    int main(){
    
      int keys[] = {0, 0, 0, 1, 2, 2, 2, 3, 3, 4};
      int ks = sizeof(keys)/sizeof(keys[0]);
      thrust::device_vector<int> d_keys(keys, keys+ks);
      thrust::device_vector<int> d_result(ks);
      thrust::sequence(d_result.begin(), d_result.end());
      int rs  = (thrust::unique_by_key(d_keys.begin(), d_keys.end(), d_result.begin())).first - d_keys.begin();
      thrust::copy_n(d_result.begin(), rs, std::ostream_iterator<int>(std::cout, ","));
      std::cout << std::endl;
    }
    
    $ nvcc -arch=sm_35 -o t3 t3.cu
    $ ./t3
    0,3,4,7,9,
    $
    

    The important activity occurring here is stream compaction and thrust provides a nice set of routines for various use-cases. For example this operation could also be done with thrust::unique_copy() and in that case, with some additional code complexity, you could eliminate the need for the thrust::sequence() call (it would be replaced by a thrust::counting_iterator zipped together with your data, and an appropriate selection functor), but it still requires an output vector of the same length.