Update raft headers

[aamijar]

Depends on NVIDIA/raft#2940. Resolves #7750.

This PR does the following:

1. Decompose usage of #include <raft/matrix/math.cuh> into necessary required headers such as:

#include <raft/linalg/matrix_vector.cuh>
#include <raft/linalg/power.cuh>
#include <raft/linalg/sqrt.cuh> 
#include <raft/matrix/threshold.cuh>
#include <raft/matrix/ratio.cuh>
#include <raft/matrix/copy.cuh>

2. Use the new mdspan apis.
3. Use a new handle per stream where necessary. This is because the mdspan apis accept handle and not stream.

[aamijar]

It seems that the reason some dask tests are failing in CI, is because the correct stream needs to be passed during these function calls? I am not sure how to resolve this yet. Do we need to have an api that accepts cuda_stream objects instead of handle?

[csadorf]

The row_reverse function uses the stream provided by the handle. I don't think we have to provide a separate stream object as long as the handle is set up correctly.

[aamijar]

Yes, we have to make sure that the correct stream is assigned to the handle.

[coderabbitai]

📝 Walkthrough

Walkthrough

Replaces deprecated RAFT matrix APIs with RAFT linalg and device-view-based calls across many modules, adds/adjusts includes and explicit handle/stream usage, updates copyright years, and changes one public function signature (dist_membership_vector adds size_t n_X_rows).

Changes

Cohort / File(s)	Summary
GLM & preprocessing cpp/src/glm/fitness.cuh, cpp/src/glm/ols.cuh, cpp/src/glm/preprocess.cuh, cpp/src/glm/preprocess_mg.cu, cpp/src/glm/qn/qn.cuh, cpp/src/glm/qn/mg/standardization.cuh, cpp/src/glm/qn_mg.cu, cpp/src/glm/ridge.cuh, cpp/src/glm/ridge_mg.cu	Replaced deprecated raft::matrix calls with raft::linalg and device-view variants (binary_mult/div_*, argmax via handle+views), updated includes, introduced explicit handle/stream usage, and bumped copyright years.
PCA & TSVD (single- & multi-GPU) cpp/src/pca/pca.cuh, cpp/src/pca/pca_mg.cu, cpp/src/tsvd/tsvd.cuh, cpp/src/tsvd/tsvd_mg.cu	Converted pointer-based matrix ops to view-based APIs (e.g., truncZeroOrigin→trunc_zero_origin, seqRoot→weighted_sqrt), replaced matrix-vector binary ops with linalg/view calls, added handle/stream management, and replaced seqRoot usages.
Solvers, linear/reg primitives & SGD cpp/src/solver/cd.cuh, cpp/src/solver/cd_mg.cu, cpp/src/solver/sgd.cuh, cpp/src_prims/functions/linearReg.cuh, cpp/src_prims/functions/logisticReg.cuh	Refactored sample-weighting, normalization and copy logic to use raft::linalg::binary_mult / binary_div_skip_zero, device views (make_device_*_view), and raft::matrix::copy_rows (view-based). Updated includes to linalg/copy headers.
OPG linear algebra & matrix math cpp/src_prims/opg/linalg/lstsq.cu, cpp/src_prims/opg/linalg/mm_aTa.cu, cpp/src_prims/opg/linalg/mv_aTb.cu, cpp/src_prims/opg/linalg/norm.cu, cpp/src_prims/opg/linalg/svd.cu, cpp/src_prims/opg/matrix/math.cu	Switched to RAFT linalg/device-view APIs for matrix-vector and elementwise ops, introduced per-block/per-stream resource handling, replaced legacy includes with targeted linalg utilities, and applied weighted/streamed variants for sqrt/power/reverse.
Prims — Misc stats & utilities cpp/src_prims/opg/stats/mean.cu, cpp/src_prims/opg/stats/stddev.cu	Removed deprecated matrix includes and updated headers; no substantive algorithmic changes.
HDBSCAN & SVM copy utilities cpp/src/hdbscan/detail/soft_clustering.cuh, cpp/src/svm/sparse_util.cuh	Replaced copyRows with raft::matrix::copy_rows using device views; updated dist_membership_vector signature to include size_t n_X_rows, adjusted call sites and view construction accordingly.
TSVD/PCA MG weighting & streaming cpp/src/pca/pca_mg.cu, cpp/src/tsvd/tsvd_mg.cu	Applied ALONG_ROWS view-based binary ops, replaced seqRoot with weighted_sqrt, and added explicit per-call stream/handle handling across MG code paths.
Prims — tests & small files python/cuml/tests/dask/test_dask_pca.py, various small include-only edits (cpp/src_prims/opg/..., cpp/src_prims/opg/linalg/...)	Bumped copyright years; updated MG PCA test n_components (30→10). Many files only had include or year updates.
Misc — genetic, hdbscan, tsvd partitioning, svd, etc. cpp/src/genetic/fitness.cuh, cpp/src/hdbscan/..., cpp/src/tsvd/..., cpp/src_prims/opg/svd.cu	Small, focused replacements: seqRoot→raft::linalg::sqrt/weighted_sqrt, and other targeted moves to linalg/view APIs and updated includes.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title 'Update raft headers' is vague and does not convey the specific nature of the changes (deprecation API migration and mdspan adoption).	Consider a more specific title like 'Migrate deprecated RAFT APIs to mdspan-based interfaces and update headers' to better reflect the substantial changes.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description clearly explains the three main objectives: decomposing raft/matrix/math.cuh, adopting mdspan APIs, and using new handles per stream.
Linked Issues check	✅ Passed	The PR successfully addresses issue #7750 by replacing all deprecated RAFT umbrella headers with specific headers and updating APIs to use mdspan-based interfaces throughout the codebase.
Out of Scope Changes check	✅ Passed	All changes are directly related to the PR objectives: replacing deprecated headers, adopting mdspan APIs, and managing RAFT handles. The single test parameter change aligns with API updates.
Merge Conflict Detection	✅ Passed	✅ No merge conflicts detected when merging into main

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 17

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)

cpp/src_prims/opg/linalg/mm_aTa.cu (1)

18-69: ⚠️ Potential issue | 🟠 Major

Add defensive validation for streams/A/local_blocks before indexing.

streams[0], A[0], and local_blocks[0] are dereferenced without guards. If n_streams == 0, streams == nullptr, A.empty(), or local_blocks doesn’t match A, this will crash. Please validate inputs upfront and fail fast.

🛡️ Suggested guard clauses

 void mm_aTa_impl(const raft::handle_t& handle,
                  Matrix::Data<math_t>& out,
                  const std::vector<Matrix::Data<math_t>*>& A,
                  const Matrix::PartDescriptor& ADesc,
                  cudaStream_t* streams,
                  int n_streams)
 {
+  if (streams == nullptr || n_streams <= 0) {
+    throw std::invalid_argument("mm_aTa_impl: streams must be non-null and n_streams > 0");
+  }
+  if (A.empty() || A[0] == nullptr) {
+    throw std::invalid_argument("mm_aTa_impl: input matrix list A must be non-empty");
+  }
   auto& comm = handle.get_comms();
   int rank   = comm.get_rank();
 
   std::vector<Matrix::RankSizePair*> local_blocks = ADesc.blocksOwnedBy(rank);
+  if (local_blocks.empty() || local_blocks.size() != A.size() || local_blocks[0] == nullptr) {
+    throw std::invalid_argument("mm_aTa_impl: local block metadata mismatch for rank");
+  }

Based on learnings: "Validate input parameters for correctness (negative dimensions, null pointers, invalid values)".

cpp/src_prims/functions/logisticReg.cuh (1)

107-108: ⚠️ Potential issue | 🟡 Minor

Pre-existing bug: logf used instead of log in double specialization.

Line 108 uses logf(1 - y_pred) (single-precision) while the first term on the same line correctly uses log(y_pred) (double-precision). This truncates intermediate precision for the second term. Not introduced by this PR, but worth fixing while you're in this file. As per coding guidelines: verify unsafe type casting between numeric types does not cause unintended precision loss.

Proposed fix

     [] __device__(double y, double y_pred) {
-      return -y * log(y_pred) - (1 - y) * logf(1 - y_pred);
+      return -y * log(y_pred) - (1 - y) * log(1 - y_pred);
     },

cpp/src/glm/qn/mg/standardization.cuh (1)

326-332: ⚠️ Potential issue | 🟡 Minor

Replace deprecated matrixVectorOp API with modern binary_mult pattern to match line 357.

Line 326 uses the legacy matrixVectorOp<false, true>(...) API with raw pointers and boolean template parameters. The analogous operation at line 357 was already migrated to the modern mdspan-based binary_mult<raft::Apply::ALONG_ROWS>(...) API. For consistency and to avoid deprecated APIs, update line 326-332 to use the same binary_mult pattern:

Suggested replacement

raft::linalg::binary_mult<raft::Apply::ALONG_ROWS>(
  handle,
  raft::make_device_matrix_view<T, int, raft::col_major>(Wweights.data, Wweights.m, Wweights.n),
  raft::make_device_vector_view<const T, int>(std_inv.data, Wweights.n));

🤖 Fix all issues with AI agents

In `@cpp/cmake/thirdparty/get_raft.cmake`:
- Around line 58-60: The find_and_configure_raft call is pointing at a personal
fork and unversioned branch; update the FORK and PINNED_TAG arguments in the
find_and_configure_raft invocation so they reference the official rapidsai
repository and a concrete release or commit tag (not "aamijar" or
"fix-const-sqrt")—use CUML_MIN_VERSION_raft for version gating as needed and
replace PINNED_TAG with the upstream release/commit identifier before merging.

In `@cpp/src_prims/opg/linalg/norm.cu`:
- Around line 73-80: The call to weighted_sqrt() breaks the multi-stream
ordering because colNorm2NoSeq_impl() ran on streams[0] and
comm.sync_stream(streams[0]) was used, while weighted_sqrt launches on
handle.get_stream(); fix by ensuring both use the same stream: if
raft::matrix::weighted_sqrt accepts a stream argument, call weighted_sqrt(...,
streams[0]) (referencing weighted_sqrt and streams[0]); otherwise insert
handle.sync_stream() (or equivalent) immediately after
comm.sync_stream(streams[0]) and before weighted_sqrt to guarantee ordering
between colNorm2NoSeq_impl and weighted_sqrt.

In `@cpp/src_prims/opg/linalg/svd.cu`:
- Around line 13-14: Remove the deprecated RAFT headers by deleting the two
include lines for <raft/matrix/math.cuh> and <raft/matrix/matrix.cuh> in this
file; verify there are no remaining usages of APIs from those headers (the
comment indicates old calls are commented out) and, if any RAFT matrix/math
symbols are actually used, replace them with the current RAFT headers that
provide the same APIs (or adjust code to use the new linalg/matrix APIs) and
rebuild to confirm no unresolved symbols remain.
- Around line 50-68: The col_reverse, row_reverse, and weighted_sqrt calls use
handle's default stream while eigDC ran on streams[0], so create and use a
temporary/local handle bound to streams[0] (the same "local-handle" pattern used
for the division loop) and call raft::matrix::col_reverse,
raft::matrix::row_reverse, and raft::matrix::weighted_sqrt with that local
handle instead of handle to ensure those ops execute on streams[0] and after
eigDC; locate the calls by the function names col_reverse, row_reverse,
weighted_sqrt and adjust the handle passed to them to a copy whose CUDA stream
has been set to streams[0].

In `@cpp/src_prims/opg/matrix/math.cu`:
- Around line 81-82: The code repeatedly constructs a raft::resources handle
inside the loop and calls raft::resource::set_cuda_stream(handle, streams[i])
without wrapping the stream index; move construction of raft::resources handle
out of the loop (hoist it once before the loop) and change the stream access to
use modulo (streams[i % n_streams]) so set_cuda_stream uses a stable handle and
a valid stream index; update the call sites where raft::resources handle and
streams[i] occur (the raft::resources handle declaration and the
raft::resource::set_cuda_stream(...) invocation) accordingly.
- Around line 45-46: The loop over localBlocks uses streams[i] and can read
out-of-bounds when localBlocks.size() > n_streams; change indexing to
round-robin using streams[i % n_streams] and avoid reconstructing
raft::resources on every iteration by creating a single raft::resources handle
before the loop and calling raft::resource::set_cuda_stream(handle, streams[i %
n_streams]) inside the loop; apply the same changes in
matrixVectorBinaryMult_impl to replace streams[i] with streams[i % n_streams]
and move raft::resources handle construction out of the per-iteration body.

In `@cpp/src/glm/qn/mg/standardization.cuh`:
- Line 15: The file is missing the header that defines
raft::linalg::matrixVectorOp used later; add an include for the correct header
so the symbol is declared. Specifically, include
<raft/linalg/matrix_vector_op.cuh> alongside the existing
<raft/linalg/matrix_vector.cuh> to ensure raft::linalg::matrixVectorOp (and
related overloads) are available for the code that calls matrixVectorOp (and
keeps binary_mult usage intact).

In `@cpp/src/glm/qn/qn.cuh`:
- Around line 267-270: The argmax call misinterprets Z's column-major layout as
row-major and also uses a floating-point output buffer: update the code around Z
and the argmax call (symbols: SimpleDenseMat<T> Z, raft::matrix::argmax,
raft::make_device_matrix_view, raft::make_device_vector_view, preds) so that Z's
memory layout and dimensions match the view you pass to argmax; either construct
Z as row-major (e.g., SimpleDenseMat<T> Z(..., ROW_MAJOR)) if you want a (X.m,
C) row-major view, or keep Z as column-major and pass the transposed dimensions
to make_device_matrix_view (C, X.m) with the proper row/col-major tag; also
change the argmax output vector type to an integer index type by using
raft::make_device_vector_view<int, int>(preds_int, static_cast<int>(X.m)) and
ensure preds is an integer device buffer (or write into an int buffer and
convert to float later) so argmax writes integer class indices into an int
container.

In `@cpp/src/glm/ridge_mg.cu`:
- Around line 50-77: The code mixes handle-based calls (zero_small_values,
binary_div_skip_zero, binary_mult) that use handle's internal stream with manual
calls using streams[0] (raft::copy, powerScalar, addScalar), causing a potential
data race; fix by ensuring the handle executes on streams[0] before any
handle-based call—either uncomment the existing
raft::resource::set_cuda_stream(handle, streams[0]) prior to zero_small_values
or create a local raft::resources (e.g., handle_local) with its CUDA stream set
to streams[0] (following the svd.cu pattern) and then use that local handle for
binary_div_skip_zero and binary_mult so all operations use the same stream.
- Line 17: Remove the deprecated include directive "#include
<raft/matrix/math.cuh>" from ridge_mg.cu; verify that no symbols from that
header are used in the file and, if any RAFT matrix helpers are still required,
replace the deprecated header with the appropriate modern RAFT headers (e.g.,
the new raft/matrix or raft::matrix APIs) and update any affected calls
accordingly.

In `@cpp/src/hdbscan/detail/soft_clustering.cuh`:
- Around line 67-73: The X_view is incorrectly created with n_exemplars rows
even though exemplar_idx indexes into the full training set (n_leaves/m); update
the code that constructs X_view (the raft::make_device_matrix_view call that
creates X_view) to use the total number of rows in X (n_leaves/m) instead of
n_exemplars, and ensure the function signature and all call sites are updated to
accept and pass that total-rows parameter so exemplar_idx_view and
raft::matrix::copy_rows operate against a matrix view whose row dimension
matches the full dataset referenced by exemplar_idx.

In `@cpp/src/pca/pca_mg.cu`:
- Around line 345-349: The matrix view for the components passed to
raft::linalg::binary_div_skip_zero is using the wrong row dimension
(prms.n_rows); update the raft::make_device_matrix_view call that wraps
components (used with singular_vals in binary_div_skip_zero) to use prms.n_cols
as the number of rows (i.e., change the view shape from (prms.n_rows,
prms.n_components) to (prms.n_cols, prms.n_components)) so the components buffer
is interpreted consistently with transform_impl and the singular_vals vector.
- Around line 183-187: binary_div_skip_zero is launched on handle.get_stream()
while scalarMultiply uses streams[0] and later gemm calls inside the loop use
streams[si] and read components that binary_div_skip_zero just modified, causing
a stream-synchronization/race; fix by running binary_div_skip_zero on the same
stream(s) as downstream ops or synchronizing before gemm: either call
binary_div_skip_zero using streams[0]/streams[si] (matching scalarMultiply/gemm)
or insert an explicit stream synchronization/event (e.g., cudaEventRecord /
streamWaitEvent or cudaStreamSynchronize) after
raft::linalg::binary_div_skip_zero and before the loop that uses gemm so
components is fully updated; update usages of
raft::linalg::binary_div_skip_zero, scalarMultiply, gemm, handle.get_stream(),
and streams to ensure consistent stream ordering.
- Around line 314-318: The components view is using incorrect dimensions and can
OOB; update the two places that create a device matrix view for components (the
scalarMultiply call and the
raft::linalg::binary_mult_skip_zero<raft::Apply::ALONG_ROWS> call) to use the
same dimensions used in transform_impl (i.e., swap the shape to match
transform_impl by passing prms.n_cols, prms.n_components), so the matrix view
matches how components was produced by the GEMM; change both
make_device_matrix_view<...>(components, prms.n_rows, prms.n_components)
occurrences to make_device_matrix_view<...>(components, prms.n_cols,
prms.n_components).
- Around line 63-69: truncCompExpVars writes to explained_var and singular_vals
on streams[0] but weighted_sqrt uses handle.get_stream(); fix by inserting an
explicit synchronization of streams[0] to the handle stream (e.g., synchronize
streams[0] or switch the handle to streams[0]) immediately after
truncCompExpVars and before calling raft::matrix::weighted_sqrt in fit_impl, and
apply the same pattern at the equivalent locations in transform_impl and
inverse_transform_impl so no handle-based operation reads/writes those buffers
without a prior stream synchronization or consistent stream assignment.

In `@cpp/src/tsvd/tsvd_mg.cu`:
- Around line 57-70: calEig launches on streams[0] and writes components_all and
explained_var_all, but trunc_zero_origin and weighted_sqrt run on
handle.get_stream() and read those buffers, and later signFlipComponents runs on
streams[0] but reads components written by trunc_zero_origin on
handle.get_stream(); add synchronization to ensure correct ordering: either
dispatch trunc_zero_origin and weighted_sqrt on streams[0] (use the same stream
as calEig) or insert an explicit stream synchronization
(cudaStreamSynchronize(streams[0]) or handle.synchronizeStream(streams[0]) /
equivalent) before calling trunc_zero_origin and weighted_sqrt, and similarly
ensure handle.get_stream() finishes before calling signFlipComponents
(synchronize handle.get_stream() or run signFlipComponents on
handle.get_stream()); update the calls around calEig, trunc_zero_origin,
weighted_sqrt, and signFlipComponents to use one consistent stream or explicit
synchronizations to eliminate the race.

In `@cpp/src/tsvd/tsvd.cuh`:
- Around line 123-130: calEig is using mixed streams: eigJacobi/eigDC write to
the explicit 'stream' while raft::matrix::col_reverse, raft::matrix::row_reverse
(and raft::linalg::transpose) operate on handle.get_stream(), causing races for
MG callers; fix by unifying streams — ensure col_reverse/row_reverse/transpose
run on the same 'stream' used by eig (either by calling the overloads that
accept an explicit cuda stream or by setting the handle's stream to 'stream'
before those calls). Update the calls manipulating 'components' and
'explained_var' in calEig so they all use the same stream variable (match
'stream' used by eigJacobi/eigDC).

🧹 Nitpick comments (2)

cpp/src/glm/qn_mg.cu (1)

23-23: Remove the commented-out include instead of leaving it as dead code.

Commented-out includes add noise and suggest an incomplete migration. Since no symbols from raft/matrix/math.hpp are used in this file, delete the line entirely.

-// `#include` <raft/matrix/math.hpp>

cpp/src/pca/pca.cuh (1)

17-17: Remove commented-out include instead of leaving it.

Suggested fix

-// `#include` <raft/matrix/matrix.cuh>

[coderabbitai]

Actionable comments posted: 3

🤖 Fix all issues with AI agents

In `@cpp/src_prims/opg/linalg/norm.cu`:
- Around line 73-82: The current code creates a new raft::resources
(handle_stream_zero) and calls raft::resource::set_cuda_stream before invoking
raft::matrix::weighted_sqrt, which discards caller allocator/comms/workspace;
instead reuse the caller's raft::handle_t by saving the caller handle's current
stream, setting the caller handle's stream to streams[0], call
raft::matrix::weighted_sqrt with that handle (not a new raft::resources), then
restore the original stream; apply this pattern wherever handle_stream_zero and
raft::resource::set_cuda_stream are used (e.g., in norm.cu, svd.cu, math.cu,
ridge_mg.cu) to ensure allocator/comms/workspace are preserved.

In `@cpp/src/glm/ridge_mg.cu`:
- Around line 50-57: handle_stream_zero is default-constructed so it does not
inherit the allocator from handle, causing temporary allocations to use the
default RMM resource; fix by creating handle_stream_zero with the same memory
resource as handle (or reuse/mutate handle's stream if safe) before calling
raft::matrix::zero_small_values: copy the workspace/allocator from handle into
handle_stream_zero (or set streams[0] on handle and restore it afterward), then
call raft::matrix::zero_small_values(handle_stream_zero,
raft::make_device_matrix_view<T,...>(S, std::size_t(1), UDesc.N), thres) to
preserve allocator affinity.

In `@cpp/src/tsvd/tsvd.cuh`:
- Around line 345-359: The trunc_zero_origin and weighted_sqrt calls currently
use handle.get_stream() which can mismatch the explicit stream used by calEig;
create and use the same stream-bound handle (the pattern used earlier:
handle_stream_zero) and pass that handle/stream when calling
raft::matrix::trunc_zero_origin and raft::matrix::weighted_sqrt so they operate
on the same stream that calEig used; ensure the device views for components_all
and explained_var_all are created/passed with that stream-bound handle (matching
symbols: calEig, handle_stream_zero, trunc_zero_origin, weighted_sqrt,
components_all, explained_var_all, singular_vals).

🧹 Nitpick comments (2)

cpp/src_prims/opg/linalg/svd.cu (1)

50-52: Remove commented-out old code.

Lines 50–52, 64, and 89–90 contain commented-out deprecated API calls that are now replaced. These should be removed to keep the file clean.

cpp/src/glm/ridge_mg.cu (1)

50-51: Remove leftover commented-out code.

Lines 50–51 contain the old deprecated calls as comments. Since the replacement is in place, these should be cleaned up to avoid clutter. The same applies to lines 62, 66–67, and 69.

[viclafargue]

Thanks, LGTM! Could not uncover any mistake. Just some old commented code should probably be removed.

[viclafargue]

Why the change?

[aamijar]

This is fixing an existing bug in the test. The dataset n_features is set to 20, but we are reducing to n_components=30 which is doesn't make sense for PCA. This bug was caught due to the new mdspan runtime validation checks.

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@cpp/src/hdbscan/detail/soft_clustering.cuh`:
- Around line 73-74: exemplar_idx is int64_t* but the code reinterpret_casts it
to size_t* for raft::make_device_vector_view (exemplar_idx_view), which is
unsafe and non-portable; instead allocate or obtain a device vector/view of the
correct integer type (e.g., int64_t or int) and perform an explicit conversion
using thrust::transform (or a similar device-side transform) to convert
exemplar_idx into the target index type before calling
raft::make_device_vector_view; update the call site that constructs
exemplar_idx_view to use the converted buffer (and prefer int consistent with
other modules like sgd.cuh/sparse_util.cuh if indices elsewhere use int) so no
reinterpret_cast is used.

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@cpp/src/tsvd/tsvd.cuh`:
- Around line 89-96: weighted_power is using handle.get_stream() while earlier
ops (rsvdFixedRank, transpose) and the following ratio call use the explicit
stream variable, risking a data race if stream != handle.get_stream(); change
the weighted_power invocation to accept and use the explicit stream (aligning
its stream parameter with the variable stream) or call a version/overload that
takes a stream so singular_vals is read on the same stream as it was written,
keeping the subsequent raft::matrix::ratio call unchanged; mirror the
stream-alignment pattern used in tsvdFit or calEig to ensure singular_vals,
explained_vars, and explained_var_ratio use the same stream.

🧹 Nitpick comments (1)

cpp/src/tsvd/tsvd.cuh (1)

345-364: Prefer creating a new handle instead of mutating the const reference's stream.

The save/mutate/restore pattern on a const raft::handle_t& is fragile:

• Not exception-safe — if trunc_zero_origin or weighted_sqrt throw, orig_stream is never restored.
• Inconsistent with the cleaner pattern used in calEig (lines 122-123) which creates a separate raft::resources.

Consider aligning with the calEig approach:

Suggested refactor

-  auto orig_stream = handle.get_stream();
-  raft::resource::set_cuda_stream(handle, stream);
-
+  raft::resources handle_stream;
+  raft::resource::set_cuda_stream(handle_stream, stream);
   raft::matrix::trunc_zero_origin(
-    handle,
+    handle_stream,
     raft::make_device_matrix_view<const math_t, std::size_t, raft::col_major>(
       components_all.data(), prms.n_cols, prms.n_cols),
     raft::make_device_matrix_view<math_t, std::size_t, raft::col_major>(
       components, n_components, prms.n_cols));

   math_t scalar = math_t(1);
   raft::matrix::weighted_sqrt(
-    handle,
+    handle_stream,
     raft::make_device_matrix_view<const math_t, std::size_t, raft::row_major>(
       explained_var_all.data(), std::size_t(1), n_components),
     raft::make_device_matrix_view<math_t, std::size_t, raft::row_major>(
       singular_vals, std::size_t(1), n_components),
     raft::make_host_scalar_view(&scalar));
-
-  raft::resource::set_cuda_stream(handle, orig_stream);

Based on learnings: "Do not create reusable CUDA resources directly in ML algorithm implementations; instead use resources from raft::handle_t" and "Avoid the default CUDA stream (NULL or 0); use raft::handle_t::get_stream() for single-stream operations and internal streams for multi-stream operations."

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@cpp/src_prims/opg/linalg/svd.cu`:
- Around line 86-93: The loop is mutating the shared const handle via
raft::resource::set_cuda_stream(handle, streams[i]); instead create a local
non-const handle copy (e.g., auto local_handle = handle;), call
raft::resource::set_cuda_stream(local_handle, streams[i]) before calling
raft::linalg::binary_div_skip_zero with local_handle, and avoid restoring the
global handle; use local_handle in the call that references U, partsToRanks,
ADesc, and S so the shared const handle is not modified. Ensure any subsequent
stream restore uses the original local copy if needed.
- Around line 53-68: The code currently mutates the caller's const handle via
raft::resource::set_cuda_stream(handle, streams[0]); instead make a local handle
copy, set the stream on that local copy, and call the matrix ops with the local
handle to preserve const-correctness and exception/thread safety (e.g. create
raft::handle_t local_handle = handle; call
raft::resource::set_cuda_stream(local_handle, streams[0]); then use local_handle
for raft::matrix::col_reverse, row_reverse, weighted_sqrt, etc.), and remove
orig_stream/restore logic; apply the same change inside the loop where
set_cuda_stream is used at the later locations (around the mentions at lines 87
and 93).

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@cpp/src/tsvd/tsvd.cuh`:
- Around line 89-100: The call to raft::matrix::weighted_power is incorrect for
squaring singular_vals (it multiplies by a scalar); replace the
weighted_power(...) call with the correct exponent API (e.g.,
raft::linalg::power_scalar or equivalent) to compute explained_vars =
singular_vals^2 into explained_vars using exponent 2 (operate on singular_vals
and write to explained_vars, keeping handle, singular_vals, explained_vars, and
prms.n_components as targets), and make the CUDA stream save/restore around this
operation exception-safe (use an RAII guard or try/finally to call
raft::resource::set_cuda_stream(handle, orig_stream) so the original stream is
always restored even if the power operation throws).

🧹 Nitpick comments (1)

cpp/src/pca/pca_mg.cu (1)

185-191: Consider hoisting raft::resources to function scope to avoid repeated construction.

Both whiten blocks (lines 185-186 and 221-222) create a separate raft::resources object configured identically with streams[0]. The same pattern repeats in fit_impl (line 63) and inverse_transform_impl (lines 320, 353). Creating a single raft::resources at function scope would reduce repetition.

♻️ Example for transform_impl

 {
   std::vector<MLCommon::Matrix::RankSizePair*> local_blocks = input_desc.partsToRanks;
+  raft::resources handle_stream_zero;
+  raft::resource::set_cuda_stream(handle_stream_zero, streams[0]);
 
   if (prms.whiten) {
     T scalar = T(sqrt(prms.n_rows - 1));
     raft::linalg::scalarMultiply(
       components, components, scalar, prms.n_cols * prms.n_components, streams[0]);
-    raft::resources handle_stream_zero;
-    raft::resource::set_cuda_stream(handle_stream_zero, streams[0]);
     raft::linalg::binary_div_skip_zero<raft::Apply::ALONG_ROWS>(

(And remove the duplicate at line 221-222.)

Also applies to: 221-227

[aamijar]

/merge