Search code examples
sortingcudathrustcub

How to sort an array of CUDA vector types


Specifically how could I sort an array of float3? Such that the .x components are the primary sort criteria, the .y components are the secondary sort criteria and the .z components are the tertiary sort criteria.

Is there a simple solution that can make a single call to cub:: DeviceRadixSort or thrust::sort_by_key?

At the moment I am thinking maybe I could create a uint32 key array where the each element has the first third of its digits taken from the first third of the input array's .x components, the second third of digits taken from the first third of the input array's .y components, and the final third of its digits taken from the first third of the input array's .z components. Or is there a better solution?


Solution

  • Using the example that Robert Crovella suggested I have crafted the following solution. Thanks again Rob.

    #include <thrust/sort.h>
    #include <thrust/device_ptr.h>
    
    struct sort_float3 {
        __host__ __device__
        bool operator()(const float3 &a, const float3 &b) const {
    
        if      (a.x <= b.x && a.y <= b.y && a.z < b.z) return true;
        else if (a.x <= b.x && a.y < b.y) return true;
        else if (a.x < b.x) return true;
        else return false;
        }
    };
    
    int main(void)
    {
        float3 *h_array;
        // Define your host array
        float3 *d_array;
        cudaMallocHost( (void**)&d_array,
                        number_of_elements * sizeof(float3) );      
        cudaMemcpy( d_array,
                    h_array, 
                    number_of_elements * sizeof(float3),
                    cudaMemcpyHostToDevice );
    
        thrust::device_ptr<float3> th_array( d_array );
        thrust::sort( th_array, 
                      th_array+number_of_elements , 
                      sort_float3() );
        return 0;
    }