Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
4858ca7
Adding Feistel network based permute
vinaydes Jul 8, 2026
0461065
Adding permutation key as a parameter
vinaydes Jul 8, 2026
aaf0fce
Formatting changes
vinaydes Jul 8, 2026
392651e
Adding randomness check for permute
vinaydes Jul 8, 2026
2f69c68
Removing redundant header inclusion
vinaydes Jul 9, 2026
40c15f0
Tidying up comments
vinaydes Jul 9, 2026
be8371d
Undoing delete
vinaydes Jul 9, 2026
6f75295
Removing redundant header
vinaydes Jul 9, 2026
27019e4
Adding permute only benchmark
vinaydes Jul 20, 2026
2b71259
Restoring the permute only kernel
vinaydes Jul 20, 2026
24bd589
Reducing the complexity of round function to achieve better bandwidth
vinaydes Jul 21, 2026
deb08b8
Removing 32-bit specialization, as it is not needed anymore
vinaydes Jul 21, 2026
3657f38
Converting to template arguments for avoiding type conversion in 32-b…
vinaydes Jul 21, 2026
a2a35be
Changing the names of functions for clarity
vinaydes Jul 21, 2026
97c01cf
Adding a test that checks for seed diversity
vinaydes Jul 21, 2026
eeb1d09
Formatting
vinaydes Jul 21, 2026
85751ea
Adding deprecated APIs for compatibility
vinaydes Jul 21, 2026
0383fbe
Skipping kernel launch if N <= 0
vinaydes Jul 21, 2026
31ac615
Adding small N test cases
vinaydes Jul 21, 2026
d19dd18
Adding changed behavior description in the deprecation message
vinaydes Jul 21, 2026
9d8224a
Adding CUDA error checking in the test
vinaydes Jul 21, 2026
108333b
Replacing Feistel logic with CCCL API for simplicity
vinaydes Jul 27, 2026
5ec3246
Fixing a typo
vinaydes Jul 9, 2026
46b9e67
Removing a narrow test
vinaydes Jul 9, 2026
f0e688d
Restoring the multi-seed diversity test
vinaydes Jul 27, 2026
14c8bf6
Updating header include list
vinaydes Jul 27, 2026
c2613ea
Adding/updating Docstrings
vinaydes Jul 27, 2026
abcde85
Deduplicating the deprecation string and changing the default behavio…
vinaydes Jul 27, 2026
f1a4283
Adding/updating Docstrings for other functions
vinaydes Jul 27, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 53 additions & 4 deletions cpp/bench/prims/random/permute.cu
Original file line number Diff line number Diff line change
@@ -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>

Expand All @@ -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),
Expand All @@ -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);
});
}

Expand Down Expand Up @@ -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
52 changes: 44 additions & 8 deletions cpp/include/raft/random/detail/make_regression.cuh
Original file line number Diff line number Diff line change
@@ -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
*/

Expand All @@ -13,7 +13,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>
Expand All @@ -29,7 +28,14 @@ namespace raft {
namespace random {
namespace detail {

/* Internal auxiliary function to help build the singular profile */
/**
* @brief Build the singular-value profile for a low-rank regression matrix.
*
* @param[out] out Generated singular values
* @param[in] n Number of singular values
* @param[in] tail_strength Relative strength of the low-rank tail
* @param[in] rank Effective matrix rank
*/
template <typename DataT, typename IdxT>
RAFT_KERNEL _singular_profile_kernel(DataT* out, IdxT n, DataT tail_strength, IdxT rank)
{
Expand All @@ -42,7 +48,18 @@ RAFT_KERNEL _singular_profile_kernel(DataT* out, IdxT n, DataT tail_strength, Id
}
}

/* Internal auxiliary function to generate a low-rank matrix */
/**
* @brief Generate a low-rank matrix with a decaying singular-value profile.
*
* @param[in] handle RAFT handle containing execution resources
* @param[out] out Generated row-major matrix
* @param[in] n_rows Number of matrix rows
* @param[in] n_cols Number of matrix columns
* @param[in] effective_rank Approximate rank of the generated matrix
* @param[in] tail_strength Relative strength of the low-rank tail
* @param[in,out] r Random number generator state
* @param[in] stream CUDA stream on which to execute
*/
template <typename DataT, typename IdxT>
static void _make_low_rank_matrix(raft::resources const& handle,
DataT* out,
Expand Down Expand Up @@ -116,8 +133,15 @@ static void _make_low_rank_matrix(raft::resources const& handle,
raft::linalg::transpose(handle, temp_out.data(), out, n_rows, n_cols, stream);
}

/* Internal auxiliary function to permute rows in the given matrix according
* to a given permutation vector */
/**
* @brief Gather matrix rows according to a permutation vector.
*
* @param[out] out Permuted output matrix
* @param[in] in Input matrix
* @param[in] perms Input row index for each output row
* @param[in] n_rows Number of matrix rows
* @param[in] n_cols Number of matrix columns
*/
template <typename DataT, typename IdxT>
RAFT_KERNEL _gather2d_kernel(
DataT* out, const DataT* in, const IdxT* perms, IdxT n_rows, IdxT n_cols)
Expand All @@ -134,6 +158,12 @@ RAFT_KERNEL _gather2d_kernel(
}
}

/**
* @brief Generate a regression data set and optionally shuffle its rows and features.
*
* When shuffling is enabled, the input seed deterministically selects distinct
* permutations for samples and features.
*/
template <typename DataT, typename IdxT>
void make_regression_caller(raft::resources const& handle,
DataT* out,
Expand Down Expand Up @@ -246,17 +276,23 @@ 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<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);
_gather2d_kernel<<<nblks_rows, Nthreads, 0, stream>>>(
values, _values, perms_samples.data(), n_rows, n_targets);
RAFT_CUDA_TRY(cudaPeekAtLastError());

// 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) {
Expand Down
Loading
Loading