Search code examples
sortingcudatupleskeythrust

unexpected behavior for thrust sort_by_key


I am trying to use sort by key to sort a set of values. However I am using two keys instead of one. The values are first sorted using key 1 followed by key2. key1 is an int and key2 is a float.

My input, expected output, actual output and my code are below. The sorting doesnot seem to have worked properly for the second key as evident from the difference in the actual and expected output in the first two rows.

INPUT:

keys1, keys2, values:

1,  0.3,    1
3,  5.1,    5
3,  3.2,    3
3,  -0.08,  8
2,  2.1,    2
2,  5.2,    8
2,  1.1,    1
1,  -0.01,  1

My expected output is

keys1, keys2, values:

1,  -0.01,  1
1,  0.3,    1
2,  1.1,    1
2,  2.1,    2
2,  5.2,    8
3, -0.08,   8
3,  3.2,    3
3,  5.1,    5

The actual output is:

keys1, keys2, values:

1,  0.3,    1
1,  -0.01,  1
2,  1.1,    1
2,  2.1,    2
2,  5.2,    8
3,  -0.08   8
3,  3.2,    3
3,  5.1,    5

My code is:

#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <thrust/execution_policy.h>

typedef thrust::tuple<int, int> Tuple;

/**************************/
/* TUPLE ORDERING FUNCTOR */
/**************************/
struct TupleComp
{
    __host__ __device__ bool operator()(const Tuple& t1, const Tuple& t2)
    {
        if (t1.get<0>() < t2.get<0>())
            return true;
        if (t1.get<0>() > t2.get<0>())
            return false;
        return t1.get<1>() < t2.get<1>();
    }
};

/********/
/* MAIN */
/********/
int main()
{
    const int N = 8;

    // --- Keys and values on the host: allocation and definition

        thrust::device_vector<int> keys1(N);
        thrust::device_vector<int> values(N);
        thrust::device_vector<float> keys2(N);

        keys1[0] = 1; keys1[1] = 3; keys1[2] = 3; keys1[3] = 3; 
        keys1[4] = 2; keys1[5] = 2; keys1[6] = 2; keys1[7] = 1; 


        values[0] = 1; values[1] = 5; values[2] = 3; values[3] = 8; 
        values[4] = 2; values[5] = 8; values[6] = 1; values[7] = 1; 

        keys2[0] = 0.3; keys2[1] = 5.1; keys2[2] = 3.2; keys2[3] = -0.08; 
        keys2[4] = 2.1; keys2[5] = 5.2; keys2[6] = 1.1; keys2[7] = -0.01; 


    std::cout << "ORIGINAL:"<< std::endl;
    std::cout << "keys1, keys2, values:" << std::endl;
    for (int i = 0; i < N; i++) {
        std::cout << keys1[i] <<'\t' << keys2[i] << '\t' << values[i] << std::endl;
    }


    thrust::sort_by_key(thrust::device, thrust::make_zip_iterator(thrust::make_tuple(keys1.begin(), keys2.begin())),
                            thrust::make_zip_iterator(thrust::make_tuple(keys1.begin() + N, keys2.begin() + N)), 
                            values.begin(), TupleComp());

    std::cout <<std::endl;

    std::cout << "ORDERED:"<< std::endl;
    std::cout << "keys1, keys2, values:" << std::endl;
    for (int i = 0; i < N; i++) {
        std::cout << keys1[i] <<'\t' << keys2[i] << '\t' << values[i] << std::endl;

    }

}

Solution

  • You've defined your tuple type incorrectly:

    typedef thrust::tuple<int, int> Tuple;
    

    it should be:

    typedef thrust::tuple<int, float> Tuple;
    

    to match the types of the keys1 and keys2, respectively.

    If you use templating for your functor types, the compiler won't make this mistake that you have made.