[FEA] permute: replace LCG with keyed Feistel permutation, add deterministic key API

cpp/bench/prims/random/permute.cu

@@ -1,13 +1,12 @@
 /*
- * 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
  */
 
 #include <common/benchmark.hpp>
 
 #include <raft/random/permute.cuh>
 #include <raft/random/rng.cuh>
-#include <raft/util/cudart_utils.hpp>
 
 #include <rmm/device_uvector.hpp>
 
@@ -20,6 +19,11 @@ struct permute_inputs {
 
 template <typename T>
 struct permute : public fixture {
+  /**
+   * @brief Construct a matrix permutation benchmark.
+   *
+   * @param[in] p Matrix dimensions, output selection, and layout
+   */
   permute(const permute_inputs& p)
     : params(p),
       perms(p.needPerms ? p.rows : 0, stream),
@@ -30,12 +34,19 @@ struct permute : public fixture {
     uniform(handle, r, in.data(), p.rows, T(-1.0), T(1.0));
   }
 
+  /** @brief Benchmark keyed permutation of a matrix and its indices. */
   void run_benchmark(::benchmark::State& state) override
   {
     raft::random::RngState r(123456ULL);
     loop_on_state(state, [this, &r]() {
-      raft::random::permute(
-        perms.data(), out.data(), in.data(), params.cols, params.rows, params.rowMajor, stream);
+      raft::random::permute(perms.data(),
+                            out.data(),
+                            in.data(),
+                            params.cols,
+                            params.rows,
+                            params.rowMajor,
+                            stream,
+                            123456ULL);
     });
   }
 
@@ -66,4 +77,42 @@ const std::vector<permute_inputs> permute_input_vecs = {
 RAFT_BENCH_REGISTER(permute<float>, "", permute_input_vecs);
 RAFT_BENCH_REGISTER(permute<double>, "", permute_input_vecs);
 
+template <typename IntType>
+struct permute_perms_only : public fixture {
+  /**
+   * @brief Construct a benchmark that generates only permutation indices.
+   *
+   * @param[in] rows Number of permutation indices to generate
+   */
+  permute_perms_only(int rows) : n_rows(rows), perms(rows, stream) {}
+
+  /** @brief Benchmark the permutation-indices-only kernel path. */
+  void run_benchmark(::benchmark::State& state) override
+  {
+    size_t bytes_processed = 0;
+    loop_on_state(state, [this, &bytes_processed]() {
+      raft::random::permute(perms.data(),
+                            (float*)nullptr,
+                            (const float*)nullptr,
+                            IntType(0),
+                            IntType(n_rows),
+                            true,
+                            stream,
+                            123456ULL);
+      bytes_processed += size_t(n_rows) * sizeof(IntType);
+    });
+    state.SetBytesProcessed(bytes_processed);
+  }
+
+ private:
+  raft::device_resources handle;
+  int n_rows;
+  rmm::device_uvector<IntType> perms;
+};
+
+RAFT_BENCH_REGISTER((permute_perms_only<int>),
+                    "",
+                    std::vector<int>({32 * 1024, 1024 * 1024, 32 * 1024 * 1024}));
+RAFT_BENCH_REGISTER((permute_perms_only<uint32_t>), "", std::vector<int>({1024 * 1024 * 1024}));
+
 }  // namespace raft::bench::random

cpp/include/raft/random/detail/make_regression.cuh

@@ -14,7 +14,6 @@
 #include <raft/core/resources.hpp>
 #include <raft/linalg/add.cuh>
 #include <raft/linalg/gemm.cuh>
-#include <raft/linalg/init.cuh>
 #include <raft/linalg/qr.cuh>
 #include <raft/linalg/transpose.cuh>
 #include <raft/matrix/diagonal.cuh>
@@ -31,7 +30,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 <typename DataT, typename IdxT>
 RAFT_KERNEL _singular_profile_kernel(DataT* out, IdxT n, DataT tail_strength, IdxT rank)
 {
@@ -44,7 +50,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 <typename DataT, typename IdxT>
 static void _make_low_rank_matrix(raft::resources const& handle,
                                   DataT* out,
@@ -130,8 +147,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 <typename DataT, typename IdxT>
 RAFT_KERNEL _gather2d_kernel(
   DataT* out, const DataT* in, const IdxT* perms, IdxT n_rows, IdxT n_cols)
@@ -148,6 +172,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 <typename DataT, typename IdxT>
 void make_regression_caller(raft::resources const& handle,
                             DataT* out,
@@ -248,9 +278,15 @@ void make_regression_caller(raft::resources const& handle,
     if (!is_dry_run) {
       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<DataT, IdxT, IdxT>(
-        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<IdxT>(n_rows, Nthreads);
       raft::launch_kernel(stream,
                           nblks_rows,
@@ -264,7 +300,7 @@ void make_regression_caller(raft::resources const& handle,
 
       // Shuffle the features from tmp_out to out
       raft::random::permute<DataT, IdxT, IdxT>(
-        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) {