Migrate RMM usage to CCCL MR design by bdice · Pull Request #98 · NVIDIA/cuCascade

bdice · 2026-04-02T21:57:58Z

Summary

Migrate all cuCascade memory resource classes from rmm::mr::device_memory_resource virtual inheritance to the CCCL-native resource concept (allocate/deallocate with cuda::stream_ref, allocate_sync/deallocate_sync, operator==, get_property).
Replace unique_ptr<device_memory_resource> owning storage with cuda::mr::any_resource<cuda::mr::device_accessible>.
Update memory_space::get_default_allocator() to return device_async_resource_ref from the any_resource.
Update pool_memory_resource usage (no longer a template in RMM 26.06; upstream passed as device_async_resource_ref).
Move memory_reservation.hpp template bodies (get_memory_resource_as<T>(), get_memory_resource_of<Tier>()) to memory_space.hpp to avoid incomplete-type errors.
Remove allocator null-checks in GPU and HOST memory_space constructors — any_resource is a value type and cannot be null.
Update internal converter helpers in representation_converter.cpp to accept rmm::device_async_resource_ref instead of raw rmm::mr::device_memory_resource*.
Update test utilities (shared_device_resource, logging_device_resource, make_mock_memory_space) for the new resource concept.
Pre-commit formatting pass on files touched by the merge with upstream/main.

Depends on NVIDIA/cudf#22008.

Affected classes

null_device_memory_resource
numa_region_pinned_host_allocator
fixed_size_host_memory_resource
reservation_aware_resource_adaptor
small_pinned_host_memory_resource
memory_space

Notes

Companion to the rapidsmpf migration PR.
cuda::stream_ref{cudaStream_t{nullptr}} in allocate_sync/deallocate_sync implementations may produce deprecation warnings. These do not break the build because the WARNINGS_AS_ERRORS CMake guard on line 80 checks the wrong variable name (WARNINGS_AS_ERRORS vs the option CUCASCADE_WARNINGS_AS_ERRORS), so -Werror is currently not enabled.
Pre-commit reformatted several files that were untouched by the migration but modified by the upstream merge (cuda_utils.hpp, test_data_batch.cpp, test_data_representation.cpp, benchmark_representation_converter.cpp, data_batch.hpp). These are whitespace/include-order changes only.

… resource concept RMM 26.06 removes rmm::mr::device_memory_resource and adopts CCCL-native memory resource concepts. This migrates all cuCascade resource classes: - null_device_memory_resource: replaced virtual override with CCCL concept methods (allocate/deallocate with cuda::stream_ref, allocate_sync, deallocate_sync, operator==, get_property). - numa_region_pinned_host_allocator: same pattern. - fixed_size_host_memory_resource: same pattern. - reservation_aware_resource_adaptor: same pattern. - small_pinned_host_memory_resource: same pattern. - memory_space: unique_ptr<device_memory_resource> replaced with cuda::mr::any_resource<cuda::mr::device_accessible> for owning storage; get_default_allocator returns device_async_resource_ref from any_resource. - memory_reservation.hpp: moved template bodies out-of-line to avoid incomplete-type errors with memory_space forward declaration. - common.hpp/cpp: updated factory typedefs and pool_memory_resource usage (no longer a template in 26.06). - pool_memory_resource upstream passed as device_async_resource_ref. - Test utilities updated for new resource concept.

copy-pr-bot · 2026-04-02T21:58:02Z

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.

mbrobbel · 2026-04-07T08:21:06Z

/ok to test a852c20

# Conflicts: # src/data/representation_converter.cpp # src/memory/memory_space.cpp # test/utils/cudf_test_utils.cpp

mbrobbel · 2026-04-15T08:19:05Z

/ok to test 9df76bf

The CCCL async_resource concept requires explicit alignment arguments and explicit cudaStream_t construction for stream_ref, unlike the old rmm::mr::device_memory_resource pointer API which had defaults. - Pass rmm::CUDA_ALLOCATION_ALIGNMENT for device resource allocate/deallocate - Pass alignof(std::max_align_t) for host resource allocate/deallocate - Use cuda::stream_ref{cudaStream_t{nullptr}} for sync allocate/deallocate - Use rmm::device_async_resource_ref in representation_converter signatures - Construct null_device_memory_resource directly (no unique_ptr indirection)

mbrobbel · 2026-04-16T13:32:33Z

/ok to test fd00654

…compliance Test resource types (shared_device_resource, logging_device_resource) now implement the synchronous allocation interface required by the CCCL resource concept. Also migrates deprecated set_current_device_resource_ref calls to set_current_device_resource and fixes any_resource construction in test setup.

… split Move all state and private methods into a non-copyable impl class (detail::reservation_aware_resource_adaptor_impl). The public wrapper now inherits from cuda::mr::shared_resource<impl>, making it copyable and movable via reference counting. This satisfies the CCCL resource concept requirements for any_resource, enabling use with device_buffer and device_uvector in RMM 26.06. Follows the same pattern used by RMM's own pool_memory_resource and tracking_resource_adaptor.

…ctations Remove extra closing brace that broke SCENARIO block structure. Fix the 'exceeds reservation' test case which previously expected overflow with two chunk_size allocations that exactly filled the reservation — add a third allocation to actually trigger overflow, and correct the expected total_allocated_bytes to reservation_size + chunk_size.

mbrobbel · 2026-04-21T12:20:12Z

/ok to test e2fa1cf

wence- · 2026-04-21T14:12:05Z

+  void* allocate_sync(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t))
+  {
+    return allocate(cuda::stream_ref{cudaStream_t{nullptr}}, bytes, alignment);
+  }


issue: needs a stream sync before returning the pointer.

wence- · 2026-04-21T14:14:46Z

-  using upstream_type = rmm::mr::device_memory_resource;
-  using type          = rmm::mr::device_memory_resource;
-  Tier tier           = TIER;
+  Tier tier = TIER;


Probably should have using type = void?

wence- · 2026-04-21T14:16:42Z

+  [[nodiscard]] rmm::device_async_resource_ref get_default_allocator() const noexcept;

  template <typename T>
-    requires std::derived_from<T, rmm::mr::device_memory_resource>


Should this requires change to being that the T is a resource_with<cuda::mr::device_accessible>?

wence- · 2026-04-21T14:19:05Z

+  //===----------------------------------------------------------------------===//
+  // Convenience allocate/deallocate with default alignment
+  //===----------------------------------------------------------------------===//

-  /// Default policy for new streams
-  std::unique_ptr<reservation_limit_policy> _default_reservation_policy;
-  std::unique_ptr<oom_handling_policy> _default_oom_policy;
+  void* allocate(cuda::stream_ref stream,
+                 std::size_t bytes,
+                 std::size_t alignment = alignof(std::max_align_t))
+  {
+    return get().allocate(stream, bytes, alignment);
+  }
+
+  void deallocate(cuda::stream_ref stream,
+                  void* ptr,
+                  std::size_t bytes,
+                  std::size_t alignment = alignof(std::max_align_t)) noexcept
+  {
+    get().deallocate(stream, ptr, bytes, alignment);
+  }


issue: I thought CCCL moved away from defaulted alignment. I think we should also do so here.

wence- · 2026-04-21T14:19:29Z

-   *
-   * For @p bytes <= MAX_SLAB_SIZE: rounds up to the next slab boundary
-   * (512 / 1 KB / 2 KB / 4 KB / 8 KB) and returns a pointer from the matching
-   * free list, expanding the pool from upstream if the list is empty.
-   *
-   * For @p bytes > MAX_SLAB_SIZE: falls back to std::malloc (pageable).
-   * The cudf::set_allocate_host_as_pinned_threshold is set to MAX_SLAB_SIZE so
-   * that cuDF's make_host_vector path uses the slab pools for metadata buffers.
-   * Larger allocations (e.g. join/sort staging buffers that call
-   * get_pinned_memory_resource() directly) are served from pageable memory.
   */


Why removing this comment?

I think this is already mentioned here
and not really an impl detail of this method.

wence- · 2026-04-21T14:21:34Z

+  return cuda::mr::any_resource<cuda::mr::device_accessible>{
+    cucascade::memory::numa_region_pinned_host_memory_resource(numa_node_id)};


Can we use brace-initializer return and just return {cucascade::...};

wence- · 2026-04-21T14:23:14Z

+//===----------------------------------------------------------------------===//
+// Allocation tracker implementations (file-local)
+//===----------------------------------------------------------------------===//



wence- · 2026-04-21T14:23:22Z

+//===----------------------------------------------------------------------===//
+// impl_type (the impl class) method definitions
+//===----------------------------------------------------------------------===//
+


Why? Also, even if you want to keep it, in the wrong place.

wence- · 2026-04-21T14:25:16Z

+                          [[maybe_unused]] std::size_t alignment)
 {
-  auto* reservation_state = _allocation_tracker->get_tracker_state(stream);
+  rmm::cuda_stream_view stream_view(stream.get());


cuda_stream_view is implicitly convertible from stream_ref, so you can just pass stream everywhere in this function.

wence- · 2026-04-21T14:26:25Z

+
+//===----------------------------------------------------------------------===//
+// Wrapper (reservation_aware_resource_adaptor) forwarding methods
+//===----------------------------------------------------------------------===//


mbrobbel

Nightly jobs are now passing. We'll have to disable the stable jobs here to merge this PR - I'll do a follow-up PR to support both stable and nightly again.

## Summary - Rewrite `RmmResourceAdaptor` as a thin shell inheriting `cuda::mr::shared_resource<detail::RmmResourceAdaptorImpl>`, with all mutable state in the impl class for copyable shared ownership. - Replace `device_memory_resource*` with `rmm::device_async_resource_ref` for non-owning references and `cuda::mr::any_resource` for owning storage. - Remove `rmm::mr::owning_wrapper` usage (removed in RMM 26.06). - Update `pool_memory_resource` usage (no longer a template; upstream passed as `device_async_resource_ref`). - Replace `set_current_device_resource(ptr)` with `set_current_device_resource_ref(ref)`. - Update test resources to satisfy CCCL resource concept (`allocate_sync`, `deallocate_sync`, `operator==`, `get_property`). - Update Cython bindings (`.pxd`/`.pyx`) to use `device_async_resource_ref` instead of `device_memory_resource`. - Point `get_cucascade.cmake` to local cuCascade source directory (companion cuCascade PR: NVIDIA/cuCascade#98). Depends on rapidsai/rmm#2361. Depends on rapidsai/ucxx#636. Depends on NVIDIA/cudf#22008. ## Notes - Depends on cuCascade migration: NVIDIA/cuCascade#98 - The `get_cucascade.cmake` change to use `SOURCE_DIR` is a development convenience and should be updated to point to the merged cuCascade commit before this PR is finalized. Authors: - Bradley Dice (https://github.com/bdice) - Niranda Perera (https://github.com/nirandaperera) Approvers: - Niranda Perera (https://github.com/nirandaperera) - Peter Andreas Entschev (https://github.com/pentschev) URL: #940

…tion - Add cudaStreamSynchronize to allocate_sync (CUCASCADE_CUDA_TRY) and deallocate_sync (CUCASCADE_ASSERT_CUDA_SUCCESS) matching RMM pattern - Consolidate cuda_utils.hpp as forwarding include to error.hpp, eliminating duplicate macro definitions across all 9 includers - Return host_accessible+device_accessible from make_default_host_memory_resource - Remove defaulted alignment from reservation_aware_resource_adaptor allocate/deallocate to match CCCL resource concept conventions - Add resource_with<device_accessible> || disk_access_limiter constraint to memory_space::get_memory_resource_as - Add using type = void to primary tier_memory_resource_trait template - Restore doc comments on small_pinned_host_memory_resource allocate/deallocate - Remove added section comments in reservation_aware_resource_adaptor.cpp - Use implicit stream_ref to cuda_stream_view conversion in impl - Use brace-initializer returns in common.cpp factory functions

nirandaperera · 2026-04-22T18:14:25Z

@mbrobbel Could you please check c8b82b8 654fde1

mbrobbel · 2026-04-22T18:18:11Z

/ok to test 654fde1

mbrobbel · 2026-04-22T18:34:25Z

@mbrobbel Could you please check c8b82b8 654fde1

Looks good to me.

mbrobbel · 2026-04-22T19:40:37Z

/merge

Updated cuCascade GIT_TAG to use the main branch and enabled shallow cloning. Cucascade was pinned to RMM refactor PR previously. Its now merged. NVIDIA/cuCascade#98 Authors: - Niranda Perera (https://github.com/nirandaperera) Approvers: - Bradley Dice (https://github.com/bdice) URL: #988

Two fixes that interact under RAPIDS 26.04 + CCCL MR design (post-NVIDIA#98): 1. Don't pass cucascade memory_capacity as cuda_async_memory_resource's initial_pool_size. Post-26.04 the constructor primes the pool by allocating that exact byte count up front and (under release_threshold=total_memory) keeps it cached. Multiple memory_space instances per device — e.g. shared_test_env constructs 3 environments upfront in unit-test main — exhaust physical GPU memory before any work runs (OOM at static-init time with allocations >= 18 GB on a 49 GB GPU). The reservation_aware_resource_adaptor enforces the cucascade-level budget independently, so the underlying RMM pool can grow lazily without losing budget enforcement. 2. Grant cudaMemPoolSetAccess(ReadWrite) on each newly-created pool for every other peer-capable device. cudaMallocAsync pools require an explicit per-pool access list for cudaMemcpyPeer* between two pools to actually transfer bytes — cudaDeviceEnablePeerAccess on its own only governs legacy cudaMalloc memory. Without this, the convert_gpu_to_gpu peer copy would silently no-op and dst would remain at its uninitialized 0xFF pattern, surfacing as deterministic FNV-1a checksum mismatches in the MGPU-04 / MGPU-06 round-trip and p2p_transfer_converter_round_trip tests. The helper `enable_pool_peer_access_for_all_visible_devices` is exposed in `<cucascade/memory/common.hpp>` so both the factory path (make_default_gpu_memory_resource) and the inline path (memory_space ctor when no factory_fn) call it. Each pool only declares what other devices may access *itself*; the call is one-sided per pool and doesn't require peer pools to exist yet. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Brings PR NVIDIA#94 (perf: reduce per-call overhead in representation converters) up to date with NVIDIA/cuCascade main (31 commits, including the CCCL MR migration NVIDIA#98 and the new disk I/O backends). Conflicts in src/data/representation_converter.cpp were resolved in favor of main's redesign, which superseded several of the perf branch's manual optimizations with cleaner RMM equivalents: - convert_gpu_to_gpu was rewritten on main to peer-copy column buffers directly (no cudf::pack), so the perf branch's removal of the pre-pack stream.synchronize() no longer applies. - Manual cudaGetDevice/cudaSetDevice guards were replaced by rmm::cuda_set_device_raii, which already skips no-op device switches (the same optimization the perf branch added by hand). - BatchCopyAccumulator placement and stream-lineage tracking taken from main. Perf optimizations that survive (merged cleanly, outside the conflicts): - shared_mutex in the converter registry (concurrent convert/has_converter no longer serialize on an exclusive lock). - Eliminated the metadata vector copy and a redundant stream.synchronize() in the packed HOST->GPU converter (convert_host_to_gpu). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

This was referenced Apr 3, 2026

Migrate RMM usage to CCCL MR design rapidsai/rapidsmpf#940

Merged

[FEA] Support memory resources from CCCL 3.2 rapidsai/rmm#2011

Closed

mbrobbel added breaking improvement labels Apr 7, 2026

Merge remote-tracking branch 'upstream/main' into rmm-cccl-migration

9df76bf

# Conflicts: # src/data/representation_converter.cpp # src/memory/memory_space.cpp # test/utils/cudf_test_utils.cpp

bdice added 2 commits April 16, 2026 08:30

Run pre-commit

fd00654

bdice marked this pull request as ready for review April 16, 2026 13:55

bdice added 5 commits April 16, 2026 10:15

Fix WARNINGS_AS_ERRORS name

9163f9d

Disable -Wnull-dereference

e2fa1cf

bdice mentioned this pull request Apr 21, 2026

CCCL Memory Resource Migration — Merge Train rapidsai/rmm#2364

Closed

mbrobbel mentioned this pull request Apr 21, 2026

ci: allow nightly jobs to fail #106

Closed

wence- suggested changes Apr 21, 2026

View reviewed changes

mbrobbel reviewed Apr 21, 2026

View reviewed changes

nirandaperera approved these changes Apr 21, 2026

View reviewed changes

nirandaperera added 2 commits April 22, 2026 11:02

temporily disabling stable CI runs

c8b82b8

disable stable

654fde1

nirandaperera mentioned this pull request Apr 22, 2026

Re-enable stable builds and CI after 26.06 release #110

Closed

mbrobbel approved these changes Apr 22, 2026

View reviewed changes

rapids-bot Bot merged commit c612588 into NVIDIA:main Apr 22, 2026
7 checks passed

nirandaperera mentioned this pull request Apr 22, 2026

Change cuCascade GIT_TAG to main rapidsai/rapidsmpf#988

Merged

		return cuda::mr::any_resource<cuda::mr::device_accessible>{
		cucascade::memory::numa_region_pinned_host_memory_resource(numa_node_id)};

Conversation

bdice commented Apr 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Affected classes

Notes

Uh oh!

copy-pr-bot Bot commented Apr 2, 2026

Uh oh!

mbrobbel commented Apr 7, 2026

Uh oh!

mbrobbel commented Apr 15, 2026

Uh oh!

mbrobbel commented Apr 16, 2026

Uh oh!

mbrobbel commented Apr 21, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mbrobbel left a comment

Choose a reason for hiding this comment

Uh oh!

nirandaperera commented Apr 22, 2026

Uh oh!

mbrobbel commented Apr 22, 2026

Uh oh!

mbrobbel commented Apr 22, 2026

Uh oh!

mbrobbel commented Apr 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bdice commented Apr 2, 2026 •

edited

Loading