Thrust: selectively copy based on another vector

Question

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

Answer 1

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}

Answer 2

Although Abator had given the right function to use. Let me try a complete example.

//~~~START:Wed, 06-Oct-2021, 21:41:22 IST
//~~~Author:Rajesh Pandian M | mrprajesh.co.in
//~~CompiledWith: nvcc a.cu -std=c++14 --expt-extended-lambda

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/copy.h>
#include <stdio.h>

int main(void) {
  const int N=5;
  int A[]={2, 3, 6, 0, 11}; //Data
  int B[]={1, 0, 1, 1, 0 }; //Stencil
  thrust::device_vector<int> dA(A, A + N);
  thrust::device_vector<int> dB(B, B + N);

  // Allocates memory based on #number of 1's
  thrust::device_vector<int> dC(thrust::count(B, B+N,1));

  //Condition on the stencil elements. If 1 is seen copy, else do not!
  thrust::copy_if( dA.begin()
                  ,dA.end()
                  ,dB.begin()
                  ,dC.begin()
                  ,[] __host__ __device__ (const int& x){
                    return 1 == x;
                  });
  //Prints
  thrust::for_each(dC.begin(), dC.end(),
                  [] __host__ __device__(const int& x){ printf("%d ",x);});
  return 0;
}