Сортировка (маленькие) массивы по ключу в CUDA

Question

Я пытаюсь написать функцию, которая принимает блок неупорядоченных пар ключ/значение, таких как

<7, 4> 
<2, 8> 
<3, 1> 
<2, 2> 
<1, 5> 
<7, 1> 
<3, 8> 
<7, 2> 

и сортирует их ключом, уменьшая значения пар с тем же ключом:

<1, 5> 
<2, 10> 
<3, 9> 
<7, 7> 

в настоящее время я использую __device__ функцию, как один ниже, который по существу bitonic вид, который будет сочетать в себе значение одного и тот же ключ и установить старые данные бесконечно большую величину (просто используя 99), так что последующий bitonic sort будет просеивать их на дно, и массив, отрезанный на значение int *, удален.

__device__ void interBitonicSortReduce(int2 *sdata, int tid, int recordNum, int *removed) { 
    int n = MIN(DEFAULT_DIMBLOCK, recordNum); 
    for (int k = 2; k <= n; k *= 2) { 
    for (int j = k/2; j > 0; j /= 2) { 
     int ixj = tid^j; 
     if (ixj > tid) { 
     if (sdata[tid].x == sdata[ixj].x && sdata[tid].x < 99) { 
      atomicAdd(&sdata[tid].y, sdata[ixj].y); 
      sdata[ixj].x = 99; 
      sdata[ixj].y = 99; 
      atomicAdd(removed, 1); 
     } 
     if ((tid & k) == 0 && sdata[tid].x > sdata[ixj].x) 
      swapData2(sdata[tid], sdata[ixj]); 
     if ((tid & k) != 0 && sdata[tid].x < sdata[ixj].x) 
      swapData2(sdata[tid], sdata[ixj]); 
     __syncthreads(); 
     } 
    } 
    } 
} 

Это прекрасно работает для небольших наборов данных, но с большими наборами (хотя по-прежнему в пределах размера одного блока) один вызов просто не будет делать это.

Можно ли попытаться объединить сортировку и сокращение в той же функции? Очевидно, что эту функцию нужно будет называть более одного раза, но можно ли точно определить, сколько раз она должна быть вызвана для исчерпания всех данных на основе ее размера?

Или я преформы сокращения отдельно с чем-то вроде этого:

__device__ int interReduce(int2 *sdata, int tid) { 
    int index = tid; 
    while (sdata[index].x == sdata[tid].x) { 
    index--; 
    if (index < 0) 
     break; 
    } 
    if (index+1 != tid) { 
    atomicAdd(&sdata[index+1].y, sdata[tid].y); 
    sdata[tid].x = 99; 
    sdata[tid].y = 99; 
    return 1; 
    } 
    return 0; 
} 

Я пытаюсь придумать наиболее эффективным решением, но мой опыт работы с CUDA и параллельными алгоритмами ограничен.

Answer 1

Вы можете использовать thrust для этого.

Использование thrust::sort_by_key с последующим thrust::reduce_by_key

Вот пример:

#include <iostream> 
#include <thrust/device_vector.h> 
#include <thrust/copy.h> 
#include <thrust/sort.h> 
#include <thrust/reduce.h> 
#include <thrust/sequence.h> 

#define N 12 
typedef thrust::device_vector<int>::iterator dintiter; 
int main(){ 

    thrust::device_vector<int> keys(N); 
    thrust::device_vector<int> values(N); 
    thrust::device_vector<int> new_keys(N); 
    thrust::device_vector<int> new_values(N); 
    thrust::sequence(keys.begin(), keys.end()); 
    thrust::sequence(values.begin(), values.end()); 

    keys[3] = 1; 
    keys[9] = 1; 
    keys[8] = 2; 
    keys[7] = 4; 

    thrust::sort_by_key(keys.begin(), keys.end(), values.begin()); 
    thrust::pair<dintiter, dintiter> new_end; 
    new_end = thrust::reduce_by_key(keys.begin(), keys.end(), values.begin(), new_keys.begin(), new_values.begin()); 

    std::cout << "results values:" << std::endl; 
    thrust::copy(new_values.begin(), new_end.second, std::ostream_iterator<int>(std::cout, " ")); 
    std::cout << std::endl << "results keys:" << std::endl; 
    thrust::copy(new_keys.begin(), new_end.first, std::ostream_iterator<int>(std::cout, " ")); 
    std::cout << std::endl; 

    return 0; 
} 

Answer 2

С вашего поста, мне кажется, что вам нужно для сортировки по ключевым многих небольших массивов. Выделите себя:

Это прекрасно работает для небольших наборов данных, но с большими наборами (хотя и по размеру одного блока) один вызов просто не будет делать этого.

Ниже вы найдете полностью обработанный пример, построенный вокруг моего ответа на Sorting many small arrays in CUDA и используя cub::BlockRadixSort.

#include <cub/cub.cuh> 
#include <stdio.h> 
#include <stdlib.h> 

#include "Utilities.cuh" 

using namespace cub; 

/**********************************/ 
/* CUB BLOCKSORT KERNEL NO SHARED */ 
/**********************************/ 
template <int BLOCK_THREADS, int ITEMS_PER_THREAD> 
__global__ void BlockSortKernel(float *d_values, int *d_keys, float *d_values_result, int *d_keys_result) 
{ 
    // --- Specialize BlockLoad, BlockStore, and BlockRadixSort collective types 
    typedef cub::BlockLoad  <int*, BLOCK_THREADS, ITEMS_PER_THREAD, BLOCK_LOAD_TRANSPOSE> BlockLoadIntT; 
    typedef cub::BlockLoad  <float*, BLOCK_THREADS, ITEMS_PER_THREAD, BLOCK_LOAD_TRANSPOSE> BlockLoadFloatT; 
    typedef cub::BlockStore  <int*, BLOCK_THREADS, ITEMS_PER_THREAD, BLOCK_STORE_TRANSPOSE> BlockStoreIntT; 
    typedef cub::BlockStore  <float*, BLOCK_THREADS, ITEMS_PER_THREAD, BLOCK_STORE_TRANSPOSE> BlockStoreFloatT; 
    typedef cub::BlockRadixSort <int , BLOCK_THREADS, ITEMS_PER_THREAD, float>     BlockRadixSortT; 

    // --- Allocate type-safe, repurposable shared memory for collectives 
    __shared__ union { 
     typename BlockLoadIntT  ::TempStorage loadInt; 
     typename BlockLoadFloatT ::TempStorage loadFloat; 
     typename BlockStoreIntT  ::TempStorage storeInt; 
     typename BlockStoreFloatT ::TempStorage storeFloat; 
     typename BlockRadixSortT ::TempStorage sort; 
    } temp_storage; 

    // --- Obtain this block's segment of consecutive keys (blocked across threads) 
    int thread_keys[ITEMS_PER_THREAD]; 
    float thread_values[ITEMS_PER_THREAD]; 
    int block_offset = blockIdx.x * (BLOCK_THREADS * ITEMS_PER_THREAD); 

    BlockLoadIntT(temp_storage.loadInt).Load(d_keys + block_offset, thread_keys); 
    BlockLoadFloatT(temp_storage.loadFloat).Load(d_values + block_offset, thread_values); 
    __syncthreads(); 

    // --- Collectively sort the keys 
    BlockRadixSortT(temp_storage.sort).SortBlockedToStriped(thread_keys, thread_values); 
    __syncthreads(); 

    // --- Store the sorted segment 
    BlockStoreIntT(temp_storage.storeInt).Store(d_keys_result + block_offset, thread_keys); 
    BlockStoreFloatT(temp_storage.storeFloat).Store(d_values_result + block_offset, thread_values); 

} 

/*******************************/ 
/* CUB BLOCKSORT KERNEL SHARED */ 
/*******************************/ 
template <int BLOCK_THREADS, int ITEMS_PER_THREAD> 
__global__ void shared_BlockSortKernel(float *d_values, int *d_keys, float *d_values_result, int *d_keys_result) 
{ 
    // --- Shared memory allocation 
    __shared__ float sharedMemoryArrayValues[BLOCK_THREADS * ITEMS_PER_THREAD]; 
    __shared__ int sharedMemoryArrayKeys[BLOCK_THREADS * ITEMS_PER_THREAD]; 

    // --- Specialize BlockStore and BlockRadixSort collective types 
    typedef cub::BlockRadixSort <int , BLOCK_THREADS, ITEMS_PER_THREAD, float> BlockRadixSortT; 

    // --- Allocate type-safe, repurposable shared memory for collectives 
    __shared__ typename BlockRadixSortT::TempStorage temp_storage; 

    int block_offset = blockIdx.x * (BLOCK_THREADS * ITEMS_PER_THREAD); 

    // --- Load data to shared memory 
    for (int k = 0; k < ITEMS_PER_THREAD; k++) { 
     sharedMemoryArrayValues[threadIdx.x * ITEMS_PER_THREAD + k] = d_values[block_offset + threadIdx.x * ITEMS_PER_THREAD + k]; 
     sharedMemoryArrayKeys[threadIdx.x * ITEMS_PER_THREAD + k] = d_keys[block_offset + threadIdx.x * ITEMS_PER_THREAD + k]; 
    } 
    __syncthreads(); 

    // --- Collectively sort the keys 
    BlockRadixSortT(temp_storage).SortBlockedToStriped(*static_cast<int(*) [ITEMS_PER_THREAD]>(static_cast<void*>(sharedMemoryArrayKeys + (threadIdx.x * ITEMS_PER_THREAD))), 
                 *static_cast<float(*)[ITEMS_PER_THREAD]>(static_cast<void*>(sharedMemoryArrayValues + (threadIdx.x * ITEMS_PER_THREAD)))); 
    __syncthreads(); 

    // --- Write data to shared memory 
    for (int k = 0; k < ITEMS_PER_THREAD; k++) { 
     d_values_result[block_offset + threadIdx.x * ITEMS_PER_THREAD + k] = sharedMemoryArrayValues[threadIdx.x * ITEMS_PER_THREAD + k]; 
     d_keys_result [block_offset + threadIdx.x * ITEMS_PER_THREAD + k] = sharedMemoryArrayKeys [threadIdx.x * ITEMS_PER_THREAD + k]; 
    } 
} 

/********/ 
/* MAIN */ 
/********/ 
int main() { 

    const int numElemsPerArray = 8; 
    const int numArrays   = 4; 
    const int N     = numArrays * numElemsPerArray; 
    const int numElemsPerThread = 4; 

    const int RANGE    = N * numElemsPerThread; 

    // --- Allocating and initializing the data on the host 
    float *h_values = (float *)malloc(N * sizeof(float)); 
    int *h_keys  = (int *) malloc(N * sizeof(int)); 
    for (int i = 0 ; i < N; i++) { 
     h_values[i] = rand() % RANGE; 
     h_keys[i] = rand() % RANGE; 
    } 

    printf("Original\n\n"); 
    for (int k = 0; k < numArrays; k++) 
     for (int i = 0; i < numElemsPerArray; i++) 
      printf("Array nr. %i; Element nr. %i; Key %i; Value %f\n", k, i, h_keys[k * numElemsPerArray + i], h_values[k * numElemsPerArray + i]); 

    // --- Allocating the results on the host 
    float *h_values_result1 = (float *)malloc(N * sizeof(float)); 
    float *h_values_result2 = (float *)malloc(N * sizeof(float)); 
    int *h_keys_result1 = (int *) malloc(N * sizeof(int)); 
    int *h_keys_result2 = (int *) malloc(N * sizeof(int)); 

    // --- Allocating space for data and results on device 
    float *d_values;   gpuErrchk(cudaMalloc((void **)&d_values,   N * sizeof(float))); 
    int *d_keys;    gpuErrchk(cudaMalloc((void **)&d_keys,   N * sizeof(int))); 
    float *d_values_result1; gpuErrchk(cudaMalloc((void **)&d_values_result1, N * sizeof(float))); 
    float *d_values_result2; gpuErrchk(cudaMalloc((void **)&d_values_result2, N * sizeof(float))); 
    int *d_keys_result1;  gpuErrchk(cudaMalloc((void **)&d_keys_result1, N * sizeof(int))); 
    int *d_keys_result2;  gpuErrchk(cudaMalloc((void **)&d_keys_result2, N * sizeof(int))); 

    // --- BlockSortKernel no shared 
    gpuErrchk(cudaMemcpy(d_values, h_values, N * sizeof(float), cudaMemcpyHostToDevice)); 
    gpuErrchk(cudaMemcpy(d_keys, h_keys, N * sizeof(int), cudaMemcpyHostToDevice)); 
    BlockSortKernel<N/numArrays/numElemsPerThread, numElemsPerThread><<<numArrays, numElemsPerArray/numElemsPerThread>>>(d_values, d_keys, d_values_result1, d_keys_result1); 
    gpuErrchk(cudaPeekAtLastError()); 
    gpuErrchk(cudaDeviceSynchronize());  
    gpuErrchk(cudaMemcpy(h_values_result1, d_values_result1, N * sizeof(float), cudaMemcpyDeviceToHost)); 
    gpuErrchk(cudaMemcpy(h_keys_result1, d_keys_result1, N * sizeof(int), cudaMemcpyDeviceToHost)); 

    printf("\n\nBlockSortKernel no shared\n\n"); 
    for (int k = 0; k < numArrays; k++) 
     for (int i = 0; i < numElemsPerArray; i++) 
      printf("Array nr. %i; Element nr. %i; Key %i; Value %f\n", k, i, h_keys_result1[k * numElemsPerArray + i], h_values_result1[k * numElemsPerArray + i]); 

    // --- BlockSortKernel with shared 
    gpuErrchk(cudaMemcpy(d_values, h_values, N * sizeof(float), cudaMemcpyHostToDevice)); 
    gpuErrchk(cudaMemcpy(d_keys, h_keys, N * sizeof(int), cudaMemcpyHostToDevice)); 
    shared_BlockSortKernel<N/numArrays/numElemsPerThread, numElemsPerThread><<<numArrays, numElemsPerArray/numElemsPerThread>>>(d_values, d_keys, d_values_result2, d_keys_result2); 
    gpuErrchk(cudaPeekAtLastError()); 
    gpuErrchk(cudaDeviceSynchronize());  
    gpuErrchk(cudaMemcpy(h_values_result2, d_values_result2, N * sizeof(float), cudaMemcpyDeviceToHost)); 
    gpuErrchk(cudaMemcpy(h_keys_result2, d_keys_result2, N * sizeof(int), cudaMemcpyDeviceToHost)); 

    printf("\n\nBlockSortKernel shared\n\n"); 
    for (int k = 0; k < numArrays; k++) 
     for (int i = 0; i < numElemsPerArray; i++) 
      printf("Array nr. %i; Element nr. %i; Key %i; Value %f\n", k, i, h_keys_result2[k * numElemsPerArray + i], h_values_result2[k * numElemsPerArray + i]); 

    return 0; 
} 

Answer 3

У меня недавно возникла проблема расширения подхода выше к случаю, когда несколько массивов должны быть заказаны в соответствии с одним и тем же ключом.

Похоже, что из-за его прототипа невозможно использовать cub::BlockRadixSort путем «упаковки» массивов с использованием итераторов и кортежей, см. C++ operating on “packed” arrays. Соответственно, я использовал подход вспомогательного индекса, предложенный в цитируемой статье.

Вот пример, который я разработал:

#include <cub/cub.cuh> 
#include <stdio.h> 
#include <stdlib.h> 

#include "Utilities.cuh" 

using namespace cub; 

/*******************************/ 
/* CUB BLOCKSORT KERNEL SHARED */ 
/*******************************/ 
template <int BLOCK_THREADS, int ITEMS_PER_THREAD> 
__global__ void shared_BlockSortKernel(float *d_valuesA, float *d_valuesB, int *d_keys, float *d_values_resultA, float *d_values_resultB, int *d_keys_result) 
{ 
    // --- Shared memory allocation 
    __shared__ float sharedMemoryArrayValuesA[BLOCK_THREADS * ITEMS_PER_THREAD]; 
    __shared__ float sharedMemoryArrayValuesB[BLOCK_THREADS * ITEMS_PER_THREAD]; 
    __shared__ int sharedMemoryArrayKeys[BLOCK_THREADS * ITEMS_PER_THREAD]; 
    __shared__ int sharedMemoryHelperIndices[BLOCK_THREADS * ITEMS_PER_THREAD]; 

    // --- Specialize BlockStore and BlockRadixSort collective types 
    typedef cub::BlockRadixSort <int , BLOCK_THREADS, ITEMS_PER_THREAD, int> BlockRadixSortT; 

    // --- Allocate type-safe, repurposable shared memory for collectives 
    __shared__ typename BlockRadixSortT::TempStorage temp_storage; 

    int block_offset = blockIdx.x * (BLOCK_THREADS * ITEMS_PER_THREAD); 

    // --- Load data to shared memory 
    for (int k = 0; k < ITEMS_PER_THREAD; k++) { 
     sharedMemoryArrayValuesA [threadIdx.x * ITEMS_PER_THREAD + k] = d_valuesA[block_offset + threadIdx.x * ITEMS_PER_THREAD + k]; 
     sharedMemoryArrayValuesB [threadIdx.x * ITEMS_PER_THREAD + k] = d_valuesB[block_offset + threadIdx.x * ITEMS_PER_THREAD + k]; 
     sharedMemoryArrayKeys [threadIdx.x * ITEMS_PER_THREAD + k] = d_keys [block_offset + threadIdx.x * ITEMS_PER_THREAD + k]; 
     sharedMemoryHelperIndices[threadIdx.x * ITEMS_PER_THREAD + k] =       threadIdx.x * ITEMS_PER_THREAD + k ; 
    } 
    __syncthreads(); 

    // --- Collectively sort the keys 
    BlockRadixSortT(temp_storage).SortBlockedToStriped(*static_cast<int(*)[ITEMS_PER_THREAD]>(static_cast<void*>(sharedMemoryArrayKeys  + (threadIdx.x * ITEMS_PER_THREAD))), 
                 *static_cast<int(*)[ITEMS_PER_THREAD]>(static_cast<void*>(sharedMemoryHelperIndices + (threadIdx.x * ITEMS_PER_THREAD)))); 
    __syncthreads(); 

    // --- Write data to shared memory 
    for (int k = 0; k < ITEMS_PER_THREAD; k++) { 
     d_values_resultA[block_offset + threadIdx.x * ITEMS_PER_THREAD + k] = sharedMemoryArrayValuesA[sharedMemoryHelperIndices[threadIdx.x * ITEMS_PER_THREAD + k]]; 
     d_values_resultB[block_offset + threadIdx.x * ITEMS_PER_THREAD + k] = sharedMemoryArrayValuesB[sharedMemoryHelperIndices[threadIdx.x * ITEMS_PER_THREAD + k]]; 
     d_keys_result [block_offset + threadIdx.x * ITEMS_PER_THREAD + k] = sharedMemoryArrayKeys        [threadIdx.x * ITEMS_PER_THREAD + k]; 
    } 
} 

/********/ 
/* MAIN */ 
/********/ 
int main() { 

    const int numElemsPerArray = 8; 
    const int numArrays   = 4; 
    const int N     = numArrays * numElemsPerArray; 
    const int numElemsPerThread = 4; 

    const int RANGE    = N * numElemsPerThread; 

    // --- Allocating and initializing the data on the host 
    float *h_valuesA = (float *)malloc(N * sizeof(float)); 
    float *h_valuesB = (float *)malloc(N * sizeof(float)); 
    int *h_keys   = (int *) malloc(N * sizeof(int)); 
    for (int i = 0 ; i < N; i++) { 
     h_valuesA[i] = rand() % RANGE; 
     h_valuesB[i] = rand() % RANGE; 
     h_keys[i] = rand() % RANGE; 
    } 

    printf("Original\n\n"); 
    for (int k = 0; k < numArrays; k++) 
     for (int i = 0; i < numElemsPerArray; i++) 
      printf("Array nr. %i; Element nr. %i; Key %i; Value A %f; Value B %f\n", k, i, h_keys[k * numElemsPerArray + i], h_valuesA[k * numElemsPerArray + i], h_valuesB[k * numElemsPerArray + i]); 

    // --- Allocating the results on the host 
    float *h_values_resultA = (float *)malloc(N * sizeof(float)); 
    float *h_values_resultB = (float *)malloc(N * sizeof(float)); 
    float *h_values_result2 = (float *)malloc(N * sizeof(float)); 
    int *h_keys_result1 = (int *) malloc(N * sizeof(int)); 
    int *h_keys_result2 = (int *) malloc(N * sizeof(int)); 

    // --- Allocating space for data and results on device 
    float *d_valuesA;   gpuErrchk(cudaMalloc((void **)&d_valuesA,  N * sizeof(float))); 
    float *d_valuesB;   gpuErrchk(cudaMalloc((void **)&d_valuesB,  N * sizeof(float))); 
    int *d_keys;    gpuErrchk(cudaMalloc((void **)&d_keys,   N * sizeof(int))); 
    float *d_values_resultA; gpuErrchk(cudaMalloc((void **)&d_values_resultA, N * sizeof(float))); 
    float *d_values_resultB; gpuErrchk(cudaMalloc((void **)&d_values_resultB, N * sizeof(float))); 
    float *d_values_result2; gpuErrchk(cudaMalloc((void **)&d_values_result2, N * sizeof(float))); 
    int *d_keys_result1;  gpuErrchk(cudaMalloc((void **)&d_keys_result1, N * sizeof(int))); 
    int *d_keys_result2;  gpuErrchk(cudaMalloc((void **)&d_keys_result2, N * sizeof(int))); 

    // --- BlockSortKernel with shared 
    gpuErrchk(cudaMemcpy(d_valuesA, h_valuesA, N * sizeof(float), cudaMemcpyHostToDevice)); 
    gpuErrchk(cudaMemcpy(d_valuesB, h_valuesB, N * sizeof(float), cudaMemcpyHostToDevice)); 
    gpuErrchk(cudaMemcpy(d_keys, h_keys, N * sizeof(int), cudaMemcpyHostToDevice)); 
    shared_BlockSortKernel<N/numArrays/numElemsPerThread, numElemsPerThread><<<numArrays, numElemsPerArray/numElemsPerThread>>>(d_valuesA, d_valuesB, d_keys, d_values_resultA, d_values_resultB, d_keys_result1); 
    gpuErrchk(cudaPeekAtLastError()); 
    gpuErrchk(cudaDeviceSynchronize());  
    gpuErrchk(cudaMemcpy(h_values_resultA, d_values_resultA, N * sizeof(float), cudaMemcpyDeviceToHost)); 
    gpuErrchk(cudaMemcpy(h_values_resultB, d_values_resultB, N * sizeof(float), cudaMemcpyDeviceToHost)); 
    gpuErrchk(cudaMemcpy(h_keys_result1, d_keys_result1, N * sizeof(int), cudaMemcpyDeviceToHost)); 

    printf("\n\nBlockSortKernel using shared memory\n\n"); 
    for (int k = 0; k < numArrays; k++) 
     for (int i = 0; i < numElemsPerArray; i++) 
      printf("Array nr. %i; Element nr. %i; Key %i; Value %f; Value %f\n", k, i, h_keys_result1[k * numElemsPerArray + i], h_values_resultA[k * numElemsPerArray + i], h_values_resultB[k * numElemsPerArray + i]); 

    return 0; 
} 

Answer 4

После моего второго ответа, я хочу, чтобы обеспечить дальнейшее расширение в случае, когда CUB используется для сортировки элементов, хранящихся в линейном общий массив памяти, который заполненный 2D сеткой потоков. Соответственно, cub::BlockRadixSort используется с 2D сеткой нитей вместо 1D сетки нитей, как в предыдущем ответе. Ниже приведен полный пример:

#include <cub/cub.cuh> 
#include <stdio.h> 
#include <stdlib.h> 

#include "Utilities.cuh" 

using namespace cub; 

/*******************************/ 
/* CUB BLOCKSORT KERNEL SHARED */ 
/*******************************/ 
template <int BLOCKSIZE_X, int BLOCKSIZE_Y, int ITEMS_PER_THREAD> 
__global__ void shared_BlockSortKernel(float *d_valuesA, float *d_valuesB, int *d_keys, float *d_values_resultA, float *d_values_resultB, int *d_keys_result) 
{ 
    // --- Shared memory allocation 
    __shared__ float sharedMemoryArrayValuesA [BLOCKSIZE_X * BLOCKSIZE_Y * ITEMS_PER_THREAD]; 
    __shared__ float sharedMemoryArrayValuesB [BLOCKSIZE_X * BLOCKSIZE_Y * ITEMS_PER_THREAD]; 
    __shared__ int sharedMemoryArrayKeys [BLOCKSIZE_X * BLOCKSIZE_Y * ITEMS_PER_THREAD]; 
    __shared__ int sharedMemoryHelperIndices[BLOCKSIZE_X * BLOCKSIZE_Y * ITEMS_PER_THREAD]; 

    // --- Specialize BlockStore and BlockRadixSort collective types 
    typedef cub::BlockRadixSort <int , BLOCKSIZE_X, ITEMS_PER_THREAD, int, 4, false, BLOCK_SCAN_WARP_SCANS, cudaSharedMemBankSizeFourByte, BLOCKSIZE_Y> BlockRadixSortT; 

    // --- Allocate type-safe, repurposable shared memory for collectives 
    __shared__ typename BlockRadixSortT::TempStorage temp_storage; 

    int block_offset = blockIdx.x * (BLOCKSIZE_X * BLOCKSIZE_Y * ITEMS_PER_THREAD); 

    // --- Load data to shared memory 
    for (int k = 0; k < ITEMS_PER_THREAD; k++) { 
     sharedMemoryArrayValuesA [(threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k] = d_valuesA[block_offset + (threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k]; 
     sharedMemoryArrayValuesB [(threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k] = d_valuesB[block_offset + (threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k]; 
     sharedMemoryArrayKeys [(threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k] = d_keys [block_offset + (threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k]; 
     sharedMemoryHelperIndices[(threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k] =       (threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k ; 
    } 
    __syncthreads(); 

    // --- Collectively sort the keys 
    BlockRadixSortT(temp_storage).SortBlockedToStriped(*static_cast<int(*)[ITEMS_PER_THREAD]>(static_cast<void*>(sharedMemoryArrayKeys  + ((threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD))), 
                 *static_cast<int(*)[ITEMS_PER_THREAD]>(static_cast<void*>(sharedMemoryHelperIndices + ((threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD)))); 
    __syncthreads(); 

    // --- Write data to shared memory 
    for (int k = 0; k < ITEMS_PER_THREAD; k++) { 
     d_values_resultA[block_offset + (threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k] = sharedMemoryArrayValuesA[sharedMemoryHelperIndices[(threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k]]; 
     d_values_resultB[block_offset + (threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k] = sharedMemoryArrayValuesB[sharedMemoryHelperIndices[(threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k]]; 
     d_keys_result [block_offset + (threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k] = sharedMemoryArrayKeys        [(threadIdx.y * BLOCKSIZE_X + threadIdx.x) * ITEMS_PER_THREAD + k]; 
    } 
} 

/********/ 
/* MAIN */ 
/********/ 
int main() { 

    const int blockSize_x  = 2; 
    const int blockSize_y  = 4; 
    const int numElemsPerArray = blockSize_x * blockSize_y; 
    const int numArrays   = 4; 
    const int N     = numArrays * numElemsPerArray; 
    const int numElemsPerThread = numElemsPerArray/(blockSize_x * blockSize_y); 

    const int RANGE    = N * numElemsPerThread; 

    // --- Allocating and initializing the data on the host 
    float *h_valuesA = (float *)malloc(N * sizeof(float)); 
    float *h_valuesB = (float *)malloc(N * sizeof(float)); 
    int *h_keys   = (int *) malloc(N * sizeof(int)); 
    for (int i = 0 ; i < N; i++) { 
     h_valuesA[i] = rand() % RANGE; 
     h_valuesB[i] = rand() % RANGE; 
     h_keys[i] = rand() % RANGE; 
    } 

    printf("Original\n\n"); 
    for (int k = 0; k < numArrays; k++) 
     for (int i = 0; i < numElemsPerArray; i++) 
      printf("Array nr. %i; Element nr. %i; Key %i; Value A %f; Value B %f\n", k, i, h_keys[k * numElemsPerArray + i], h_valuesA[k * numElemsPerArray + i], h_valuesB[k * numElemsPerArray + i]); 

    // --- Allocating the results on the host 
    float *h_values_resultA = (float *)malloc(N * sizeof(float)); 
    float *h_values_resultB = (float *)malloc(N * sizeof(float)); 
    float *h_values_result2 = (float *)malloc(N * sizeof(float)); 
    int *h_keys_result1 = (int *) malloc(N * sizeof(int)); 
    int *h_keys_result2 = (int *) malloc(N * sizeof(int)); 

    // --- Allocating space for data and results on device 
    float *d_valuesA;   gpuErrchk(cudaMalloc((void **)&d_valuesA,  N * sizeof(float))); 
    float *d_valuesB;   gpuErrchk(cudaMalloc((void **)&d_valuesB,  N * sizeof(float))); 
    int *d_keys;    gpuErrchk(cudaMalloc((void **)&d_keys,   N * sizeof(int))); 
    float *d_values_resultA; gpuErrchk(cudaMalloc((void **)&d_values_resultA, N * sizeof(float))); 
    float *d_values_resultB; gpuErrchk(cudaMalloc((void **)&d_values_resultB, N * sizeof(float))); 
    float *d_values_result2; gpuErrchk(cudaMalloc((void **)&d_values_result2, N * sizeof(float))); 
    int *d_keys_result1;  gpuErrchk(cudaMalloc((void **)&d_keys_result1, N * sizeof(int))); 
    int *d_keys_result2;  gpuErrchk(cudaMalloc((void **)&d_keys_result2, N * sizeof(int))); 

    // --- BlockSortKernel with shared 
    gpuErrchk(cudaMemcpy(d_valuesA, h_valuesA, N * sizeof(float), cudaMemcpyHostToDevice)); 
    gpuErrchk(cudaMemcpy(d_valuesB, h_valuesB, N * sizeof(float), cudaMemcpyHostToDevice)); 
    gpuErrchk(cudaMemcpy(d_keys, h_keys, N * sizeof(int), cudaMemcpyHostToDevice)); 
    shared_BlockSortKernel<blockSize_x, blockSize_y, numElemsPerThread><<<numArrays, numElemsPerArray/numElemsPerThread>>>(d_valuesA, d_valuesB, d_keys, d_values_resultA, d_values_resultB, d_keys_result1); 
    gpuErrchk(cudaPeekAtLastError()); 
    gpuErrchk(cudaDeviceSynchronize());  
    gpuErrchk(cudaMemcpy(h_values_resultA, d_values_resultA, N * sizeof(float), cudaMemcpyDeviceToHost)); 
    gpuErrchk(cudaMemcpy(h_values_resultB, d_values_resultB, N * sizeof(float), cudaMemcpyDeviceToHost)); 
    gpuErrchk(cudaMemcpy(h_keys_result1, d_keys_result1, N * sizeof(int), cudaMemcpyDeviceToHost)); 

    printf("\n\nBlockSortKernel using shared memory\n\n"); 
    for (int k = 0; k < numArrays; k++) 
     for (int i = 0; i < numElemsPerArray; i++) 
      printf("Array nr. %i; Element nr. %i; Key %i; Value %f; Value %f\n", k, i, h_keys_result1[k * numElemsPerArray + i], h_values_resultA[k * numElemsPerArray + i], h_values_resultB[k * numElemsPerArray + i]); 

    return 0; 
}