Can thrust support device segmented sort?

[zlwu92]

Hi,

I want to know if there's support of thrust or cub library that can support segmented sort, and also allow me to pass my self-defined compare functor into it as param. Thanks.

[gevtushenko]

Hello @zlwu92!

There's segmented sort in CUB, but it doesn't allow you to pass a custom comparison operator. You could try thrust::sort or cub::DeviceMergeSort to achieve this functionality:

#include <iostream>
#include <thrust/device_vector.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/sort.h>

struct custom_compare_t {
    __device__ bool operator()(thrust::tuple<int, int> lhs, thrust::tuple<int, int> rhs) {
        const int lhs_segment = thrust::get<0>(lhs);
        const int rhs_segment = thrust::get<0>(rhs);

        if (lhs_segment != rhs_segment) {
            return lhs_segment < rhs_segment;
        }

        const int lhs_value = thrust::get<1>(lhs);
        const int rhs_value = thrust::get<1>(rhs);
        return lhs_value > rhs_value;
    }
};

int main() {
    thrust::device_vector<int> values(6);
    thrust::device_vector<int> segments(6);

    segments[0] = 0; segments[1] =  0; segments[2] =  0; segments[3] =  1; segments[4] =  1; segments[5] =  1;
    values[0] = 10;    values[1] = 40;   values[2] = 30;   values[3] = 20;   values[4] = 30;   values[5] = 10;
    
    std::cout << "before: " << std::endl;
    for (int i = 0; i < 6; i++) {
        std::cout << segments[i] << " " << values[i] << std::endl;
    }

    auto begin = thrust::make_zip_iterator(thrust::make_tuple(segments.begin(), values.begin()));
    auto end = thrust::make_zip_iterator(thrust::make_tuple(segments.end(), values.end()));

    thrust::sort(begin, end, custom_compare_t{});

    std::cout << "after: " << std::endl;
    for (int i = 0; i < 6; i++) {
        std::cout << segments[i] << " " << values[i] << std::endl;
    }

    return 0;
}

/*
before: 
0 10
0 40
0 30
1 20
1 30
1 10
after: 
0 40
0 30
0 10
1 30
1 20
1 10
*/