diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index e740b27a3d..2e9ab70b9b 100644 --- a/cpp/bench/prims/random/permute.cu +++ b/cpp/bench/prims/random/permute.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,7 +7,6 @@ #include #include -#include #include @@ -20,6 +19,11 @@ struct permute_inputs { template struct permute : public fixture { + /** + * @brief Construct a matrix permutation benchmark. + * + * @param[in] p Matrix dimensions, output selection, and layout + */ permute(const permute_inputs& p) : params(p), perms(p.needPerms ? p.rows : 0, stream), @@ -30,12 +34,19 @@ struct permute : public fixture { uniform(handle, r, in.data(), p.rows, T(-1.0), T(1.0)); } + /** @brief Benchmark keyed permutation of a matrix and its indices. */ void run_benchmark(::benchmark::State& state) override { raft::random::RngState r(123456ULL); loop_on_state(state, [this, &r]() { - raft::random::permute( - perms.data(), out.data(), in.data(), params.cols, params.rows, params.rowMajor, stream); + raft::random::permute(perms.data(), + out.data(), + in.data(), + params.cols, + params.rows, + params.rowMajor, + stream, + 123456ULL); }); } @@ -66,4 +77,42 @@ const std::vector permute_input_vecs = { RAFT_BENCH_REGISTER(permute, "", permute_input_vecs); RAFT_BENCH_REGISTER(permute, "", permute_input_vecs); +template +struct permute_perms_only : public fixture { + /** + * @brief Construct a benchmark that generates only permutation indices. + * + * @param[in] rows Number of permutation indices to generate + */ + permute_perms_only(int rows) : n_rows(rows), perms(rows, stream) {} + + /** @brief Benchmark the permutation-indices-only kernel path. */ + void run_benchmark(::benchmark::State& state) override + { + size_t bytes_processed = 0; + loop_on_state(state, [this, &bytes_processed]() { + raft::random::permute(perms.data(), + (float*)nullptr, + (const float*)nullptr, + IntType(0), + IntType(n_rows), + true, + stream, + 123456ULL); + bytes_processed += size_t(n_rows) * sizeof(IntType); + }); + state.SetBytesProcessed(bytes_processed); + } + + private: + raft::device_resources handle; + int n_rows; + rmm::device_uvector perms; +}; + +RAFT_BENCH_REGISTER((permute_perms_only), + "", + std::vector({32 * 1024, 1024 * 1024, 32 * 1024 * 1024})); +RAFT_BENCH_REGISTER((permute_perms_only), "", std::vector({1024 * 1024 * 1024})); + } // namespace raft::bench::random diff --git a/cpp/include/raft/random/detail/make_regression.cuh b/cpp/include/raft/random/detail/make_regression.cuh index 773eae7b39..b122692ea0 100644 --- a/cpp/include/raft/random/detail/make_regression.cuh +++ b/cpp/include/raft/random/detail/make_regression.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -13,7 +13,6 @@ #include #include #include -#include #include #include #include @@ -29,7 +28,14 @@ namespace raft { namespace random { namespace detail { -/* Internal auxiliary function to help build the singular profile */ +/** + * @brief Build the singular-value profile for a low-rank regression matrix. + * + * @param[out] out Generated singular values + * @param[in] n Number of singular values + * @param[in] tail_strength Relative strength of the low-rank tail + * @param[in] rank Effective matrix rank + */ template RAFT_KERNEL _singular_profile_kernel(DataT* out, IdxT n, DataT tail_strength, IdxT rank) { @@ -42,7 +48,18 @@ RAFT_KERNEL _singular_profile_kernel(DataT* out, IdxT n, DataT tail_strength, Id } } -/* Internal auxiliary function to generate a low-rank matrix */ +/** + * @brief Generate a low-rank matrix with a decaying singular-value profile. + * + * @param[in] handle RAFT handle containing execution resources + * @param[out] out Generated row-major matrix + * @param[in] n_rows Number of matrix rows + * @param[in] n_cols Number of matrix columns + * @param[in] effective_rank Approximate rank of the generated matrix + * @param[in] tail_strength Relative strength of the low-rank tail + * @param[in,out] r Random number generator state + * @param[in] stream CUDA stream on which to execute + */ template static void _make_low_rank_matrix(raft::resources const& handle, DataT* out, @@ -116,8 +133,15 @@ static void _make_low_rank_matrix(raft::resources const& handle, raft::linalg::transpose(handle, temp_out.data(), out, n_rows, n_cols, stream); } -/* Internal auxiliary function to permute rows in the given matrix according - * to a given permutation vector */ +/** + * @brief Gather matrix rows according to a permutation vector. + * + * @param[out] out Permuted output matrix + * @param[in] in Input matrix + * @param[in] perms Input row index for each output row + * @param[in] n_rows Number of matrix rows + * @param[in] n_cols Number of matrix columns + */ template RAFT_KERNEL _gather2d_kernel( DataT* out, const DataT* in, const IdxT* perms, IdxT n_rows, IdxT n_cols) @@ -134,6 +158,12 @@ RAFT_KERNEL _gather2d_kernel( } } +/** + * @brief Generate a regression data set and optionally shuffle its rows and features. + * + * When shuffling is enabled, the input seed deterministically selects distinct + * permutations for samples and features. + */ template void make_regression_caller(raft::resources const& handle, DataT* out, @@ -246,9 +276,15 @@ void make_regression_caller(raft::resources const& handle, constexpr IdxT Nthreads = 256; + // Derive two distinct permutation keys from the seed so the shuffle stays + // reproducible for a given seed while the samples and features get + // independent permutations. + const uint64_t samples_key = seed; + const uint64_t features_key = seed ^ 0x9e3779b97f4a7c15ULL; + // Shuffle the samples from out to tmp_out raft::random::permute( - perms_samples.data(), tmp_out.data(), out, n_cols, n_rows, true, stream); + perms_samples.data(), tmp_out.data(), out, n_cols, n_rows, true, stream, samples_key); IdxT nblks_rows = raft::ceildiv(n_rows, Nthreads); _gather2d_kernel<<>>( values, _values, perms_samples.data(), n_rows, n_targets); @@ -256,7 +292,7 @@ void make_regression_caller(raft::resources const& handle, // Shuffle the features from tmp_out to out raft::random::permute( - perms_features.data(), out, tmp_out.data(), n_rows, n_cols, false, stream); + perms_features.data(), out, tmp_out.data(), n_rows, n_cols, false, stream, features_key); // Shuffle the coefficients accordingly if (coef != nullptr) { diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index d991d6fc50..c503fa6e70 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,26 +11,74 @@ #include #include - -#include +#include +#include namespace raft { namespace random { namespace detail { -template +/** + * @brief Generate permutation indices without copying input data. + * + * @tparam IntType Output permutation index type + * @tparam IdxType Input index and extent type + * @tparam TPB Threads per block + * @tparam ITEMS_PER_THREAD Permutation indices generated by each thread + * @tparam ShuffleIterator Iterator that maps output indices to input indices + * + * @param[out] perms Generated permutation indices + * @param[in] shuffled_indices Iterator that defines the permutation + * @param[in] N Number of indices to generate + */ +template +RAFT_KERNEL permsOnlyKernel(IntType* perms, ShuffleIterator shuffled_indices, IdxType N) +{ + IdxType base = IdxType(blockIdx.x) * IdxType(TPB * ITEMS_PER_THREAD) + threadIdx.x; +#pragma unroll + for (int i = 0; i < ITEMS_PER_THREAD; i++) { + IdxType idx = base + IdxType(i * TPB); + if (idx < N) { perms[idx] = IntType(shuffled_indices[idx]); } + } +} + +/** + * @brief Generate permutation indices and optionally copy permuted matrix rows. + * + * @tparam Type Input and output element type + * @tparam IntType Output permutation index type + * @tparam IdxType Input index and extent type + * @tparam TPB Threads per block + * @tparam rowMajor Whether the matrices use row-major layout + * @tparam ShuffleIterator Iterator that maps output indices to input indices + * + * @param[out] perms Generated permutation indices, or nullptr + * @param[out] out Permuted output matrix, or nullptr + * @param[in] in Input matrix, or nullptr + * @param[in] shuffled_indices Iterator that defines the permutation + * @param[in] N Number of matrix rows + * @param[in] D Number of matrix columns + */ +template RAFT_KERNEL permuteKernel( - IntType* perms, Type* out, const Type* in, IdxType a, IdxType b, IdxType N, IdxType D) + IntType* perms, Type* out, const Type* in, ShuffleIterator shuffled_indices, IdxType N, IdxType D) { namespace cg = cooperative_groups; const int WARP_SIZE = 32; int tid = threadIdx.x + blockIdx.x * blockDim.x; - // having shuffled input indices and coalesced output indices appears - // to be preferable to the reverse, especially for column major - IntType inIdx = ((a * int64_t(tid)) + b) % N; IntType outIdx = tid; + IntType inIdx = (tid < N) ? IntType(shuffled_indices[tid]) : IntType(0); if (perms != nullptr && tid < N) { perms[outIdx] = inIdx; } @@ -39,19 +87,17 @@ RAFT_KERNEL permuteKernel( if (rowMajor) { cg::thread_block_tile warp = cg::tiled_partition(cg::this_thread_block()); - __shared__ IntType inIdxShm[TPB]; - __shared__ IntType outIdxShm[TPB]; - inIdxShm[threadIdx.x] = inIdx; - outIdxShm[threadIdx.x] = outIdx; - warp.sync(); - - int warpID = threadIdx.x / WARP_SIZE; + // The warp cooperatively copies its 32 rows one at a time so that the 32 + // lanes stride together along D, giving coalesced global loads and stores. + // Copying row i needs lane i's (inIdx, outIdx); int laneID = threadIdx.x % WARP_SIZE; - for (int i = warpID * WARP_SIZE; i < warpID * WARP_SIZE + WARP_SIZE; ++i) { - if (outIdxShm[i] < N) { + for (int i = 0; i < WARP_SIZE; ++i) { + IntType inIdxI = warp.shfl(inIdx, i); + IntType outIdxI = warp.shfl(outIdx, i); + if (outIdxI < N) { #pragma unroll for (int j = laneID; j < D; j += WARP_SIZE) { - out[outIdxShm[i] * D + j] = in[inIdxShm[i] * D + j]; + out[outIdxI * D + j] = in[inIdxI * D + j]; } } } @@ -66,14 +112,26 @@ RAFT_KERNEL permuteKernel( // This is wrapped in a type to allow for partial template specialization template struct permute_impl_t { + /** + * @brief Launch the permutation kernel with the widest valid vector type. + * + * @param[out] perms Generated permutation indices, or nullptr + * @param[out] out Permuted output matrix + * @param[in] in Input matrix + * @param[in] N Number of matrix rows + * @param[in] D Number of matrix columns + * @param[in] nblks Number of thread blocks to launch + * @param[in] shuffled_indices Iterator that defines the permutation + * @param[in] stream CUDA stream on which to launch the kernel + */ + template static void permuteImpl(IntType* perms, Type* out, const Type* in, IdxType N, IdxType D, int nblks, - IdxType a, - IdxType b, + ShuffleIterator shuffled_indices, cudaStream_t stream) { // determine vector type and set new pointers @@ -85,11 +143,11 @@ struct permute_impl_t { if (D % VLen == 0 && raft::is_aligned(vout, sizeof(VType)) && raft::is_aligned(vin, sizeof(VType))) { permuteKernel - <<>>(perms, vout, vin, a, b, N, D / VLen); + <<>>(perms, vout, vin, shuffled_indices, N, D / VLen); RAFT_CUDA_TRY(cudaPeekAtLastError()); } else { // otherwise try the next lower vector length permute_impl_t::permuteImpl( - perms, out, in, N, D, nblks, a, b, stream); + perms, out, in, N, D, nblks, shuffled_indices, stream); } } }; @@ -97,22 +155,51 @@ struct permute_impl_t { // at vector length 1 we just execute a scalar version to break the recursion template struct permute_impl_t { + /** + * @brief Launch the scalar permutation kernel. + * + * @param[out] perms Generated permutation indices, or nullptr + * @param[out] out Permuted output matrix + * @param[in] in Input matrix + * @param[in] N Number of matrix rows + * @param[in] D Number of matrix columns + * @param[in] nblks Number of thread blocks to launch + * @param[in] shuffled_indices Iterator that defines the permutation + * @param[in] stream CUDA stream on which to launch the kernel + */ + template static void permuteImpl(IntType* perms, Type* out, const Type* in, IdxType N, IdxType D, int nblks, - IdxType a, - IdxType b, + ShuffleIterator shuffled_indices, cudaStream_t stream) { permuteKernel - <<>>(perms, out, in, a, b, N, D); + <<>>(perms, out, in, shuffled_indices, N, D); RAFT_CUDA_TRY(cudaPeekAtLastError()); } }; +/** + * @brief Generate a keyed permutation and optionally apply it to a matrix. + * + * @tparam Type Input and output element type + * @tparam IntType Output permutation index type + * @tparam IdxType Input index and extent type + * @tparam TPB Threads per block + * + * @param[out] perms Generated permutation indices, or nullptr + * @param[out] out Permuted output matrix, or nullptr + * @param[in] in Input matrix, or nullptr + * @param[in] D Number of matrix columns + * @param[in] N Number of matrix rows + * @param[in] rowMajor Whether the matrices use row-major layout + * @param[in] stream CUDA stream on which to launch kernels + * @param[in] key Key used to construct the random bijection + */ template void permute(IntType* perms, Type* out, @@ -120,15 +207,23 @@ void permute(IntType* perms, IntType D, IntType N, bool rowMajor, - cudaStream_t stream) + cudaStream_t stream, + uint64_t key) { - auto nblks = raft::ceildiv(N, (IntType)TPB); + if (N <= 0) { return; } + + cuda::shuffle_iterator shuffled_indices{cuda::random_bijection{N, cuda::std::minstd_rand{key}}}; - // always keep 'a' to be coprime to N - IdxType a = rand() % N; - while (raft::gcd(a, N) != 1) - a = (a + 1) % N; - IdxType b = rand() % N; + if (out == nullptr) { + constexpr int ITEMS_PER_THREAD = 8; + auto nblks = raft::ceildiv(N, IntType(TPB * ITEMS_PER_THREAD)); + permsOnlyKernel + <<>>(perms, shuffled_indices, N); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + return; + } + + auto nblks = raft::ceildiv(N, (IntType)TPB); if (rowMajor) { permute_impl_t::permuteImpl( - perms, out, in, N, D, nblks, a, b, stream); + perms, out, in, N, D, nblks, shuffled_indices, stream); } } diff --git a/cpp/include/raft/random/permute.cuh b/cpp/include/raft/random/permute.cuh index 6a308d1fd4..41ee16571d 100644 --- a/cpp/include/raft/random/permute.cuh +++ b/cpp/include/raft/random/permute.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -74,6 +74,8 @@ using perms_out_view_t = typename perms_out_view in, std::optional> permsOut, - std::optional> out) + std::optional> out, + uint64_t key) { static_assert(std::is_integral_v, "permute: The type of each element " @@ -126,13 +129,26 @@ void permute(raft::resources const& handle, D, N, is_row_major, - resource::get_cuda_stream(handle)); + resource::get_cuda_stream(handle), + key); } } /** * @brief Overload of `permute` that compiles if users pass in `std::nullopt` * for either or both of `permsOut` and `out`. + * + * @tparam InputOutputValueType Input and output matrix element type + * @tparam IdxType Matrix index and extent type + * @tparam Layout Matrix layout type + * @tparam PermsOutType Optional permutation view type or `std::nullopt_t` + * @tparam OutType Optional output matrix view type or `std::nullopt_t` + * + * @param[in] handle RAFT handle containing the CUDA stream + * @param[in] in Input matrix + * @param[out] permsOut Optional generated permutation indices + * @param[out] out Optional permuted output matrix + * @param[in] key Key that selects the permutation */ template in, PermsOutType&& permsOut, - OutType&& out) + OutType&& out, + uint64_t key) { // If PermsOutType is std::optional> // for some T, then that type T need not be related to any of the @@ -159,7 +176,7 @@ void permute(raft::resources const& handle, std::optional permsOut_arg = std::forward(permsOut); std::optional out_arg = std::forward(out); - permute(handle, in, permsOut_arg, out_arg); + permute(handle, in, permsOut_arg, out_arg, key); } /** @} */ @@ -182,8 +199,82 @@ void permute(raft::resources const& handle, * @param[in] rowMajor true if the matrices are row major, * false if they are column major * @param[in] stream CUDA stream on which to run + * @param[in] key 64-bit key that selects the permutation. The same key + * (with the same @c N) always produces the same permutation. + */ +template +void permute(IntType* perms, + Type* out, + const Type* in, + IntType D, + IntType N, + bool rowMajor, + cudaStream_t stream, + uint64_t key) +{ + detail::permute(perms, out, in, D, N, rowMajor, stream, key); +} + +#define KEYLESS_PERMUTE_DEPRECATED_WARNING \ + "permute() now requires an explicit key (uint64_t). " \ + "This deprecated shim uses rand() to preserve per-call variation. " \ + "Pass an explicit key to make the permutation reproducible. " \ + "This overload will be removed in a future release." + +/** + * @brief Deprecated keyless mdspan overload that preserves per-call variation. + * + * This overload draws a key from `rand()`. Use the keyed overload when + * reproducibility or explicit control of the permutation is required. + * + * @deprecated Use the overload that takes an explicit @c uint64_t key. + */ +template +[[deprecated(KEYLESS_PERMUTE_DEPRECATED_WARNING)]] +void permute(raft::resources const& handle, + raft::device_matrix_view in, + std::optional> permsOut, + std::optional> out) +{ + permute(handle, in, permsOut, out, static_cast(rand())); +} + +/** + * @brief Deprecated keyless `std::nullopt` overload that preserves per-call variation. + * + * This overload draws a key from `rand()`. Use the keyed overload when + * reproducibility or explicit control of the permutation is required. + * + * @deprecated Use the overload that takes an explicit @c uint64_t key. + */ +template +[[deprecated(KEYLESS_PERMUTE_DEPRECATED_WARNING)]] +void permute(raft::resources const& handle, + raft::device_matrix_view in, + PermsOutType&& permsOut, + OutType&& out) +{ + permute(handle, + in, + std::forward(permsOut), + std::forward(out), + static_cast(rand())); +} + +/** + * @brief Deprecated keyless raw-pointer overload that preserves per-call variation. + * + * This overload draws a key from `rand()`. Use the keyed overload when + * reproducibility or explicit control of the permutation is required. + * + * @deprecated Use the overload that takes an explicit @c uint64_t key. */ template +[[deprecated(KEYLESS_PERMUTE_DEPRECATED_WARNING)]] void permute(IntType* perms, Type* out, const Type* in, @@ -192,7 +283,8 @@ void permute(IntType* perms, bool rowMajor, cudaStream_t stream) { - detail::permute(perms, out, in, D, N, rowMajor, stream); + detail::permute( + perms, out, in, D, N, rowMajor, stream, static_cast(rand())); } }; // namespace random diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index 44b5b9c66c..2f798371d7 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -25,6 +26,13 @@ struct PermInputs { unsigned long long int seed; }; +/** + * @brief Print permutation test parameters for GoogleTest diagnostics. + * + * @param[in,out] os Output stream + * @param[in] dims Test parameters + * @return The output stream + */ template ::std::ostream& operator<<(::std::ostream& os, const PermInputs& dims) { @@ -37,6 +45,7 @@ class PermTest : public ::testing::TestWithParam> { using test_data_type = T; protected: + /** @brief Construct an empty raw-pointer permutation test fixture. */ PermTest() : in(0, resource::get_cuda_stream(handle)), out(0, resource::get_cuda_stream(handle)), @@ -44,6 +53,7 @@ class PermTest : public ::testing::TestWithParam> { { } + /** @brief Allocate test inputs and run the keyed raw-pointer overload. */ void SetUp() override { auto stream = resource::get_cuda_stream(handle); @@ -65,7 +75,7 @@ class PermTest : public ::testing::TestWithParam> { out_ptr = out.data(); uniform(handle, r, in_ptr, len, T(-1.0), T(1.0)); } - permute(outPerms_ptr, out_ptr, in_ptr, D, N, params.rowMajor, stream); + permute(outPerms_ptr, out_ptr, in_ptr, D, N, params.rowMajor, stream, params.seed); resource::sync_stream(handle); } @@ -85,6 +95,7 @@ class PermMdspanTest : public ::testing::TestWithParam> { using test_data_type = T; protected: + /** @brief Construct an empty mdspan permutation test fixture. */ PermMdspanTest() : in(0, resource::get_cuda_stream(handle)), out(0, resource::get_cuda_stream(handle)), @@ -102,6 +113,7 @@ class PermMdspanTest : public ::testing::TestWithParam> { using vector_view_t = raft::device_vector_view; protected: + /** @brief Allocate test inputs and run the keyed mdspan overloads. */ void SetUp() override { auto stream = resource::get_cuda_stream(handle); @@ -133,16 +145,16 @@ class PermMdspanTest : public ::testing::TestWithParam> { std::optional> outPerms_view; if (outPerms_ptr != nullptr) { outPerms_view.emplace(outPerms_ptr, N); } - permute(handle, in_view, outPerms_view, out_view); + permute(handle, in_view, outPerms_view, out_view, params.seed); // None of these three permute calls should have an effect. // The point is to test whether the function can deduce the // element type of outPerms if given nullopt. std::optional> out_view_empty; std::optional> outPerms_view_empty; - permute(handle, in_view, std::nullopt, out_view_empty); - permute(handle, in_view, outPerms_view_empty, std::nullopt); - permute(handle, in_view, std::nullopt, std::nullopt); + permute(handle, in_view, std::nullopt, out_view_empty, params.seed); + permute(handle, in_view, outPerms_view_empty, std::nullopt, params.seed); + permute(handle, in_view, std::nullopt, std::nullopt, params.seed); }; if (params.rowMajor) { @@ -164,6 +176,17 @@ class PermMdspanTest : public ::testing::TestWithParam> { int* outPerms_ptr = nullptr; }; +/** + * @brief Compare a device array with a contiguous range. + * + * @param[in] actual Device array to compare + * @param[in] size Number of elements + * @param[in] start First expected value + * @param[in] eq_compare Equality comparison function + * @param[in] doSort Whether to sort the actual values before comparison + * @param[in] stream CUDA stream used for the device-to-host copy + * @return GoogleTest assertion result + */ template ::testing::AssertionResult devArrMatchRange( const T* actual, size_t size, T start, L eq_compare, bool doSort = true, cudaStream_t stream = 0) @@ -183,6 +206,19 @@ template return ::testing::AssertionSuccess(); } +/** + * @brief Verify that output rows match the input rows selected by a permutation. + * + * @param[in] perms Device permutation indices + * @param[in] out Device output matrix + * @param[in] in Device input matrix + * @param[in] D Number of matrix columns + * @param[in] N Number of matrix rows + * @param[in] rowMajor Whether the matrices use row-major layout + * @param[in] eq_compare Equality comparison function + * @param[in] stream CUDA stream used for device-to-host copies + * @return GoogleTest assertion result + */ template ::testing::AssertionResult devArrMatchShuffle(const int* perms, const T* out, @@ -216,6 +252,13 @@ template const std::vector> inputsf = { // only generate permutations + // small-N: identity path (N=1), 2-bit minimum domain (N=2), non-power-of-two + // cycle-walk cases (N=3, N=7), and sub-warp size (N=15) + {1, 8, true, false, true, 1234ULL}, + {2, 8, true, false, true, 1234ULL}, + {3, 8, true, false, true, 1234ULL}, + {7, 8, true, false, true, 1234ULL}, + {15, 8, true, false, true, 1234ULL}, {32, 8, true, false, true, 1234ULL}, {32, 8, true, false, true, 1234567890ULL}, {1024, 32, true, false, true, 1234ULL}, @@ -228,6 +271,11 @@ const std::vector> inputsf = { {100000, 32, true, false, true, 1234567890ULL}, {100001, 33, true, false, true, 1234567890ULL}, // permute and shuffle the data row major + {1, 8, true, true, true, 1234ULL}, + {2, 8, true, true, true, 1234ULL}, + {3, 8, true, true, true, 1234ULL}, + {7, 8, true, true, true, 1234ULL}, + {15, 8, true, true, true, 1234ULL}, {32, 8, true, true, true, 1234ULL}, {32, 8, true, true, true, 1234567890ULL}, {1024, 32, true, true, true, 1234ULL}, @@ -240,6 +288,11 @@ const std::vector> inputsf = { {100000, 32, true, true, true, 1234567890ULL}, {100001, 31, true, true, true, 1234567890ULL}, // permute and shuffle the data column major + {1, 8, true, true, false, 1234ULL}, + {2, 8, true, true, false, 1234ULL}, + {3, 8, true, true, false, 1234ULL}, + {7, 8, true, true, false, 1234ULL}, + {15, 8, true, true, false, 1234ULL}, {32, 8, true, true, false, 1234ULL}, {32, 8, true, true, false, 1234567890ULL}, {1024, 32, true, true, false, 1234ULL}, @@ -269,6 +322,7 @@ const std::vector> inputsf = { } while (false) using PermTestF = PermTest; +/** @brief Validate raw-pointer permutation output for single-precision inputs. */ TEST_P(PermTestF, Result) { using test_data_type = PermTestF::test_data_type; @@ -277,6 +331,7 @@ TEST_P(PermTestF, Result) INSTANTIATE_TEST_CASE_P(PermTests, PermTestF, ::testing::ValuesIn(inputsf)); using PermMdspanTestF = PermMdspanTest; +/** @brief Validate mdspan permutation output for single-precision inputs. */ TEST_P(PermMdspanTestF, Result) { using test_data_type = PermTestF::test_data_type; @@ -286,6 +341,11 @@ INSTANTIATE_TEST_CASE_P(PermMdspanTests, PermMdspanTestF, ::testing::ValuesIn(in const std::vector> inputsd = { // only generate permutations + {1, 8, true, false, true, 1234ULL}, + {2, 8, true, false, true, 1234ULL}, + {3, 8, true, false, true, 1234ULL}, + {7, 8, true, false, true, 1234ULL}, + {15, 8, true, false, true, 1234ULL}, {32, 8, true, false, true, 1234ULL}, {32, 8, true, false, true, 1234567890ULL}, {1024, 32, true, false, true, 1234ULL}, @@ -298,6 +358,11 @@ const std::vector> inputsd = { {100000, 32, true, false, true, 1234567890ULL}, {100001, 33, true, false, true, 1234567890ULL}, // permute and shuffle the data row major + {1, 8, true, true, true, 1234ULL}, + {2, 8, true, true, true, 1234ULL}, + {3, 8, true, true, true, 1234ULL}, + {7, 8, true, true, true, 1234ULL}, + {15, 8, true, true, true, 1234ULL}, {32, 8, true, true, true, 1234ULL}, {32, 8, true, true, true, 1234567890ULL}, {1024, 32, true, true, true, 1234ULL}, @@ -310,6 +375,11 @@ const std::vector> inputsd = { {100000, 32, true, true, true, 1234567890ULL}, {100001, 31, true, true, true, 1234567890ULL}, // permute and shuffle the data column major + {1, 8, true, true, false, 1234ULL}, + {2, 8, true, true, false, 1234ULL}, + {3, 8, true, true, false, 1234ULL}, + {7, 8, true, true, false, 1234ULL}, + {15, 8, true, true, false, 1234ULL}, {32, 8, true, true, false, 1234ULL}, {32, 8, true, true, false, 1234567890ULL}, {1024, 32, true, true, false, 1234ULL}, @@ -323,6 +393,7 @@ const std::vector> inputsd = { {100001, 33, true, true, false, 1234567890ULL}}; using PermTestD = PermTest; +/** @brief Validate raw-pointer permutation output for double-precision inputs. */ TEST_P(PermTestD, Result) { using test_data_type = PermTestF::test_data_type; @@ -331,6 +402,7 @@ TEST_P(PermTestD, Result) INSTANTIATE_TEST_CASE_P(PermTests, PermTestD, ::testing::ValuesIn(inputsd)); using PermMdspanTestD = PermMdspanTest; +/** @brief Validate mdspan permutation output for double-precision inputs. */ TEST_P(PermMdspanTestD, Result) { using test_data_type = PermTestF::test_data_type; @@ -338,5 +410,67 @@ TEST_P(PermMdspanTestD, Result) } INSTANTIATE_TEST_CASE_P(PermMdspanTests, PermMdspanTestD, ::testing::ValuesIn(inputsd)); +/** + * @brief Count matches between a reference permutation and many keyed permutations. + * + * Each thread uses `(base_seed + tid + 1)` and records how many generated + * indices match either the reference permutation or the original index. + * + * @param[in] ref_perm Reference permutation generated from `base_seed` + * @param[in] N Number of permutation indices + * @param[in] base_seed Seed used for the reference permutation + * @param[out] match_counts Match count generated by each thread + */ +__global__ void seed_diversity_kernel(const uint32_t* ref_perm, + uint32_t N, + uint64_t base_seed, + int* match_counts) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + cuda::shuffle_iterator shuffled_indices{ + cuda::random_bijection{N, cuda::std::minstd_rand{base_seed + uint64_t(tid) + 1}}}; + int count = 0; + for (uint32_t i = 0; i < N; i++) { + uint32_t val = shuffled_indices[i]; + if (val == ref_perm[i] || val == i) { count++; } + } + match_counts[tid] = count; +} + +/** @brief Verify that many consecutive seeds produce diverse permutations. */ +TEST(PermTest, SeedDiversity) +{ + constexpr uint32_t N = 1000; + constexpr uint64_t base_seed = 42ULL; + constexpr int TPB = 256; + constexpr int nblocks = 64; + constexpr int total_threads = nblocks * TPB; + constexpr float max_match_frac = 0.05f; + + raft::resources handle; + auto stream = resource::get_cuda_stream(handle); + + rmm::device_uvector d_ref(N, stream); + rmm::device_uvector d_matches(total_threads, stream); + detail::permute( + d_ref.data(), nullptr, nullptr, 0, N, true, stream, base_seed); + + seed_diversity_kernel<<>>(d_ref.data(), N, base_seed, d_matches.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + std::vector h_matches(total_threads); + raft::update_host(h_matches.data(), d_matches.data(), total_threads, stream); + resource::sync_stream(handle); + + int total_matches = 0; + for (int count : h_matches) { + total_matches += count; + } + + int max_allowed = static_cast(max_match_frac * float(N) * float(total_threads)); + EXPECT_LT(total_matches, max_allowed) + << "Too many index matches across seeds: " << total_matches << " >= " << max_allowed; +} + } // end namespace random } // end namespace raft