Sorting and scan a large number of arrays of varying sizes

[wanghongyu2001]

I want to use the CUB library for sorting, with different blocks handling arrays of different sizes. However, I need to specify the parameters of cub::BlockRadixSort at compile time. But I can only determine the sizes of these arrays at runtime, as each block handles arrays of different sizes. So, it seems that I cannot use CUB in this case. Is there any solution? Or should I use batch Thrust sort instead?

For example, if I need to sort and compute prefix sums for hundreds of thousands of arrays, with sizes ranging from 0 to several hundred thousand, is there any way to efficiently perform these operations on the GPU using the CUB library?

Any suggestions are helpful.

[elstehle]

Thank you for reaching out with your question, @wanghongyu2001.

For example, if I need to sort and compute prefix sums for hundreds of thousands of arrays, with sizes ranging from 0 to several hundred thousand, is there any way to efficiently perform these operations on the GPU using the CUB library?

If your arrays reside at different offsets within the same memory allocation, you can use cub::DeviceSegmentedSort.

Here's an godbolt example that you can play around with.

Usage example

#include <cub/cub.cuh>

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

int main(int argc, char **argv)
{
    using data_t = int;
    // Sample data
    int  num_items = 7;
    int  num_segments = 2;
    std::vector<int> offsets = {0, 3, 7};
    std::vector<int> keys_in = {8, 6, 7, 5, 3, 0, 9};
    
    // Prepare device-side data
    thrust::device_vector<data_t> d_offsets{offsets};
    thrust::device_vector<data_t> d_keys_in{keys_in};
    thrust::device_vector<data_t> d_keys_out{keys_in.size()};

    data_t *d_keys_in_ptr = thrust::raw_pointer_cast(d_keys_in.data());
    data_t *d_keys_out_ptr = thrust::raw_pointer_cast(d_keys_out.data());

    // Determine temporary device storage requirements
    void     *d_temp_storage = NULL;
    size_t   temp_storage_bytes = 0;
    cub::DeviceSegmentedSort::SortKeys(
        d_temp_storage, temp_storage_bytes,
        d_keys_in_ptr, d_keys_out_ptr, 
        num_items, num_segments, d_offsets.cbegin(), d_offsets.cbegin() + 1);

    // Allocate temporary storage
    cudaMalloc(&d_temp_storage, temp_storage_bytes);

    // Run sorting operation
    cub::DeviceSegmentedSort::SortKeys(
        d_temp_storage, temp_storage_bytes,
        d_keys_in_ptr, d_keys_out_ptr, 
        num_items, num_segments, d_offsets.cbegin(), d_offsets.cbegin() + 1);
    
    // Get results & print
    auto next_segment_it = offsets.cbegin();
    thrust::host_vector<data_t> h_sorted_out{d_keys_out};
    for(std::size_t i=0; i < h_sorted_out.size(); i++){
        if(i == *next_segment_it){
            std::cout << "\n --- Segment beginning at offset " << *next_segment_it << " --- \n";
            next_segment_it++;
        }
        std::cout << i << ": " << h_sorted_out[i] << "\n";
    }
    return 0;
}

Output:

 --- Segment beginning at offset 0 --- 
0: 6
1: 7
2: 8

 --- Segment beginning at offset 3 --- 
3: 0
4: 3
5: 5
6: 9

Since you also asked about prefix sum, could you elaborate what your to-be-prefix-summed input looks like?

[wanghongyu2001]

Thank you.

It's not like that; each segment is independent and does not affect others. The returned value is just the index of this segment, so it has high parallelism.
E. g.

0, 3, 8                 <= offset
0  1  2 | 3  4  5  6  7 <= indexes
0, 4, 0 | 1, 4, 2, 2, 1 <= input
0, 1, 0 | 0, 2, 1, 1, 0 <= output

If using cub::DeviceSelect::UniqueByKey, does it not support segments?

[elstehle]

I am not sure whether the output of the second segment in your example would is correct. Wouldn't it be:

0, 3, 8                 <= offset
0  1  2 | 3  4  5  6  7 <= indexes
0, 4, 0 | 1, 4, 2, 2, 1 <= input
0, 1, 0 | 0, 1, 2, 2, 0 <= output

If using cub::DeviceSelect::UniqueByKey, does it not support segments?

No, unfortunately, cub::DeviceSelect::UniqueByKey does not support a segmented version of the algorithm yet.

[wanghongyu2001]

Yes, you are correct. I made a mistake.

0, 3, 8                 <= offset
0  1  2 | 3  4  5  6  7 <= indexes
0, 4, 0 | 1, 4, 2, 2, 1 <= input
0, 1, 0 | 0, 2, 1, 1, 0 <= output

So, is there an appropriate solution for this?

Thank you very much for your assistance.

[elstehle]

Aha, if I understand correctly, what you are trying to implement is a segmented order-preserving dictionary encoding?

There is an option to use cub::DeviceSelect::UniqueByKey but it will require a few extra steps, thrust::zip_iterators during cub::DeviceUniqueByKey, etc. For simplicity, I am generating the absolute indexes into a "global" dictionary (similar to my original output suggestion example). It will only require some extra work to replace it with relative dictionary offsets. Here's a quick rundown of the data transformation steps:

0, 3, 8                 <= offset
0  1  2 | 3  4  5  6  7 <= indexes
0, 4, 0 | 1, 4, 2, 2, 1 <= input
---- Sort
0, 1, 2 | 3, 4, 5, 6, 7 <= sort_indexes (indexes of sorted items)
0, 0, 4 | 1, 1, 2, 2, 4 <= post_sort_items (items after segmented sort)
0, 2, 1 | 3, 7, 5, 6, 4 <= original_item_indexes (indexes after segmented sort)
---- Generate auxiliary data from offsets
0, 0, 0 | 1, 1, 1, 1, 1 <= array_id (generated from offsets, needed for subsequent unique_by_key)
---- UniqueByKey, will use a zip_iterator of array_id and post_sort_items
0, 2 | 3, 5, 7 <= unique_sort_indexes (sort_indexes of the unique items)
0, 1 | 2, 3, 4 <= dictionary_entry
---- Scatter: for i=0..total_unique_items: prepared_output[unique_sort_indexes[i]] = dictionary_entry[i];
0  1  2 | 3  4  5  6  7 <= indexes
0, -, 1 | 2, -, 3, -, 4 <= prepared_output (within the sorted order, what's the index of the )
---- Prefix-scan that fills in the blanks, sweeping the latest non-blank item to the right
0, 0, 1 | 2, 2, 3, 3, 4 <= filled_prepared_output
---- Now we have this:
0, 0, 1 | 2, 2, 3, 3, 4 <= filled_prepared_output
0, 2, 1 | 3, 7, 5, 6, 4 <= original_item_indexes (from the indexes after segmented sort)
---- Then Scatter back: for i=0..len: output[original_item_indexes[i]] = filled_prepared_output[i]
0  1  2 | 3  4  5  6  7 <= indexes
0, 1, 0 | 2, 4, 3, 3, 2 <=  output

As you can see, it requires a couple of cub algorithms. If this is an approach you'd be interested in, let me know and I can see if I can support in translating this to the respective cub algorithm invocations. However, I am not sure when I will get to that.

[wanghongyu2001]

Yes, this is exactly what I want. Thank you!

For convenience, let me describe the problem again. Please don't feel discouraged by my verbosity.

I have hundreds of thousands of segments, each with a size ranging from 0 to several hundred thousand. Each segment is independent and not related to others. They are only adjacent in storage. So they have high parallelism.

E.g.

0, 3, 8                 <= offset
0  1  2 | 3  4  5  6  7 <= indexes
0, 4, 0 | 1, 4, 2, 2, 1 <= input
0, 1, 0 | 0, 2, 1, 1, 0 <= output