Code Yarns ‍👨‍💻
Tech BlogPersonal Blog

Thrust: Compact Multiple Vectors Using Predicate

📅 2011-Jun-27 ⬩ ✍️ Ashwin Nanjappa ⬩ 📚 Archive

From this earlier post, we can see that the zip_iterator makes it really easy to compact multiple vectors (of same size) based on the duplicate values in one of those vectors.

Another scenario that arises frequently is the need to compact multiple vectors (of same size) based on testing the value in one of those vectors (using a predicate). We may want to remove all the elements for which the predicate is true.

For example, say I have 2 vectors. A previous kernel might have invalidated some of the values in the first vector by setting them to -1. Now, I want to compact these 2 vectors such that the elements corresponding to -1 in the first vector is removed from both vectors.

The code required to the above compaction is very similar to that in the earlier blog post:

// Many vectors
thrust::device_vector< int > vec0;
thrust::device_vector< int > vec1;

// Make zip_iterator easy to use
typedef thrust::device_vector< int >::iterator  IntDIter;
typedef thrust::tuple< IntDIter, IntDIter >     IntDIterTuple2;
typedef thrust::zip_iterator< IntDIterTuple2 >  ZipDIter;

// Remove elements in many vectors if element in vec0 is negative
ZipDIter newEnd = thrust::remove_if(    thrust::make_zip_iterator( thrust::make_tuple( vec0.begin(), vec1.begin() ) ),
                                        thrust::make_zip_iterator( thrust::make_tuple( vec0.end(), vec1.end() ) ),
                                        isTuple2Negative() );

// Erase the removed elements from the vectors
IntDIterTuple2 endTuple = newEnd.get_iterator_tuple();
vec0.erase( thrust::get<0>( endTuple ), vec0.end() );
vec1.erase( thrust::get<1>( endTuple ), vec1.end() );

The only extra work needed is to carefully write a predicate that does what we want:

// Make predicate easy to write
typedef thrust::tuple< int, int > IntTuple2;

// Predicate
struct isTuple2Negative
    __host__ __device__ bool operator() ( const IntTuple2& tup )
        const int x = thrust::get<0>( tup );
        return ( x < 0 );

That is it, the compaction works like magic! 😊

Tried with: CUDA 4.0 and Thrust 1.4.0