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;
}

[0x5459]

Thanks for your answer. The sample you gave meets my needs. However, zip will cause a large number of copies and sorting of large data structures will cause performance degradation. Will cub consider directly supporting sorting of multiple keywords?

[pauleonix]

@0x5459 What exactly do you mean? A zip iterator doesn't require a copy of the data, it is an iterator accessing the original buffers d_keys_in1 and d_keys_in2 from your example under the hood, i.e. it implements an SoA memory access pattern while giving an AoS-like interface to the algorithm.

[0x5459]

Sorry, I'm new to cuda and cccl. Now I understand. thank you for your answers. 😄

[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.