 template<
         typename DataType,
         unsigned int ItemsPerThread,
-        unsigned int ThreadsInBlock,
+        unsigned int BlockDimX,
+        unsigned int BlockDimY,
+        unsigned int BlockDimZ,
         typename ActionType>
 void BlockShuffleTest(DataType *data, ActionType action)
 {
-    BlockShuffleTestKernel<DataType, ThreadsInBlock, ItemsPerThread><<<1, ThreadsInBlock>>> (data, action);
+    dim3 block(BlockDimX, BlockDimY, BlockDimZ);
+    BlockShuffleTestKernel<DataType, BlockDimX, BlockDimY, BlockDimZ, ItemsPerThread><<<1, block>>> (data, action);
 
     CubDebugExit(cudaPeekAtLastError());
     CubDebugExit(cudaDeviceSynchronize());


template <
        typename DataType,
        unsigned int ThreadsInBlock,
        unsigned int ItemsPerThread,
        typename ActionType>
__global__ void BlockShuffleTestKernel(
        DataType *data,
        ActionType action)
{
    typedef cub::BlockShuffle<DataType, ThreadsInBlock> BlockShuffle;

    __shared__ typename BlockShuffle::TempStorage temp_storage_shuffle;

    DataType thread_data[ItemsPerThread];

    data += threadIdx.x * ItemsPerThread;
    for (unsigned int item = 0; item < ItemsPerThread; item++)
    {
        thread_data[item] = data[item];

    Test<2> (g_allocator);
    Test<15> (g_allocator);

    Test<int32_t, 1, 64, 2, 2>(g_allocator);

    return 0;
}
