From add8e6205f2a950a2485dab3f4363241f3ad7768 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 8 Jul 2026 15:38:28 +0100 Subject: [PATCH 01/23] 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 02dc19ecc324a3af44f2aab64ea38e82fe2b8561 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 8 Jul 2026 15:48:18 +0100 Subject: [PATCH 02/23] 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 ca3104ec905d7f8b10fd17b8ad086424b15d1d40 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 8 Jul 2026 16:14:20 +0100 Subject: [PATCH 03/23] 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 2e65b956e4a53faae3779021c4bf93dc7410b2a7 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Wed, 8 Jul 2026 16:31:40 +0100 Subject: [PATCH 04/23] 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 83cef8fb348c3d519f1ec7c56a1e90c67d2751cb Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 09:36:57 +0100 Subject: [PATCH 05/23] 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 9f1f05af91c3338d73babcbd3b520bd1da537ddf Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 10:29:15 +0100 Subject: [PATCH 06/23] 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 91e88bec86caaf7cb93f69ae616da67599044357 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 14:41:26 +0100 Subject: [PATCH 07/23] 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 f3198fbbf272f8b614d51135eba116f042118f36 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 14:41:47 +0100 Subject: [PATCH 08/23] 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 c9e7ca2fcefcd1feba0ab2e19116db4f034e86a9 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 14:55:35 +0100 Subject: [PATCH 09/23] 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 607c2720f206f987f751d1c0a9c7683e6d027315 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Thu, 9 Jul 2026 15:02:38 +0100 Subject: [PATCH 10/23] 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 412df93494d0533ee79800d13cfd160deec0cc3c Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 20 Jul 2026 16:52:48 +0100 Subject: [PATCH 11/23] 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 5e6bd89031175b419938cbb9168f7e454162b9c7 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Mon, 20 Jul 2026 17:46:02 +0100 Subject: [PATCH 12/23] 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 03f9c04cc4316d4f269bc195eb44b07365ee82ba Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 11:55:11 +0100 Subject: [PATCH 13/23] 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 9ea5f8f48c56bde0184bc347168f3ff427446fef Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 12:02:11 +0100 Subject: [PATCH 14/23] 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 14c58d190dbe2313ad802275f682e4d69c4be4e4 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 13:31:48 +0100 Subject: [PATCH 15/23] 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 128b1ce250a8c68fd2d62e1e1d3b972f8bb79309 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 13:49:08 +0100 Subject: [PATCH 16/23] 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/23] 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 d77909fb3f11f630465e19215ef9437f9fe4926d Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 14:56:28 +0100 Subject: [PATCH 18/23] 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 f08ebaa9db980e3ad4d51e90594bbf6f97d56933 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 20:04:34 +0100 Subject: [PATCH 19/23] 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 067fd5340f5ba634ed165e3caf4ebf39329bd52e Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 20:14:37 +0100 Subject: [PATCH 20/23] 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 24dcaa24f2bbcec1abefa5cb33909041f965e79d Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 20:18:45 +0100 Subject: [PATCH 21/23] 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 ffdffb7826da74555317d560ba09828be8381472 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 22:32:18 +0100 Subject: [PATCH 22/23] 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 80a3471640cdef1be5a205d4d0678d6efb2892c2 Mon Sep 17 00:00:00 2001 From: Vinay D Date: Tue, 21 Jul 2026 22:32:42 +0100 Subject: [PATCH 23/23] 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);