Block-wide primitives like cub::BlockRadixSort, cub::BlockLoad, and cub::BlockStore are used within CUDA kernels to perform collective operations across a thread block.
Specialization
These classes must be specialized by providing:
- The data type being processed.
- The number of threads per block.
- The number of keys/items per thread.
- (Implicitly) The targeted compilation architecture.
Memory Management
Once specialized, these classes expose an opaque TempStorage member type. To manage memory efficiently, thread blocks should use a union of these TempStorage types to statically allocate the required shared memory. This ensures that the total shared memory allocated is only as large as the largest required buffer among the used primitives.
Data Access Patterns
cub::BlockLoad and cub::BlockStore are designed to provide coalesced access to device memory. They use a striped access pattern (consecutive threads accessing consecutive items) and then transpose the keys into a blocked arrangement across threads.
#include <cub/cub.cuh>
// Block-sorting CUDA kernel
__global__ void BlockSortKernel(int *d_in, int *d_out)
{
using namespace cub;
// Specialize BlockRadixSort, BlockLoad, and BlockStore for 128 threads
// owning 16 integer items each
typedef BlockRadixSort<int, 128, 16> BlockRadixSort;
typedef BlockLoad<int, 128, 16, BLOCK_LOAD_TRANSPOSE> BlockLoad;
typedef BlockStore<int, 128, 16, BLOCK_STORE_TRANSPOSE> BlockStore;
// Allocate shared memory using a union of TempStorage types
__shared__ union {
typename BlockRadixSort::TempStorage sort;
typename BlockLoad::TempStorage load;
typename BlockStore::TempStorage store;
} temp_storage;
int block_offset = blockIdx.x * (128 * 16);
// 1. Load data from device memory to thread-local arrays
int thread_keys[16];
BlockLoad(temp_storage.load).Load(d_in + block_offset, thread_keys);
__syncthreads();
// 2. Collectively sort the keys within the block
BlockRadixSort(temp_storage.sort).Sort(thread_keys);
__syncthreads();
// 3. Store the sorted segment back to device memory
BlockStore(temp_storage.store).Store(d_out + block_offset, thread_keys);
}