Search code examples
cudathrust

Thrust: selectively copy based on another vector


I'd like to use a set of thrust operations to selectively copy the elements of one vector A into a new vector B based on a predicate on elements in a third vector C.

Here's an example case: I want to copy elements (in order) from A when the corresponding element in B is 1 to C and don't if it is 0. I want |C| < |A| if there are 0s in B. We can pre-determine the size of C by a reduction on B. e.g:

A = [2, 3, 6, 0, 11]

B = [1, 0, 1, 1, 0]

C = [2, 6, 0]

Any help is greatly appreciated


Solution

  • This algorithm is known as stream compaction. It is implemented in thrust::copy_if.

    The following example is taken from the Thrust documentation.

    #include <thrust/copy.h>
    ...
    struct is_even
    {
      __host__ __device__
      bool operator()(const int x)
      {
        return (x % 2) == 0;
      }
    };
    ...
    int N = 6;
    int data[N]    = { 0, 1,  2, 3, 4, 5};
    int stencil[N] = {-2, 0, -1, 0, 1, 2};
    int result[4];
    thrust::copy_if(data, data + N, stencil, result, is_even());
    // data remains    = { 0, 1,  2, 3, 4, 5};
    // stencil remains = {-2, 0, -1, 0, 1, 2};
    // result is now     { 0, 1,  3, 5}