diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index e740b27a3d..43362fa389 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 */ @@ -34,8 +34,14 @@ struct permute : public fixture { { 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 +72,36 @@ 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 { + permute_perms_only(int rows) : n_rows(rows), perms(rows, stream) {} + + 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..9fb0b55e51 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 */ @@ -246,9 +246,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 +262,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..33faef4fd2 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 */ @@ -12,25 +12,165 @@ #include +#include #include +#include namespace raft { namespace random { namespace detail { +/* + * Keyed index permutation for permute(), built from a small Feistel network + * whose round function is the MurmurHash3 32-bit finalizer (fmix32). + * + * permute() needs a bijection f: [0, N) -> [0, N) so that mapping every output + * index to a distinct input index yields a true permutation (no duplicated or + * dropped rows). + * + * Because N is usually not a power of two, we use *cycle walking*: choose the + * smallest n with 2^n >= N, permute over the 2^n domain, and if the result is + * >= N, permute again until it lands in [0, N). Restricting a permutation of a + * finite set to a subset by cycle walking is itself a permutation of that + * subset, and it always terminates. With a good mixer the expected number of + * applications is 2^n / N < 2. + * + * Reference: https://en.wikipedia.org/wiki/Feistel_cipher + */ + +// Simplified version of MurmurHash3 32-bit finalizer. +// Diffuses all input bits into all output bits; +// used as the Feistel round function. Reference: +// https://github.com/aappleby/smhasher/blob/07bb4de10a63e8cc2e1724865454eba635742383/src/MurmurHash3.cpp#L68 +HDI uint32_t fmix32(uint32_t h) +{ + // h ^= h >> 16; + h *= 0x85ebca6bu; + h ^= h >> 13; + // h *= 0xc2b2ae35u; + // h ^= h >> 16; + return h; +} + +// Precomputed schedule for the keyed Feistel permutation of [0, N). Every field +// here depends only on (N, key), so it is pre-computed on the host. +// +// The round count is fixed at ROUNDS (4), which is enough for near-ideal +// avalanche at these widths. +struct kperm_params { + static constexpr int ROUNDS = 4; + + uint64_t N; // N (domain size); N <= 1 selects the identity permutation + uint64_t mask_n; // low-n-bit mask, n = ceil(log2(N)) + uint32_t b_bits; // width of the low part (shift to extract H = m >> b_bits) + uint32_t a; // low-bit mask for the high part: (1 << a_bits) - 1 + uint32_t b; // low-bit mask for the low part: (1 << b_bits) - 1 + uint32_t prefix[ROUNDS]; // per-round key-derived mixing constant +}; + +// Build the Feistel schedule for permuting [0, N) with the given 64-bit key. +// Runs once on the host per permute() call. +HDI kperm_params make_kperm_params(uint64_t N, uint64_t key) +{ + kperm_params p{}; + p.N = N; + if (N <= 1) return p; // identity domain; remaining fields go unused + + // n = ceil(log2(N)): smallest n with 2^n >= N (the cycle-walking domain). + int n = 0; + uint64_t pow2 = 1; + while (pow2 < N && n < 64) { + pow2 <<= 1; + ++n; + } + + // Clamp the width to at least 2 bits so the low half always has >= 1 bit + // This makes b == 0 impossible, so the device mixer needs no per-round + // guard against a zero-width low half. + if (n < 2) n = 2; + + const uint32_t a_bits = uint32_t((n + 1) / 2); // high part width + const uint32_t b_bits = uint32_t(n / 2); // low part width + p.b_bits = b_bits; + p.a = (a_bits >= 32) ? ~uint32_t(0) : ((uint32_t(1) << a_bits) - 1); + p.b = (b_bits >= 32) ? ~uint32_t(0) : ((uint32_t(1) << b_bits) - 1); // b_bits >= 1 here + p.mask_n = (n >= 64) ? ~uint64_t(0) : ((uint64_t(1) << n) - 1); + + // Per-round key schedule: diffuse the 64-bit key once, then derive a distinct + // prefix per round with a golden-ratio round spread. Identical across threads. + const uint32_t klo = uint32_t(key); + const uint32_t khi = uint32_t(key >> 32); + const uint32_t golden_ratio = 0x9e3779b9u; + + const uint32_t kbase = fmix32(fmix32(klo) ^ khi); + for (int i = 0; i < kperm_params::ROUNDS; i++) { + p.prefix[i] = fmix32(kbase ^ (uint32_t(i) + golden_ratio)); + } + return p; +} + +// Feistel permutation on m. +// The 4 rounds are unrolled by hand: even rounds mix the low part into the high +// part (a bits), odd rounds mix the high part into the low part (b bits). +// WordType is uint32_t for 32-bit domains and uint64_t for larger ones. +template +HDI WordType kperm_mix_nbits(WordType m, const kperm_params& p) +{ + m &= WordType(p.mask_n); + uint32_t L = uint32_t(m & p.b); // low b_bits + uint32_t H = uint32_t(m >> p.b_bits); // high a_bits + + H ^= fmix32(L ^ p.prefix[0]) & p.a; // round 0 (even): H from L + L ^= fmix32(H ^ p.prefix[1]) & p.b; // round 1 (odd): L from H + H ^= fmix32(L ^ p.prefix[2]) & p.a; // round 2 (even): H from L + L ^= fmix32(H ^ p.prefix[3]) & p.b; // round 3 (odd): L from H + return (WordType(H) << p.b_bits) | WordType(L); +} + +// Keyed permutation of [0, N): returns the index that output slot idx pulls +// from, using the precomputed schedule p. A bijection over [0, N) for any key. +// +// idx MUST lie in [0, N) for the cycle walk to halt. +template +HDI IdxType kperm_index(IdxType idx, const kperm_params& p) +{ + if (p.N <= 1) return idx; // 0- or 1-element domain: identity + + // Use a 32-bit word for 32-bit (or narrower) index types to avoid unnecessary + // 32->64->32 promotion. For larger types, stay in 64 bits. + using WordType = std::conditional_t; + + // Cycle walk: re-permute over [0, 2^n) until the result falls in [0, N). When + // N is a power of two this is a single application (the loop body runs once). + WordType x = WordType(idx); + do { + x = kperm_mix_nbits(x, p); + } while (x >= WordType(p.N)); + return IdxType(x); +} + +template +RAFT_KERNEL permsOnlyKernel(IntType* perms, kperm_params fp, 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(kperm_index(idx, fp)); } + } +} + 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, kperm_params fp, 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) ? kperm_index(IntType(tid), fp) : IntType(0); if (perms != nullptr && tid < N) { perms[outIdx] = inIdx; } @@ -39,19 +179,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]; } } } @@ -72,8 +210,7 @@ struct permute_impl_t { IdxType N, IdxType D, int nblks, - IdxType a, - IdxType b, + kperm_params fp, cudaStream_t stream) { // determine vector type and set new pointers @@ -85,11 +222,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, fp, 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, fp, stream); } } }; @@ -103,12 +240,11 @@ struct permute_impl_t { IdxType N, IdxType D, int nblks, - IdxType a, - IdxType b, + kperm_params fp, cudaStream_t stream) { permuteKernel - <<>>(perms, out, in, a, b, N, D); + <<>>(perms, out, in, fp, N, D); RAFT_CUDA_TRY(cudaPeekAtLastError()); } }; @@ -120,15 +256,26 @@ void permute(IntType* perms, IntType D, IntType N, bool rowMajor, - cudaStream_t stream) + cudaStream_t stream, + uint64_t key) { + if (N <= 0) { return; } + + if (out == nullptr) { + constexpr int ITEMS_PER_THREAD = 8; + kperm_params fp = make_kperm_params(uint64_t(N), key); + auto nblks = raft::ceildiv(N, IntType(TPB * ITEMS_PER_THREAD)); + permsOnlyKernel + <<>>(perms, fp, N); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + return; + } + auto nblks = raft::ceildiv(N, (IntType)TPB); - // 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; + // build the keyed Feistel schedule for [0, N) once on the host from the + // caller-supplied key; the same schedule is passed as an argument + kperm_params fp = make_kperm_params(uint64_t(N), key); if (rowMajor) { permute_impl_t::permuteImpl( - perms, out, in, N, D, nblks, a, b, stream); + perms, out, in, N, D, nblks, fp, stream); } } diff --git a/cpp/include/raft/random/permute.cuh b/cpp/include/raft/random/permute.cuh index 6a308d1fd4..2d7a4268bb 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,7 +129,8 @@ void permute(raft::resources const& handle, D, N, is_row_major, - resource::get_cuda_stream(handle)); + resource::get_cuda_stream(handle), + key); } } @@ -142,7 +146,8 @@ 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 +164,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 +187,85 @@ 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); +} + +/** + * @deprecated Use the overload that takes an explicit @c uint64_t key. + * + * @warning BEHAVIOR CHANGE: this shim always uses @c key=0 and therefore + * returns the same fixed permutation for every call with the same + * @c in.extent(0). The old implementation drew the permutation from + * @c rand() and produced a different result on each call. Pass a varying + * key to the keyed overload to restore per-call variation. + */ +template +[[deprecated( + "permute() now requires an explicit key (uint64_t). " + "BEHAVIOR CHANGE: this shim uses key=0 (same fixed permutation every call). " + "The old overload used rand() -- pass a varying key to restore per-call variation.")]] +void permute(raft::resources const& handle, + raft::device_matrix_view in, + std::optional> permsOut, + std::optional> out) +{ + permute(handle, in, permsOut, out, uint64_t{0}); +} + +/** + * @deprecated Use the overload that takes an explicit @c uint64_t key. + * + * @warning BEHAVIOR CHANGE: this shim always uses @c key=0 and therefore + * returns the same fixed permutation for every call with the same + * @c in.extent(0). The old implementation drew the permutation from + * @c rand() and produced a different result on each call. Pass a varying + * key to the keyed overload to restore per-call variation. + */ +template +[[deprecated( + "permute() now requires an explicit key (uint64_t). " + "BEHAVIOR CHANGE: this shim uses key=0 (same fixed permutation every call). " + "The old overload used rand() -- pass a varying key to restore per-call variation.")]] +void permute(raft::resources const& handle, + raft::device_matrix_view in, + PermsOutType&& permsOut, + OutType&& out) +{ + permute( + handle, in, std::forward(permsOut), std::forward(out), uint64_t{0}); +} + +/** + * @deprecated Use the overload that takes an explicit @c uint64_t key. + * + * @warning BEHAVIOR CHANGE: this shim always uses @c key=0 and therefore + * returns the same fixed permutation for every call with the same @c N. + * The old implementation drew the permutation from @c rand() and produced + * a different result on each call. Pass a varying key to the keyed overload + * to restore per-call variation. */ template +[[deprecated( + "permute() now requires an explicit key (uint64_t). " + "BEHAVIOR CHANGE: this shim uses key=0 (same fixed permutation every call). " + "The old overload used rand() -- pass a varying key to restore per-call variation.")]] void permute(IntType* perms, Type* out, const Type* in, @@ -192,7 +274,7 @@ 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, uint64_t{0}); } }; // namespace random diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index 44b5b9c66c..fb2fc1c23d 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 @@ -65,7 +66,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); } @@ -133,16 +134,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) { @@ -216,6 +217,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 +236,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 +253,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}, @@ -286,6 +304,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 +321,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 +338,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}, @@ -338,5 +371,65 @@ TEST_P(PermMdspanTestD, Result) } INSTANTIATE_TEST_CASE_P(PermMdspanTests, PermMdspanTestD, ::testing::ValuesIn(inputsd)); +// Each thread generates a permutation keyed by (base_seed + tid + 1) and counts +// how many elements coincide with the reference. Two independent uniform random +// permutations agree on exactly 1 position in expectation, so the total across +// all threads should be far below the 5% threshold used here. +__global__ void kperm_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; + detail::kperm_params p = detail::make_kperm_params(uint64_t(N), base_seed + uint64_t(tid) + 1); + int count = 0; + for (uint32_t i = 0; i < N; i++) { + uint32_t val = detail::kperm_index(i, p); + if (val == ref_perm[i] || val == i) { count++; } + } + match_counts[tid] = count; +} + +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; + + // Build reference permutation on the host. + detail::kperm_params ref_p = detail::make_kperm_params(uint64_t(N), base_seed); + std::vector h_ref(N); + for (uint32_t i = 0; i < N; i++) { + h_ref[i] = detail::kperm_index(i, ref_p); + } + + 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); + raft::update_device(d_ref.data(), h_ref.data(), N, stream); + + kperm_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 c : h_matches) { + total_matches += c; + } + + 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