CUB select if with returned indexes

Question

I have recently been running into performance issues when using the Thrust library. These come from thrust allocating memory in the base of a large nested loop structure. This is obviously unwanted, with ideal execution using a pre-allocated slab of global memory. I would like to remove or improve the offending code through one of three ways:

1. Implementing a custom thrust memory allocator
2. Replacing the thrust code with CUB code (with pre-allocated temp storage)
3. Write a custom kernel to do what I want

Although the third option would be my normal preferred choice, the operation that I want to perform is a copy_if/select_if type operation where both the data and indexes are returned. Writing a custom kernel would likely be reinventing the wheel and so I would prefer to go with one of the other two options.

I have been hearing great things about CUB, and so I see this as an ideal chance to use it in anger. What I would like to know is:

How would one implement a CUB select_if with returned indexes?

Can this be done with an ArgIndexInputIterator and a functor like so?

struct GreaterThan
{
    int compare;

    __host__ __device__ __forceinline__
    GreaterThan(int compare) : compare(compare) {}

    __host__ __device__ __forceinline__
    bool operator()(const cub::ArgIndexInputIterator<int> &a) const {
        return (a.value > compare);
    }
};

with the following in the main body of the code:

//d_in = device int array
//d_temp_storage = some preallocated block


int threshold_value;
GreaterThan select_op(threshold_value);

cub::ArgIndexInputIterator<int> input_itr(d_in);
cub::ArgIndexInputIterator<int> output_itr(d_out); //????


CubDebugExit(DeviceSelect::If(d_temp_storage, temp_storage_bytes, input_itr, output_itr, d_num_selected, num_items, select_op));

Will this try and do any memory allocation under the hood?

EDIT:

So going off Robert Crovella's comment, the functor should take the product of dereferencing a cub::ArgIndexInputIterator<int>, which should be a cub::ItemOffsetPair<int> making the functor now:

struct GreaterThan
{
    int compare;

    __host__ __device__ __forceinline__
    GreaterThan(int compare) : compare(compare) {}

    __host__ __device__ __forceinline__
    bool operator()(const cub::ItemOffsetPair<int,int> &a) const {
        return (a.value > compare);
    }
};

and in the code, d_out should be a device array of cub::ItemOffsetPair<int,int>:

//d_in = device int array
//d_temp_storage = some preallocated block

cub::ItemOffsetPair<int,int> * d_out;
//allocate d_out

int threshold_value;
GreaterThan select_op(threshold_value);

cub::ArgIndexInputIterator<int,int> input_itr(d_in);
CubDebugExit(DeviceSelect::If(d_temp_storage, temp_storage_bytes, input_itr, d_out, d_num_selected, num_items, select_op));

Answer 1

After some fiddling and asking around, I was able to get a simple code along the lines of what you suggest working:

$ cat t348.cu
#include <cub/cub.cuh>
#include <stdio.h>
#define DSIZE 6

struct GreaterThan
{

    __host__ __device__ __forceinline__
    bool operator()(const cub::ItemOffsetPair<int, ptrdiff_t> &a) const {
        return (a.value > DSIZE/2);
    }
};

int main(){

  int num_items = DSIZE;
  int *d_in;
  cub::ItemOffsetPair<int,ptrdiff_t> * d_out;
  int *d_num_selected;
  int *d_temp_storage = NULL;
  size_t temp_storage_bytes = 0;

  cudaMalloc((void **)&d_in, num_items*sizeof(int));
  cudaMalloc((void **)&d_num_selected, sizeof(int));
  cudaMalloc((void **)&d_out, num_items*sizeof(cub::ItemOffsetPair<int,ptrdiff_t>));

  int h_in[DSIZE] = {5, 4, 3, 2, 1, 0};
  cudaMemcpy(d_in, h_in, num_items*sizeof(int), cudaMemcpyHostToDevice);

  cub::ArgIndexInputIterator<int *> input_itr(d_in);


  cub::DeviceSelect::If(d_temp_storage, temp_storage_bytes, input_itr, d_out, d_num_selected, num_items, GreaterThan());

  cudaMalloc(&d_temp_storage, temp_storage_bytes);

  cub::DeviceSelect::If(d_temp_storage, temp_storage_bytes, input_itr, d_out, d_num_selected, num_items, GreaterThan());
  int h_num_selected = 0;
  cudaMemcpy(&h_num_selected, d_num_selected, sizeof(int), cudaMemcpyDeviceToHost);
  cub::ItemOffsetPair<int, ptrdiff_t> h_out[h_num_selected];
  cudaMemcpy(h_out, d_out, h_num_selected*sizeof(cub::ItemOffsetPair<int, ptrdiff_t>), cudaMemcpyDeviceToHost);
  for (int i =0 ; i < h_num_selected; i++)
    printf("index: %d, offset: %d, value: %d\n", i, h_out[i].offset, h_out[i].value);

  return 0;
}
$ nvcc -arch=sm_20 -o t348 t348.cu
$ ./t348
index: 0, offset: 0, value: 5
index: 1, offset: 1, value: 4
$

RHEL 6.2, cub v1.2.2, CUDA 5.5