Сортировка 3-х массивов по ключу в CUDA (возможно, с использованием Thrust)

Question

У меня есть 3 массивы одинакового размера (более 300.000 элементов). Один массив чисел с плавающей точкой и два массива индексов. Итак, для каждого номера у меня есть 2 идентификатора.

Все массивы 3 уже находятся в глобальной памяти графического процессора. Я хочу отсортировать все номера по их идентификаторам соответственно.

Есть ли способ использовать библиотеку Thrust для выполнения этой задачи? Есть ли лучший способ, чем библиотека Thrust?

Конечно, я предпочитаю не копировать их в память хоста пару раз. Кстати, это массивы, а не векторы.

Спасибо за вашу помощь заранее.

---

Предварительное решение , но это очень медленно. Это занимает почти 4 секунд, а мой размер массива порядка 300000

thrust::device_ptr<float> keys(afterSum);
thrust::device_ptr<int> vals0(d_index);
thrust::device_ptr<int> vals1(blockId); 

thrust::device_vector<int> sortedIndex(numElements);
thrust::device_vector<int> sortedBlockId(numElements);

thrust::counting_iterator<int> iter(0);
thrust::device_vector<int> indices(numElements);
thrust::copy(iter, iter + indices.size(), indices.begin()); 

thrust::sort_by_key(keys, keys + numElements , indices.begin());    

thrust::gather(indices.begin(), indices.end(), vals0, sortedIndex.begin());
thrust::gather(indices.begin(), indices.end(), vals1, sortedBlockId.begin());

thrust::host_vector<int> h_sortedIndex=sortedIndex;
thrust::host_vector<int> h_sortedBlockId=sortedBlockId;

Answer 1

Конечно, вы можете использовать Thrust.Во-первых, вам нужно обернуть необработанные указатели устройства CUDA thrust::device_ptr.Предполагая, что ваши значения с плавающей точкой находятся в массиве pkeys, а идентификаторы - в массивах pvals0 и pvals1, а numElements - это длина массивов, примерно так должно работать:

#include <thrust/device_ptr.h>
#include <thrust/sort.h>
#include <thrust/gather.h>
#include <thrust/iterator/counting_iterator.h>

cudaEvent_t start, stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);

cudaEventRecord(start);

thrust::device_ptr<float> keys(pkeys);
thrust::device_ptr<int> vals0(pvals0);
thrust::device_ptr<int> vals1(pvals1);

// allocate space for the output
thrust::device_vector<int> sortedVals0(numElements);
thrust::device_vector<int> sortedVals1(numElements);

// initialize indices vector to [0,1,2,..]
thrust::counting_iterator<int> iter(0);
thrust::device_vector<int> indices(numElements);
thrust::copy(iter, iter + indices.size(), indices.begin());

// first sort the keys and indices by the keys
thrust::sort_by_key(keys.begin(), keys.end(), indices.begin());

// Now reorder the ID arrays using the sorted indices
thrust::gather(indices.begin(), indices.end(), vals0.begin(), sortedVals0.begin());
thrust::gather(indices.begin(), indices.end(), vals1.begin(), sortedVals1.begin());

cudaEventRecord(stop);
cudaEventSynchronize(stop);
float milliseconds = 0;
cudaEventElapsedTime(&milliseconds, start, stop);
printf("Took %f milliseconds for %d elements\n", milliseconds, numElements);

Answer 2

Я сравнил два подхода, предложенных выше, а именно: использование thrust::zip_iterator и использование thrust::gather. Я проверил их в случае сортировки двух массивов по ключу или трех массивов в соответствии с просьбой автора. Во всех двух случаях подход с использованием thrust::gather оказался более быстрым.

Дело 2 Массив

#include <time.h>       // --- time
#include <stdlib.h>     // --- srand, rand

#include <thrust\host_vector.h>
#include <thrust\device_vector.h>
#include <thrust\sort.h>
#include <thrust\iterator\zip_iterator.h>

#include "TimingGPU.cuh"

//#define VERBOSE
//#define COMPACT

int main() {

    const int N = 1048576;
    //const int N = 10;

    TimingGPU timerGPU;

    // --- Initialize random seed
    srand(time(NULL));

    thrust::host_vector<int> h_code(N);
    thrust::host_vector<double> h_x(N);
    thrust::host_vector<double> h_y(N);

    for (int k = 0; k < N; k++) {       
        // --- Generate random numbers between 0 and 9
        h_code[k] = rand() % 10 + 1;
        h_x[k] = ((double)rand() / (RAND_MAX));
        h_y[k] = ((double)rand() / (RAND_MAX));
    }

    thrust::device_vector<int> d_code(h_code);

    thrust::device_vector<double> d_x(h_x);
    thrust::device_vector<double> d_y(h_y);

#ifdef VERBOSE
    printf("Before\n");
    for (int k = 0; k < N; k++) printf("code = %i; x = %f; y = %f\n", h_code[k], h_x[k], h_y[k]);
#endif

    timerGPU.StartCounter();
#ifdef COMPACT
    thrust::sort_by_key(d_code.begin(), d_code.end(), thrust::make_zip_iterator(thrust::make_tuple(d_x.begin(), d_y.begin())));
#else

    // --- Initialize indices vector to [0,1,2,..]
    thrust::counting_iterator<int> iter(0);
    thrust::device_vector<int> indices(N);
    thrust::copy(iter, iter + indices.size(), indices.begin());

    // --- First, sort the keys and indices by the keys
    thrust::sort_by_key(d_code.begin(), d_code.end(), indices.begin());

    // Now reorder the ID arrays using the sorted indices
    thrust::gather(indices.begin(), indices.end(), d_x.begin(), d_x.begin());
    thrust::gather(indices.begin(), indices.end(), d_y.begin(), d_y.begin());
#endif

    printf("Timing GPU = %f\n", timerGPU.GetCounter());

#ifdef VERBOSE
    h_code = d_code;
    h_x = d_x;
    h_y = d_y;

    printf("After\n");
    for (int k = 0; k < N; k++) printf("code = %i; x = %f; y = %f\n", h_code[k], h_x[k], h_y[k]);
#endif
}

Дело 3 Массив

#include <time.h>       // --- time
#include <stdlib.h>     // --- srand, rand

#include <thrust\host_vector.h>
#include <thrust\device_vector.h>
#include <thrust\sort.h>
#include <thrust\iterator\zip_iterator.h>

#include "TimingGPU.cuh"

//#define VERBOSE
//#define COMPACT

int main() {

    const int N = 1048576;
    //const int N = 10;

    TimingGPU timerGPU;

    // --- Initialize random seed
    srand(time(NULL));

    thrust::host_vector<int> h_code(N);
    thrust::host_vector<double> h_x(N);
    thrust::host_vector<double> h_y(N);
    thrust::host_vector<double> h_z(N);

    for (int k = 0; k < N; k++) {
        // --- Generate random numbers between 0 and 9
        h_code[k] = rand() % 10 + 1;
        h_x[k] = ((double)rand() / (RAND_MAX));
        h_y[k] = ((double)rand() / (RAND_MAX));
        h_z[k] = ((double)rand() / (RAND_MAX));
    }

    thrust::device_vector<int> d_code(h_code);

    thrust::device_vector<double> d_x(h_x);
    thrust::device_vector<double> d_y(h_y);
    thrust::device_vector<double> d_z(h_z);

#ifdef VERBOSE
    printf("Before\n");
    for (int k = 0; k < N; k++) printf("code = %i; x = %f; y = %f\n", h_code[k], h_x[k], h_y[k]);
#endif

    timerGPU.StartCounter();
#ifdef COMPACT
    thrust::sort_by_key(d_code.begin(), d_code.end(), thrust::make_zip_iterator(thrust::make_tuple(d_x.begin(), d_y.begin(), d_z.begin())));
#else

    // --- Initialize indices vector to [0,1,2,..]
    thrust::counting_iterator<int> iter(0);
    thrust::device_vector<int> indices(N);
    thrust::copy(iter, iter + indices.size(), indices.begin());

    // --- First, sort the keys and indices by the keys
    thrust::sort_by_key(d_code.begin(), d_code.end(), indices.begin());

    // Now reorder the ID arrays using the sorted indices
    thrust::gather(indices.begin(), indices.end(), d_x.begin(), d_x.begin());
    thrust::gather(indices.begin(), indices.end(), d_y.begin(), d_y.begin());
    thrust::gather(indices.begin(), indices.end(), d_z.begin(), d_z.begin());
#endif

    printf("Timing GPU = %f\n", timerGPU.GetCounter());

#ifdef VERBOSE
    h_code = d_code;
    h_x = d_x;
    h_y = d_y;

    printf("After\n");
    for (int k = 0; k < N; k++) printf("code = %i; x = %f; y = %f\n", h_code[k], h_x[k], h_y[k]);
#endif
}

Время в случае массивов 2 для N = 1048576

zip_iterator  = 7.34ms
gather        = 4.27ms

Сроки в случае 3 массивов для N = 1048576

zip_iterator  = 9.64ms
gather        = 4.22ms

Тесты, выполненные на карте NVIDIA GTX 960.

Answer 3

Я бы использовал zip_iterator для выполнения одного sort_by_key для обоих векторов индексов одновременно.

Это будет выглядеть так:

    typedef typename thrust::tuple<thrust::device_vector<int>::iterator, thrust::device_vector<int>::iterator> IteratorTuple;
    typedef typename thrust::zip_iterator<IteratorTuple> ZipIterator;   

    // here I suppose your 3 arrays are pointed to by device_ptr as suggested by @harrism
    thrust::device_vector<float> key(pKey, pKey + numElements);
    thrust::device_vector<int> val0(pVal0, pVal0 + numElements);
    thrust::device_vector<int> val1(pVal1, pVal1 + numElements);

    ZipIterator iterBegin(thrust::make_tuple(val0.begin(), val1.begin()));  
    thrust::sort_by_key(key.begin(), key.end(), iterBegin);