Thrust: selectively move elements to another vector

Question

I'm trying to figure out the best way to do the following using Thrust: vector A has a million floats, they have some particular order. I want to move to vector B every element x in A for which x>7.0 such that the order of elements is maintain in both vectors A and B. Importantly, only a tiny fraction of elements need be moved. Efficiency is more important for my code than elegance.

My idea was to use thrust::copy_if from A to B and then thrust::remove_if on A. But I don't know the exact number of elements to be copy, and since apparently the memory for B must be allocated in advance, another counting operation is necessary. An inelegant way to skip the counting operation is to pre-allocate "enough" memory for vector B.

Using thrust::remove_copy_if has much the same problems: you need to allocate memory for B in advance, and also it doesn't actually remove anything from A so another thrust::remove_if is required anyway.

Another idea I had was to use thrust::stable_sort with some custom-made comparison functor, to push all elements I want out to the end of A, and then somehow figure out how many there are and thrust::copy them to B. This also looks pretty inelegant...

Answer 1

You're on the right track with thrust::copy_if. Just allocate two more buffers of the same size as the first one. Then copy_if > 7.0f to the first one and copy_if <= 7.0f to the second one. Allocating buffers of the same size as the original buffer is fine as long as you know there's room, and 1 million floats only takes up 4MB.

Edit:

I did a performance comparison of the copy_if and stable_partition approaches. On my card, a GTX660, stable_partition took around 150% as long as copy_if for "split" values of 0.1f, 0.5f and 0.9f. I added tests to ensure that both methods are stable (maintain the order of the values).

#include <cuda.h>
#include <curand.h>

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

#include <iostream>
#include <cassert>

#define CHECK_CUDA_CALL(x) do { if((x)!=cudaSuccess) { \
    printf("Error at %s:%d\n",__FILE__,__LINE__);\
    return EXIT_FAILURE;}} while(0)


#define CHECK_CURAND_CALL(x) do { if((x)!=CURAND_STATUS_SUCCESS) { \
    printf("Error at %s:%d\n",__FILE__,__LINE__);\
    return EXIT_FAILURE;}} while(0)


#define SPLIT  0.1f

struct is_low
{
  __host__ __device__ bool operator()(const float x)
  {
    return x <= SPLIT;
  }
};


struct is_high
{
  __host__ __device__ bool operator()(const float x)
  {
    return x > SPLIT;
  }
};


class EventTimer {
public:
  EventTimer() : mStarted(false), mStopped(false) {
    cudaEventCreate(&mStart);
    cudaEventCreate(&mStop);
  }
  ~EventTimer() {
    cudaEventDestroy(mStart);
    cudaEventDestroy(mStop);
  }
  void start(cudaStream_t s = 0) {
    cudaEventRecord(mStart, s); 
    mStarted = true;
    mStopped = false;
  }
  void stop(cudaStream_t s = 0)  {
    assert(mStarted);
    cudaEventRecord(mStop, s); 
    mStarted = false;
    mStopped = true;
  }
  float elapsed() {
    assert(mStopped);
    if (!mStopped) return 0; 
    cudaEventSynchronize(mStop);
    float elapsed = 0;
    cudaEventElapsedTime(&elapsed, mStart, mStop);
    return elapsed;
  }

private:
  bool mStarted, mStopped;
  cudaEvent_t mStart, mStop;
};


int main(int argc, char *argv[])
{
    const size_t n = 1024 * 1024 * 50;

    // Create prng
    curandGenerator_t gen;
    CHECK_CURAND_CALL(curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT));

    // Set seed
    CHECK_CURAND_CALL(curandSetPseudoRandomGeneratorSeed(gen, 1234ULL));

    // Generate n floats on device 
    thrust::device_vector<float> vec_rnd_d(n);
    float* ptr_rnd_d = thrust::raw_pointer_cast(vec_rnd_d.data());
    CHECK_CURAND_CALL(curandGenerateUniform(gen, ptr_rnd_d, n));

    thrust::device_vector<float> vec_low_d(n);
    thrust::device_vector<float> vec_high_d(n);

    for (int i = 0; i < 5; ++i) {
      EventTimer timer;
      timer.start();
      thrust::device_vector<float>::iterator iter_end;
      iter_end = thrust::copy_if(vec_rnd_d.begin(), vec_rnd_d.end(), vec_low_d.begin(), is_low());
      thrust::copy_if(vec_rnd_d.begin(), vec_rnd_d.end(), vec_high_d.begin(), is_high());
      timer.stop();
      std::cout << "copy_if: " << timer.elapsed() << "ms" << std::endl;

      // check result
      thrust::host_vector<float> vec_rnd_h = vec_rnd_d;
      thrust::host_vector<float> vec_low_h = vec_low_d;
      thrust::host_vector<float> vec_high_h = vec_high_d;
      thrust::host_vector<float>::iterator low_iter_h = vec_low_h.begin();
      thrust::host_vector<float>::iterator high_iter_h = vec_high_h.begin();
      for (thrust::host_vector<float>::iterator rnd_iter_h = vec_rnd_h.begin();
          rnd_iter_h != vec_rnd_h.end(); ++rnd_iter_h) {
        if (*rnd_iter_h <= SPLIT) {
          assert(*low_iter_h == *rnd_iter_h);
          ++low_iter_h;
        }
        else {
          assert(*high_iter_h == *rnd_iter_h);
          ++high_iter_h;
        }
      }
    }

    for (int i = 0; i < 5; ++i) {
      thrust::device_vector<float> vec_rnd_copy = vec_rnd_d;
      EventTimer timer;
      timer.start();
      thrust::device_vector<float>::iterator iter_split = 
        thrust::stable_partition(vec_rnd_copy.begin(), vec_rnd_copy.end(), is_low());
      timer.stop();
      size_t n_low = iter_split - vec_rnd_copy.begin();
      std::cout << "stable_partition: " << timer.elapsed() << "ms" << std::endl;

      // check result
      thrust::host_vector<float> vec_rnd_h = vec_rnd_d;
      thrust::host_vector<float> vec_partitioned_h = vec_rnd_copy;
      thrust::host_vector<float>::iterator low_iter_h = vec_partitioned_h.begin();
      thrust::host_vector<float>::iterator high_iter_h = vec_partitioned_h.begin() + n_low;
      for (thrust::host_vector<float>::iterator rnd_iter_h = vec_rnd_h.begin();
          rnd_iter_h != vec_rnd_h.end(); ++rnd_iter_h) {
        if (*rnd_iter_h <= SPLIT) {
          assert(*low_iter_h == *rnd_iter_h);
          ++low_iter_h;
        }
        else {
          assert(*high_iter_h == *rnd_iter_h);
          ++high_iter_h;
        }
      }
    }

    CHECK_CURAND_CALL(curandDestroyGenerator(gen));

    return EXIT_SUCCESS;
}

Output:

C:\rd\projects\cpp\test_cuda\Release>test_cuda.exe
copy_if: 40.2919ms
copy_if: 38.0157ms
copy_if: 38.5036ms
copy_if: 37.6751ms
copy_if: 38.1054ms
stable_partition: 59.5473ms
stable_partition: 61.4016ms
stable_partition: 59.1854ms
stable_partition: 61.3195ms
stable_partition: 59.1205ms

Answer 2

To answer my own question, I finally found thrust::stable_partition, which is more efficient and elegant than all "copy_if"-alternatives. It just moves all elements that fail to satisfy a predicate to the end of the array and returns the start of the second sequence. Pointer arithmetic gives the size of B, but in fact it's not necessary anymore:

thrust::device_vector<float>::iterator iter = thrust::stable_partition(A.begin(), A.end(), pred)
thrust::device_vector<float> B(iter, A.end())
A.erase(iter, A.end());