[INFRA] Import NVIDIA/CCCL upstream as optimization reference library
CCCL (CUDA C++ Core Libraries) provides: - CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk) - Thrust: high-level parallel algorithms (transform_reduce, sort, scan) - libcudacxx: CUDA C++ standard library (atomics, barriers, memory) - cudax: experimental features (memory resources, allocators) - Tuning policies: per-SM hardware-specific algorithm parameters Competition optimization vectors mapped to CCCL: - Output TPS (83% weight): warp_reduce, block_reduce, device_topk - Input TPS (14% weight): device_scan, block_load, prefetch - Cache TPS (3% weight): prefix caching strategy patterns - Memory (0.9 util): pooled/cached/buddy allocators Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only) License: Apache-2.0
This commit is contained in:
EngineX CI
326
cccl_upstream/cub/test/catch2_test_block_shuffle.cu
Normal file
326
cccl_upstream/cub/test/catch2_test_block_shuffle.cu
Normal file
@@ -0,0 +1,326 @@
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
|
||||
// SPDX-License-Identifier: BSD-3
|
||||
|
||||
/******************************************************************************
|
||||
* Test of BlockMergeSort utilities
|
||||
******************************************************************************/
|
||||
|
||||
// Ensure printing of CUDA runtime errors to console
|
||||
#define CUB_STDERR
|
||||
|
||||
#include <cub/block/block_shuffle.cuh>
|
||||
|
||||
#include <thrust/sort.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#include <c2h/catch2_test_helper.h>
|
||||
|
||||
template <int BlockDimX, int BlockDimY, int BlockDimZ, int ItemsPerThread, class T, class ActionT>
|
||||
__global__ void block_shuffle_kernel(T* data, ActionT action)
|
||||
{
|
||||
using block_shuffle_t = cub::BlockShuffle<T, BlockDimX, BlockDimY, BlockDimZ>;
|
||||
using temp_storage_t = typename block_shuffle_t::TempStorage;
|
||||
|
||||
__shared__ temp_storage_t temp_storage;
|
||||
|
||||
T thread_data[ItemsPerThread];
|
||||
|
||||
data += cub::RowMajorTid(BlockDimX, BlockDimY, BlockDimZ) * ItemsPerThread;
|
||||
for (int item = 0; item < ItemsPerThread; item++)
|
||||
{
|
||||
thread_data[item] = data[item];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
block_shuffle_t block_shuffle(temp_storage);
|
||||
action(block_shuffle, thread_data);
|
||||
|
||||
for (int item = 0; item < ItemsPerThread; item++)
|
||||
{
|
||||
data[item] = thread_data[item];
|
||||
}
|
||||
}
|
||||
|
||||
struct up_op_t
|
||||
{
|
||||
template <class BlockShuffleT, class T, int ItemsPerThread>
|
||||
__device__ void operator()(BlockShuffleT& block_shuffle, T (&thread_data)[ItemsPerThread]) const
|
||||
{
|
||||
block_shuffle.Up(thread_data, thread_data);
|
||||
}
|
||||
};
|
||||
|
||||
struct offset_op_t
|
||||
{
|
||||
int m_distance;
|
||||
|
||||
__host__ offset_op_t(int distance)
|
||||
: m_distance(distance)
|
||||
{}
|
||||
|
||||
template <class BlockShuffleT, class T, int ItemsPerThread>
|
||||
__device__ void operator()(BlockShuffleT& block_shuffle, T (&thread_data)[ItemsPerThread]) const
|
||||
{
|
||||
block_shuffle.Offset(thread_data[0], thread_data[0], m_distance);
|
||||
}
|
||||
};
|
||||
|
||||
struct rotate_op_t
|
||||
{
|
||||
unsigned int m_distance;
|
||||
|
||||
__host__ rotate_op_t(unsigned int distance)
|
||||
: m_distance(distance)
|
||||
{}
|
||||
|
||||
template <class BlockShuffleT, class T, int ItemsPerThread>
|
||||
__device__ void operator()(BlockShuffleT& block_shuffle, T (&thread_data)[ItemsPerThread]) const
|
||||
{
|
||||
block_shuffle.Rotate(thread_data[0], thread_data[0], m_distance);
|
||||
}
|
||||
};
|
||||
|
||||
template <class T>
|
||||
struct up_with_suffix_op_t
|
||||
{
|
||||
int m_target_thread_id;
|
||||
T* m_d_suffix_ptr;
|
||||
|
||||
__host__ up_with_suffix_op_t(int target_thread_id, T* d_suffix_ptr)
|
||||
: m_target_thread_id(target_thread_id)
|
||||
, m_d_suffix_ptr(d_suffix_ptr)
|
||||
{}
|
||||
|
||||
template <class BlockShuffleT, int ItemsPerThread>
|
||||
__device__ void operator()(BlockShuffleT& block_shuffle, T (&thread_data)[ItemsPerThread]) const
|
||||
{
|
||||
T suffix{};
|
||||
|
||||
block_shuffle.Up(thread_data, thread_data, suffix);
|
||||
|
||||
if (cub::RowMajorTid(static_cast<int>(blockDim.x), static_cast<int>(blockDim.y), static_cast<int>(blockDim.z))
|
||||
== m_target_thread_id)
|
||||
{
|
||||
m_d_suffix_ptr[0] = suffix;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct down_op_t
|
||||
{
|
||||
template <class BlockShuffleT, class T, int ItemsPerThread>
|
||||
__device__ void operator()(BlockShuffleT& block_shuffle, T (&thread_data)[ItemsPerThread]) const
|
||||
{
|
||||
block_shuffle.Down(thread_data, thread_data);
|
||||
}
|
||||
};
|
||||
|
||||
template <class T>
|
||||
struct down_with_prefix_op_t
|
||||
{
|
||||
int m_target_thread_id;
|
||||
T* m_d_prefix_ptr;
|
||||
|
||||
__host__ down_with_prefix_op_t(int target_thread_id, T* d_prefix_ptr)
|
||||
: m_target_thread_id(target_thread_id)
|
||||
, m_d_prefix_ptr(d_prefix_ptr)
|
||||
{}
|
||||
|
||||
template <class BlockShuffleT, int ItemsPerThread>
|
||||
__device__ void operator()(BlockShuffleT& block_shuffle, T (&thread_data)[ItemsPerThread]) const
|
||||
{
|
||||
T prefix{};
|
||||
|
||||
block_shuffle.Down(thread_data, thread_data, prefix);
|
||||
|
||||
if (cub::RowMajorTid(static_cast<int>(blockDim.x), static_cast<int>(blockDim.y), static_cast<int>(blockDim.z))
|
||||
== m_target_thread_id)
|
||||
{
|
||||
m_d_prefix_ptr[0] = prefix;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <int ItemsPerThread, int BlockDimX, int BlockDimY, int BlockDimZ, class T, class ActionT>
|
||||
void block_shuffle(c2h::device_vector<T>& data, ActionT action)
|
||||
{
|
||||
dim3 block(BlockDimX, BlockDimY, BlockDimZ);
|
||||
block_shuffle_kernel<BlockDimX, BlockDimY, BlockDimZ, ItemsPerThread>
|
||||
<<<1, block>>>(thrust::raw_pointer_cast(data.data()), action);
|
||||
|
||||
REQUIRE(cudaSuccess == cudaPeekAtLastError());
|
||||
REQUIRE(cudaSuccess == cudaDeviceSynchronize());
|
||||
}
|
||||
|
||||
// %PARAM% MULTI_DIM mdim 0:1
|
||||
// %PARAM% DIM_IDX dim_idx 0:1:2
|
||||
|
||||
#if MULTI_DIM
|
||||
using block_dim_xs = c2h::enum_type_list<int, 7, 32, 64>;
|
||||
using block_dim_yz = c2h::enum_type_list<int, 2>;
|
||||
#else
|
||||
using block_dim_xs = c2h::enum_type_list<int, 64, 512, 1024>;
|
||||
using block_dim_yz = c2h::enum_type_list<int, 1>;
|
||||
#endif
|
||||
|
||||
using block_dim_x = c2h::enum_type_list<int, c2h::get<DIM_IDX, block_dim_xs>::value>;
|
||||
|
||||
using types = c2h::type_list<std::int32_t, std::int64_t>;
|
||||
using items_per_thread = c2h::enum_type_list<int, 1, 2, 15>;
|
||||
using single_item_per_thread = c2h::enum_type_list<int, 1>;
|
||||
|
||||
template <class TestType>
|
||||
struct params_t
|
||||
{
|
||||
using type = typename c2h::get<0, TestType>;
|
||||
|
||||
static constexpr int items_per_thread = c2h::get<1, TestType>::value;
|
||||
static constexpr int block_dim_x = c2h::get<2, TestType>::value;
|
||||
static constexpr int block_dim_y = c2h::get<3, TestType>::value;
|
||||
static constexpr int block_dim_z = block_dim_y;
|
||||
static constexpr int threads_in_block = block_dim_x * block_dim_y * block_dim_z;
|
||||
static constexpr int tile_size = items_per_thread * threads_in_block;
|
||||
};
|
||||
|
||||
C2H_TEST("Block shuffle offset works", "[shuffle][block]", types, single_item_per_thread, block_dim_x, block_dim_yz)
|
||||
{
|
||||
using params = params_t<TestType>;
|
||||
using type = typename params::type;
|
||||
|
||||
c2h::device_vector<type> d_data(params::tile_size);
|
||||
c2h::gen(C2H_SEED(10), d_data);
|
||||
|
||||
const int distance = GENERATE_COPY(take(4, random(1 - params::tile_size, params::tile_size - 1)));
|
||||
|
||||
c2h::host_vector<type> h_data = d_data;
|
||||
c2h::host_vector<type> h_ref(params::tile_size);
|
||||
|
||||
for (int i = 0; i < static_cast<int>(h_data.size()); i++)
|
||||
{
|
||||
const int source = i + distance;
|
||||
h_ref[i] = (source >= 0) && (source < params::tile_size) ? h_data[source] : h_data[i];
|
||||
}
|
||||
|
||||
block_shuffle<params::items_per_thread, params::block_dim_x, params::block_dim_y, params::block_dim_z>(
|
||||
d_data, offset_op_t{distance});
|
||||
|
||||
REQUIRE(h_ref == d_data);
|
||||
}
|
||||
|
||||
C2H_TEST("Block shuffle rotate works", "[shuffle][block]", types, single_item_per_thread, block_dim_x, block_dim_yz)
|
||||
{
|
||||
using params = params_t<TestType>;
|
||||
using type = typename params::type;
|
||||
|
||||
c2h::device_vector<type> d_data(params::tile_size);
|
||||
c2h::gen(C2H_SEED(10), d_data);
|
||||
|
||||
c2h::device_vector<type> d_ref = d_data;
|
||||
|
||||
const unsigned int distance = GENERATE_COPY(take(4, random(0, params::tile_size - 1)));
|
||||
|
||||
c2h::host_vector<type> h_ref = d_data;
|
||||
std::rotate(h_ref.begin(), h_ref.begin() + distance, h_ref.end());
|
||||
|
||||
block_shuffle<params::items_per_thread, params::block_dim_x, params::block_dim_y, params::block_dim_z>(
|
||||
d_data, rotate_op_t{distance});
|
||||
|
||||
REQUIRE(h_ref == d_data);
|
||||
}
|
||||
|
||||
C2H_TEST("Block shuffle up works", "[shuffle][block]", types, items_per_thread, block_dim_x, block_dim_yz)
|
||||
{
|
||||
using params = params_t<TestType>;
|
||||
using type = typename params::type;
|
||||
|
||||
c2h::device_vector<type> d_data(params::tile_size);
|
||||
c2h::gen(C2H_SEED(10), d_data);
|
||||
|
||||
c2h::device_vector<type> d_ref(params::tile_size);
|
||||
thrust::copy(d_data.begin(), d_data.end() - 1, d_ref.begin() + 1);
|
||||
thrust::copy(d_data.begin(), d_data.begin() + 1, d_ref.begin());
|
||||
|
||||
block_shuffle<params::items_per_thread, params::block_dim_x, params::block_dim_y, params::block_dim_z>(
|
||||
d_data, up_op_t{});
|
||||
|
||||
REQUIRE(d_ref == d_data);
|
||||
}
|
||||
|
||||
C2H_TEST("Block shuffle up works when suffix is required",
|
||||
"[shuffle][block]",
|
||||
types,
|
||||
items_per_thread,
|
||||
block_dim_x,
|
||||
block_dim_yz)
|
||||
{
|
||||
using params = params_t<TestType>;
|
||||
using type = typename params::type;
|
||||
|
||||
c2h::device_vector<type> d_data(params::tile_size);
|
||||
c2h::gen(C2H_SEED(10), d_data);
|
||||
|
||||
const int target_thread_id = GENERATE_COPY(take(2, random(0, params::threads_in_block - 1)));
|
||||
|
||||
c2h::device_vector<type> d_ref(params::tile_size);
|
||||
thrust::copy(d_data.begin(), d_data.end() - 1, d_ref.begin() + 1);
|
||||
thrust::copy(d_data.begin(), d_data.begin() + 1, d_ref.begin());
|
||||
|
||||
c2h::device_vector<type> d_suffix(1);
|
||||
c2h::device_vector<type> d_suffix_ref(1);
|
||||
thrust::copy(d_data.end() - 1, d_data.end(), d_suffix_ref.begin());
|
||||
|
||||
block_shuffle<params::items_per_thread, params::block_dim_x, params::block_dim_y, params::block_dim_z>(
|
||||
d_data, up_with_suffix_op_t<type>{target_thread_id, thrust::raw_pointer_cast(d_suffix.data())});
|
||||
|
||||
REQUIRE(d_ref == d_data);
|
||||
REQUIRE(d_suffix_ref == d_suffix);
|
||||
}
|
||||
|
||||
C2H_TEST("Block shuffle down works", "[shuffle][block]", types, items_per_thread, block_dim_x, block_dim_yz)
|
||||
{
|
||||
using params = params_t<TestType>;
|
||||
using type = typename params::type;
|
||||
|
||||
c2h::device_vector<type> d_data(params::tile_size);
|
||||
c2h::gen(C2H_SEED(10), d_data);
|
||||
|
||||
c2h::device_vector<type> d_ref(params::tile_size);
|
||||
thrust::copy(d_data.begin() + 1, d_data.end(), d_ref.begin());
|
||||
thrust::copy(d_data.end() - 1, d_data.end(), d_ref.end() - 1);
|
||||
|
||||
block_shuffle<params::items_per_thread, params::block_dim_x, params::block_dim_y, params::block_dim_z>(
|
||||
d_data, down_op_t{});
|
||||
|
||||
REQUIRE(d_ref == d_data);
|
||||
}
|
||||
|
||||
C2H_TEST("Block shuffle down works when prefix is required",
|
||||
"[shuffle][block]",
|
||||
types,
|
||||
items_per_thread,
|
||||
block_dim_x,
|
||||
block_dim_yz)
|
||||
{
|
||||
using params = params_t<TestType>;
|
||||
using type = typename params::type;
|
||||
|
||||
c2h::device_vector<type> d_data(params::tile_size);
|
||||
c2h::gen(C2H_SEED(10), d_data);
|
||||
|
||||
const int target_thread_id = GENERATE_COPY(take(2, random(0, params::threads_in_block - 1)));
|
||||
|
||||
c2h::device_vector<type> d_ref(params::tile_size);
|
||||
thrust::copy(d_data.begin() + 1, d_data.end(), d_ref.begin());
|
||||
thrust::copy(d_data.end() - 1, d_data.end(), d_ref.end() - 1);
|
||||
|
||||
c2h::device_vector<type> d_prefix(1);
|
||||
c2h::device_vector<type> d_prefix_ref(1);
|
||||
thrust::copy(d_data.begin(), d_data.begin() + 1, d_prefix_ref.begin());
|
||||
|
||||
block_shuffle<params::items_per_thread, params::block_dim_x, params::block_dim_y, params::block_dim_z>(
|
||||
d_data, down_with_prefix_op_t<type>{target_thread_id, thrust::raw_pointer_cast(d_prefix.data())});
|
||||
|
||||
REQUIRE(d_ref == d_data);
|
||||
REQUIRE(d_prefix_ref == d_prefix);
|
||||
}
|
||||
Reference in New Issue
Block a user