Compatibility with CCCL 3.2 (#755)

This PR introduces some fixes needed for CCCL 3.2.

The noteworthy changes are:
- Pinned memory resource must be copyable (required by CCCL 3.2's memory resource concepts), which means the "PIMPL" idiom must use a `shared_ptr` to make the resource copyable
- Header/namespace updates (pinned MR is no longer experimental)
- Resource adaptor comparison does not work with `optional` types due to constraints in the CCCL resource type machinery. This requires a workaround where optionals are not compared directly, but instead checked for `has_value()` before comparing their underlying objects. It may be possible to improve this as I continue refactoring RMM's bindings around CCCL memory resources.

Authors:
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Niranda Perera (https://github.com/nirandaperera)

URL: #755

Author: bdice

cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp

@@ -149,7 +149,7 @@ class PinnedMemoryResource final : public HostMemoryResource {
     // using PImpl idiom to hide cudax .cuh headers from rapidsmpf. cudax cuh headers will
     // only be used by the impl in .cu file.
     struct PinnedMemoryResourceImpl;
-    std::unique_ptr<PinnedMemoryResourceImpl> impl_;
+    std::shared_ptr<PinnedMemoryResourceImpl> impl_;
 };
 
 static_assert(cuda::mr::resource<PinnedMemoryResource>);

cpp/src/memory/pinned_memory_resource.cu

@@ -7,23 +7,41 @@
 
 #include <cuda_runtime_api.h>
 
-#include <cuda/memory_resource>
-
 #include <rmm/resource_ref.hpp>
 
 #include <rapidsmpf/error.hpp>
 #include <rapidsmpf/memory/pinned_memory_resource.hpp>
 #include <rapidsmpf/utils.hpp>
 
 #if RAPIDSMPF_CUDA_VERSION_AT_LEAST(RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION)
+#if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 2)
+#include <cuda/memory_resource>
+#else
 #include <cuda/experimental/memory_resource.cuh>
 #endif
+#endif
 
 namespace rapidsmpf {
 
 #if RAPIDSMPF_CUDA_VERSION_AT_LEAST(RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION)
 
 namespace {
+#if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 2)
+cuda::memory_pool_properties get_memory_pool_properties() {
+    return cuda::memory_pool_properties{
+        // It was observed that priming async pools have little effect for performance.
+        // See <https://github.com/rapidsai/rmm/issues/1931>.
+        .initial_pool_size = 0,
+        // Before <https://github.com/NVIDIA/cccl/pull/6718>, the default
+        // `release_threshold` was 0, which defeats the purpose of having a pool. We
+        // now set it so the pool never releases unused pinned memory.
+        .release_threshold = std::numeric_limits<size_t>::max(),
+        // This defines how the allocations can be exported (IPC). See the docs of
+        // `cudaMemPoolCreate` in <https://docs.nvidia.com/cuda/cuda-runtime-api>.
+        .allocation_handle_type = ::cudaMemAllocationHandleType::cudaMemHandleTypeNone
+    };
+}
+#else
 cuda::experimental::memory_pool_properties get_memory_pool_properties() {
     return cuda::experimental::memory_pool_properties{
         // It was observed that priming async pools have little effect for performance.
@@ -39,6 +57,7 @@ cuda::experimental::memory_pool_properties get_memory_pool_properties() {
             cuda::experimental::cudaMemAllocationHandleType::cudaMemHandleTypeNone
     };
 }
+#endif
 }  // namespace
 
 struct PinnedMemoryResource::PinnedMemoryResourceImpl {
@@ -55,8 +74,13 @@ struct PinnedMemoryResource::PinnedMemoryResourceImpl {
         resource.deallocate(stream, ptr, bytes, alignment);
     }
 
+#if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 2)
+    cuda::pinned_memory_pool pool;
+    cuda::pinned_memory_pool_ref resource;
+#else
     cuda::experimental::pinned_memory_pool pool;
     cuda::experimental::pinned_memory_resource resource;
+#endif
 };
 #else  // CUDA_VERSION < RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION
 
@@ -77,7 +101,7 @@ struct PinnedMemoryResource::PinnedMemoryResourceImpl {
 #endif
 
 PinnedMemoryResource::PinnedMemoryResource(int numa_id)
-    : impl_(std::make_unique<PinnedMemoryResourceImpl>(numa_id)) {
+    : impl_(std::make_shared<PinnedMemoryResourceImpl>(numa_id)) {
     RAPIDSMPF_EXPECTS(
         is_pinned_memory_resources_supported(),
         "PinnedMemoryResource is not supported for CUDA versions "

cpp/src/rmm_resource_adaptor.cpp

@@ -126,8 +126,19 @@ bool RmmResourceAdaptor::do_is_equal(
     if (cast == nullptr) {
         return false;
     }
-    return get_upstream_resource() == cast->get_upstream_resource()
-           && get_fallback_resource() == cast->get_fallback_resource();
+    // Manual comparison of optionals to avoid recursive constraint satisfaction in
+    // CCCL 3.2. std::optional::operator== triggers infinite concept checking when the
+    // wrapped type (rmm::device_async_resource_ref) inherits from CCCL's concept-based
+    // resource_ref.
+    // TODO: Revert this after the RMM resource ref types are replaced with
+    // plain cuda::mr ref types. This depends on
+    // https://github.com/rapidsai/rmm/issues/2011.
+    auto this_fallback = get_fallback_resource();
+    auto other_fallback = cast->get_fallback_resource();
+    bool fallbacks_equal =
+        (this_fallback.has_value() == other_fallback.has_value())
+        && (!this_fallback.has_value() || (*this_fallback == *other_fallback));
+    return get_upstream_resource() == cast->get_upstream_resource() && fallbacks_equal;
 }
 
 }  // namespace rapidsmpf