Is there any way to use cub for multi-keyword sorting?

[0x5459]

Hi, Is there any way to use cub for multi-keyword sorting? Or is there any plan to support multi-keyword sorting?

Pseudo code:

// Declare, allocate, and initialize device-accessible pointers for sorting data
int  num_items;          // e.g., 7
int  *d_keys_in1;        // e.g., [8, 8, 7, 5, 3, 3, 9]
int  *d_keys_in2;        // e.g., [1, 2, 0, 0, 4, 3, 0]
int  *d_keys_out;        // e.g., [        ...        ]
int  *d_values_in;       // e.g., [0, 1, 2, 3, 4, 5, 6]
int  *d_values_out;      // e.g., [        ...        ]

// Run sorting operation
cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, 
    d_keys_in1, d_keys_in2, d_keys_out, d_values_in, d_values_out, num_items);

// d_keys_out                 <-- [3, 3, 5, 7, 8, 8, 9]
// d_values_out               <-- [5, 4, 3, 2, 0, 1, 6]

[pauleonix]

You can zip the two keys together using thrust::zip_iterator and sort by them as tuples. But while this in theory should be radix-sortable, you might have to put it through a thrust::transform_iterator to get a single 64b type that is radix sortable. Or you use merge-sort.

[pauleonix]

Seems like the supported way of radix-sorting custom types is a "decomposer object" according to the docs. Not sure how that is done though.

[pauleonix]

See cub/examples/device/example_device_radix_sort_custom.cu for an example using the decomposer object.

[elstehle]

@pauleonix beat me to it 🚀

Seems like the supported way of radix-sorting custom types is a "decomposer object" according to the docs. Not sure how that is done though.

This is indeed the preferred way of sorting composed sort keys.

I've put together a code sample on how you can use the decomposer object to implement your use case. Let me know if you have any questions about it.

Link to code sample:
https://godbolt.org/z/5dGexEhYo

Code Sample

#include <cstdint>
#include <cuda/std/tuple>
#include <cub/device/device_radix_sort.cuh>
#include <thrust/copy.h>
#include <thrust/sort.h>
#include <thrust/tuple.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/iterator/zip_iterator.h>

struct decomposer2_t
{
    template<typename key_t, typename key2_t>
  __host__ __device__ ::cuda::std::tuple<key_t&, key2_t&>
  operator()(thrust::tuple<key_t, key2_t> &key) const
  {
    return {thrust::get<0>(key), thrust::get<1>(key)};
  }
};

int main(int argc, char **argv)
{
    using key1_t = int;
    using key2_t = unsigned int;
    using value_t = int;

    // Sample data
    std::vector<key1_t> keys1 = {8, 8, 7, 5, 3, 3, 9};
    std::vector<key2_t> keys2 = {1, 2, 0, 0, 4, 3, 0};
    std::vector<value_t> values = {0, 1, 2, 3, 4, 5, 6};
    
    // Run exclusive scan using composition as custom scan operator
    std::size_t num_items = keys1.size();
    thrust::device_vector<key1_t> d_keys1{keys1};
    thrust::device_vector<key2_t> d_keys2{keys2};
    thrust::device_vector<value_t> d_values_in{values};
    thrust::device_vector<value_t> d_values_out{values};

    // Temporarily, these are the sort-keys input and output arrays used within the radix sort
    thrust::device_vector<thrust::tuple<key1_t, key2_t>> d_zipped_sort_keys_in{num_items};
    thrust::device_vector<thrust::tuple<key1_t, key2_t>> d_zipped_sort_keys_out{num_items};

    // Compose the sort keys from two decomposed arrays: (key1_t*, key2_t*) -> tuple<key1_t,key2_t>*
    auto zipped_sort_keys_it = thrust::make_zip_iterator(d_keys1.cbegin(), d_keys2.cbegin());
    thrust::copy(zipped_sort_keys_it, zipped_sort_keys_it + num_items, d_zipped_sort_keys_in.begin());

    std::size_t temp_storage_bytes{};
    cub::DeviceRadixSort::SortPairs(
      nullptr,
      temp_storage_bytes,
      thrust::raw_pointer_cast(d_zipped_sort_keys_in.data()),
      thrust::raw_pointer_cast(d_zipped_sort_keys_out.data()),
      thrust::raw_pointer_cast(d_values_in.data()),
      thrust::raw_pointer_cast(d_values_out.data()),
      num_items,
      decomposer2_t{});

    void *temp_storage{};
    cudaMalloc(&temp_storage, temp_storage_bytes);

    cub::DeviceRadixSort::SortPairs(
      temp_storage,
      temp_storage_bytes,
      thrust::raw_pointer_cast(d_zipped_sort_keys_in.data()),
      thrust::raw_pointer_cast(d_zipped_sort_keys_out.data()),
      thrust::raw_pointer_cast(d_values_in.data()),
      thrust::raw_pointer_cast(d_values_out.data()),
      num_items,
      decomposer2_t{});

    // Get results & print
    thrust::host_vector<thrust::tuple<key1_t, key2_t>> h_keys{d_zipped_sort_keys_out};
    thrust::host_vector<value_t> h_values{d_values_out};
    std::cout << "d_keys   <-- [";
    for(std::size_t i = 0; i<h_keys.size(); i++){
        std::cout << "(" << thrust::get<0>(h_keys[i]) << ", " << thrust::get<1>(h_keys[i]) << ") " << ((i<h_keys.size()-1) ? ", " : "]\n");
    }
    std::cout << "d_values <-- [";
    for(std::size_t i = 0; i<h_values.size(); i++){
        std::cout << h_values[i] << ((i<h_values.size()-1) ? ", " : "]\n");
    }

    return 0;
}

[elstehle]

No problem, @0x5459, we are here to help 😉

And you are right, DeviceRadixSort is one of the very few algorithms that does not work with fancy iterators like the zip_iterator directly but works with pointers. Hence, we need to materialize the zipped sort keys once before we're sorting them using thrust::copy and pass a pointer to the materialized zipped sort keys, d_zipped_sort_keys_in, to cub::DeviceRadixSort.

auto zipped_sort_keys_it = thrust::make_zip_iterator(d_keys1.cbegin(), d_keys2.cbegin());
    thrust::copy(zipped_sort_keys_it, zipped_sort_keys_it + num_items, d_zipped_sort_keys_in.begin());
 ...
 cub::DeviceRadixSort::SortPairs(
      temp_storage,
      temp_storage_bytes,
      thrust::raw_pointer_cast(d_zipped_sort_keys_in.data()),
      ...

This is a one time copy required for cub::DeviceRadixSort. Other algorithms, like cub::DeviceScan can work with arbitrary iterators and therefore you could pass zipped_sort_keys_it as the input (d_in) without requiring extra materialization before invoking the algorithm.

[pauleonix]

Sorry for missing that. Maybe measure if DeviceRadixSort with that copy is still cheaper than DeviceMergeSort which takes a general iterator like DeviceScan.

[kevkrist]

Forgive me if this is a silly question, but I don't understand exactly why DeviceRadixSort couldn't tolerate, say, a transform iterator, whose dereferenced type is a primitive type like uint64_t (the use case I have in mind is hashing some keys). Because you'd just BlockLoad the iterator and only operate on the uint64_t values. Am I missing something? Also, it doesn't seem to me to be simple to hack a fork of cub to make this possible, if it is possible. Please correct me if I'm wrong (I'd find it very useful).

[elstehle]

Yes, it is possible to make DeviceRadixSort work with arbitrary iterators and it will likely be implemented once we make progress on adding support for efficient partial sorting.

There's one caveat with DeviceRadixSort that makes dealing with arbitrary iterators a bit more intricate. DeviceRadixSort is a multi-pass algorithm. That is, it needs to repeatedly process its input to bring it into sorted order. This is achieved using double-buffering. If we assume, say, a 32-bit radix key, 8-bit radix digits, and, hence, four sorting passes, the radix sorting passes would look as follows:

1. Read the input from the user-provided iterator, write it to the temporary buffer
2. Read the input from the temporary buffer, write it to the user-provided output buffer
3. Read the input from user-provided output buffer , write it to the temporary buffer
4. Read the input from the temporary buffer, write it to the user-provided output buffer

Now, if the user-provided iterator is of a different type than the temporary buffer that results in two kernel template instantiations: (1) one that has the user-provided iterator as input iterator and (2) one that has the temporary buffer as input iterator (i.e., a raw pointer). This results in increased compile time. This, for instance, isn't the case for (almost) any other CUB algorithms, which are single-pass.

If the user provides an output iterator instead of an output buffer, this would obviously further complicate the situation, requiring further temporary storage buffer allocations.

[kevkrist]

Great, thanks for the clarification @elstehle !