Dry Run Protocol - #2961
Conversation
…mory Introduce a dry-run execution framework that replaces device and host memory resources with lightweight fake allocators to measure peak memory usage without holding real memory. New files: - dry_run_memory_resource.hpp: dry_run_allocator (lock-free bump allocator), dry_run_device_memory_resource, dry_run_host_memory_resource, dry_run_resource_manager (RAII), and dry_run_execute() helper. - dry_run_flag.hpp: boolean dry-run flag as a raft resource, allowing algorithms to skip kernel execution during profiling. - tests/util/dry_run_memory_resource.cpp: unit tests. The dry_run_allocator probes the upstream once to obtain a base address, then atomically bumps a pointer for each allocation — no mutex, no map, no real memory held after the initial probe.
…pinned_memory_resource Add pinned and managed resources to the raft::resources handle to make it possible to customize / temporarily replace these resources
…aking change due to transitive includes in downstream libraries
Merges Remove deprecated headers (NVIDIA#2939). Conflict resolutions: - rsvd.cuh: Use new mdspan-based raft::matrix::sqrt and reciprocal APIs (they have internal dry-run guards); kept cudaMemsetAsync guard - svd.cuh: Use raft::matrix::weighted_sqrt (has internal dry-run guard) - matrix.cuh: Accept deletion (deprecated, removed in main) Co-authored-by: Cursor <cursoragent@cursor.com>
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughIntroduces dry-run resource tracking for allocation profiling. Adds resource adaptors, RAII activation, container reallocation, and dry-run guards across RAFT APIs. Updates tests and documentation for allocation prediction and guard behavior. ChangesDry-Run Memory Profiling
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (10)
cpp/include/raft/matrix/slice.cuh (1)
52-60: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Keep the public input validation before the dry-run exit.
Line 52 returns before the layout and bounds checks, so dry-run silently accepts slices that the real call would reject. That makes dry-run behavior diverge at the public API boundary.
Move the dry-run return after the
RAFT_EXPECTS(...)block and beforedetail::sliceMatrix(...).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/matrix/slice.cuh` around lines 52 - 60, The public validation in sliceMatrix is being skipped during dry-run because the early resource::get_dry_run_flag(handle) return comes before the RAFT_EXPECTS checks. Move the dry-run exit to after the existing layout and bounds validation in sliceMatrix, so is_row_or_column_major and the row/col extent checks always run before any return, and only detail::sliceMatrix is bypassed in dry-run mode.cpp/include/raft/sparse/convert/detail/bitset_to_csr.cuh (1)
123-147: 🩺 Stability & Availability | 🔴 CriticalCRITICAL: Wrap both CUB scan calls in
RAFT_CUDA_TRY.
cub::DeviceScan::ExclusiveSumreturnscudaError_t; leaving either call unchecked can hide a failed workspace query or scan until later stream work.Suggested fix
- cub::DeviceScan::ExclusiveSum( - nullptr, scan_ws_bytes, sub_nnz.data(), sub_nnz.data(), sub_nnz_size + 1, stream); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + nullptr, scan_ws_bytes, sub_nnz.data(), sub_nnz.data(), sub_nnz_size + 1, stream)); @@ - cub::DeviceScan::ExclusiveSum( - scan_ws.data(), scan_ws_bytes, sub_nnz.data(), sub_nnz.data(), sub_nnz_size + 1, stream); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + scan_ws.data(), scan_ws_bytes, sub_nnz.data(), sub_nnz.data(), sub_nnz_size + 1, stream));🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/convert/detail/bitset_to_csr.cuh` around lines 123 - 147, Both cub::DeviceScan::ExclusiveSum calls in bitset_to_csr need to be checked with RAFT_CUDA_TRY so CUDA failures are surfaced immediately. Wrap the workspace-size query call and the actual scan call in RAFT_CUDA_TRY, keeping the existing scan_ws_bytes and scan_ws usage intact, and make sure the fix is applied in the bitset-to-CSR conversion path near calc_nnz_by_rows and csr initialization.Sources: Coding guidelines, Path instructions
cpp/include/raft/spectral/detail/matrix_wrappers.hpp (1)
106-118: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftHIGH:
nrm1()now undercounts dry-run memory.Returning
0beforethrust::reduceskips any temporary storage that the allocator-backed Thrust policy would request, so dry-run can underestimate memory duringmodularity_matrix_tconstruction. Have you considered keeping a temp-storage sizing/probe path here instead of short-circuiting before the reduce setup? As per path instructions, "Allocations through rmm/raft memory resources must NOT be guarded—allocation attempts must still be tracked to produce accurate memory statistics."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/spectral/detail/matrix_wrappers.hpp` around lines 106 - 118, `nrm1()` is short-circuiting on dry-run before Thrust sets up its reduction, which skips allocator-backed temporary storage accounting and underestimates memory during `modularity_matrix_t` construction. Update `raft::spectral::detail::matrix_wrappers::nrm1()` to preserve the Thrust reduce/setup path for dry-run memory probing, and only avoid the actual computation if needed after temp-storage sizing has been recorded; keep the `thrust::reduce`/`thrust_policy` path reachable so rmm/raft allocations are still tracked.Source: Path instructions
cpp/include/raft/random/detail/rmat_rectangular_generator.cuh (1)
209-224: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Dry-run bypasses
thetashape validationThis
returnnow happens before the host-sidetheta.extent(0)check, so dry-run accepts inputs that the real call would reject. That makes dry-run mask integration bugs instead of preserving the API contract.As per coding guidelines, public APIs should “validate their metadata wherever possible,” and the dry-run guard should only skip the CUDA work.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/random/detail/rmat_rectangular_generator.cuh` around lines 209 - 224, The dry-run early return in rmat_rectangular_gen is bypassing the host-side theta shape validation, so invalid metadata can slip through. Move the theta.extent(0) check so it runs before the resource::get_dry_run_flag(handle) return, while still keeping the CUDA work skipped in dry-run. Keep the existing validation logic in the rmat_rectangular_gen path and ensure dry-run preserves the same API contract as a real call.Sources: Coding guidelines, Path instructions
cpp/include/raft/stats/contingency_matrix.cuh (1)
115-128: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftHIGH:
get_input_class_cardinality()leaves host outputs indeterminate in dry-runThis early return skips both pointer validation and the writes to
minLabel/maxLabel. Callers use this helper to sizeout_matbeforecontingency_matrix(), so the documented workflow is no longer safe under dry-run.As per path instructions, dry-run support should keep
raft::resourcesentry points safe to call, and the developer guide still expects metadata/host-side contract checks where possible.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/contingency_matrix.cuh` around lines 115 - 128, `get_input_class_cardinality()` currently returns immediately on dry-run, leaving `minLabel` and `maxLabel` unset and skipping their validation. Update this helper so the host outputs are always initialized to safe values (or otherwise deterministically populated) even when `resource::get_dry_run_flag(handle)` is true, while keeping the existing `RAFT_EXPECTS` checks and the `detail::getInputClassCardinality` path intact for normal execution. Reference the `get_input_class_cardinality`, `resource::get_dry_run_flag`, and `detail::getInputClassCardinality` symbols when adjusting the control flow.Sources: Coding guidelines, Path instructions
cpp/include/raft/stats/detail/homogeneity_score.cuh (1)
46-56: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Dry-run now reports perfect homogeneity for non-empty inputs.
When
dry_runis true,mutual_info_score(...)andentropy(...)both short-circuit to0, so the existingcomputedEntropy == 0 ? 1.0 : ...fallback returns a perfect score without doing any work. That makes dry-run observably wrong for any caller that reads the result during profiling. Consider short-circuitingdry_runin this wrapper before the entropy-zero special case and returning a neutral sentinel such as0.0.Suggested fix
double homogeneity_score(bool dry_run, const T* truthClusterArray, const T* predClusterArray, int size, T lowerLabelRange, T upperLabelRange, cudaStream_t stream) { if (size == 0) return 1.0; + if (dry_run) return 0.0; double computedMI, computedEntropy; computedMI = mutual_info_score( dry_run, truthClusterArray, predClusterArray, size, lowerLabelRange, upperLabelRange, stream); computedEntropy = entropy(dry_run, truthClusterArray, size, lowerLabelRange, upperLabelRange, stream);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/detail/homogeneity_score.cuh` around lines 46 - 56, The homogeneity wrapper is returning a perfect score during dry-run because both `mutual_info_score(...)` and `entropy(...)` can short-circuit to zero, which then triggers the `computedEntropy == 0` fallback in `homogeneity_score(...)`. Update this logic so `dry_run` is handled explicitly before the entropy-zero branch, and return a neutral sentinel such as `0.0` instead of `1.0` when no real computation is performed. Keep the change localized to the `homogeneity_score` path in this header and preserve the existing non-dry-run behavior for normal inputs.cpp/include/raft/stats/detail/meanvar.cuh (1)
203-207: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winHIGH: Dry-run still queries CUDA occupancy in the row-major path.
Line 204 and Line 207 execute CUDA/device queries before the
if (!dry_run)guard, someanvar(..., dry_run=true, rowMajor=true)still touches the CUDA runtime. That breaks the dry-run contract and can fail on setups using dry-run specifically to avoid GPU work.Suggested fix
- int occupancy; - RAFT_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, meanvar_kernel_rowmajor<T, I, BlockSize>, BlockSize, 0)); - gs.y = - std::min(gs.y, raft::ceildiv<decltype(gs.y)>(occupancy * getMultiProcessorCount(), gs.x)); + if (!dry_run) { + int occupancy; + RAFT_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &occupancy, meanvar_kernel_rowmajor<T, I, BlockSize>, BlockSize, 0)); + gs.y = std::min( + gs.y, raft::ceildiv<decltype(gs.y)>(occupancy * getMultiProcessorCount(), gs.x)); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/detail/meanvar.cuh` around lines 203 - 207, The row-major dry-run path in meanvar still performs CUDA runtime queries before the dry_run check, so move the occupancy and multiprocessor-count logic behind the existing dry_run guard. Update the meanvar row-major branch in meanvar.cuh so meanvar(..., dry_run=true, rowMajor=true) skips cudaOccupancyMaxActiveBlocksPerMultiprocessor and getMultiProcessorCount entirely, only computing gs.y when not in dry-run mode.Source: Path instructions
cpp/include/raft/core/host_mdarray.hpp (1)
219-228: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Dry-run should not suppress host-scalar initialization.
This write is host-side only, so skipping it leaves
make_host_scalar(res, v)uninitialized under dry-run and can make host-side branching or workspace sizing diverge from real execution. Consider keeping the assignment unconditional here.Suggested fix
- if (!resource::get_dry_run_flag(res)) { scalar(0) = v; } + scalar(0) = v;As per path instructions, "Any expensive function or any function involving CUDA-calls must be guarded" and dry-run APIs taking
raft::resourcesshould remain safe to call.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/core/host_mdarray.hpp` around lines 219 - 228, The host-side initialization in make_host_scalar is being skipped when resource::get_dry_run_flag(res) is set, which leaves the returned host_scalar uninitialized during dry-run. Update make_host_scalar in host_mdarray.hpp so the assignment to scalar(0) = v happens unconditionally after construction, while keeping any CUDA or expensive work guarded elsewhere; this preserves correct host-side behavior without breaking dry-run safety.Source: Path instructions
cpp/include/raft/core/pinned_mdarray.hpp (1)
115-122: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Dry-run should not skip pinned-scalar initialization.
This is still a host-side write. Guarding it leaves
make_pinned_scalar(handle, v)undefined in dry-run and can change host-visible control flow without improving allocation tracking.Suggested fix
- if (!resource::get_dry_run_flag(handle)) { scalar(0) = v; } + scalar(0) = v;As per path instructions, "Any expensive function or any function involving CUDA-calls must be guarded" and dry-run APIs taking
raft::resourcesshould remain safe to call.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/core/pinned_mdarray.hpp` around lines 115 - 122, The dry-run guard in make_pinned_scalar is incorrectly skipping the host-side assignment to the pinned_scalar, leaving the returned scalar uninitialized in dry-run mode. Remove the conditional around the scalar initialization in make_pinned_scalar so the value is always written after constructing pinned_scalar<ElementType>, while keeping any CUDA/allocation-related behavior safe through the existing resource/policy path.Source: Path instructions
cpp/include/raft/core/managed_mdarray.hpp (1)
115-122: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Dry-run should still initialize managed scalars.
managed_scalaris host-accessible, so guarding this write can leave the value undefined during dry-run and make any host-side logic that inspects it diverge from the real path. The allocation is already preserved, so this assignment should stay unconditional.Suggested fix
- if (!resource::get_dry_run_flag(handle)) { scalar(0) = v; } + scalar(0) = v;As per path instructions, "Any expensive function or any function involving CUDA-calls must be guarded" and dry-run APIs taking
raft::resourcesshould remain safe to call.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/core/managed_mdarray.hpp` around lines 115 - 122, The make_managed_scalar helper is conditionally skipping initialization during dry-run, which can leave managed_scalar<ElementType> uninitialized even though it is host-accessible. Update make_managed_scalar so the scalar(0) = v assignment is unconditional, while keeping the allocation path safe for dry-run; use the existing make_managed_scalar and managed_scalar<ElementType> symbols to locate the write and remove the resource::get_dry_run_flag(handle) guard around it.Source: Path instructions
🟠 Major comments (26)
cpp/include/raft/linalg/detail/norm.cuh-65-67 (1)
65-67: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH:
colNormCaller<LinfNorm>bypasses dry-run.Line 67 passes a hardcoded
false, so column-wiseLinfNormstill runs the reduction path during dry-run instead of only tracking allocations.Suggested fix
} else if constexpr (norm_type == LinfNorm) { reduce<rowMajor, false, Type, OutType, IdxType>( - false, dots, data, D, N, (OutType)0, stream, false, raft::abs_op(), raft::max_op(), fin_op); + dry_run, dots, data, D, N, (OutType)0, stream, false, raft::abs_op(), raft::max_op(), fin_op); } else {As per path instructions, “Any expensive function or any function involving CUDA-calls must be guarded via
resource::get_dry_run_flag(res).”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/linalg/detail/norm.cuh` around lines 65 - 67, `colNormCaller<LinfNorm>` is still invoking the reduction path during dry-run because it passes a hardcoded false into `reduce`; update this branch to follow the same dry-run guard used elsewhere by checking `resource::get_dry_run_flag(res)` before calling `reduce`. Keep the fix localized in `colNormCaller` within norm.cuh, and ensure the `LinfNorm` path only performs allocation tracking when dry-run is enabled rather than executing the CUDA reduction.Source: Path instructions
cpp/include/raft/linalg/strided_reduction.cuh-132-132 (1)
132-132: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Dry-run masks invalid
strided_reductioncalls.Line 132 returns before the output-size checks and before the unsupported-type guard inside
stridedReduction(...), so a dry-run can succeed for a call that the real path rejects immediately. Have you considered leaving the validation/type gate in place and skipping only the final launch?As per coding guidelines, Public Interface says to “validate their metadata wherever possible.” Based on path instructions, dry-run should guard expensive/CUDA work via
resource::get_dry_run_flag(res), not bypass cheap host-side validation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/linalg/strided_reduction.cuh` at line 132, The dry-run early return in stridedReduction is skipping host-side validation, so invalid calls can appear to succeed. Keep the metadata/output-size checks and unsupported-type guard in place in stridedReduction, and use resource::get_dry_run_flag(handle) only to skip the expensive CUDA launch path after validation has already run.Sources: Coding guidelines, Path instructions
cpp/include/raft/linalg/add.cuh-105-107 (1)
105-107: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Keep public mdspan validation outside the dry-run fast path.
Lines 107, 145, and 181 return before the contiguity/size
RAFT_EXPECTSchecks, so dry-run can report success for malformed inputs that the real execution path still rejects. Have you considered moving the dry-run guard to just before the actual kernel dispatch instead?As per coding guidelines, Public Interface says to “validate their metadata wherever possible.” Based on path instructions, dry-run should guard expensive/CUDA work via
resource::get_dry_run_flag(res), not bypass cheap host-side validation.Also applies to: 145-145, 181-181
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/linalg/add.cuh` around lines 105 - 107, The dry-run early return in add and the other public overloads is bypassing the mdspan contiguity and size RAFT_EXPECTS checks, so malformed inputs can incorrectly succeed. Move the resource::get_dry_run_flag(handle) guard in add, and the equivalent guards in the other overloads, to just before the actual kernel dispatch so host-side validation always runs while only the expensive CUDA work is skipped.Sources: Coding guidelines, Path instructions
cpp/include/raft/matrix/sqrt.cuh-37-40 (1)
37-40: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Preserve the shape checks in dry-run mode.
Both out-of-place overloads return before
RAFT_EXPECTS(in.size() == out.size()), so dry-run hides mismatched shapes that real execution would fail on.Validate sizes first, then short-circuit before launching the math kernels.
Also applies to: 76-83
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/matrix/sqrt.cuh` around lines 37 - 40, The out-of-place sqrt overloads in the sqrt.cuh API are returning early on resource::get_dry_run_flag(handle) before the RAFT_EXPECTS size check, which lets mismatched input/output shapes slip through in dry-run mode. Move the shape validation in the affected sqrt overloads (the out-of-place paths around the seqRoot call) so RAFT_EXPECTS(in.size() == out.size()) runs before any dry-run short-circuit, then keep the dry-run return only after validation and before launching the kernel.cpp/include/raft/matrix/detail/select_k-inl.cuh-131-132 (1)
131-132: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: This dry-run guard is too deep for the sorted radix path.
When
sorted == true, the caller still launchesraft::linalg::map_offset(...)to populateoffsetsbefore reaching this return, so dry-run still performs CUDA work on the radix+sort path. That breaks the dry-run contract even though the actual segmented sort is skipped.As per path instructions, “Any expensive function or any function involving CUDA-calls must be guarded via
resource::get_dry_run_flag(res).”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/matrix/detail/select_k-inl.cuh` around lines 131 - 132, The dry-run guard in the select-k flow is placed too late, so the sorted radix path still executes CUDA work through raft::linalg::map_offset before returning. Move the resource::get_dry_run_flag(handle) check earlier in the select-k implementation, before any path that can launch CUDA work, and ensure the guard covers the sorted==true branch as well as the unsorted path. Use the select-k logic around raft::linalg::map_offset and the surrounding sorted radix handling to locate the fix.Source: Path instructions
cpp/include/raft/matrix/detail/select_warpsort.cuh-1193-1193 (1)
1193-1193: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: This early return under-reports dry-run memory for warpsort.
select_k_()now does the right thing by allocatingtmp_val/tmp_idxbefore checkingdry_run, but Line 1193 returns before that helper is reached. Any caller using this overload in dry-run will miss those allocations and get a smaller peak-memory estimate.As per path instructions, “Allocations through rmm/raft memory resources must NOT be guarded—allocation attempts must still be tracked to produce accurate memory statistics.”
Suggested fix
- if (resource::get_dry_run_flag(res)) { return; } ASSERT(k <= kMaxCapacity, "Current max k is %d (requested %d)", kMaxCapacity, k);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/matrix/detail/select_warpsort.cuh` at line 1193, The dry-run guard in the warpsort path is skipping allocation accounting, so peak-memory estimates are too low. Update the overload in select_warpsort.cuh around the select_k_() call site to avoid returning before the helper runs, and let the tmp_val/tmp_idx allocation attempts be reached even when resource::get_dry_run_flag(res) is true. Keep the dry-run behavior in the downstream logic, but do not guard the rmm/raft-backed allocations that must be tracked for memory statistics.Source: Path instructions
cpp/include/raft/sparse/solver/detail/lanczos.cuh-198-227 (1)
198-227: 🩺 Stability & Availability | 🟠 MajorHIGH: Use the RAFT Thrust policy here.
thrust::sequenceandthrust::sortonthrust::devicecan run outsidehandle’s stream, so the laterraft::matrix::gather(...)calls may read unsynchronized indices. Useraft::resource::get_thrust_policy(handle)here instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/solver/detail/lanczos.cuh` around lines 198 - 227, The index initialization and sorting in the lanczos selection path currently uses the generic thrust::device policy, which can bypass the handle stream and leave later gather operations racing unsynchronized data. Update the thrust::sequence and both thrust::sort calls in this selection block to use the RAFT Thrust policy from raft::resource::get_thrust_policy(handle), keeping the existing indices/selected_indices logic intact so all work stays ordered on the handle’s stream.Sources: Coding guidelines, Path instructions
cpp/include/raft/solver/detail/lap_functions.cuh-267-276 (1)
267-276: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Zero-cover dry-run undercounts its peak workspace.
executeZeroCover()now returns beforezeroCoverIteration()is reached, and the dry-run branch insidezeroCoverIteration()only sizes the CSR buffers. The real path also allocatespredicates_vandaddresses_vatSP * N, so the dry-run report will miss the dominant memory in this phase.As per path instructions, "Allocations through rmm/raft memory resources must NOT be guarded—allocation attempts must still be tracked to produce accurate memory statistics," and "early-return dry-run guards don’t accidentally skip required resource setup/teardown..."
Also applies to: 353-353
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/solver/detail/lap_functions.cuh` around lines 267 - 276, The dry-run path in zero-cover is undercounting workspace because `executeZeroCover()` now exits before `zeroCoverIteration()`, and `zeroCoverIteration()`’s dry-run branch only allocates the CSR buffers while skipping the `predicates_v` and `addresses_v` allocations that dominate memory use. Update `executeZeroCover`/`zeroCoverIteration` so dry-run still performs the same allocation tracking as the real path, including the `predicates_v` and `addresses_v` `rmm::device_uvector` allocations, while still returning early before any compute; keep the behavior aligned with the existing `csr_ptrs_v`/`csr_neighbors_v` setup so memory statistics remain accurate.Source: Path instructions
cpp/include/raft/sparse/linalg/detail/spmm.hpp-90-95 (1)
90-95: 🩺 Stability & Availability | 🟠 MajorHIGH:
bufferSizeis a byte count, buttmpis allocated asValueTypeelements.
cusparsespmm_bufferSizereturns bytes, so this over-allocates the workspace bysizeof(ValueType)and skews dry-run memory accounting. Use a byte buffer or convert the count to elements before allocating.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/linalg/detail/spmm.hpp` around lines 90 - 95, The workspace allocation in the sparse SpMM path is using `bufferSize` as if it were a count of `ValueType` elements, but `cusparsespmm_bufferSize` returns bytes. Update the allocation in the `spmm`/workspace setup around `tmp` to treat the buffer as raw bytes or to convert bytes to element count before constructing the `rmm::device_uvector`, and make sure the dry-run path still accounts for the same byte-sized workspace.cpp/include/raft/linalg/detail/lstsq.cuh-213-216 (1)
213-216: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winHIGH: Clean up
gesvdj_paramson the non-dry-run path too.Issue: the dry-run branch destroys
gesvdj_params, but the normal path exits after Line 244 withoutcusolverDnDestroyGesvdjInfo.
Why: repeatedlstsqSvdJacobicalls leak cuSOLVER resources.As per path instructions,
cpp/REVIEW_GUIDELINES.mdrequires checking resource cleanup on all paths.Suggested fix
raft::linalg::gemv(handle, U, n_rows, minmn, b, Ub, true, stream); raft::linalg::binaryOp(Ub, Ub, S, minmn, DivideByNonZero<math_t>(), stream); raft::linalg::gemv(handle, V, n_cols, minmn, Ub, w, false, stream); + RAFT_CUSOLVER_TRY(cusolverDnDestroyGesvdjInfo(gesvdj_params)); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/linalg/detail/lstsq.cuh` around lines 213 - 216, The `lstsqSvdJacobi` flow cleans up `gesvdj_params` only in the `resource::get_dry_run_flag(handle)` branch, so add matching `cusolverDnDestroyGesvdjInfo(gesvdj_params)` cleanup on the normal execution path before the function returns. Make sure every exit path in `lstsq.cuh` that creates `gesvdj_params` also destroys it, using the same `gesvdj_params` symbol so repeated calls do not leak cuSOLVER resources.Source: Path instructions
cpp/include/raft/stats/detail/scores.cuh-162-184 (1)
162-184: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winHIGH: Use element counts for
device_uvectorand avoid untracked host allocation in dry-run.Issue:
device_uvector<double>sizes are element counts, son * sizeof(double)allocates 8× too many doubles; dry-run also allocatesh_sorted_abs_diffs(n)viastd::vectorbefore returning.
Why: this over-reports dry-run memory and can cause avoidable OOM.As per path instructions, dry-run must keep rmm/raft allocation tracking accurate while skipping unnecessary work.
Suggested fix
- int array_size = n * sizeof(double); - rmm::device_uvector<double> abs_diffs_array(array_size, stream); - rmm::device_uvector<double> sorted_abs_diffs(array_size, stream); - rmm::device_uvector<double> tmp_sums(2 * sizeof(double), stream); + rmm::device_uvector<double> abs_diffs_array(n, stream); + rmm::device_uvector<double> sorted_abs_diffs(n, stream); + rmm::device_uvector<double> tmp_sums(2, stream); @@ - std::vector<double> mean_errors(2); - std::vector<double> h_sorted_abs_diffs(n); int thread_cnt = 256; int block_cnt = raft::ceildiv(n, thread_cnt); if (dry_run) { return; } + + std::vector<double> mean_errors(2); + std::vector<double> h_sorted_abs_diffs(n);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/detail/scores.cuh` around lines 162 - 184, The allocation logic in the scores computation is using byte counts where `rmm::device_uvector<double>` expects element counts, so `abs_diffs_array`, `sorted_abs_diffs`, and `tmp_sums` should be sized in doubles rather than `sizeof(double)` multiples; fix this in the relevant scoring routine before the CUB `SortKeys` call. Also move the host-side `std::vector<double> h_sorted_abs_diffs` allocation behind the `dry_run` early return (or otherwise skip it during dry-run) so dry-run only accounts for tracked rmm/raft device allocations and does not create untracked host memory.Source: Path instructions
cpp/include/raft/random/detail/multi_variable_gaussian.cuh-188-190 (1)
188-190: 🩺 Stability & Availability | 🟠 MajorHIGH: Guard MVG setup before the constructor does CUDA work
cpp/include/raft/random/detail/multi_variable_gaussian.cuh:146-190Issue:
give_gaussian()exits in dry-run, butbuild_multi_variable_gaussian_token_impl()still constructsmulti_variable_gaussian_impl, which creates the CURAND generator and runs cuSOLVER buffer-size calls first.
Impact: dry-run paths still execute CUDA/library setup and can fail before reaching the guard.Have you considered making the token/constructor setup dry-run-aware and deferring only the workspace accounting?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/random/detail/multi_variable_gaussian.cuh` around lines 188 - 190, The dry-run guard in give_gaussian() is too late because multi_variable_gaussian_impl is still constructed earlier by build_multi_variable_gaussian_token_impl(), which triggers CURAND and cuSOLVER setup before the early return. Make the token/constructor path dry-run-aware by checking resource::get_dry_run_flag(handle) before creating multi_variable_gaussian_impl, and defer only the workspace accounting needed for the token so dry-run execution avoids all CUDA/library initialization.Source: Path instructions
cpp/include/raft/stats/detail/adjusted_rand_index.cuh-125-150 (1)
125-150: 🩺 Stability & Availability | 🟠 MajorUse a wider type for contingency sizing
Issue:
nUniqClasses * nUniqClassesis evaluated inMathT; with the defaultMathT=int, it overflows past ~46k classes/samples.
Why: dry-run setsnUniqClasses = size, so large inputs can turn the allocation size into a bogus value before any work runs. The same risk exists in the non-dry-run path when the label range is large.
Consider computing the extent insize_t(or another wider type) before allocating the contingency buffers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/detail/adjusted_rand_index.cuh` around lines 125 - 150, The contingency buffer sizing in adjusted_rand_index.cuh is computed in MathT, which can overflow when nUniqClasses is large, especially in the dry_run path where it can equal size. Update the sizing logic around the dContingencyMatrix allocation and the nUniqClasses calculation in the adjusted rand index helper to use a wider extent type such as size_t before multiplying, then cast safely for the device allocation. Keep the existing dry_run and countUnique flow intact, but ensure the final matrix element count is computed in an overflow-safe type.Sources: Coding guidelines, Path instructions
cpp/tests/core/temporary_device_buffer.cu-67-68 (1)
67-68: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Raw-stream copy bypasses dry-run here.
This
raft::copy(..., resource::get_cuda_stream(h))call still performs real CUDA work during the dry-run pass, so the test can mutateresultwhile claiming the path is dry-run compliant. Use theraft::resourcesoverload (or guard the raw copy) inside the lambda.Suggested fix
- raft::copy( - result.data(), d_view.data_handle(), d_view.extent(0), resource::get_cuda_stream(h)); + auto result_view = raft::make_device_vector_view<int>(result.data(), d_view.extent(0)); + raft::copy(h, result_view, raft::make_const_mdspan(d_view));As per path instructions, use
raft::resourcesinstead of raw streams/handles, and guard CUDA work in dry-run mode viaresource::get_dry_run_flag(res).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/core/temporary_device_buffer.cu` around lines 67 - 68, The copy inside the temporary device buffer test is using a raw CUDA stream, which bypasses dry-run handling and can still perform real device work. Update the lambda around the raft::copy call to use the raft::resources-based overload tied to the existing resource handle, and gate the copy with resource::get_dry_run_flag(res) so the dry-run pass does not mutate result. Keep the fix localized to the temporary_device_buffer test and the lambda containing raft::copy.Source: Path instructions
cpp/tests/core/bitset.cu-331-332 (1)
331-332: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Guard raw CUDA work inside the dry-run lambda.
Issue: Line 331 and Line 387 still launch GPU work during the dry-run pass.
Why: This violates dry-run semantics and can mask non-compliant kernels.Suggested fix
- RAFT_CUDA_TRY(cudaMemsetAsync( - repeat_device.data_handle(), 0, eval_n_elements * sizeof(bitset_t), stream)); + if (!resource::get_dry_run_flag(h)) { + RAFT_CUDA_TRY(cudaMemsetAsync(repeat_device.data_handle(), + 0, + eval_n_elements * sizeof(bitset_t), + resource::get_cuda_stream(h))); + } ... - raft::linalg::range(query_device.data_handle(), query_device.size(), stream); + if (!resource::get_dry_run_flag(h)) { + raft::linalg::range( + query_device.data_handle(), query_device.size(), resource::get_cuda_stream(h)); + }As per path instructions, “Any expensive function or any function involving CUDA-calls must be guarded via
resource::get_dry_run_flag(res).”Also applies to: 387-387
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/core/bitset.cu` around lines 331 - 332, Guard the raw CUDA work in the dry-run path so it only runs when resource::get_dry_run_flag(res) is false. Update the bitset test logic around the dry-run lambda in the relevant test helpers (including the cudaMemsetAsync call and the other CUDA launch at the referenced location) to skip these GPU operations during dry-run, keeping the lambda’s non-CUDA validation separate from device work.Source: Path instructions
cpp/tests/linalg/rsvd.cu-128-149 (1)
128-149: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: Percentage-based RSVD cases never hit the percentage API path.
Line 131 is dead here:
params.kis already rewritten from0to a concrete rank earlier inSetUp(), so the ratio-based fixtures now run the fixed-rank Jacobi variants instead ofrsvd_perc*. That drops coverage for the public percentage APIs and can hide dry-run regressions there.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/linalg/rsvd.cu` around lines 128 - 149, The percentage-based RSVD branch in the RSVD test is unreachable because `SetUp()` rewrites `params.k` before this lambda runs, so the test never exercises `rsvd_perc` or `rsvd_perc_symmetric`. Update the selection logic in the `rsvd` test around `raft::execute_with_dry_run_check` to branch on the original test mode or an explicit percentage/fixed-rank flag instead of `params.k`, so the ratio-based fixtures actually call the percentage APIs while fixed-rank cases continue to use `rsvd_fixed_rank_jacobi` and `rsvd_fixed_rank_symmetric_jacobi`.cpp/tests/sparse/csr_transpose.cu-97-113 (1)
97-113: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winHIGH: This dry-run wrapper breaks stream ordering with the test’s prepared inputs.
make_data()uploads into buffers onraft_handle/stream, but this lambda launchescsr_transposeon a differentraft::resourcesstream viaresource::get_cuda_stream(h). Without an explicit dependency, transpose can observe partially populated inputs. Have you considered runningexecute_with_dry_run_checkonraft_handleinstead so setup and compute stay ordered on the same resource? As per path instructions, verifycpp/**/*changes againstdocs/source/developer_guide.md, especially “Asynchronous operations and stream ordering”.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/sparse/csr_transpose.cu` around lines 97 - 113, The dry-run wrapper in csr_transpose test is launching work on a different stream than the one used by make_data(), which can break ordering and let csr_transpose read inputs before uploads complete. Update the execute_with_dry_run_check call to use the same raft_handle/resource context as the test setup so the prepared buffers and the transpose kernel share stream ordering, and keep the relevant resource/stream selection aligned through resource::get_cuda_stream and the surrounding test harness.Source: Path instructions
docs/source/dry_run_protocol.md-43-49 (1)
43-49: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winThe example teaches the wrong stream ownership model.
This snippet introduces a public-style API that takes a raw
cudaStream_talongsideraft::resources, which contradicts RAFT’s dry-run/resource guidance and can mislead users away from resource-ordered streams. Showauto stream = resource::get_cuda_stream(handle);instead of accepting a separate stream parameter. As per coding guidelines, prefer “raft::resourcesover raw handles/streams.” As per path instructions, docs changes should prioritize accuracy and consistency with current APIs.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/source/dry_run_protocol.md` around lines 43 - 49, The dry-run example in algo currently models an incorrect ownership pattern by accepting a raw cudaStream_t alongside raft::resources, which conflicts with the resource-first API guidance. Update the example to derive the stream from raft::resources using resource::get_cuda_stream(handle) inside algo, and remove the separate stream parameter so the signature and call sites reflect the preferred resource-ordered model. Keep the rest of the snippet aligned with this RAFT API style and ensure the doc example stays consistent with current dry-run guidance.Sources: Coding guidelines, Path instructions
cpp/tests/stats/meanvar.cu-84-85 (1)
84-85: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winHIGH: The allocation assertion now depends on a private
detailtype.Using
sizeof(raft::stats::detail::mean_var<T>)bakes private workspace layout into the test. That makes this suite fail on benign internal refactors instead of public dry-run regressions.As per coding guidelines, "Public API over
detail. When re-using functionality or writing tests, call the public API rather than functions in thedetailnamespace." As per path instructions, verify the code follows the idioms indocs/source/developer_guide.md, especially the "Preferred APIs and idioms" section.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/stats/meanvar.cu` around lines 84 - 85, The allocation assertion in the mean/var test is tied to the private mean_var detail type, so update the check to use the public stats API or a public dry-run/size query instead of sizeof(raft::stats::detail::mean_var<T>). Keep the existing rowMajor and params.cols logic in the test, but replace the private workspace-size dependency with the public symbol used by the mean/variance path so the test only validates supported API behavior.Sources: Coding guidelines, Path instructions
cpp/tests/stats/cov.cu-67-91 (1)
67-91: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winHIGH: This still leaves one
coventry point unverified under dry-run.Lines 67-91 wrap the mdspan overload, but Line 106 still calls
cov<false>(handle, ..., stream)directly. That means a regression in that publicraft::resourcespath can slip through while this file appears dry-run covered.As per path instructions, "Any function taking a
raft::resourceshandle is safe to call in dry-run mode."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/stats/cov.cu` around lines 67 - 91, The dry-run coverage in cov.cu only exercises the mdspan overload, while the public raft::resources entry point cov<false>(handle, ..., stream) remains unverified. Update the test to also invoke the raft::resources-based cov path inside the dry-run check, using the existing cov and raft::execute_with_dry_run_check symbols, so both overloads are covered under dry-run and regressions in the handle-taking API are caught.Source: Path instructions
cpp/tests/util/bitonic_sort.cu-146-152 (1)
146-152: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: this dry-run callback still launches the kernel.
The lambda ignores
hand calls the raw-streambitonic_launchdirectly, so the dry-run leg still executesbitonic_kernel<<<...>>>. That means this test never actually validates “skip GPU work in dry-run mode”; it only re-runs the kernel. Have you considered guarding the launch withif (!resource::get_dry_run_flag(h))or routing through araft::resources-aware helper? As per path instructions, dry-run CUDA work must be guarded viaresource::get_dry_run_flag(res)and RAFT code should preferraft::resourcesover raw streams.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/util/bitonic_sort.cu` around lines 146 - 152, The dry-run test callback in bitonic sort is still invoking the GPU launch unconditionally, so it never actually verifies that dry-run skips work. Update the lambda passed to execute_with_dry_run_check to use the provided raft::resources argument instead of ignoring it, and guard the bitonic_launch invocation with resource::get_dry_run_flag(h) (or equivalent raft::resources-aware helper) so the kernel only runs when not in dry-run mode. Prefer threading raft::resources through the launch path rather than calling the raw-stream bitonic_launch directly.Source: Path instructions
cpp/tests/sparse/symmetrize.cu-95-104 (1)
95-104: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftHIGH:
outescapes the dry-run resource scope.
outis an owningCOOcreated beforeexecute_with_dry_run_check, then populated inside the callback and dereferenced afterward. If the dry-run leg resizes/fills those buffers against the temporary dry-run resources, the laterassert_symmetrylaunch can end up reading storage whose allocator/resource context has already been torn down. Have you considered keeping the owningCOOentirely inside the callback (and asserting only on the non-dry-run leg), or splitting the dry-run accounting from the functional validation? As per path instructions, early-return dry-run guards must not break required resource setup/teardown, and the developer guide requires dry-run allocation tracking to stay within the dry-run resource scope.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/sparse/symmetrize.cu` around lines 95 - 104, `out` is being created outside `execute_with_dry_run_check`, so the owning COO can outlive the dry-run resource scope and be used after the temporary allocator is torn down. Move the `raft::sparse::COO<value_t, value_idx, nnz_t> out` setup and the `symmetrize`/`assert_symmetry` validation into the non-dry-run path, or otherwise keep the owning storage entirely within the callback used by `execute_with_dry_run_check`. Use the `execute_with_dry_run_check` and `raft::sparse::linalg::symmetrize` sites as the key locations to refactor so dry-run accounting stays separate from functional validation.Source: Path instructions
cpp/include/raft/core/detail/copy.hpp-403-405 (1)
403-405: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH: This guard drops copy-side workspace allocations from dry-run.
detail::copyallocates an intermediatedevice_mdarrayin theuse_intermediate_src/use_intermediate_dstpaths. Returning here skips those allocations entirely, so dry-run under-reports peak memory for layout/dtype conversion copies.Suggested fix
- if (resource::get_dry_run_flag(res)) { return; } + if constexpr (!(config::use_intermediate_src || config::use_intermediate_dst)) { + if (resource::get_dry_run_flag(res)) { return; } + }As per path instructions, "Allocations through rmm/raft memory resources must NOT be guarded—allocation attempts must still be tracked" and dry-run guards must not skip required setup/teardown.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/core/detail/copy.hpp` around lines 403 - 405, The dry-run early return in detail::copy is skipping workspace allocations needed by the use_intermediate_src and use_intermediate_dst paths, so peak memory is under-reported. Remove the blanket guard at the start of copy and ensure the intermediate device_mdarray setup still runs under dry-run so allocation attempts through the memory resource are tracked; only skip the actual data movement if needed. Keep the fix localized to detail::copy and the intermediate-copy branches that allocate temporary buffers.Source: Path instructions
cpp/include/raft/core/sparse_types.hpp-181-181 (1)
181-181: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winHIGH:
initialize_sparsity()no longer updates the sparse structure.Line 181 only reallocates
c_elements_.get_elements()/view()still size fromstructure_.get_nnz(), and anyStructureType::initialize_sparsity()override is skipped, so the matrix can keep reporting the oldnnzwhile the backing storage changed.Suggested fix
- void initialize_sparsity(nnz_type nnz) { c_elements_.reallocate(nnz); }; + void initialize_sparsity(nnz_type nnz) + { + structure_.initialize_sparsity(nnz); + c_elements_.reallocate(nnz); + };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/core/sparse_types.hpp` at line 181, `initialize_sparsity()` in `sparse_types.hpp` only reallocates `c_elements_`, so the sparse metadata stays stale and any `StructureType::initialize_sparsity()` override is bypassed. Update `initialize_sparsity(nnz_type nnz)` to also initialize the underlying `structure_` and ensure `structure_.get_nnz()` reflects the new sparsity before `get_elements()` and `view()` are used. Preserve the structure-specific initialization path by invoking the structure’s own `initialize_sparsity` behavior instead of replacing it with a bare reallocate.Source: Path instructions
cpp/include/raft/core/dry_run_resources.hpp-54-74 (1)
54-74: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftHIGH: Serialize process-global resource swaps.
Issue: each active
dry_run_resourcesinstance snapshots and replaces process-global host/device resources without a lifetime lock.
Why: overlapping dry-run scopes can restore globals out of order, making unrelated allocations hit the wrong or destroyed adaptor.Suggested fix
`#include` <cstddef> `#include` <cstdint> `#include` <memory> +#include <mutex> `#include` <utility> @@ class dry_run_resources : public resources { public: explicit dry_run_resources(const resources& existing) : resources(existing), + global_resource_lock_(global_resource_mutex()), active_(!resource::get_dry_run_flag(existing)), old_host_(raft::mr::get_default_host_resource()), old_device_(rmm::mr::get_current_device_resource_ref()) @@ private: + static auto global_resource_mutex() -> std::recursive_mutex& + { + static std::recursive_mutex mutex; + return mutex; + } + + // Hold while process-global memory resources are swapped. + std::unique_lock<std::recursive_mutex> global_resource_lock_; + // Declaration order determines destruction order.As per coding guidelines, RAFT algorithms should remain thread-safe with different
raft::resourcesinstances.Also applies to: 159-192
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/core/dry_run_resources.hpp` around lines 54 - 74, Serialize the process-global host/device resource swap performed by dry_run_resources so overlapping active scopes cannot restore resources out of order. Update dry_run_resources’ constructor/destructor path to take and hold a shared lifetime lock around the existing snapshots and the calls to raft::mr::set_default_host_resource and rmm::mr::set_current_device_resource, and keep the lock until the base-class cleanup completes. Use the dry_run_resources symbols and the init()/destructor flow to ensure only one active instance can replace the global adaptors at a time while still allowing separate raft::resources instances to remain thread-safe.Source: Coding guidelines
cpp/include/raft/mr/dry_run_resource.hpp-80-98 (1)
80-98: 🩺 Stability & Availability | 🟠 MajorHIGH: Keep probe deallocation on the allocation stream.
probe_containerallocates the async probe with the caller’scuda::stream_ref, but its destructor always frees it oncudaStreamPerThread. For stream-ordered memory resources, that can break ordering and race with pending work on the original stream.Suggested fix
class probe_container { MR mr_; void* ptr_; std::size_t size_; std::size_t alignment_; + cuda::stream_ref stream_{cudaStreamPerThread}; @@ probe_container(MR mr, cuda::stream_ref stream, std::size_t size, std::size_t alignment = alignof(std::max_align_t)) - : mr_(std::move(mr)), ptr_(nullptr), size_(size), alignment_(alignment) + : mr_(std::move(mr)), ptr_(nullptr), size_(size), alignment_(alignment), stream_(stream) { ptr_ = mr_.allocate(stream, size_, alignment_); } @@ if (ptr_ == nullptr) return; if constexpr (cuda::mr::resource<MR>) { - mr_.deallocate(cuda::stream_ref{cudaStreamPerThread}, ptr_, size_, alignment_); + mr_.deallocate(stream_, ptr_, size_, alignment_); } else { mr_.deallocate_sync(ptr_, size_, alignment_); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/mr/dry_run_resource.hpp` around lines 80 - 98, probe_container in dry_run_resource.hpp allocates using the caller-provided cuda::stream_ref but currently deallocates on cudaStreamPerThread in the destructor, which breaks stream ordering for cuda::mr::resource types. Update probe_container to retain the allocation stream as state and use that same stream when calling mr_.deallocate, while preserving the existing synchronous deallocation path for non-resource MR implementations.Source: Path instructions
A non-breaking src-only changes to modernize the use of raft primitives across cuVS source code. The general rule applied here is to prefer raft helpers taking `raft::resources` as an argument over other raft helpers over third-party libraries. - thrust::fill / thrust::fill_n → raft::matrix::fill - thrust::transform → raft::linalg::map - thrust::sequence / thrust::tabulate → raft::linalg::map_offset - raft::linalg::unaryOp / raft::linalg::binaryOp → raft::linalg::map - raft::linalg::add (pointer-based) → raft::linalg::add (mdspan-based) - raft::copy (pointer-based) → raft::copy (mdspan-based) - raft::update_device / raft::update_host → raft::copy (mdspan-based) - raft::linalg::rowNorm → raft::linalg::norm - raft::linalg::reduce (pointer-based) → raft::linalg::reduce (mdspan-based) - cudaMemsetAsync → raft::matrix::fill The purpose of this PR is to improve the consistency in using the library code (even though sometimes at the cost of a bit more auxiliary code). This is also a prerequisite to achieving dry run compliance in cuVS if we choose to merge that in NVIDIA/raft#2961 Authors: - Artem M. Chirkin (https://github.com/achirkin) Approvers: - Dante Gama Dessavre (https://github.com/dantegd) URL: NVIDIA#1837
huuanhhuyn
left a comment
huuanhhuyn
left a comment
There was a problem hiding this comment.
Thanks @achirkin for the heavy-lifting! I have went through this PR halfway.
| */ | ||
| inline void sync_stream(const resources& res, rmm::cuda_stream_view stream) | ||
| { | ||
| if (raft::resource::get_dry_run_flag(res)) { return; } |
There was a problem hiding this comment.
What is the motivation for guarding sync_stream?
Your dry_run_protocol.md suggests the opposite that sync_stream is safe without guard.
There was a problem hiding this comment.
The guide says it's safe to call this function, because it is dry-run compliant (takes the raft::resources as an argument) - hence it guards from any CUDA work inside.
cudaStreamSynchronize is cuda CUDA and it takes time, so we must guard it.
There was a problem hiding this comment.
If I understand dry-run correctly, when res is in dry-run mode, the stream contains no CUDA work and contains only light-weight fake memory operations. Therefore it would be fast to sync.
If user intentionally adds a CUDA/memory work to the stream without guarding, he wants the sync_stream to work I suppose.
I mean this guarding on sync stream wouldn't bring any benefit but limit its usage a little.
There was a problem hiding this comment.
It is still CUDA work and a call to cuda context, which may briefly lock. In dry run mode, nobody should ever call CUDA, that is a hard constraint.
|
|
||
| 1. **Allocations must not be guarded.** Every `rmm::device_uvector`, `rmm::device_scalar`, `rmm::device_buffer`, `raft::make_(device|host|pinned|managed)_(mdarray|matrix|vector|scalar)` allocation must execute in both modes so the tracker sees it. | ||
|
|
||
| 2. **CUDA work must be guarded.** Kernel launches, Thrust algorithms, cuBLAS/cuSOLVER/cuSPARSE compute calls, `cudaMemcpyAsync`, `cudaMemsetAsync`, and `raft::interruptible::synchronize` must not run in dry-run mode. |
There was a problem hiding this comment.
if all CUDA work is guarded / skipped in dry-run. How do we measure allocations whose allocation sizes are produced by a kernel?
If that kernel is dependent on earlier kernels and so on, dry-run must then execute all of them without guarding?
There was a problem hiding this comment.
The case you're describing falls under "allocation size depends on values/data". This is indeed generally a problem and we resort to providing safe upper bounds on allocations. Have a look at the sparse namespace - we have a lot cases like this.
In tests, this is covered by DATA_DRIVEN provenance: we compare that dry run produced stats are not smaller than the real produced stats.
There was a problem hiding this comment.
Is finding the upper bound always possible? And what if the upper bound is 3x, 10x?
| * This ensures the memory is released promptly. | ||
| */ | ||
| void reallocate(size_type size) | ||
| { |
There was a problem hiding this comment.
Although your implementation avoids peak data (old + new) when the new size is larger than the current size, it has two issues:
- when the new size is the same -> it does NOT reallocate while its name is "reallocate"
- when the new size is smaller -> it DOES allocate while it could more efficiently shrink the current buffer (similar to what
resize()already does)
I would suggest to do the current implementation only when the new size is larger. When the new size is less or equal, simply call resize() and clear the buffer.
I am not sure about clearing the buffer. It adds some overhead, but it sounds safer to me.
There was a problem hiding this comment.
Funny enough, both of these are pretty much intentional features. Here's a bit more context for this.
By design, mdarrays and mdspans are not supposed to be resizeable at all. However, at some point we decided to implement sparse matrices using the same container policies as mdarrays. For sparse data, one needs the resize functionality to allow changing sparsity pattern. So the resize function is exclusively used for sparse data and only to allocate a new sparse matrix storage of the required size (initialize_sparsity).
Back to reallocate addressing your concerns. The intention here is to always free up the memory as soon as requested: the released buffer may be huge (e.g. half of GPU memory); the user may want to use the memory as soon as possible, so we want it to be available.
Essentially reallocate is almost the same as running resize + shrink_to_fit on an underlying RMM container. But there's a catch: reallocate guarantees to NOT copy the data. This catch is important for dry run: we fake the memory allocations, so any memory copy would cause OOM / break the program.
There was a problem hiding this comment.
Thank you for sharing the context with sparse matrix.
I still don't understand why when the new size is less than or equal to the current size, we don't simply free the redundant memory and keep the remaining new size part. This will give back immediately current_size - new_size to the user. Or do you mean we have to free everything at the current address and then allocate the new size some where else? If so, why?
There was a problem hiding this comment.
The current implementation makes sense to me only when the new size is greater than the current size
| policy_t policy{}; | ||
| auto scalar = device_scalar<ElementType, IndexType>{handle, extents, policy}; | ||
| scalar(0) = v; | ||
| if (!resource::get_dry_run_flag(handle)) { scalar(0) = v; } |
There was a problem hiding this comment.
I guess we don't need to guard this scalar operation (same for other resources). It is lightweight to run and keep the code a little bit simpler.
There was a problem hiding this comment.
The scalar is on device - in the default per-device memory resource, which is faked in dry run mode. Therefore we have to guard it (and it falls under "GPU work" category anyway).
| @@ -0,0 +1,92 @@ | |||
| # Dry Run Protocol | |||
|
|
|||
| The dry run protocol lets callers estimate an algorithm's memory footprint without executing it. When enabled, the runtime swaps memory resources for lightweight trackers that record every allocation and deallocation, producing peak-usage statistics at the end. | |||
There was a problem hiding this comment.
I suppose we are not able to track non-raft allocations with dry-run. It is probably fine but we should be aware and document it somewhere.
Are we using third-party libraries which do allocations without RAFT somewhere? I am not aware about it.
There was a problem hiding this comment.
Plenty of rapids org libraries use the rmm memory resources directly. They are tracked though.
There was a problem hiding this comment.
And non-rapidsai libraries like Cutlass, CUB, Thrust, etc.? Are their allocations tracked?
There was a problem hiding this comment.
This is covered in the same file below in the section "What Needs Guarding". In general, I don't think it's sustainable to maintain a list of all possible libraries to tell whether they use rmm or not. But it's normally clear from the API or documentation of those libraries.
Please also refer to the presentation for the relevant overview https://drive.google.com/file/d/1uMkM-Xzhi-2ZAKe_z3zvCfHoefn0bAAx/view?usp=drive_link
divyegala
left a comment
divyegala
left a comment
There was a problem hiding this comment.
Would you consider guarding the kernel launch behind a dispatch function/macro where the dry run flag is checked, like we talked in your offline brown bag session?
|
Hi @divyegala , I definitely want to do that, but maybe in a separate PR in parallel, where we could focus on discussing the pros and cons of different approaches to do that? The main issue I see is that for being useful for dry run, the dispatcher should take |
|
@achirkin thank you, that makes sense to me. I will start reviewing the PR in its current form. 👍 |
There was a problem hiding this comment.
Actionable comments posted: 1
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/include/raft/matrix/detail/select_radix.cuh (1)
1280-1337: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftDo not query CUDA occupancy during dry-run.
Lines 1302-1337 dispatch to
radix_topk_one_blockorradix_topk. Both functions callcalc_chunk_sizebefore theirif (dry_run)return.calc_chunk_sizecallscudaOccupancyMaxActiveBlocksPerMultiprocessor. The multi-block path can also callcalc_grid_dim, which makes the same CUDA Runtime call.A dry-run invocation therefore enters CUDA Runtime before it returns. Guard the occupancy queries. Use a dry-run workspace sizing path that does not call CUDA Runtime APIs.
#!/bin/bash set -euo pipefail file="cpp/include/raft/matrix/detail/select_radix.cuh" # Inspect the dry-run return placement relative to occupancy queries. sed -n '806,837p;880,920p;1153,1190p;1275,1340p' "$file" # Confirm all occupancy-query call sites and dry-run guards. rg -n -C 3 'cudaOccupancyMaxActiveBlocksPerMultiprocessor|calc_chunk_size|calc_grid_dim|if \(dry_run\)' "$file"As per path instructions, “functions using raw streams or CUDA work must guard execution with
resource::get_dry_run_flag(res).”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/matrix/detail/select_radix.cuh` around lines 1280 - 1337, Ensure dry-run execution in the radix top-k dispatch does not invoke CUDA Runtime occupancy APIs. Update calc_chunk_size and calc_grid_dim, and the radix_topk_one_block/radix_topk paths that call them, to use a deterministic workspace-sizing path when dry_run is true and only query CUDA occupancy for real execution; preserve normal dispatch and sizing behavior otherwise.cpp/include/raft/linalg/power.cuh (1)
79-95: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winMove the dry-run guard below the input validation.
In both
powerandpower_scalar, the guard returns before theRAFT_EXPECTSchecks. In dry-run mode, contiguity and size mismatches are therefore not reported, and the caller receives a memory estimate for an input that will fail in the real run.Neither function allocates memory, so moving the guard after the validation block does not change allocation tracking. It only skips the kernel launch, which is the intent.
cpp/include/raft/stats/stddev.cuhandcpp/include/raft/stats/dispersion.cuhin this PR already keep theirRAFT_EXPECTSchecks ahead of the dry-run dispatch.🛡️ Proposed fix for both overloads
void power(raft::resources const& handle, InType in1, InType in2, OutType out) { - if (resource::get_dry_run_flag(handle)) { return; } using in_value_t = typename InType::value_type; using out_value_t = typename OutType::value_type; RAFT_EXPECTS(raft::is_row_or_column_major(out), "Output must be contiguous"); RAFT_EXPECTS(raft::is_row_or_column_major(in1), "Input 1 must be contiguous"); RAFT_EXPECTS(raft::is_row_or_column_major(in2), "Input 2 must be contiguous"); RAFT_EXPECTS(out.size() == in1.size() && in1.size() == in2.size(), "Size mismatch between Output and Inputs"); + if (resource::get_dry_run_flag(handle)) { return; } power<in_value_t, out_value_t, typename OutType::index_type>({ - if (resource::get_dry_run_flag(handle)) { return; } using in_value_t = typename InType::value_type; using out_value_t = typename OutType::value_type; RAFT_EXPECTS(raft::is_row_or_column_major(out), "Output must be contiguous"); RAFT_EXPECTS(raft::is_row_or_column_major(in), "Input must be contiguous"); RAFT_EXPECTS(out.size() == in.size(), "Size mismatch between Output and Input"); + if (resource::get_dry_run_flag(handle)) { return; } powerScalar<in_value_t, out_value_t, typename OutType::index_type>(One caveat:
power_scalardereferences*scalar.data_handle()after the guard. That host read is safe only if the scalar is a real host allocation. Confirm thatmake_host_scalarin dry-run mode still yields a readable pointer, becausecpp/include/raft/core/host_mdarray.hppskips value initialization in dry-run mode.Also applies to: 118-132
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/linalg/power.cuh` around lines 79 - 95, Move the dry-run guard below all RAFT_EXPECTS validation in both power and power_scalar, so contiguity and size checks always execute before skipping the kernel launch. Preserve the existing validation and dry-run behavior, and verify that power_scalar’s dereference of scalar.data_handle() remains safe with make_host_scalar in dry-run mode; adjust only if that pointer is not readable.
🟠 Major comments (20)
cpp/include/raft/sparse/op/detail/reduce.cuh-137-149 (1)
137-149: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReject scan sizes that do not fit CUB's item count.
Lines 137 and 149 narrow
diff.size()toint. IfnnzexceedsINT_MAX, the CUB scan receives a truncated or invalid item count. The later read fromdiff.back()then does not represent a scan over all nonzeros. Validate the bound before the workspace query, or use a supported CUB interface with a wider item-count type.As per coding guidelines, add input validation when invalid dimensions can cause incorrect behavior.
#!/bin/bash set -euo pipefail rg -n -C 4 'max_duplicates|static_cast<int>\(diff.size\(\)\)|nnz_t' \ cpp/include/raft/sparse/op/detail/reduce.cuh cpp/include/raft/sparse🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/op/detail/reduce.cuh` around lines 137 - 149, Validate before both CUB scan calls that diff.size() is no greater than the maximum representable int, and reject the input with the established validation mechanism when it exceeds that bound. Update the workspace query and scan in the surrounding reduction flow to use only a validated item count, preserving full coverage of all nonzeros and preventing narrowing overflow.cpp/include/raft/sparse/op/detail/reduce.cuh-136-149 (1)
136-149: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winCheck both CUB scan calls.
Lines 136-137 and 148-149 ignore the CUDA status from
cub::DeviceScan::ExclusiveSum. A failed workspace query can leavescan_ws_bytesinvalid. A failed scan can continue to the host copy and use invalid scan data. Wrap both calls withRAFT_CUDA_TRY.Proposed fix
- cub::DeviceScan::ExclusiveSum( - nullptr, scan_ws_bytes, diff.data(), diff.data(), static_cast<int>(diff.size()), stream); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + nullptr, scan_ws_bytes, diff.data(), diff.data(), static_cast<int>(diff.size()), stream)); ... - cub::DeviceScan::ExclusiveSum( - scan_ws.data(), scan_ws_bytes, diff.data(), diff.data(), static_cast<int>(diff.size()), stream); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + scan_ws.data(), scan_ws_bytes, diff.data(), diff.data(), static_cast<int>(diff.size()), stream));As per coding guidelines, “kernel launches, memory operations, and synchronization must use RAFT/CUDA error checking.” As per path instructions, use checked CUDA macros.
#!/bin/bash set -euo pipefail rg -n -C 3 'DeviceScan::ExclusiveSum|RAFT_CUDA_TRY' \ cpp/include/raft/sparse/op/detail/reduce.cuh🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/op/detail/reduce.cuh` around lines 136 - 149, Wrap both cub::DeviceScan::ExclusiveSum calls in RAFT_CUDA_TRY, including the workspace-size query and the actual scan, so CUDA failures are propagated before scan_ws allocation or subsequent data use.cpp/include/raft/linalg/detail/lstsq.cuh-213-216 (1)
213-216: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winUse exception-safe cleanup for
gesvdj_params.The dry-run branch destroys
gesvdj_params, but the normal path does not destroy it. A later allocation failure or checked cuSOLVER failure also bypasses this cleanup. Use an RAII owner or a scope guard immediately aftercusolverDnCreateGesvdjInfosucceeds.As per path instructions, review CUDA call sites for “leaks or exception-path cleanup failures.”
#!/bin/bash set -euo pipefail rg -n -C 8 \ 'lstsqSvdJacobi|cusolverDnCreateGesvdjInfo|cusolverDnDestroyGesvdjInfo' \ cpp/include/raft/linalg/detail/lstsq.cuh🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/linalg/detail/lstsq.cuh` around lines 213 - 216, Make cleanup of gesvdj_params exception-safe in the code surrounding cusolverDnCreateGesvdjInfo: establish an RAII owner or scope guard immediately after successful creation that calls cusolverDnDestroyGesvdjInfo exactly once on every exit path, including dry-run returns, allocation failures, and checked cuSOLVER errors. Remove the manual dry-run-only destruction in the lstsqSvdJacobi flow to avoid double cleanup.cpp/include/raft/sparse/convert/detail/bitset_to_csr.cuh-123-126 (1)
123-126: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winCheck both CUB scan calls.
cub::DeviceScan::ExclusiveSumreturns a CUDA status. If either call fails, the code can continue with invalid workspace or CSR row-offset data. Wrap both calls withRAFT_CUDA_TRY.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/convert/detail/bitset_to_csr.cuh` around lines 123 - 126, Wrap both cub::DeviceScan::ExclusiveSum calls in RAFT_CUDA_TRY, including the workspace-size query and the actual scan execution. Ensure any CUDA failure is propagated before using scan_ws or the generated CSR row offsets.Sources: Coding guidelines, Path instructions
cpp/include/raft/stats/detail/scores.cuh-183-183 (1)
183-183: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMove host-vector allocations after the dry-run return.
std::vector<double> h_sorted_abs_diffs(n)allocates host memory before this check. A large dry-run request can consume or exhaust host memory without contributing to RAFT allocation statistics. Keep the RMM allocations before this check, but constructmean_errorsandh_sorted_abs_diffsonly after dry-run mode returns.Proposed fix
- std::vector<double> mean_errors(2); - std::vector<double> h_sorted_abs_diffs(n); int thread_cnt = 256; int block_cnt = raft::ceildiv(n, thread_cnt); if (dry_run) { return; } + std::vector<double> mean_errors(2); + std::vector<double> h_sorted_abs_diffs(n); RAFT_CUDA_TRY(cudaMemsetAsync(tmp_sums.data(), 0, 2 * sizeof(double), stream));As per path instructions, dry-run implementations must suppress real allocations while leaving RMM/RAFT allocation attempts unguarded.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/detail/scores.cuh` at line 183, Move the dry_run return in the surrounding score computation before constructing the host vectors mean_errors and h_sorted_abs_diffs, while keeping the existing RMM/RAFT allocation attempts before that return. Ensure dry-run mode performs no host-vector allocations and exits immediately after those allocation attempts.Source: Path instructions
cpp/tests/sparse/reduce.cu-65-82 (1)
65-82: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftThe COO output
outis allocated inside the dry-run scope but declared outside it.
outis constructed at Line 65, outside the wrapper. During the dry-run pass,max_duplicatesreaches the dry-run branch incpp/include/raft/sparse/op/detail/reduce.cuhand callsout.allocate(nnz, m, n, false, stream). That allocation uses the temporary dry-run memory resource installed byexecute_with_dry_run_check. When the wrapper exits,outstill owns those probe-backed buffers, and they are later freed or reallocated against the restored resource.This is the same hazard that
cpp/tests/core/bitset.cudocuments at Lines 201-203, where the whole object lifetime is kept inside the dry-run scope. Declareoutinside the callback, or confirm that the dry-run resource outlives every object it allocated.The inline comment at Lines 66-67 states that the COO output is not tracked. That statement does not match the implementation, which allocates the output during the dry-run pass.
#!/bin/bash # Check the lifetime semantics of the dry-run resource wrapper and the COO allocate path. set -euo pipefail fd -t f 'dry_run_resources.hpp' cpp/include | while IFS= read -r f; do echo "=== $f ===" cat -n "$f" done echo "=== execute_with_dry_run_check definition ===" rg -n -C 25 'execute_with_dry_run_check' cpp/include cpp/tests --glob '*.{hpp,cuh,h}' echo "=== COO::allocate ===" rg -n -C 15 'void allocate' cpp/include/raft/sparse/coo.hpp🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/sparse/reduce.cu` around lines 65 - 82, Move construction of the COO variable out of the surrounding scope and into the dry-run callback containing max_duplicates, keeping its entire lifetime within execute_with_dry_run_check. Remove or update the inaccurate comments claiming the COO output is not tracked, and ensure the callback’s output remains available to the subsequent test logic as required.cpp/tests/core/bitset.cu-388-399 (1)
388-399: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winGate the
anyandnoneassertions on the dry-run flag.Lines 390-391 assert values without checking
resource::get_dry_run_flag(h), but every other value-dependent assertion in this callback is gated. In the dry-run pass,my_bitset.reset(h, false)performs no CUDA work, soany(h)andnone(h)read probe-backed memory whose contents are undefined. These two assertions can fail during the dry-run pass.The assertions at Lines 395-399 already cover the same two calls correctly.
💚 Proposed fix
// Test count() operations my_bitset.reset(h, false); - ASSERT_EQ(my_bitset.any(h), false); - ASSERT_EQ(my_bitset.none(h), true); + if (!resource::get_dry_run_flag(h)) { + ASSERT_EQ(my_bitset.any(h), false); + ASSERT_EQ(my_bitset.none(h), true); + } raft::linalg::map_offset(h, query_device.view(), raft::cast_op<index_t>{});🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/core/bitset.cu` around lines 388 - 399, Gate the initial my_bitset.any(h) and my_bitset.none(h) assertions after reset on resource::get_dry_run_flag(h), matching the existing guarded value-dependent assertions below. Keep the reset and subsequent setup unchanged, and preserve the assertions during non-dry-run execution.cpp/include/raft/solver/detail/lap_functions.cuh-269-271 (1)
269-271: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winCalculate dry-run allocation counts in
size_t.Both sites evaluate
SP * Nbeforedevice_uvectorreceives the value. If the product exceeds the source integer type, overflow produces an invalid allocation size and an incorrect peak-memory estimate.
cpp/include/raft/solver/detail/lap_functions.cuh#L269-L271: derive a checkedsize_t matrix_sizebefore allocatingpredicates_v,addresses_v, andcsr_neighbors_v.cpp/include/raft/solver/detail/lap_functions.cuh#L432-L436: derive the same checkedsize_t matrix_sizebefore allocatingpredicates_v,addresses_v, andelements_v.#!/bin/bash set -euo pipefail rg -n -C 3 'device_uvector<.*>\([^;]*SP \* N' cpp/include/raft/solver/detail/lap_functions.cuh🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/solver/detail/lap_functions.cuh` around lines 269 - 271, In cpp/include/raft/solver/detail/lap_functions.cuh at lines 269-271 and 432-436, compute a checked size_t matrix_size from SP and N before constructing the device_uvector instances; use that variable for predicates_v, addresses_v, csr_neighbors_v, and elements_v at the respective sites, preserving the existing allocation behavior while preventing source-type overflow.cpp/include/raft/stats/detail/contingencyMatrix.cuh-209-219 (1)
209-219: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPrevent overflow in the dry-run workspace estimate.
Line 217 evaluates
4 * nSamplesasint. IfnSamples > INT_MAX / 4, signed overflow occurs before conversion tosize_t. The dry-run estimate can then under-report or corrupt the required workspace size.Validate
nSamplesand promote it before every size calculation.Proposed fix
+ RAFT_EXPECTS(nSamples >= 0, "nSamples must be non-negative"); + auto sample_count = static_cast<size_t>(nSamples); if (dry_run) { - auto tmpStagingMemorySize = raft::alignTo<size_t>(nSamples * sizeof(T), 256); + auto tmpStagingMemorySize = raft::alignTo<size_t>(sample_count * sizeof(T), 256); tmpStagingMemorySize *= 2; - size_t cubWorkspaceUpperBound = 4 * nSamples * sizeof(T); + size_t cubWorkspaceUpperBound = size_t{4} * sample_count * sizeof(T);#!/bin/bash set -euo pipefail python3 - <<'PY' import ctypes max_int = (1 << (ctypes.sizeof(ctypes.c_int) * 8 - 1)) - 1 print(f"INT_MAX={max_int}") print(f"First overflowing nSamples for 4 * nSamples: {max_int // 4 + 1}") PY🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/detail/contingencyMatrix.cuh` around lines 209 - 219, Prevent signed overflow in the dry-run calculations within the contingency-matrix workspace sizing block by validating nSamples against the supported range and promoting it to size_t before every multiplication, including the tmpStagingMemorySize and cubWorkspaceUpperBound expressions. Preserve the existing alignment and workspace estimate behavior for valid sample counts, and reject or otherwise handle values that cannot be represented safely.cpp/include/raft/stats/detail/meanvar.cuh-212-223 (1)
212-223: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard the occupancy CUDA call during dry runs.
Line 212 guards the memset and kernel launches, but lines 203-207 still call
cudaOccupancyMaxActiveBlocksPerMultiprocessorwhendry_runis true. This makes a memory-only probe issue a CUDA runtime query. Move the occupancy calculation andgs.yadjustment into the!dry_runbranch.As per path instructions, dry-run implementations must suppress CUDA work while preserving allocation tracking.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/detail/meanvar.cuh` around lines 212 - 223, Move the cudaOccupancyMaxActiveBlocksPerMultiprocessor call and related gs.y adjustment into the existing !dry_run branch around meanvar_kernel_rowmajor and meanvar_kernel_fill. Ensure dry_run performs no CUDA runtime queries or kernel work while retaining the existing buffer allocation and tracking behavior.Source: Path instructions
cpp/include/raft/sparse/solver/detail/lanczos.cuh-158-169 (1)
158-169: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winCheck all five CUDA kernel launches with
RAFT_CUDA_TRY(cudaPeekAtLastError()).Add the check immediately after
kernel_triangular_populate,kernel_triangular_beta_k,kernel_clamp_down,kernel_clamp_down_vector, andkernel_normalize.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/solver/detail/lanczos.cuh` around lines 158 - 169, Add RAFT_CUDA_TRY(cudaPeekAtLastError()) immediately after each of the five CUDA launches in lanczos.cuh: kernel_triangular_populate and kernel_triangular_beta_k in lines 158-169, plus kernel_clamp_down, kernel_clamp_down_vector, and kernel_normalize in lines 386-415. Keep the checks directly adjacent to their respective launches.Sources: Coding guidelines, Path instructions
cpp/include/raft/sparse/convert/detail/bitmap_to_csr.cuh-309-315 (1)
309-315: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAllocate
scan_wsfromdevice_memory. The current constructor uses the default device resource, so custom workspace-resource accounting and isolation do not include the CUB scan workspace.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/convert/detail/bitmap_to_csr.cuh` around lines 309 - 315, Construct scan_ws with device_memory, matching the resource used for sub_nnz, so the CUB scan workspace is accounted for by the configured workspace resource.Sources: Coding guidelines, Path instructions
cpp/include/raft/sparse/solver/detail/lanczos.cuh-198-227 (1)
198-227: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse the handle-bound Thrust policy.
thrust::deviceruns on the default CUDA stream, while neighboring operations use the handle stream. This can break ordering for non-default handle streams. Useauto thrust_policy = resource::get_thrust_policy(handle)forthrust::sequenceand boththrust::sortcalls.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/sparse/solver/detail/lanczos.cuh` around lines 198 - 227, Use a handle-bound Thrust execution policy for all operations in this selection block: define thrust_policy via resource::get_thrust_policy(handle), then pass it to thrust::sequence and both thrust::sort calls instead of thrust::device. Preserve the existing index ordering and eigenvalue selection logic.Sources: Coding guidelines, Path instructions
cpp/include/raft/matrix/detail/select_k-inl.cuh-131-132 (1)
131-132: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winCheck each CUB sort status before the dry-run return.
SortPairsandSortPairsDescendingreturn CUDA status. The dry-run path returns after an unchecked workspace query. If that query fails, dry-run can report memory usage instead of returning the CUDA error. Wrap both workspace queries and both execution calls withRAFT_CUDA_TRY.#!/bin/bash set -euo pipefail sed -n '90,170p' cpp/include/raft/matrix/detail/select_k-inl.cuh rg -n -C 2 'RAFT_CUDA_TRY' cpp/include/raft/core/detail/macros.hpp cpp/include/raft🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/matrix/detail/select_k-inl.cuh` around lines 131 - 132, Update the CUB sort calls in the select-k implementation, including both SortPairs and SortPairsDescending workspace-size queries and execution calls, to use RAFT_CUDA_TRY. Ensure each status is checked before the resource::get_dry_run_flag(handle) early return so failed dry-run workspace queries propagate the CUDA error.cpp/include/raft/stats/detail/silhouette_score.cuh-229-230 (1)
229-230: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftUse one ordered stream for dependent operations.
pairwise_distance(handle, ...)andmatrix_vector_op(handle, ...)use the stream inhandle. The dependent operations usestream. If these streams differ, the reductions can readdistanceMatrixoraverageDistanceBetweenSampleAndClusterbefore the producer work completes.Use
resource::get_cuda_stream(handle)for all operations in this function, or add explicit event ordering between the two streams. As per path instructions, maintain asynchronous stream ordering and use RAFT-managed streams for internal work.Also applies to: 268-276
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/stats/detail/silhouette_score.cuh` around lines 229 - 230, Update the operations in the silhouette-score function, including pairwise_distance and the dependent matrix_vector_op/reduction calls, to use the same RAFT-managed stream obtained via resource::get_cuda_stream(handle). Preserve asynchronous ordering and ensure every producer and consumer of distanceMatrix and averageDistanceBetweenSampleAndCluster is submitted to that stream.Source: Path instructions
cpp/include/raft/matrix/threshold.cuh-38-38 (1)
38-38: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winKeep the dimension check outside the dry-run guard.
RAFT_EXPECTS(in.size() == out.size(), ...)is skipped when dry-run mode is enabled. Dry-run then accepts an input/output shape that normal execution rejects. This can produce a memory estimate for a call that cannot run.Move the dry-run guard after the size check, while still skipping
detail::setSmallValuesZero.Proposed fix
- if (resource::get_dry_run_flag(handle)) { return; } RAFT_EXPECTS(in.size() == out.size(), "Input and output matrices must have same size"); + if (resource::get_dry_run_flag(handle)) { return; }As per coding guidelines, validate invalid dimensions where they can cause incorrect behavior. As per path instructions, dry-run guards should suppress CUDA work without bypassing required precondition checks.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/matrix/threshold.cuh` at line 38, Move the resource::get_dry_run_flag(handle) guard in the threshold operation after the RAFT_EXPECTS input/output size validation, so dimension mismatches are rejected in both modes while dry-run still skips detail::setSmallValuesZero and other CUDA work.Sources: Coding guidelines, Path instructions
cpp/include/raft/mr/dry_run_resource.hpp-181-207 (1)
181-207: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe probe ignores the alignment of later allocations.
std::call_oncecreates the probe with the alignment of the first request only. A later request with a larger alignment receives the same pointer, which can be under-aligned for that request. Aligned vector loads and stores on that pointer are then invalid, and the returned pointer breaks the alignment contract of the memory-resource API.The probe is also 256 bytes for every request size. Any code path that is not dry-run guarded and writes to the returned pointer overruns the probe.
Allocate the probe with the maximum alignment RAFT can request, or reallocate the probe when a request needs a stricter alignment.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/mr/dry_run_resource.hpp` around lines 181 - 207, Update allocate_sync and allocate so the shared state_->probe always satisfies the strictest alignment requested by RAFT, rather than preserving the first allocation’s alignment; reallocate or otherwise replace the probe when a later request requires greater alignment. Also size the probe to accommodate the requested bytes, or otherwise ensure returned storage cannot be overrun by an unguarded request, while preserving allocation/deallocation accounting.Source: Coding guidelines
cpp/include/raft/core/host_container_policy.hpp-118-127 (1)
118-127: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReset
bytesize_when you free the buffer.Line 122 frees
data_and Line 123 setsdata_tonullptr, butbytesize_keeps the old value. If the allocation on Line 125 throws, the container is left withdata_ == nullptrand a non-zerobytesize_. Any lateroperator[]on that object dereferences a null pointer.The destructor is safe because it tests both fields, so this is an invariant defect and not a double free.
🛠️ Proposed fix
if (data_ != nullptr) { mr_.deallocate_sync(data_, bytesize_); data_ = nullptr; + bytesize_ = 0; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/core/host_container_policy.hpp` around lines 118 - 127, Update host_container::reallocate so bytesize_ is reset to zero immediately after deallocating data_ and clearing data_. Preserve the existing allocation and swap flow, ensuring an allocation failure cannot leave a null data_ pointer paired with a stale nonzero size.Source: Coding guidelines
cpp/include/raft/random/detail/rng_impl.cuh-16-23 (1)
16-23: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winAdd the
dry_run_flag.hppinclude.Line 396 calls
resource::get_dry_run_flag(res)and Line 397 callsresource::get_cuda_stream(res), but this header does not include<raft/core/resource/dry_run_flag.hpp>or<raft/core/resource/cuda_stream.hpp>. The build then depends on transitive includes and can break when an upstream header changes.🛠️ Proposed fix
`#include` <raft/core/detail/macros.hpp> `#include` <raft/core/device_mdarray.hpp> `#include` <raft/core/math.hpp> `#include` <raft/core/operators.cuh> +#include <raft/core/resource/cuda_stream.hpp> +#include <raft/core/resource/dry_run_flag.hpp> `#include` <raft/linalg/map.cuh>🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/random/detail/rng_impl.cuh` around lines 16 - 23, Add direct includes for the declarations used by the RNG implementation: include the dry-run flag and CUDA stream resource headers before the code calling resource::get_dry_run_flag and resource::get_cuda_stream. Keep the existing scatter and CUB includes unchanged.cpp/include/raft/spectral/detail/matrix_wrappers.hpp-393-396 (1)
393-396: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
auto handle = get_handle()copies theraft::resourcesobject in bothmvoverrides.sparse_matrix_t::get_handle()returnsresources const&(Line 323), soautodeducesraft::resourcesby value and each call constructs and destroys a full handle copy. Both overrides run once per solver iteration, so the copy is repeated on a hot path.
cpp/include/raft/spectral/detail/matrix_wrappers.hpp#L393-L396: changeauto handletoauto const& handleinlaplacian_matrix_t::mv.cpp/include/raft/spectral/detail/matrix_wrappers.hpp#L459-L462: changeauto handletoauto const& handleinmodularity_matrix_t::mv.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/raft/spectral/detail/matrix_wrappers.hpp` around lines 393 - 396, Update the handle declarations in both laplacian_matrix_t::mv (cpp/include/raft/spectral/detail/matrix_wrappers.hpp:393-396) and modularity_matrix_t::mv (cpp/include/raft/spectral/detail/matrix_wrappers.hpp:459-462) to bind the resources returned by sparse_matrix_t::get_handle() as a const reference instead of copying them; leave the cublas, stream, and dry-run retrieval unchanged.
🟡 Minor comments (1)
docs/source/developer_guide.md-325-329 (1)
325-329: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winLine 329 renders as part of the last nested bullet.
All attempted allocations in the above resources are tracked...immediately follows a nested list item with no blank line. CommonMark and MyST treat that as a lazy continuation of the list item, so the sentence renders inside the bullet instead of as a standalone paragraph. Insert a blank line.📝 Proposed fix
- workspace memory resources managed by `raft::resources`. + All attempted allocations in the above resources are tracked and reported, thus enabling planning of the memory usage with a relatively small overhead of simulated execution.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docs/source/developer_guide.md` around lines 325 - 329, Insert a blank line after the final nested bullet under the listed memory resources and before “All attempted allocations...” so that sentence renders as a standalone paragraph rather than as continuation text within the bullet.
🧹 Nitpick comments (3)
cpp/tests/sparse/preprocess.cu (1)
166-180: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign the TF-IDF dry-run expectation with the BM25 expectation.
encode_tfidfandencode_bm25follow the same allocation pattern in this test, but the BM25 branches usealloc_behavior::DATA_DRIVENwithsizeof(float) * coo_a.nnz, while the TF-IDF branches usealloc_behavior::ARGUMENT_DRIVENwith1. An expected allocation of1byte makes the check almost unconditional, so a regression in TF-IDF allocation estimation would not fail the test.If
encode_tfidfreally allocates independently ofnnz, add a short comment that states why. Otherwise use the same data-driven bound as the BM25 branches.🧪 Proposed alignment with the BM25 branches
if (coo_on) { raft::execute_with_dry_run_check( handle, [&](raft::resources const& h) { raft::sparse::matrix::encode_tfidf<float, int>(h, coo_a_matrix, result.view()); }, - raft::alloc_behavior::ARGUMENT_DRIVEN, - 1); + raft::alloc_behavior::DATA_DRIVEN, + sizeof(float) * coo_a.nnz); } else { raft::execute_with_dry_run_check( handle, [&](raft::resources const& h) { raft::sparse::matrix::encode_tfidf<float, int>(h, csr_matrix, result.view()); }, - raft::alloc_behavior::ARGUMENT_DRIVEN, - 1); + raft::alloc_behavior::DATA_DRIVEN, + sizeof(float) * coo_a.nnz); }#!/bin/bash # Inspect the encode_tfidf / encode_bm25 implementations to compare their allocation patterns. set -euo pipefail fd -t f 'preprocessing.cuh' cpp/include | while IFS= read -r f; do echo "=== $f ===" rg -n -C 12 'encode_tfidf|encode_bm25|get_dry_run_flag|device_uvector|make_device_' "$f" done🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/sparse/preprocess.cu` around lines 166 - 180, Align both encode_tfidf dry-run checks in the preprocessing test with the BM25 branches by using DATA_DRIVEN allocation behavior and a bound based on sizeof(float) multiplied by coo_a.nnz. If encode_tfidf intentionally allocates independently of nnz, retain the current expectation and add a concise comment explaining that behavior.cpp/tests/util/dry_run_resources.cpp (1)
364-387: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for
dry_run_resourcesnested insidedry_run_resources.
dry_run_resourcescomputesactive_(!resource::get_dry_run_flag(existing))and returns an emptymemory_statsfromget_bytes_peak()when inactive. This test covers dry-run nested in stats, but no test covers dry-run nested in dry-run. That inactive path controls whether the inner object skipsinit()and whether the destructor restores globals, so a regression there would silently produce zero statistics.💚 Proposed additional test
TEST(DryRunResources, NestedDryRunIsInactive) { raft::resources res; dry_run_resources outer(res); { dry_run_resources inner(outer); EXPECT_TRUE(resource::get_dry_run_flag(inner)); // Inner is inactive: it reports no statistics of its own. EXPECT_EQ(inner.get_bytes_peak().total(), 0UL); } // The outer handle must still be in dry-run mode after the inner one is destroyed. EXPECT_TRUE(resource::get_dry_run_flag(outer)); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/util/dry_run_resources.cpp` around lines 364 - 387, Add a test alongside NestedDryRunInStats that constructs an outer dry_run_resources and an inner dry_run_resources using the outer resource. Verify the inner remains marked as dry-run, reports zero from get_bytes_peak().total(), and that the outer still reports the dry-run flag after the inner is destroyed.cpp/tests/test_utils.cuh (1)
368-389: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe diagnostic
printffires for expectedDATA_DRIVENmismatches.In
DATA_DRIVENmodedry >= actualis the expected outcome, not a failure. The condition at Lines 368-371 triggers on any inequality, so every conformingDATA_DRIVENtest prints a full six-category dump. Gate the diagnostic on the assertion actually failing, or restrict it to theARGUMENT_DRIVENandNO_ALLOCATIONSmodes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/test_utils.cuh` around lines 368 - 389, Update the mismatch diagnostic around the dry/actual workspace comparison so it does not print for expected DATA_DRIVEN cases where dry exceeds actual. Gate the printf using the same assertion-failure condition, or limit it to ARGUMENT_DRIVEN and NO_ALLOCATIONS modes, while preserving diagnostics for genuine mismatches.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/include/raft/mr/dry_run_resource.hpp`:
- Around line 90-105: Update probe_container’s three-argument construction and
destructor to track whether allocation used allocate_sync or the asynchronous
allocate API. Use that stored allocation mode to call the corresponding
deallocate_sync or deallocate function, rather than selecting solely via
cuda::mr::resource<MR>, while preserving the existing stream choice for
asynchronous deallocation.
---
Outside diff comments:
In `@cpp/include/raft/linalg/power.cuh`:
- Around line 79-95: Move the dry-run guard below all RAFT_EXPECTS validation in
both power and power_scalar, so contiguity and size checks always execute before
skipping the kernel launch. Preserve the existing validation and dry-run
behavior, and verify that power_scalar’s dereference of scalar.data_handle()
remains safe with make_host_scalar in dry-run mode; adjust only if that pointer
is not readable.
In `@cpp/include/raft/matrix/detail/select_radix.cuh`:
- Around line 1280-1337: Ensure dry-run execution in the radix top-k dispatch
does not invoke CUDA Runtime occupancy APIs. Update calc_chunk_size and
calc_grid_dim, and the radix_topk_one_block/radix_topk paths that call them, to
use a deterministic workspace-sizing path when dry_run is true and only query
CUDA occupancy for real execution; preserve normal dispatch and sizing behavior
otherwise.
---
Major comments:
In `@cpp/include/raft/core/host_container_policy.hpp`:
- Around line 118-127: Update host_container::reallocate so bytesize_ is reset
to zero immediately after deallocating data_ and clearing data_. Preserve the
existing allocation and swap flow, ensuring an allocation failure cannot leave a
null data_ pointer paired with a stale nonzero size.
In `@cpp/include/raft/linalg/detail/lstsq.cuh`:
- Around line 213-216: Make cleanup of gesvdj_params exception-safe in the code
surrounding cusolverDnCreateGesvdjInfo: establish an RAII owner or scope guard
immediately after successful creation that calls cusolverDnDestroyGesvdjInfo
exactly once on every exit path, including dry-run returns, allocation failures,
and checked cuSOLVER errors. Remove the manual dry-run-only destruction in the
lstsqSvdJacobi flow to avoid double cleanup.
In `@cpp/include/raft/matrix/detail/select_k-inl.cuh`:
- Around line 131-132: Update the CUB sort calls in the select-k implementation,
including both SortPairs and SortPairsDescending workspace-size queries and
execution calls, to use RAFT_CUDA_TRY. Ensure each status is checked before the
resource::get_dry_run_flag(handle) early return so failed dry-run workspace
queries propagate the CUDA error.
In `@cpp/include/raft/matrix/threshold.cuh`:
- Line 38: Move the resource::get_dry_run_flag(handle) guard in the threshold
operation after the RAFT_EXPECTS input/output size validation, so dimension
mismatches are rejected in both modes while dry-run still skips
detail::setSmallValuesZero and other CUDA work.
In `@cpp/include/raft/mr/dry_run_resource.hpp`:
- Around line 181-207: Update allocate_sync and allocate so the shared
state_->probe always satisfies the strictest alignment requested by RAFT, rather
than preserving the first allocation’s alignment; reallocate or otherwise
replace the probe when a later request requires greater alignment. Also size the
probe to accommodate the requested bytes, or otherwise ensure returned storage
cannot be overrun by an unguarded request, while preserving
allocation/deallocation accounting.
In `@cpp/include/raft/random/detail/rng_impl.cuh`:
- Around line 16-23: Add direct includes for the declarations used by the RNG
implementation: include the dry-run flag and CUDA stream resource headers before
the code calling resource::get_dry_run_flag and resource::get_cuda_stream. Keep
the existing scatter and CUB includes unchanged.
In `@cpp/include/raft/solver/detail/lap_functions.cuh`:
- Around line 269-271: In cpp/include/raft/solver/detail/lap_functions.cuh at
lines 269-271 and 432-436, compute a checked size_t matrix_size from SP and N
before constructing the device_uvector instances; use that variable for
predicates_v, addresses_v, csr_neighbors_v, and elements_v at the respective
sites, preserving the existing allocation behavior while preventing source-type
overflow.
In `@cpp/include/raft/sparse/convert/detail/bitmap_to_csr.cuh`:
- Around line 309-315: Construct scan_ws with device_memory, matching the
resource used for sub_nnz, so the CUB scan workspace is accounted for by the
configured workspace resource.
In `@cpp/include/raft/sparse/convert/detail/bitset_to_csr.cuh`:
- Around line 123-126: Wrap both cub::DeviceScan::ExclusiveSum calls in
RAFT_CUDA_TRY, including the workspace-size query and the actual scan execution.
Ensure any CUDA failure is propagated before using scan_ws or the generated CSR
row offsets.
In `@cpp/include/raft/sparse/op/detail/reduce.cuh`:
- Around line 137-149: Validate before both CUB scan calls that diff.size() is
no greater than the maximum representable int, and reject the input with the
established validation mechanism when it exceeds that bound. Update the
workspace query and scan in the surrounding reduction flow to use only a
validated item count, preserving full coverage of all nonzeros and preventing
narrowing overflow.
- Around line 136-149: Wrap both cub::DeviceScan::ExclusiveSum calls in
RAFT_CUDA_TRY, including the workspace-size query and the actual scan, so CUDA
failures are propagated before scan_ws allocation or subsequent data use.
In `@cpp/include/raft/sparse/solver/detail/lanczos.cuh`:
- Around line 158-169: Add RAFT_CUDA_TRY(cudaPeekAtLastError()) immediately
after each of the five CUDA launches in lanczos.cuh: kernel_triangular_populate
and kernel_triangular_beta_k in lines 158-169, plus kernel_clamp_down,
kernel_clamp_down_vector, and kernel_normalize in lines 386-415. Keep the checks
directly adjacent to their respective launches.
- Around line 198-227: Use a handle-bound Thrust execution policy for all
operations in this selection block: define thrust_policy via
resource::get_thrust_policy(handle), then pass it to thrust::sequence and both
thrust::sort calls instead of thrust::device. Preserve the existing index
ordering and eigenvalue selection logic.
In `@cpp/include/raft/spectral/detail/matrix_wrappers.hpp`:
- Around line 393-396: Update the handle declarations in both
laplacian_matrix_t::mv
(cpp/include/raft/spectral/detail/matrix_wrappers.hpp:393-396) and
modularity_matrix_t::mv
(cpp/include/raft/spectral/detail/matrix_wrappers.hpp:459-462) to bind the
resources returned by sparse_matrix_t::get_handle() as a const reference instead
of copying them; leave the cublas, stream, and dry-run retrieval unchanged.
In `@cpp/include/raft/stats/detail/contingencyMatrix.cuh`:
- Around line 209-219: Prevent signed overflow in the dry-run calculations
within the contingency-matrix workspace sizing block by validating nSamples
against the supported range and promoting it to size_t before every
multiplication, including the tmpStagingMemorySize and cubWorkspaceUpperBound
expressions. Preserve the existing alignment and workspace estimate behavior for
valid sample counts, and reject or otherwise handle values that cannot be
represented safely.
In `@cpp/include/raft/stats/detail/meanvar.cuh`:
- Around line 212-223: Move the cudaOccupancyMaxActiveBlocksPerMultiprocessor
call and related gs.y adjustment into the existing !dry_run branch around
meanvar_kernel_rowmajor and meanvar_kernel_fill. Ensure dry_run performs no CUDA
runtime queries or kernel work while retaining the existing buffer allocation
and tracking behavior.
In `@cpp/include/raft/stats/detail/scores.cuh`:
- Line 183: Move the dry_run return in the surrounding score computation before
constructing the host vectors mean_errors and h_sorted_abs_diffs, while keeping
the existing RMM/RAFT allocation attempts before that return. Ensure dry-run
mode performs no host-vector allocations and exits immediately after those
allocation attempts.
In `@cpp/include/raft/stats/detail/silhouette_score.cuh`:
- Around line 229-230: Update the operations in the silhouette-score function,
including pairwise_distance and the dependent matrix_vector_op/reduction calls,
to use the same RAFT-managed stream obtained via
resource::get_cuda_stream(handle). Preserve asynchronous ordering and ensure
every producer and consumer of distanceMatrix and
averageDistanceBetweenSampleAndCluster is submitted to that stream.
In `@cpp/tests/core/bitset.cu`:
- Around line 388-399: Gate the initial my_bitset.any(h) and my_bitset.none(h)
assertions after reset on resource::get_dry_run_flag(h), matching the existing
guarded value-dependent assertions below. Keep the reset and subsequent setup
unchanged, and preserve the assertions during non-dry-run execution.
In `@cpp/tests/sparse/reduce.cu`:
- Around line 65-82: Move construction of the COO variable out of the
surrounding scope and into the dry-run callback containing max_duplicates,
keeping its entire lifetime within execute_with_dry_run_check. Remove or update
the inaccurate comments claiming the COO output is not tracked, and ensure the
callback’s output remains available to the subsequent test logic as required.
---
Minor comments:
In `@docs/source/developer_guide.md`:
- Around line 325-329: Insert a blank line after the final nested bullet under
the listed memory resources and before “All attempted allocations...” so that
sentence renders as a standalone paragraph rather than as continuation text
within the bullet.
---
Nitpick comments:
In `@cpp/tests/sparse/preprocess.cu`:
- Around line 166-180: Align both encode_tfidf dry-run checks in the
preprocessing test with the BM25 branches by using DATA_DRIVEN allocation
behavior and a bound based on sizeof(float) multiplied by coo_a.nnz. If
encode_tfidf intentionally allocates independently of nnz, retain the current
expectation and add a concise comment explaining that behavior.
In `@cpp/tests/test_utils.cuh`:
- Around line 368-389: Update the mismatch diagnostic around the dry/actual
workspace comparison so it does not print for expected DATA_DRIVEN cases where
dry exceeds actual. Gate the printf using the same assertion-failure condition,
or limit it to ARGUMENT_DRIVEN and NO_ALLOCATIONS modes, while preserving
diagnostics for genuine mismatches.
In `@cpp/tests/util/dry_run_resources.cpp`:
- Around line 364-387: Add a test alongside NestedDryRunInStats that constructs
an outer dry_run_resources and an inner dry_run_resources using the outer
resource. Verify the inner remains marked as dry-run, reports zero from
get_bytes_peak().total(), and that the outer still reports the dry-run flag
after the inner is destroyed.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
5 participants
The dry run protocol defines a mechanism to simulate the execution of algorithms to get a precise estimate of the memory requirements for a real execution with the same parameters.
This PR:
raft::util::dry_run_execute,raft::dry_run_resources, andresource::get_dry_run_flagto let callers estimate peak memory usage of any RAFT algorithm without executing GPU work.Note for reviewers
The PR contains a lot of small tedious changes to cover all of raft library and the tests components.
Please start reading at
docs/sourceupdates to learn more about the topic and the principles guiding these changes.