How to perform a DeviceSelect on indexes?

[dennisorlando]

Hi, I want to perform a DeviceSelect on an array based on the index and not the value of every element but I have no idea how.
E.g. I would like to define an "is_even" function such that given int A[] = {13,4,7,8,9,10, 11} produces [13, 7,9, 11] (even indexes) and not [4,7,10] (even numbers)

[gonidelis]

DeviceSelect::FlaggedIf is your friend.
https://nvidia.github.io/cccl/cub/api/structcub_1_1DeviceSelect.html#_CPPv4I00000EN3cub12DeviceSelect9FlaggedIfE11cudaError_tPvR6size_t14InputIteratorT12FlagIterator15OutputIteratorT20NumSelectedIteratorTi8SelectOp12cudaStream_t

[gonidelis]

Actually you don't need the whole FlaggedIf, you can just use DeviceSelect::Flagged(). See reproducer below :))

[gonidelis]

#include <cub/device/device_select.cuh>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>


thrust::device_vector<int> create_flags(int num_elements) {
    thrust::device_vector<int> vec(num_elements);
    for(int i = 0; i < num_elements; ++i) {
        vec[i] = !(i % 2); // even indexes
    }
    return vec;
}

int main()
{
    thrust::device_vector<int> d_in    = {13,4,7,8,9,10,11};
    int num_items            = d_in.size();
    thrust::device_vector<int> d_flags = create_flags(num_items);
    thrust::device_vector<int> d_out(num_items / 2 + 1);  // output cannot be more than half the input size + 1 
    thrust::device_vector<int> d_num_selected_out(num_items);

    // Determine temporary device storage requirements
    void* d_temp_storage      = nullptr;
    size_t temp_storage_bytes = 0;
    cub::DeviceSelect::Flagged(
    d_temp_storage,
    temp_storage_bytes,
    d_in.begin(),
    d_flags.begin(),
    d_out.begin(),
    d_num_selected_out.data(),
    num_items);

    cudaMalloc(&d_temp_storage, temp_storage_bytes);

    cub::DeviceSelect::Flagged(
    d_temp_storage,
    temp_storage_bytes,
    d_in.begin(),
    d_flags.begin(),
    d_out.begin(),
    d_num_selected_out.data(),
    num_items);

    thrust::host_vector<int> h_vec = d_out;

    for(int i = 0; i < h_vec.size(); ++i) {
        printf("%d ", h_vec[i]);
    }
}

[gonidelis]

So in general think of the flags vector as a selector. If you can form it in the way you desire to "pick on items" from the input vector you are golden. Also don't get confused with the FlaggedIf version as it extends the functionality to enable the user to perform further operations on top of your flags vector before applying it on the input vector (seems like an overkill in your case). I hope I was clear enough. If not, don't hesitate asking.

[dennisorlando]

Hi! thank you very much for the help :)
I was thinking about your solution but I would need to create a bit vector each time I would want to perform this operation, which kind of defeats the purpose of thrust as a parallelization library.
Anyways, to add more details: I'm working on an existing algorithm and I'm trying to optimize it; the algorithm is in itself very fast which means I'm (sadly) unable to add pretty much anything without increasing the program's latency.

In particular, I have a vector A of size S split into N chunks and another array which tells me how many elements are there in each chunk. As an example, with S=9 and N=3 I would have something like:

data = [1,2, nil, 7, nil, nil, 7, 8, 10 ];
n_data = [2, 1, 3];

....which I then want to compact ("push everything left") into

out = [1, 2, 7, 7, 8, 10, nil, nil, nil]

using whatever operation / library works; I thought a DeviceSelect would fit best. The current implementation relies on cudaMemcpies but I was checking whether or not cub/thrust would perorm better

[pauleonix]

@dennisorlando The trick is using fancy iterators, in particular counting_iterator. Both Thrust and CUB provide one and you can use them interchangeably (the Thrust versions of fancy iterators provide factory functions which avoid some boilerplate but that mostly matters for more complicated ones like transform_iterator with a device lambda).

[pauleonix]

@dennisorlando I.e. I would propose using

#include <cuda/functional>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>

// ...

auto is_even_idx_iter = thrust::make_transform_iterator(
    thrust::make_counting_iterator(0),
    cuda::proclaim_return_type<bool>([] __device__ (int idx) {
        return idx % 2 == 0;
    })); 

instead of d_flags.begin().

[pauleonix]

@dennisorlando These fancy iterators are the non-rangeified precursors of the std::views::iota view factory and the std::views::transform view adaptor in C++20.