Repository navigation
Migrate RMM usage to CCCL MR design - #7951
Merged
Merged
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
51 tasks done
bdice
force-pushed
the
rmm-cccl-migration
branch
4 times, most recently
from
April 17, 2026 07:09
48202c6 to
4c9b17f
Compare
bdice
added
breaking
bdice
force-pushed
the
rmm-cccl-migration
branch
from
April 17, 2026 07:29
4c9b17f to
67ee0f1
Compare
bdice
changed the title
bdice
force-pushed
the
rmm-cccl-migration
branch
5 times, most recently
from
April 18, 2026 15:52
8fba1a6 to
b55ccf5
Compare
bdice
marked this pull request as ready for review
April 19, 2026 19:16
📝 WalkthroughWalkthroughThe pull request updates device memory allocation and deallocation calls across C++ source files and examples to include explicit alignment parameters using Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related issues
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 5
🧹 Nitpick comments (2)
ci/test_python_common.sh (1)
6-6: Source helper script relative to the script location.Line 6 uses a repo-root relative path, so invocation from another working directory can fail. Resolve from
BASH_SOURCE(and add a ShellCheck source hint).Proposed patch
set -euo pipefail -source ./ci/use_conda_packages_from_prs.sh +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=ci/use_conda_packages_from_prs.sh +source "${SCRIPT_DIR}/use_conda_packages_from_prs.sh"Based on learnings: For code reviews in rapidsai/cuml, whenever reviewing shell scripts, use ShellCheck to validate syntax and catch potential issues before posting comments.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ci/test_python_common.sh` at line 6, The script sources ./ci/use_conda_packages_from_prs.sh using a repo-root relative path which breaks when the script is executed from another working directory; update the sourcing in ci/test_python_common.sh to resolve the helper script relative to the current script location using ${BASH_SOURCE[0]} (or a derived DIR variable) so the path is stable regardless of CWD, and add a ShellCheck source directive (e.g. # shellcheck source=ci/use_conda_packages_from_prs.sh) above the source line to satisfy static checks.ci/use_conda_packages_from_prs.sh (1)
28-31: Validate channel values before mutating conda config.Add a non-empty check for each channel so failures are explicit and early.
Proposed patch
for _channel in "${RAPIDS_PREPENDED_CONDA_CHANNELS[@]}" do + if [[ -z "${_channel}" ]]; then + echo "ERROR: empty channel value in RAPIDS_PREPENDED_CONDA_CHANNELS" >&2 + return 1 + fi conda config --system --add channels "${_channel}" RAPIDS_EXTRA_CONDA_CHANNEL_ARGS+=(--prepend-channel "${_channel}") doneAs per coding guidelines:
ci/**/*: Verify proper error handling and meaningful error messages.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ci/use_conda_packages_from_prs.sh` around lines 28 - 31, The loop that mutates conda config iterates RAPIDS_PREPENDED_CONDA_CHANNELS without validating entries; update the loop in use_conda_packages_from_prs.sh to skip or fail fast on empty/whitespace channel values by checking each _channel for non-empty content (e.g., trim and test -n) before calling conda config and appending to RAPIDS_EXTRA_CONDA_CHANNEL_ARGS; on an invalid value emit a clear error message naming the variable/loop (RAPIDS_PREPENDED_CONDA_CHANNELS / _channel) and exit with a non-zero status so failures are explicit.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ci/test_cpp.sh`:
- Line 6: The script currently does "source ./ci/use_conda_packages_from_prs.sh"
which breaks out-of-repo invocation; replace that relative source with a source
that uses a computed SCRIPT_DIR (e.g., derived from "${BASH_SOURCE[0]}" or
similar) so the script always sources ci/use_conda_packages_from_prs.sh relative
to the script location, and then perform the existing cd to the repo root;
update the occurrence of "source ./ci/use_conda_packages_from_prs.sh" to "source
$SCRIPT_DIR/use_conda_packages_from_prs.sh" (and ensure SCRIPT_DIR is set before
this source and the later cd still runs).
In `@ci/use_wheels_from_prs.sh`:
- Around line 32-41: Add a guard before the heredoc that validates the
PIP_CONSTRAINT variable and fails with a clear message if it is empty or unset:
test that PIP_CONSTRAINT is non-empty (e.g. [ -n "${PIP_CONSTRAINT:-}" ] || {
echo "PIP_CONSTRAINT must be set" >&2; exit 1; }), then proceed to the existing
cat >> "${PIP_CONSTRAINT}" <<EOF block; reference the PIP_CONSTRAINT variable
and the heredoc block when locating where to insert the check.
In `@cpp/examples/symreg/symreg_example.cpp`:
- Around line 239-240: The allocation for d_finalprogs is using element count
and rmm::CUDA_ALLOCATION_ALIGNMENT, but allocate()/deallocate() expect bytes and
other code (e.g. symFit allocation in genetic.cu) uses alignof(program); change
the curr_mr.allocate call for d_finalprogs to request params.population_size *
sizeof(cg::program) bytes and use the same alignment as the symFit buffer
(alignof(cg::program) or the exact alignof(program) used in genetic.cu) so
allocation and subsequent deallocation/swapping use identical size-in-bytes and
alignment.
In `@cpp/src/genetic/genetic.cu`:
- Around line 231-232: The deallocation call currently uses h_nextprogs[i].len
but should use the old program length stored in tmp after raft::copy overwrote
it; update the rmm::mr::get_current_device_resource_ref().deallocate(...)
invocation to pass tmp.len * sizeof(node) (and alignof(node) as before) instead
of h_nextprogs[i].len so the freed size matches the original allocation
(references: tmp, h_nextprogs, d_oldprogs, raft::copy, and the deallocate call).
In `@python/cuml/cuml/manifold/umap/umap.pyx`:
- Line 1290: The call to convert an RMM device resource for
make_any_device_resource is using get_current_device_resource().get_mr(), but
the newer RMM API requires extracting the C pointer via
get_current_device_resource().c_ref.value(); update the invocation that
currently passes get_current_device_resource().get_mr() (around
make_any_device_resource usage) to instead pass
get_current_device_resource().c_ref.value() so make_any_device_resource receives
the proper device resource handle.
---
Nitpick comments:
In `@ci/test_python_common.sh`:
- Line 6: The script sources ./ci/use_conda_packages_from_prs.sh using a
repo-root relative path which breaks when the script is executed from another
working directory; update the sourcing in ci/test_python_common.sh to resolve
the helper script relative to the current script location using
${BASH_SOURCE[0]} (or a derived DIR variable) so the path is stable regardless
of CWD, and add a ShellCheck source directive (e.g. # shellcheck
source=ci/use_conda_packages_from_prs.sh) above the source line to satisfy
static checks.
In `@ci/use_conda_packages_from_prs.sh`:
- Around line 28-31: The loop that mutates conda config iterates
RAPIDS_PREPENDED_CONDA_CHANNELS without validating entries; update the loop in
use_conda_packages_from_prs.sh to skip or fail fast on empty/whitespace channel
values by checking each _channel for non-empty content (e.g., trim and test -n)
before calling conda config and appending to RAPIDS_EXTRA_CONDA_CHANNEL_ARGS; on
an invalid value emit a clear error message naming the variable/loop
(RAPIDS_PREPENDED_CONDA_CHANNELS / _channel) and exit with a non-zero status so
failures are explicit.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: c86a26e5-791f-4cd7-bc39-414050dde04a
📒 Files selected for processing (19)
ci/build_cpp.shci/build_docs.shci/build_python.shci/build_wheel_cuml.shci/build_wheel_libcuml.shci/run_clang_tidy.shci/test_cpp.shci/test_notebooks.shci/test_python_common.shci/test_wheel.shci/test_wheel_dask.shci/test_wheel_integrations.shci/use_conda_packages_from_prs.shci/use_wheels_from_prs.shcpp/examples/symreg/symreg_example.cppcpp/src/genetic/genetic.cucpp/tests/sg/genetic/evolution_test.cucpp/tests/sg/genetic/program_test.cupython/cuml/cuml/manifold/umap/umap.pyx
Contributor
Author
divyegala
approved these changes
Apr 20, 2026
viclafargue
approved these changes
Apr 21, 2026
| auto curr_mr = rmm::mr::get_current_device_resource_ref(); | ||
| d_finalprogs = static_cast<cg::program_t>(curr_mr.allocate(stream, params.population_size)); | ||
| d_finalprogs = static_cast<cg::program_t>( | ||
| curr_mr.allocate(stream, params.population_size * sizeof(cg::program), alignof(cg::program))); |
Contributor
There was a problem hiding this comment.
It looks like the PR description states adding rmm::CUDA_ALLOCATION_ALIGNMENT to allocate/deallocate. Could you update the it?
Also, could we also state the allocation bugfix in genetic code in the PR description?
symreg_example: pass byte count (population_size * sizeof(program)) instead of element count to allocate/deallocate, and use alignof(program) to match genetic.cu. Remove unused rmm/aligned.hpp include. parallel_evolve: use tmp.len (the old program length from d_oldprogs) instead of h_nextprogs[i].len (the new program length) in the deallocate call so the freed size matches the original allocation.
bdice
force-pushed
the
rmm-cccl-migration
branch
from
April 21, 2026 11:57
a23a300 to
ba9cd2b
Compare
Contributor
Author
|
/merge |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
allocate/deallocatecalls to the new CCCL 3-argument API that requires explicit alignmentrmm.librmm.per_device_resourceCython import withrmm.pylibrmm.memory_resourceand usemake_any_device_resourceto obtain the resource fordevice_bufferconstructionDepends on rapidsai/rmm#2361.
Depends on rapidsai/ucxx#636.
Depends on NVIDIA/raft#2996.
Depends on NVIDIA/cuvs#1990.
Changes
cpp/src/genetic/genetic.cu: Add explicitalignof(node)/alignof(program)to allallocateanddeallocatecalls inparallel_evolveandsymFit; fix deallocation bug inparallel_evolvewhereh_nextprogs[i].lenwas incorrectly used instead oftmp.lento compute the buffer size being freedcpp/examples/symreg/symreg_example.cpp: Useparams.population_size * sizeof(cg::program)andalignof(cg::program)forallocate/deallocatecalls, fixing incorrect byte-size computation; remove unused<rmm/aligned.hpp>includecpp/tests/sg/genetic/evolution_test.cu: Add alignment arguments to allocate/deallocate inSymRegtestcpp/tests/sg/genetic/program_test.cu: Add alignment arguments toSetUp/TearDownallocate/deallocate callspython/cuml/cuml/manifold/umap/umap.pyx: Replaceget_current_device_resource()withmake_any_device_resource(get_current_device_resource().get_mr())fordevice_bufferconstruction