From 4858ca7b8fbe61d566ee74d16017c301224574be Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 8 Jul 2026 15:38:28 +0100 Subject: [PATCH 01/29] Adding Feistel network based permute --- cpp/include/raft/random/detail/permute.cuh | 212 +++++++++++++++++---- 1 file changed, 177 insertions(+), 35 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index d991d6fc50..24e5d767dc 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -12,15 +12,164 @@ #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). The previous affine map out = (a*x + b) mod N is a valid + * bijection but is *linear*: a single input-bit flip reaches only the higher + * output bits (poor avalanche), so nearby indices map to structured, correlated + * destinations. The Feistel construction below is a non-linear bijection with + * near-ideal strict-avalanche behavior, so the permutation looks random while + * remaining exactly collision-free. + * + * A balanced Feistel network is a permutation over n-bit words for *any* round + * function -- invertibility does not depend on the round function being a + * bijection, which is why an fmix32-keyed mixer is a valid round function. We + * run 4 rounds (empirically enough for near-ideal avalanche at these widths). + * + * 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. + */ + +// MurmurHash3 32-bit finalizer. Diffuses all input bits into all output bits; +// used as the (non-invertible, but that is fine) Feistel round function. +HDI uint32_t feistel_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) and not on the index being permuted, so it is +// derived once on the host and passed by value to every thread. This avoids +// re-deriving the key schedule and masks in each thread -- and on each +// cycle-walk retry -- which would otherwise be identical redundant work. +// +// The round count is fixed at ROUNDS (4), which is enough for near-ideal +// avalanche at these widths. Keeping it a compile-time constant lets the round +// loop in feistel_mix_nbits unroll fully with the even/odd branch folded out. +struct feistel_permute_params { + static constexpr int ROUNDS = 4; + + uint64_t U; // N (domain size); U <= 1 selects the identity permutation + uint64_t mask_n; // low-n-bit mask, n = ceil(log2(N)) (the cycle-walk domain) + uint32_t mask_b; // low-half mask (b bits); high half needs no mask (see below) + uint32_t a; // high part width, ceil(n/2), 1..32 + uint32_t b; // low part width, floor(n/2), 1..32 (n is clamped to >= 2) + 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. +inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t key) +{ + feistel_permute_params p{}; + p.U = 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) { + pow2 <<= 1; + ++n; + } + + // Clamp the width to at least 2 bits so the low half always has >= 1 bit + // (b = floor(n/2) >= 1). Only N == 2 is affected (n = 1 -> 2): it now permutes + // over the 4-element domain [0, 4) and cycle-walks the two dead values, still a + // valid, terminating permutation of {0, 1}. 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 = uint32_t((n + 1) / 2); // high part width, 1..32 + const uint32_t b = uint32_t(n / 2); // low part width, 1..32 + p.a = a; + p.b = b; + p.mask_n = (n >= 64) ? ~uint64_t(0) : ((uint64_t(1) << n) - 1); + p.mask_b = (b >= 32) ? ~uint32_t(0) : ((uint32_t(1) << b) - 1); // b >= 1 here + + // 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 kbase = feistel_fmix32(feistel_fmix32(klo) ^ khi); + for (int i = 0; i < feistel_permute_params::ROUNDS; i++) { + p.prefix[i] = feistel_fmix32(kbase ^ (uint32_t(i) + 0x9e3779b9u)); + } + return p; +} + +// One balanced Feistel permutation over the low n bits of m, using the schedule +// precomputed in p. The n bits split into a high part of ceil(n/2) bits and a +// low part of floor(n/2) bits (each <= 32, so each half fits fmix32). Each round +// XORs an fmix32 mix of one half, keyed per round, into the other half, +// alternating which half is updated. Every step is invertible, so the whole map +// is a bijection over [0, 2^n). +// +// 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). Since +// n is clamped to >= 2, both halves always have >= 1 bit, so every round is +// unconditional -- no branch and no loop overhead. +HDI uint64_t feistel_mix_nbits(uint64_t m, const feistel_permute_params& p) +{ + m &= p.mask_n; + uint32_t L = uint32_t(m & p.mask_b); // low floor(n/2) bits + uint32_t H = uint32_t(m >> p.b); // high ceil(n/2) bits (n - b = a bits, already masked) + + // Each round's mix is fmix32(...) >> (32 - width), so it already has <= a (or + // <= b) bits; XORing into a same-width half keeps it within that width, so no + // per-round masking is needed. H stays a bits and L stays b bits throughout. + H ^= feistel_fmix32(L ^ p.prefix[0]) >> (32 - p.a); // round 0 (even): H from L + L ^= feistel_fmix32(H ^ p.prefix[1]) >> (32 - p.b); // round 1 (odd): L from H + H ^= feistel_fmix32(L ^ p.prefix[2]) >> (32 - p.a); // round 2 (even): H from L + L ^= feistel_fmix32(H ^ p.prefix[3]) >> (32 - p.b); // round 3 (odd): L from H + return (uint64_t(H) << p.b) | uint64_t(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). Cycle walking is guaranteed to terminate only when it +// starts inside the target set: an in-range idx sits on an orbit that returns to +// it (in range), so the walk always halts. An out-of-range idx can land on a +// cycle contained entirely in the dead zone [N, 2^n) and loop forever, so callers +// must not pass idx >= N (e.g. gate padding threads whose tid >= N). +template +HDI IdxType feistel_permute_index(IdxType idx, const feistel_permute_params& p) +{ + if (p.U <= 1) return idx; // 0- or 1-element domain: identity + + // 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). + uint64_t x = uint64_t(idx); + do { + x = feistel_mix_nbits(x, p); + } while (x >= p.U); + return IdxType(x); +} + 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, feistel_permute_params fp, IdxType N, IdxType D) { namespace cg = cooperative_groups; const int WARP_SIZE = 32; @@ -28,9 +177,12 @@ RAFT_KERNEL permuteKernel( 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; + // to be preferable to the reverse, especially for column major. + // Only in-range slots get a permuted source index: feistel_permute_index + // requires idx in [0, N) to guarantee the cycle walk terminates, and padding + // threads (tid >= N) never have their inIdx read (their row is skipped below). IntType outIdx = tid; + IntType inIdx = (tid < N) ? feistel_permute_index(IntType(tid), fp) : IntType(0); if (perms != nullptr && tid < N) { perms[outIdx] = inIdx; } @@ -39,19 +191,19 @@ 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); that exchange is purely + // intra-warp, so a warp shuffle broadcasts it register-to-register -- no + // shared memory or block barrier required. 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 +224,7 @@ struct permute_impl_t { IdxType N, IdxType D, int nblks, - IdxType a, - IdxType b, + feistel_permute_params fp, cudaStream_t stream) { // determine vector type and set new pointers @@ -85,11 +236,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 +254,11 @@ struct permute_impl_t { IdxType N, IdxType D, int nblks, - IdxType a, - IdxType b, + feistel_permute_params fp, cudaStream_t stream) { permuteKernel - <<>>(perms, out, in, a, b, N, D); + <<>>(perms, out, in, fp, N, D); RAFT_CUDA_TRY(cudaPeekAtLastError()); } }; @@ -124,11 +274,10 @@ void permute(IntType* perms, { 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; + // draw one 64-bit key and build the keyed Feistel schedule for [0, N) once on + // the host; the same schedule is broadcast (by value) to every thread + uint64_t key = (uint64_t(rand()) << 32) ^ uint64_t(rand()); + feistel_permute_params fp = make_feistel_permute_params(uint64_t(N), key); if (rowMajor) { permute_impl_t 0) ? 16 / sizeof(Type) : 1>::permuteImpl(perms, - out, - in, - N, - D, - nblks, - a, - b, - stream); + (16 / sizeof(Type) > 0) ? 16 / sizeof(Type) : 1>::permuteImpl( + perms, out, in, N, D, nblks, fp, stream); } else { permute_impl_t::permuteImpl( - perms, out, in, N, D, nblks, a, b, stream); + perms, out, in, N, D, nblks, fp, stream); } } From 0461065aa271db0b84ecca44c10d076eb7c55e77 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 8 Jul 2026 15:48:18 +0100 Subject: [PATCH 02/29] Adding permutation key as a parameter --- cpp/bench/prims/random/permute.cu | 10 +++++++-- .../raft/random/detail/make_regression.cuh | 10 +++++++-- cpp/include/raft/random/detail/permute.cuh | 9 ++++---- cpp/include/raft/random/permute.cuh | 21 +++++++++++++------ cpp/tests/random/permute.cu | 10 ++++----- 5 files changed, 41 insertions(+), 19 deletions(-) diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index e740b27a3d..937939473a 100644 --- a/cpp/bench/prims/random/permute.cu +++ b/cpp/bench/prims/random/permute.cu @@ -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); }); } diff --git a/cpp/include/raft/random/detail/make_regression.cuh b/cpp/include/raft/random/detail/make_regression.cuh index 773eae7b39..07bd11a5e7 100644 --- a/cpp/include/raft/random/detail/make_regression.cuh +++ b/cpp/include/raft/random/detail/make_regression.cuh @@ -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 24e5d767dc..038706f663 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -270,13 +270,14 @@ 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); - // draw one 64-bit key and build the keyed Feistel schedule for [0, N) once on - // the host; the same schedule is broadcast (by value) to every thread - uint64_t key = (uint64_t(rand()) << 32) ^ uint64_t(rand()); + // build the keyed Feistel schedule for [0, N) once on the host from the + // caller-supplied key; the same schedule is broadcast (by value) to every + // thread feistel_permute_params fp = make_feistel_permute_params(uint64_t(N), key); if (rowMajor) { diff --git a/cpp/include/raft/random/permute.cuh b/cpp/include/raft/random/permute.cuh index 6a308d1fd4..c0b0c22312 100644 --- a/cpp/include/raft/random/permute.cuh +++ b/cpp/include/raft/random/permute.cuh @@ -15,6 +15,7 @@ #include #include +#include #include #include @@ -74,6 +75,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 +130,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 +147,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 +165,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,6 +188,8 @@ 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, @@ -190,9 +198,10 @@ void permute(IntType* perms, IntType D, IntType N, bool rowMajor, - cudaStream_t stream) + cudaStream_t stream, + uint64_t key) { - detail::permute(perms, out, in, D, N, rowMajor, stream); + detail::permute(perms, out, in, D, N, rowMajor, stream, key); } }; // namespace random diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index 44b5b9c66c..e7f262f526 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -65,7 +65,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 +133,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) { From aaf0fce68879e041a633ee4883c409302cf92817 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 8 Jul 2026 16:14:20 +0100 Subject: [PATCH 03/29] Formatting changes --- cpp/bench/prims/random/permute.cu | 2 +- .../raft/random/detail/make_regression.cuh | 2 +- cpp/include/raft/random/detail/permute.cuh | 22 ++++++++++++------- cpp/include/raft/random/permute.cuh | 2 +- 4 files changed, 17 insertions(+), 11 deletions(-) diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index 937939473a..09d7374a5c 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 */ diff --git a/cpp/include/raft/random/detail/make_regression.cuh b/cpp/include/raft/random/detail/make_regression.cuh index 07bd11a5e7..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 */ diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 038706f663..3342d4d311 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 */ @@ -69,11 +69,11 @@ HDI uint32_t feistel_fmix32(uint32_t h) struct feistel_permute_params { static constexpr int ROUNDS = 4; - uint64_t U; // N (domain size); U <= 1 selects the identity permutation - uint64_t mask_n; // low-n-bit mask, n = ceil(log2(N)) (the cycle-walk domain) - uint32_t mask_b; // low-half mask (b bits); high half needs no mask (see below) - uint32_t a; // high part width, ceil(n/2), 1..32 - uint32_t b; // low part width, floor(n/2), 1..32 (n is clamped to >= 2) + uint64_t U; // N (domain size); U <= 1 selects the identity permutation + uint64_t mask_n; // low-n-bit mask, n = ceil(log2(N)) (the cycle-walk domain) + uint32_t mask_b; // low-half mask (b bits); high half needs no mask (see below) + uint32_t a; // high part width, ceil(n/2), 1..32 + uint32_t b; // low part width, floor(n/2), 1..32 (n is clamped to >= 2) uint32_t prefix[ROUNDS]; // per-round key-derived mixing constant }; @@ -286,8 +286,14 @@ void permute(IntType* perms, IdxType, TPB, true, - (16 / sizeof(Type) > 0) ? 16 / sizeof(Type) : 1>::permuteImpl( - perms, out, in, N, D, nblks, fp, stream); + (16 / sizeof(Type) > 0) ? 16 / sizeof(Type) : 1>::permuteImpl(perms, + out, + in, + N, + D, + nblks, + fp, + stream); } else { permute_impl_t::permuteImpl( 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 c0b0c22312..fdecfcc62f 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 */ From 392651ed5d1c8b97c018794d55d411e982b6e15f Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 8 Jul 2026 16:31:40 +0100 Subject: [PATCH 04/29] Adding randomness check for permute --- cpp/tests/random/permute.cu | 72 ++++++++++++++++++++++++++++++++++++- 1 file changed, 71 insertions(+), 1 deletion(-) diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index e7f262f526..b9b4c3707b 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 */ @@ -13,6 +13,7 @@ #include #include +#include #include namespace raft { @@ -338,5 +339,74 @@ TEST_P(PermMdspanTestD, Result) } INSTANTIATE_TEST_CASE_P(PermMdspanTests, PermMdspanTestD, ::testing::ValuesIn(inputsd)); +/* + * Randomness test for the Feistel-based permutation via a chi-square + * goodness-of-fit test. + * + * For a random permutation of [0, N), the consecutive differences + * d[i] = (perm[i+1] - perm[i]) mod N are approximately uniformly distributed + * over [0, N). We bin them into B equal-width buckets and compare the observed + * counts against the flat expected count E = (N-1)/B using Pearson's statistic + * + * chi2 = sum_b (O_b - E)^2 / E. + * + * Under the uniform (random) hypothesis chi2 follows a chi-square distribution + * with dof = B-1, so it concentrates around its mean (dof) with std sqrt(2*dof) + * -- for B = 1000 that is 999 +/- ~45. Structure in the permutation (e.g. a + * predictable relationship between adjacent outputs) piles the differences into + * a few buckets and inflates chi2 far above dof. + * + * The key is an explicit, fixed parameter, so the permutation and the resulting + * statistic are fully deterministic (not flaky). We assert chi2 lands in a + * generous central band around dof; a well-mixing permutation clears this + * easily, while a non-uniform one falls outside it. + */ +TEST(PermuteRandomness, ConsecutiveDifferencesChiSquareUniform) +{ + // N is a non-power-of-two (prime) so the permutation also cycle-walks. + const int N = 100003; + const uint64_t key = 1234567890ULL; + const int B = 1000; // number of histogram buckets + + raft::resources handle; + auto stream = resource::get_cuda_stream(handle); + rmm::device_uvector perms(N, stream); + // null in/out: only the permutation indices are produced. + permute(perms.data(), + static_cast(nullptr), + static_cast(nullptr), + /* D = */ 1, + N, + /* rowMajor = */ true, + stream, + key); + std::vector p(N); + raft::update_host(p.data(), perms.data(), N, stream); + resource::sync_stream(handle); + + std::vector hist(B, 0); + const int total = N - 1; + for (int i = 0; i + 1 < N; ++i) { + int d = static_cast(((static_cast(p[i + 1]) - p[i]) % N + N) % N); + int b = static_cast((static_cast(d) * B) / N); // map [0, N) -> [0, B) + if (b >= B) b = B - 1; + hist[b]++; + } + + const double expected = static_cast(total) / B; + double chi2 = 0.0; + for (int b = 0; b < B; ++b) { + double diff = static_cast(hist[b]) - expected; + chi2 += diff * diff / expected; + } + + const double dof = B - 1; // expected mean of chi2 under uniformity + // Generous central band (mean +/- ~11 std) to demonstrate uniformity without + // being flaky. A random permutation sits near dof; structure pushes chi2 high. + EXPECT_GT(chi2, 0.5 * dof) << "chi2=" << chi2 << " suspiciously low (dof=" << dof << ")"; + EXPECT_LT(chi2, 1.5 * dof) << "chi2=" << chi2 << " too high -- consecutive differences are " + << "not uniform, permutation is not random (dof=" << dof << ")"; +} + } // end namespace random } // end namespace raft From 2f69c68fcee9b8a4c38a9d228309993096259df1 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 09:36:57 +0100 Subject: [PATCH 05/29] Removing redundant header inclusion --- cpp/include/raft/random/permute.cuh | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/include/raft/random/permute.cuh b/cpp/include/raft/random/permute.cuh index fdecfcc62f..de459a4245 100644 --- a/cpp/include/raft/random/permute.cuh +++ b/cpp/include/raft/random/permute.cuh @@ -15,7 +15,6 @@ #include #include -#include #include #include From 40c15f040b92292788215642c937a34be6765151 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 10:29:15 +0100 Subject: [PATCH 06/29] Tidying up comments --- cpp/include/raft/random/detail/permute.cuh | 96 +++++++--------------- 1 file changed, 31 insertions(+), 65 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 3342d4d311..6fa428443c 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -25,17 +25,7 @@ namespace detail { * * 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). The previous affine map out = (a*x + b) mod N is a valid - * bijection but is *linear*: a single input-bit flip reaches only the higher - * output bits (poor avalanche), so nearby indices map to structured, correlated - * destinations. The Feistel construction below is a non-linear bijection with - * near-ideal strict-avalanche behavior, so the permutation looks random while - * remaining exactly collision-free. - * - * A balanced Feistel network is a permutation over n-bit words for *any* round - * function -- invertibility does not depend on the round function being a - * bijection, which is why an fmix32-keyed mixer is a valid round function. We - * run 4 rounds (empirically enough for near-ideal avalanche at these widths). + * 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 @@ -43,10 +33,13 @@ namespace detail { * 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 */ // MurmurHash3 32-bit finalizer. Diffuses all input bits into all output bits; -// used as the (non-invertible, but that is fine) Feistel round function. +// used as the Feistel round function. Reference: +// https://github.com/aappleby/smhasher/blob/07bb4de10a63e8cc2e1724865454eba635742383/src/MurmurHash3.cpp#L68 HDI uint32_t feistel_fmix32(uint32_t h) { h ^= h >> 16; @@ -58,22 +51,18 @@ HDI uint32_t feistel_fmix32(uint32_t h) } // Precomputed schedule for the keyed Feistel permutation of [0, N). Every field -// here depends only on (N, key) and not on the index being permuted, so it is -// derived once on the host and passed by value to every thread. This avoids -// re-deriving the key schedule and masks in each thread -- and on each -// cycle-walk retry -- which would otherwise be identical redundant work. +// 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. Keeping it a compile-time constant lets the round -// loop in feistel_mix_nbits unroll fully with the even/odd branch folded out. +// avalanche at these widths. struct feistel_permute_params { static constexpr int ROUNDS = 4; - uint64_t U; // N (domain size); U <= 1 selects the identity permutation - uint64_t mask_n; // low-n-bit mask, n = ceil(log2(N)) (the cycle-walk domain) - uint32_t mask_b; // low-half mask (b bits); high half needs no mask (see below) - uint32_t a; // high part width, ceil(n/2), 1..32 - uint32_t b; // low part width, floor(n/2), 1..32 (n is clamped to >= 2) + 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 mask_b; // low-half mask (b bits); + uint32_t a; // widht of the high part + uint32_t b; // width of the low part uint32_t prefix[ROUNDS]; // per-round key-derived mixing constant }; @@ -82,7 +71,7 @@ struct feistel_permute_params { inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t key) { feistel_permute_params p{}; - p.U = N; + 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). @@ -94,14 +83,12 @@ inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t k } // Clamp the width to at least 2 bits so the low half always has >= 1 bit - // (b = floor(n/2) >= 1). Only N == 2 is affected (n = 1 -> 2): it now permutes - // over the 4-element domain [0, 4) and cycle-walks the two dead values, still a - // valid, terminating permutation of {0, 1}. This makes b == 0 impossible, so - // the device mixer needs no per-round guard against a zero-width low half. + // 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 = uint32_t((n + 1) / 2); // high part width, 1..32 - const uint32_t b = uint32_t(n / 2); // low part width, 1..32 + const uint32_t a = uint32_t((n + 1) / 2); // high part width + const uint32_t b = uint32_t(n / 2); // low part width p.a = a; p.b = b; p.mask_n = (n >= 64) ? ~uint64_t(0) : ((uint64_t(1) << n) - 1); @@ -109,35 +96,26 @@ inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t k // 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 klo = uint32_t(key); + const uint32_t khi = uint32_t(key >> 32); + const uint32_t golden_ration = 0x9e3779b9u; + const uint32_t kbase = feistel_fmix32(feistel_fmix32(klo) ^ khi); for (int i = 0; i < feistel_permute_params::ROUNDS; i++) { - p.prefix[i] = feistel_fmix32(kbase ^ (uint32_t(i) + 0x9e3779b9u)); + p.prefix[i] = feistel_fmix32(kbase ^ (uint32_t(i) + golden_ratio)); } return p; } -// One balanced Feistel permutation over the low n bits of m, using the schedule -// precomputed in p. The n bits split into a high part of ceil(n/2) bits and a -// low part of floor(n/2) bits (each <= 32, so each half fits fmix32). Each round -// XORs an fmix32 mix of one half, keyed per round, into the other half, -// alternating which half is updated. Every step is invertible, so the whole map -// is a bijection over [0, 2^n). -// +// 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). Since -// n is clamped to >= 2, both halves always have >= 1 bit, so every round is -// unconditional -- no branch and no loop overhead. +// part (a bits), odd rounds mix the high part into the low part (b bits). HDI uint64_t feistel_mix_nbits(uint64_t m, const feistel_permute_params& p) { m &= p.mask_n; - uint32_t L = uint32_t(m & p.mask_b); // low floor(n/2) bits - uint32_t H = uint32_t(m >> p.b); // high ceil(n/2) bits (n - b = a bits, already masked) + uint32_t L = uint32_t(m & p.mask_b); // low + uint32_t H = uint32_t(m >> p.b); // high - // Each round's mix is fmix32(...) >> (32 - width), so it already has <= a (or - // <= b) bits; XORing into a same-width half keeps it within that width, so no - // per-round masking is needed. H stays a bits and L stays b bits throughout. H ^= feistel_fmix32(L ^ p.prefix[0]) >> (32 - p.a); // round 0 (even): H from L L ^= feistel_fmix32(H ^ p.prefix[1]) >> (32 - p.b); // round 1 (odd): L from H H ^= feistel_fmix32(L ^ p.prefix[2]) >> (32 - p.a); // round 2 (even): H from L @@ -148,22 +126,18 @@ HDI uint64_t feistel_mix_nbits(uint64_t m, const feistel_permute_params& p) // 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). Cycle walking is guaranteed to terminate only when it -// starts inside the target set: an in-range idx sits on an orbit that returns to -// it (in range), so the walk always halts. An out-of-range idx can land on a -// cycle contained entirely in the dead zone [N, 2^n) and loop forever, so callers -// must not pass idx >= N (e.g. gate padding threads whose tid >= N). +// idx MUST lie in [0, N) for the cycle walk to halt. template HDI IdxType feistel_permute_index(IdxType idx, const feistel_permute_params& p) { - if (p.U <= 1) return idx; // 0- or 1-element domain: identity + if (p.N <= 1) return idx; // 0- or 1-element domain: identity // 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). uint64_t x = uint64_t(idx); do { x = feistel_mix_nbits(x, p); - } while (x >= p.U); + } while (x >= p.N); return IdxType(x); } @@ -176,11 +150,6 @@ RAFT_KERNEL permuteKernel( 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. - // Only in-range slots get a permuted source index: feistel_permute_index - // requires idx in [0, N) to guarantee the cycle walk terminates, and padding - // threads (tid >= N) never have their inIdx read (their row is skipped below). IntType outIdx = tid; IntType inIdx = (tid < N) ? feistel_permute_index(IntType(tid), fp) : IntType(0); @@ -193,9 +162,7 @@ RAFT_KERNEL permuteKernel( // 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); that exchange is purely - // intra-warp, so a warp shuffle broadcasts it register-to-register -- no - // shared memory or block barrier required. + // Copying row i needs lane i's (inIdx, outIdx); int laneID = threadIdx.x % WARP_SIZE; for (int i = 0; i < WARP_SIZE; ++i) { IntType inIdxI = warp.shfl(inIdx, i); @@ -276,8 +243,7 @@ void permute(IntType* perms, auto nblks = raft::ceildiv(N, (IntType)TPB); // build the keyed Feistel schedule for [0, N) once on the host from the - // caller-supplied key; the same schedule is broadcast (by value) to every - // thread + // caller-supplied key; the same schedule is passed as an argument feistel_permute_params fp = make_feistel_permute_params(uint64_t(N), key); if (rowMajor) { From 5ec3246263809cd16d832f18a7c7edb727317ef4 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 14:41:26 +0100 Subject: [PATCH 07/29] Fixing a typo --- cpp/include/raft/random/detail/permute.cuh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 6fa428443c..cb649680d2 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -96,9 +96,9 @@ inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t k // 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_ration = 0x9e3779b9u; + 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 = feistel_fmix32(feistel_fmix32(klo) ^ khi); for (int i = 0; i < feistel_permute_params::ROUNDS; i++) { From 46b9e67bde749fd12924dbe71f766def2d30df3e Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 14:41:47 +0100 Subject: [PATCH 08/29] Removing a narrow test --- cpp/tests/random/permute.cu | 72 ------------------------------------- 1 file changed, 72 deletions(-) diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index b9b4c3707b..15d32146cb 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -338,75 +338,3 @@ TEST_P(PermMdspanTestD, Result) _PERMTEST_BODY(test_data_type); } INSTANTIATE_TEST_CASE_P(PermMdspanTests, PermMdspanTestD, ::testing::ValuesIn(inputsd)); - -/* - * Randomness test for the Feistel-based permutation via a chi-square - * goodness-of-fit test. - * - * For a random permutation of [0, N), the consecutive differences - * d[i] = (perm[i+1] - perm[i]) mod N are approximately uniformly distributed - * over [0, N). We bin them into B equal-width buckets and compare the observed - * counts against the flat expected count E = (N-1)/B using Pearson's statistic - * - * chi2 = sum_b (O_b - E)^2 / E. - * - * Under the uniform (random) hypothesis chi2 follows a chi-square distribution - * with dof = B-1, so it concentrates around its mean (dof) with std sqrt(2*dof) - * -- for B = 1000 that is 999 +/- ~45. Structure in the permutation (e.g. a - * predictable relationship between adjacent outputs) piles the differences into - * a few buckets and inflates chi2 far above dof. - * - * The key is an explicit, fixed parameter, so the permutation and the resulting - * statistic are fully deterministic (not flaky). We assert chi2 lands in a - * generous central band around dof; a well-mixing permutation clears this - * easily, while a non-uniform one falls outside it. - */ -TEST(PermuteRandomness, ConsecutiveDifferencesChiSquareUniform) -{ - // N is a non-power-of-two (prime) so the permutation also cycle-walks. - const int N = 100003; - const uint64_t key = 1234567890ULL; - const int B = 1000; // number of histogram buckets - - raft::resources handle; - auto stream = resource::get_cuda_stream(handle); - rmm::device_uvector perms(N, stream); - // null in/out: only the permutation indices are produced. - permute(perms.data(), - static_cast(nullptr), - static_cast(nullptr), - /* D = */ 1, - N, - /* rowMajor = */ true, - stream, - key); - std::vector p(N); - raft::update_host(p.data(), perms.data(), N, stream); - resource::sync_stream(handle); - - std::vector hist(B, 0); - const int total = N - 1; - for (int i = 0; i + 1 < N; ++i) { - int d = static_cast(((static_cast(p[i + 1]) - p[i]) % N + N) % N); - int b = static_cast((static_cast(d) * B) / N); // map [0, N) -> [0, B) - if (b >= B) b = B - 1; - hist[b]++; - } - - const double expected = static_cast(total) / B; - double chi2 = 0.0; - for (int b = 0; b < B; ++b) { - double diff = static_cast(hist[b]) - expected; - chi2 += diff * diff / expected; - } - - const double dof = B - 1; // expected mean of chi2 under uniformity - // Generous central band (mean +/- ~11 std) to demonstrate uniformity without - // being flaky. A random permutation sits near dof; structure pushes chi2 high. - EXPECT_GT(chi2, 0.5 * dof) << "chi2=" << chi2 << " suspiciously low (dof=" << dof << ")"; - EXPECT_LT(chi2, 1.5 * dof) << "chi2=" << chi2 << " too high -- consecutive differences are " - << "not uniform, permutation is not random (dof=" << dof << ")"; -} - -} // end namespace random -} // end namespace raft From be8371db67a992e915648e85d6a2ec22cce84cd2 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 14:55:35 +0100 Subject: [PATCH 09/29] Undoing delete --- cpp/tests/random/permute.cu | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index 15d32146cb..a8dab1ed13 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -338,3 +338,6 @@ TEST_P(PermMdspanTestD, Result) _PERMTEST_BODY(test_data_type); } INSTANTIATE_TEST_CASE_P(PermMdspanTests, PermMdspanTestD, ::testing::ValuesIn(inputsd)); + +} // end namespace random +} // end namespace raft From 6f75295b7c66bfd09a5a098165f8760d7ae8ed1c Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 15:02:38 +0100 Subject: [PATCH 10/29] Removing redundant header --- cpp/tests/random/permute.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index a8dab1ed13..a90f209236 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -13,7 +13,6 @@ #include #include -#include #include namespace raft { From 27019e4ed145de1d8c1f429d4d5d0036257747a1 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 20 Jul 2026 16:52:48 +0100 Subject: [PATCH 11/29] Adding permute only benchmark --- cpp/bench/prims/random/permute.cu | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index 09d7374a5c..208482ed12 100644 --- a/cpp/bench/prims/random/permute.cu +++ b/cpp/bench/prims/random/permute.cu @@ -72,4 +72,34 @@ 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 From 2b71259de294ba6717558824314d91d3ba172acf Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 20 Jul 2026 17:46:02 +0100 Subject: [PATCH 12/29] Restoring the permute only kernel --- cpp/include/raft/random/detail/permute.cuh | 38 ++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index cb649680d2..b843d45bbd 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -14,6 +14,7 @@ #include #include +#include namespace raft { namespace random { @@ -141,6 +142,28 @@ HDI IdxType feistel_permute_index(IdxType idx, const feistel_permute_params& p) return IdxType(x); } +template +RAFT_KERNEL permsOnlyKernel32(uint32_t* perms, feistel_permute_params fp, uint32_t N) +{ + uint32_t base = uint32_t(blockIdx.x) * uint32_t(TPB * ITEMS_PER_THREAD) + threadIdx.x; +#pragma unroll + for (int i = 0; i < ITEMS_PER_THREAD; i++) { + uint32_t idx = base + uint32_t(i * TPB); + if (idx < N) { perms[idx] = feistel_permute_index(idx, fp); } + } +} + +template +RAFT_KERNEL permsOnlyKernel(IntType* perms, feistel_permute_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(feistel_permute_index(idx, fp)); } + } +} + template RAFT_KERNEL permuteKernel( IntType* perms, Type* out, const Type* in, feistel_permute_params fp, IdxType N, IdxType D) @@ -240,6 +263,21 @@ void permute(IntType* perms, cudaStream_t stream, uint64_t key) { + if (out == nullptr) { + constexpr int ITEMS_PER_THREAD = 8; + feistel_permute_params fp = make_feistel_permute_params(uint64_t(N), key); + if constexpr (std::is_same_v) { + auto nblks = raft::ceildiv(N, IntType(TPB * ITEMS_PER_THREAD)); + permsOnlyKernel32<<>>(perms, fp, N); + } else { + 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); // build the keyed Feistel schedule for [0, N) once on the host from the From 24bd5896fb8d2094d157b9f4f60f76901eeb57a8 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 11:55:11 +0100 Subject: [PATCH 13/29] Reducing the complexity of round function to achieve better bandwidth --- cpp/include/raft/random/detail/permute.cuh | 40 +++++++++++----------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index b843d45bbd..2cc3fe6384 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -43,11 +43,11 @@ namespace detail { // https://github.com/aappleby/smhasher/blob/07bb4de10a63e8cc2e1724865454eba635742383/src/MurmurHash3.cpp#L68 HDI uint32_t feistel_fmix32(uint32_t h) { - h ^= h >> 16; + // h ^= h >> 16; h *= 0x85ebca6bu; h ^= h >> 13; - h *= 0xc2b2ae35u; - h ^= h >> 16; + // h *= 0xc2b2ae35u; + // h ^= h >> 16; return h; } @@ -61,9 +61,9 @@ struct feistel_permute_params { 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 mask_b; // low-half mask (b bits); - uint32_t a; // widht of the high part - uint32_t b; // width of the low part + 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 }; @@ -88,12 +88,12 @@ inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t k // guard against a zero-width low half. if (n < 2) n = 2; - const uint32_t a = uint32_t((n + 1) / 2); // high part width - const uint32_t b = uint32_t(n / 2); // low part width - p.a = a; - p.b = b; - p.mask_n = (n >= 64) ? ~uint64_t(0) : ((uint64_t(1) << n) - 1); - p.mask_b = (b >= 32) ? ~uint32_t(0) : ((uint32_t(1) << b) - 1); // b >= 1 here + 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. @@ -114,14 +114,14 @@ inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t k HDI uint64_t feistel_mix_nbits(uint64_t m, const feistel_permute_params& p) { m &= p.mask_n; - uint32_t L = uint32_t(m & p.mask_b); // low - uint32_t H = uint32_t(m >> p.b); // high - - H ^= feistel_fmix32(L ^ p.prefix[0]) >> (32 - p.a); // round 0 (even): H from L - L ^= feistel_fmix32(H ^ p.prefix[1]) >> (32 - p.b); // round 1 (odd): L from H - H ^= feistel_fmix32(L ^ p.prefix[2]) >> (32 - p.a); // round 2 (even): H from L - L ^= feistel_fmix32(H ^ p.prefix[3]) >> (32 - p.b); // round 3 (odd): L from H - return (uint64_t(H) << p.b) | uint64_t(L); + uint32_t L = uint32_t(m & p.b); // low b_bits + uint32_t H = uint32_t(m >> p.b_bits); // high a_bits + + H ^= feistel_fmix32(L ^ p.prefix[0]) & p.a; // round 0 (even): H from L + L ^= feistel_fmix32(H ^ p.prefix[1]) & p.b; // round 1 (odd): L from H + H ^= feistel_fmix32(L ^ p.prefix[2]) & p.a; // round 2 (even): H from L + L ^= feistel_fmix32(H ^ p.prefix[3]) & p.b; // round 3 (odd): L from H + return (uint64_t(H) << p.b_bits) | uint64_t(L); } // Keyed permutation of [0, N): returns the index that output slot idx pulls From deb08b81313803daf3ef65214eb9e7b80c40e7f2 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 12:02:11 +0100 Subject: [PATCH 14/29] Removing 32-bit specialization, as it is not needed anymore --- cpp/include/raft/random/detail/permute.cuh | 26 +++++----------------- 1 file changed, 5 insertions(+), 21 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 2cc3fe6384..3b2760c1a5 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -142,17 +142,6 @@ HDI IdxType feistel_permute_index(IdxType idx, const feistel_permute_params& p) return IdxType(x); } -template -RAFT_KERNEL permsOnlyKernel32(uint32_t* perms, feistel_permute_params fp, uint32_t N) -{ - uint32_t base = uint32_t(blockIdx.x) * uint32_t(TPB * ITEMS_PER_THREAD) + threadIdx.x; -#pragma unroll - for (int i = 0; i < ITEMS_PER_THREAD; i++) { - uint32_t idx = base + uint32_t(i * TPB); - if (idx < N) { perms[idx] = feistel_permute_index(idx, fp); } - } -} - template RAFT_KERNEL permsOnlyKernel(IntType* perms, feistel_permute_params fp, IdxType N) { @@ -264,16 +253,11 @@ void permute(IntType* perms, uint64_t key) { if (out == nullptr) { - constexpr int ITEMS_PER_THREAD = 8; - feistel_permute_params fp = make_feistel_permute_params(uint64_t(N), key); - if constexpr (std::is_same_v) { - auto nblks = raft::ceildiv(N, IntType(TPB * ITEMS_PER_THREAD)); - permsOnlyKernel32<<>>(perms, fp, N); - } else { - auto nblks = raft::ceildiv(N, (IntType)(TPB * ITEMS_PER_THREAD)); - permsOnlyKernel - <<>>(perms, fp, N); - } + constexpr int ITEMS_PER_THREAD = 8; + feistel_permute_params fp = make_feistel_permute_params(uint64_t(N), key); + auto nblks = raft::ceildiv(N, IntType(TPB * ITEMS_PER_THREAD)); + permsOnlyKernel + <<>>(perms, fp, N); RAFT_CUDA_TRY(cudaPeekAtLastError()); return; } From 3657f384e03d59b5095fe315e01429b2ab431aed Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 13:31:48 +0100 Subject: [PATCH 15/29] Converting to template arguments for avoiding type conversion in 32-bit case --- cpp/include/raft/random/detail/permute.cuh | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 3b2760c1a5..1e5df64df2 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -111,17 +111,19 @@ inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t k // 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). -HDI uint64_t feistel_mix_nbits(uint64_t m, const feistel_permute_params& p) +// WordType is uint32_t for 32-bit domains and uint64_t for larger ones. +template +HDI WordType feistel_mix_nbits(WordType m, const feistel_permute_params& p) { - m &= 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 + 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 ^= feistel_fmix32(L ^ p.prefix[0]) & p.a; // round 0 (even): H from L L ^= feistel_fmix32(H ^ p.prefix[1]) & p.b; // round 1 (odd): L from H H ^= feistel_fmix32(L ^ p.prefix[2]) & p.a; // round 2 (even): H from L L ^= feistel_fmix32(H ^ p.prefix[3]) & p.b; // round 3 (odd): L from H - return (uint64_t(H) << p.b_bits) | uint64_t(L); + return (WordType(H) << p.b_bits) | WordType(L); } // Keyed permutation of [0, N): returns the index that output slot idx pulls @@ -133,12 +135,16 @@ HDI IdxType feistel_permute_index(IdxType idx, const feistel_permute_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). - uint64_t x = uint64_t(idx); + WordType x = WordType(idx); do { x = feistel_mix_nbits(x, p); - } while (x >= p.N); + } while (x >= WordType(p.N)); return IdxType(x); } From a2a35bec73a227e95bf4dde106835a760277763a Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 13:49:08 +0100 Subject: [PATCH 16/29] Changing the names of functions for clarity --- cpp/include/raft/random/detail/permute.cuh | 47 +++++++++++----------- 1 file changed, 24 insertions(+), 23 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 1e5df64df2..dc7b806892 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -38,10 +38,11 @@ namespace detail { * Reference: https://en.wikipedia.org/wiki/Feistel_cipher */ -// MurmurHash3 32-bit finalizer. Diffuses all input bits into all output bits; +// 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 feistel_fmix32(uint32_t h) +HDI uint32_t fmix32(uint32_t h) { // h ^= h >> 16; h *= 0x85ebca6bu; @@ -56,7 +57,7 @@ HDI uint32_t feistel_fmix32(uint32_t h) // // The round count is fixed at ROUNDS (4), which is enough for near-ideal // avalanche at these widths. -struct feistel_permute_params { +struct kperm_params { static constexpr int ROUNDS = 4; uint64_t N; // N (domain size); N <= 1 selects the identity permutation @@ -69,9 +70,9 @@ struct feistel_permute_params { // Build the Feistel schedule for permuting [0, N) with the given 64-bit key. // Runs once on the host per permute() call. -inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t key) +inline kperm_params make_kperm_params(uint64_t N, uint64_t key) { - feistel_permute_params p{}; + kperm_params p{}; p.N = N; if (N <= 1) return p; // identity domain; remaining fields go unused @@ -101,9 +102,9 @@ inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t k const uint32_t khi = uint32_t(key >> 32); const uint32_t golden_ratio = 0x9e3779b9u; - const uint32_t kbase = feistel_fmix32(feistel_fmix32(klo) ^ khi); - for (int i = 0; i < feistel_permute_params::ROUNDS; i++) { - p.prefix[i] = feistel_fmix32(kbase ^ (uint32_t(i) + golden_ratio)); + 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; } @@ -113,16 +114,16 @@ inline feistel_permute_params make_feistel_permute_params(uint64_t N, uint64_t k // 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 feistel_mix_nbits(WordType m, const feistel_permute_params& p) +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 ^= feistel_fmix32(L ^ p.prefix[0]) & p.a; // round 0 (even): H from L - L ^= feistel_fmix32(H ^ p.prefix[1]) & p.b; // round 1 (odd): L from H - H ^= feistel_fmix32(L ^ p.prefix[2]) & p.a; // round 2 (even): H from L - L ^= feistel_fmix32(H ^ p.prefix[3]) & p.b; // round 3 (odd): L from H + 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); } @@ -131,7 +132,7 @@ HDI WordType feistel_mix_nbits(WordType m, const feistel_permute_params& p) // // idx MUST lie in [0, N) for the cycle walk to halt. template -HDI IdxType feistel_permute_index(IdxType idx, const feistel_permute_params& p) +HDI IdxType kperm_index(IdxType idx, const kperm_params& p) { if (p.N <= 1) return idx; // 0- or 1-element domain: identity @@ -143,25 +144,25 @@ HDI IdxType feistel_permute_index(IdxType idx, const feistel_permute_params& p) // N is a power of two this is a single application (the loop body runs once). WordType x = WordType(idx); do { - x = feistel_mix_nbits(x, p); + x = kperm_mix_nbits(x, p); } while (x >= WordType(p.N)); return IdxType(x); } template -RAFT_KERNEL permsOnlyKernel(IntType* perms, feistel_permute_params fp, IdxType N) +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(feistel_permute_index(idx, fp)); } + if (idx < N) { perms[idx] = IntType(kperm_index(idx, fp)); } } } template RAFT_KERNEL permuteKernel( - IntType* perms, Type* out, const Type* in, feistel_permute_params fp, 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; @@ -169,7 +170,7 @@ RAFT_KERNEL permuteKernel( int tid = threadIdx.x + blockIdx.x * blockDim.x; IntType outIdx = tid; - IntType inIdx = (tid < N) ? feistel_permute_index(IntType(tid), fp) : IntType(0); + IntType inIdx = (tid < N) ? kperm_index(IntType(tid), fp) : IntType(0); if (perms != nullptr && tid < N) { perms[outIdx] = inIdx; } @@ -209,7 +210,7 @@ struct permute_impl_t { IdxType N, IdxType D, int nblks, - feistel_permute_params fp, + kperm_params fp, cudaStream_t stream) { // determine vector type and set new pointers @@ -239,7 +240,7 @@ struct permute_impl_t { IdxType N, IdxType D, int nblks, - feistel_permute_params fp, + kperm_params fp, cudaStream_t stream) { permuteKernel @@ -260,7 +261,7 @@ void permute(IntType* perms, { if (out == nullptr) { constexpr int ITEMS_PER_THREAD = 8; - feistel_permute_params fp = make_feistel_permute_params(uint64_t(N), key); + kperm_params fp = make_kperm_params(uint64_t(N), key); auto nblks = raft::ceildiv(N, IntType(TPB * ITEMS_PER_THREAD)); permsOnlyKernel <<>>(perms, fp, N); @@ -272,7 +273,7 @@ void permute(IntType* perms, // 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 - feistel_permute_params fp = make_feistel_permute_params(uint64_t(N), key); + kperm_params fp = make_kperm_params(uint64_t(N), key); if (rowMajor) { permute_impl_t Date: Tue, 21 Jul 2026 14:53:34 +0100 Subject: [PATCH 17/29] Adding a test that checks for seed diversity --- cpp/include/raft/random/detail/permute.cuh | 4 +- cpp/tests/random/permute.cu | 58 ++++++++++++++++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index dc7b806892..6cd1389f8e 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -70,7 +70,7 @@ struct kperm_params { // Build the Feistel schedule for permuting [0, N) with the given 64-bit key. // Runs once on the host per permute() call. -inline kperm_params make_kperm_params(uint64_t N, uint64_t key) +HDI kperm_params make_kperm_params(uint64_t N, uint64_t key) { kperm_params p{}; p.N = N; @@ -79,7 +79,7 @@ inline kperm_params make_kperm_params(uint64_t N, uint64_t key) // n = ceil(log2(N)): smallest n with 2^n >= N (the cycle-walking domain). int n = 0; uint64_t pow2 = 1; - while (pow2 < N) { + while (pow2 < N && n < 64) { pow2 <<= 1; ++n; } diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index a90f209236..a252f1ca3b 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -338,5 +339,62 @@ 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()); + + 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 From eeb1d094a3ef8ebe6348702379974603516e9f3b Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 14:56:28 +0100 Subject: [PATCH 18/29] Formatting --- cpp/bench/prims/random/permute.cu | 4 +++- cpp/include/raft/random/detail/permute.cuh | 6 +++--- cpp/tests/random/permute.cu | 10 ++++++---- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index 208482ed12..43362fa389 100644 --- a/cpp/bench/prims/random/permute.cu +++ b/cpp/bench/prims/random/permute.cu @@ -99,7 +99,9 @@ struct permute_perms_only : public fixture { 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({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/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 6cd1389f8e..aead8b5500 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -91,8 +91,8 @@ HDI kperm_params make_kperm_params(uint64_t N, uint64_t key) 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_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); @@ -261,7 +261,7 @@ void permute(IntType* perms, { if (out == nullptr) { constexpr int ITEMS_PER_THREAD = 8; - kperm_params fp = make_kperm_params(uint64_t(N), key); + kperm_params fp = make_kperm_params(uint64_t(N), key); auto nblks = raft::ceildiv(N, IntType(TPB * ITEMS_PER_THREAD)); permsOnlyKernel <<>>(perms, fp, N); diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index a252f1ca3b..e6ccfcab16 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -348,9 +348,9 @@ __global__ void kperm_seed_diversity_kernel(const uint32_t* ref_perm, 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; + 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++; } @@ -389,7 +389,9 @@ TEST(PermTest, SeedDiversity) resource::sync_stream(handle); int total_matches = 0; - for (int c : h_matches) { total_matches += c; } + 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) From 85751ea78ca608e6e9aa1edcc00aaa4a19089f38 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 20:04:34 +0100 Subject: [PATCH 19/29] Adding deprecated APIs for compatibility --- cpp/include/raft/random/permute.cuh | 47 +++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/cpp/include/raft/random/permute.cuh b/cpp/include/raft/random/permute.cuh index de459a4245..b59628e77c 100644 --- a/cpp/include/raft/random/permute.cuh +++ b/cpp/include/raft/random/permute.cuh @@ -203,6 +203,53 @@ void permute(IntType* perms, detail::permute(perms, out, in, D, N, rowMajor, stream, key); } +/** @deprecated Pass an explicit uint64_t key. This overload uses key=0. */ +template +[[deprecated( + "permute() now requires an explicit key argument (uint64_t). " + "This overload forwards with key=0 and will be removed in a future release.")]] +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 Pass an explicit uint64_t key. This overload uses key=0. */ +template +[[deprecated( + "permute() now requires an explicit key argument (uint64_t). " + "This overload forwards with key=0 and will be removed in a future release.")]] +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 Pass an explicit uint64_t key. This overload uses key=0. */ +template +[[deprecated( + "permute() now requires an explicit key argument (uint64_t). " + "This overload forwards with key=0 and will be removed in a future release.")]] +void permute(IntType* perms, + Type* out, + const Type* in, + IntType D, + IntType N, + bool rowMajor, + cudaStream_t stream) +{ + detail::permute(perms, out, in, D, N, rowMajor, stream, uint64_t{0}); +} + }; // namespace random } // namespace raft #endif From 0383fbeffc824ab5cc07e3277c03e7d7e17e8a7a Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 20:14:37 +0100 Subject: [PATCH 20/29] Skipping kernel launch if N <= 0 --- cpp/include/raft/random/detail/permute.cuh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index aead8b5500..33faef4fd2 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -259,6 +259,8 @@ void permute(IntType* perms, 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); From 31ac615cd2916650b4ff2aa090861d6643b7c612 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 20:18:45 +0100 Subject: [PATCH 21/29] Adding small N test cases --- cpp/tests/random/permute.cu | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index e6ccfcab16..6f4f99b9b3 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -217,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}, @@ -229,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}, @@ -241,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}, @@ -287,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}, @@ -299,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}, @@ -311,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}, From d19dd18b2f6a23c92ed5614893a7bfa0b9335794 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 22:32:18 +0100 Subject: [PATCH 22/29] Adding changed behavior description in the deprecation message --- cpp/include/raft/random/permute.cuh | 45 +++++++++++++++++++++++------ 1 file changed, 36 insertions(+), 9 deletions(-) diff --git a/cpp/include/raft/random/permute.cuh b/cpp/include/raft/random/permute.cuh index b59628e77c..2d7a4268bb 100644 --- a/cpp/include/raft/random/permute.cuh +++ b/cpp/include/raft/random/permute.cuh @@ -203,11 +203,20 @@ void permute(IntType* perms, detail::permute(perms, out, in, D, N, rowMajor, stream, key); } -/** @deprecated Pass an explicit uint64_t key. This overload uses key=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 argument (uint64_t). " - "This overload forwards with key=0 and will be removed in a future release.")]] + "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, @@ -216,15 +225,24 @@ void permute(raft::resources const& handle, permute(handle, in, permsOut, out, uint64_t{0}); } -/** @deprecated Pass an explicit uint64_t key. This overload uses key=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 argument (uint64_t). " - "This overload forwards with key=0 and will be removed in a future release.")]] + "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, @@ -234,11 +252,20 @@ void permute(raft::resources const& handle, handle, in, std::forward(permsOut), std::forward(out), uint64_t{0}); } -/** @deprecated Pass an explicit uint64_t key. This overload uses key=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 argument (uint64_t). " - "This overload forwards with key=0 and will be removed in a future release.")]] + "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, From 9d8224a7fbd3fd350afbcd61c8d7ad92017a0d96 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 22:32:42 +0100 Subject: [PATCH 23/29] Adding CUDA error checking in the test --- cpp/tests/random/permute.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index 6f4f99b9b3..fb2fc1c23d 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -415,6 +415,7 @@ TEST(PermTest, SeedDiversity) 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); From 108333bc593a00afd3ac03d2af3b23649bbf524f Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 27 Jul 2026 15:00:00 +0100 Subject: [PATCH 24/29] Replacing Feistel logic with CCCL API for simplicity --- cpp/include/raft/random/detail/permute.cuh | 179 ++++----------------- cpp/tests/random/permute.cu | 58 ++----- 2 files changed, 44 insertions(+), 193 deletions(-) diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 33faef4fd2..2d76ee7546 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -11,158 +11,38 @@ #include #include +#include +#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) +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(kperm_index(idx, fp)); } + if (idx < N) { perms[idx] = IntType(shuffled_indices[idx]); } } } -template +template RAFT_KERNEL permuteKernel( - IntType* perms, Type* out, const Type* in, kperm_params fp, 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; @@ -170,7 +50,7 @@ RAFT_KERNEL permuteKernel( int tid = threadIdx.x + blockIdx.x * blockDim.x; IntType outIdx = tid; - IntType inIdx = (tid < N) ? kperm_index(IntType(tid), fp) : IntType(0); + IntType inIdx = (tid < N) ? IntType(shuffled_indices[tid]) : IntType(0); if (perms != nullptr && tid < N) { perms[outIdx] = inIdx; } @@ -204,13 +84,14 @@ RAFT_KERNEL permuteKernel( // This is wrapped in a type to allow for partial template specialization template struct permute_impl_t { + template static void permuteImpl(IntType* perms, Type* out, const Type* in, IdxType N, IdxType D, int nblks, - kperm_params fp, + ShuffleIterator shuffled_indices, cudaStream_t stream) { // determine vector type and set new pointers @@ -222,11 +103,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, fp, 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, fp, stream); + perms, out, in, N, D, nblks, shuffled_indices, stream); } } }; @@ -234,17 +115,18 @@ struct permute_impl_t { // at vector length 1 we just execute a scalar version to break the recursion template struct permute_impl_t { + template static void permuteImpl(IntType* perms, Type* out, const Type* in, IdxType N, IdxType D, int nblks, - kperm_params fp, + ShuffleIterator shuffled_indices, cudaStream_t stream) { permuteKernel - <<>>(perms, out, in, fp, N, D); + <<>>(perms, out, in, shuffled_indices, N, D); RAFT_CUDA_TRY(cudaPeekAtLastError()); } }; @@ -261,22 +143,19 @@ void permute(IntType* perms, { if (N <= 0) { return; } + cuda::shuffle_iterator shuffled_indices{cuda::random_bijection{N, cuda::std::minstd_rand{key}}}; + 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); + <<>>(perms, shuffled_indices, N); RAFT_CUDA_TRY(cudaPeekAtLastError()); return; } auto nblks = raft::ceildiv(N, (IntType)TPB); - // 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, fp, stream); + perms, out, in, N, D, nblks, shuffled_indices, stream); } } diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index fb2fc1c23d..9fb7eaef3a 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -371,64 +371,36 @@ 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); + rmm::device_uvector d_other(N, stream); + detail::permute( + d_ref.data(), nullptr, nullptr, 0, N, true, stream, base_seed); + detail::permute( + d_other.data(), nullptr, nullptr, 0, N, true, stream, base_seed + 1); - 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); + std::vector h_ref(N); + std::vector h_other(N); + raft::update_host(h_ref.data(), d_ref.data(), N, stream); + raft::update_host(h_other.data(), d_other.data(), N, stream); resource::sync_stream(handle); - int total_matches = 0; - for (int c : h_matches) { - total_matches += c; + int matches = 0; + for (uint32_t i = 0; i < N; i++) { + if (h_other[i] == h_ref[i] || h_other[i] == i) { matches++; } } - 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; + int max_allowed = static_cast(max_match_frac * float(N)); + EXPECT_LT(matches, max_allowed) << "Too many index matches across seeds: " << matches + << " >= " << max_allowed; } } // end namespace random From f0e688d56ecbcbc4be961bc623ed020fd2e728ce Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 27 Jul 2026 15:28:49 +0100 Subject: [PATCH 25/29] Restoring the multi-seed diversity test --- cpp/tests/random/permute.cu | 48 +++++++++++++++++++++++++++---------- 1 file changed, 35 insertions(+), 13 deletions(-) diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index 9fb7eaef3a..7fdd96a62a 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -371,36 +371,58 @@ 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 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; +} + 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_other(N, stream); + rmm::device_uvector d_matches(total_threads, stream); detail::permute( d_ref.data(), nullptr, nullptr, 0, N, true, stream, base_seed); - detail::permute( - d_other.data(), nullptr, nullptr, 0, N, true, stream, base_seed + 1); - std::vector h_ref(N); - std::vector h_other(N); - raft::update_host(h_ref.data(), d_ref.data(), N, stream); - raft::update_host(h_other.data(), d_other.data(), N, stream); + 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 matches = 0; - for (uint32_t i = 0; i < N; i++) { - if (h_other[i] == h_ref[i] || h_other[i] == i) { matches++; } + int total_matches = 0; + for (int count : h_matches) { + total_matches += count; } - int max_allowed = static_cast(max_match_frac * float(N)); - EXPECT_LT(matches, max_allowed) << "Too many index matches across seeds: " << matches - << " >= " << max_allowed; + 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 From 14c8bf6bf0d9711c18a4e6fabea81251ac87cbad Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 27 Jul 2026 16:37:27 +0100 Subject: [PATCH 26/29] Updating header include list --- cpp/bench/prims/random/permute.cu | 1 - cpp/include/raft/random/detail/make_regression.cuh | 1 - cpp/include/raft/random/detail/permute.cuh | 2 -- 3 files changed, 4 deletions(-) diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index 43362fa389..e3209e6d63 100644 --- a/cpp/bench/prims/random/permute.cu +++ b/cpp/bench/prims/random/permute.cu @@ -7,7 +7,6 @@ #include #include -#include #include diff --git a/cpp/include/raft/random/detail/make_regression.cuh b/cpp/include/raft/random/detail/make_regression.cuh index 9fb0b55e51..7daf0e45a5 100644 --- a/cpp/include/raft/random/detail/make_regression.cuh +++ b/cpp/include/raft/random/detail/make_regression.cuh @@ -13,7 +13,6 @@ #include #include #include -#include #include #include #include diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 2d76ee7546..5ba827436e 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -14,8 +14,6 @@ #include #include -#include - namespace raft { namespace random { namespace detail { From c2613ea5ba937fc1cff26b53f1b025dae45a9cdd Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 27 Jul 2026 17:32:43 +0100 Subject: [PATCH 27/29] Adding/updating Docstrings --- cpp/bench/prims/random/permute.cu | 7 ++ .../raft/random/detail/make_regression.cuh | 6 ++ cpp/include/raft/random/detail/permute.cuh | 71 +++++++++++++++++++ cpp/include/raft/random/permute.cuh | 12 ++++ cpp/tests/random/permute.cu | 18 +++-- 5 files changed, 110 insertions(+), 4 deletions(-) diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index e3209e6d63..08f2920f7a 100644 --- a/cpp/bench/prims/random/permute.cu +++ b/cpp/bench/prims/random/permute.cu @@ -29,6 +29,7 @@ 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); @@ -73,8 +74,14 @@ 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; diff --git a/cpp/include/raft/random/detail/make_regression.cuh b/cpp/include/raft/random/detail/make_regression.cuh index 7daf0e45a5..d67f0252bf 100644 --- a/cpp/include/raft/random/detail/make_regression.cuh +++ b/cpp/include/raft/random/detail/make_regression.cuh @@ -133,6 +133,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, diff --git a/cpp/include/raft/random/detail/permute.cuh b/cpp/include/raft/random/detail/permute.cuh index 5ba827436e..c503fa6e70 100644 --- a/cpp/include/raft/random/detail/permute.cuh +++ b/cpp/include/raft/random/detail/permute.cuh @@ -18,6 +18,19 @@ namespace raft { namespace random { namespace detail { +/** + * @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 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, @@ -113,6 +155,18 @@ 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, @@ -129,6 +183,23 @@ struct permute_impl_t { } }; +/** + * @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, diff --git a/cpp/include/raft/random/permute.cuh b/cpp/include/raft/random/permute.cuh index 2d7a4268bb..b5557783e5 100644 --- a/cpp/include/raft/random/permute.cuh +++ b/cpp/include/raft/random/permute.cuh @@ -137,6 +137,18 @@ void permute(raft::resources const& handle, /** * @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 > { { } + /** @brief Allocate test inputs and run the keyed raw-pointer overload. */ void SetUp() override { auto stream = resource::get_cuda_stream(handle); @@ -103,6 +104,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); @@ -371,10 +373,17 @@ 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. +/** + * @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, @@ -391,6 +400,7 @@ __global__ void seed_diversity_kernel(const uint32_t* ref_perm, match_counts[tid] = count; } +/** @brief Verify that many consecutive seeds produce diverse permutations. */ TEST(PermTest, SeedDiversity) { constexpr uint32_t N = 1000; From abcde850c5d14b122d87283e247024e0f7a61258 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 27 Jul 2026 17:33:41 +0100 Subject: [PATCH 28/29] Deduplicating the deprecation string and changing the default behavior from key = 0 to key = rand() --- cpp/include/raft/random/permute.cuh | 66 ++++++++++++++--------------- 1 file changed, 32 insertions(+), 34 deletions(-) diff --git a/cpp/include/raft/random/permute.cuh b/cpp/include/raft/random/permute.cuh index b5557783e5..41ee16571d 100644 --- a/cpp/include/raft/random/permute.cuh +++ b/cpp/include/raft/random/permute.cuh @@ -215,69 +215,66 @@ void permute(IntType* perms, 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." + /** - * @deprecated Use the overload that takes an explicit @c uint64_t key. + * @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. * - * @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. + * @deprecated Use the overload that takes an explicit @c uint64_t key. */ 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.")]] +[[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, uint64_t{0}); + permute(handle, in, permsOut, out, static_cast(rand())); } /** - * @deprecated Use the overload that takes an explicit @c uint64_t key. + * @brief Deprecated keyless `std::nullopt` overload that preserves per-call variation. * - * @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. + * 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( - "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.")]] +[[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), uint64_t{0}); + permute(handle, + in, + std::forward(permsOut), + std::forward(out), + static_cast(rand())); } /** - * @deprecated Use the overload that takes an explicit @c uint64_t key. + * @brief Deprecated keyless raw-pointer overload that preserves per-call variation. * - * @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. + * 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( - "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.")]] +[[deprecated(KEYLESS_PERMUTE_DEPRECATED_WARNING)]] void permute(IntType* perms, Type* out, const Type* in, @@ -286,7 +283,8 @@ void permute(IntType* perms, bool rowMajor, cudaStream_t stream) { - detail::permute(perms, out, in, D, N, rowMajor, stream, uint64_t{0}); + detail::permute( + perms, out, in, D, N, rowMajor, stream, static_cast(rand())); } }; // namespace random From f1a428369cd7e36d4040605f2d1a94489bf54d14 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 27 Jul 2026 17:39:35 +0100 Subject: [PATCH 29/29] Adding/updating Docstrings for other functions --- cpp/bench/prims/random/permute.cu | 5 +++ .../raft/random/detail/make_regression.cuh | 33 +++++++++++++++-- cpp/tests/random/permute.cu | 37 +++++++++++++++++++ 3 files changed, 71 insertions(+), 4 deletions(-) diff --git a/cpp/bench/prims/random/permute.cu b/cpp/bench/prims/random/permute.cu index 08f2920f7a..2e9ab70b9b 100644 --- a/cpp/bench/prims/random/permute.cu +++ b/cpp/bench/prims/random/permute.cu @@ -19,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), diff --git a/cpp/include/raft/random/detail/make_regression.cuh b/cpp/include/raft/random/detail/make_regression.cuh index d67f0252bf..b122692ea0 100644 --- a/cpp/include/raft/random/detail/make_regression.cuh +++ b/cpp/include/raft/random/detail/make_regression.cuh @@ -28,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) { @@ -41,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, @@ -115,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) diff --git a/cpp/tests/random/permute.cu b/cpp/tests/random/permute.cu index 79b46e6505..2f798371d7 100644 --- a/cpp/tests/random/permute.cu +++ b/cpp/tests/random/permute.cu @@ -26,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) { @@ -38,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)), @@ -87,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)), @@ -167,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) @@ -186,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, @@ -289,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; @@ -297,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; @@ -358,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; @@ -366,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;