Add CMake option for disabling legacy MR interface

cpp/CMakeLists.txt

@@ -44,6 +44,7 @@ option(BUILD_BENCHMARKS "Configure CMake to build (google) benchmarks" OFF)
 # This is mostly so that dependent libraries are configured in shared mode for downstream dependents
 # of RMM that get their common dependencies transitively.
 option(BUILD_SHARED_LIBS "Build RMM shared libraries" ON)
+option(RMM_ENABLE_LEGACY_MR_INTERFACE "Enable legacy memory resource interface" ON)
 set(RMM_LOGGING_LEVEL
     "INFO"
     CACHE STRING "Choose the logging level.")
@@ -54,6 +55,7 @@ message(VERBOSE "RMM: Build with NVTX support: ${RMM_NVTX}")
 # Set logging level. Must go before including gtests and benchmarks. Set the possible values of
 # build type for cmake-gui.
 message(STATUS "RMM: RMM_LOGGING_LEVEL = '${RMM_LOGGING_LEVEL}'")
+message(STATUS "RMM: Legacy MR interface enabled: ${RMM_ENABLE_LEGACY_MR_INTERFACE}")
 
 # cudart can be linked statically or dynamically
 option(CUDA_STATIC_RUNTIME "Statically link the CUDA runtime" OFF)
@@ -132,6 +134,11 @@ if(RMM_NVTX)
   target_compile_definitions(rmm PUBLIC RMM_NVTX)
 endif()
 
+# Control legacy MR interface visibility
+if(RMM_ENABLE_LEGACY_MR_INTERFACE)
+  target_compile_definitions(rmm PUBLIC RMM_ENABLE_LEGACY_MR_INTERFACE)
+endif()
+
 # ##################################################################################################
 # * tests and benchmarks ---------------------------------------------------------------------------
 

cpp/include/rmm/detail/cccl_adaptors.hpp

@@ -34,6 +34,7 @@ class cccl_resource_ref : public ResourceType {
 
   cccl_resource_ref(base&& other) : base(std::move(other)) {}
 
+#ifdef RMM_ENABLE_LEGACY_MR_INTERFACE
   void* allocate(std::size_t bytes) { return this->allocate_sync(bytes); }
 
   void* allocate(std::size_t bytes, std::size_t alignment)
@@ -50,8 +51,8 @@ class cccl_resource_ref : public ResourceType {
   {
     return this->deallocate_sync(ptr, bytes, alignment);
   }
+#endif  // RMM_ENABLE_LEGACY_MR_INTERFACE
 
-#if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 1)
   void* allocate_sync(std::size_t bytes) { return base::allocate_sync(bytes); }
 
   void* allocate_sync(std::size_t bytes, std::size_t alignment)
@@ -68,24 +69,6 @@ class cccl_resource_ref : public ResourceType {
   {
     return base::deallocate_sync(ptr, bytes, alignment);
   }
-#else
-  void* allocate_sync(std::size_t bytes) { return base::allocate(bytes); }
-
-  void* allocate_sync(std::size_t bytes, std::size_t alignment)
-  {
-    return base::allocate(bytes, alignment);
-  }
-
-  void deallocate_sync(void* ptr, std::size_t bytes) noexcept
-  {
-    return base::deallocate(ptr, bytes);
-  }
-
-  void deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept
-  {
-    return base::deallocate(ptr, bytes, alignment);
-  }
-#endif
 };
 
 template <typename ResourceType>
@@ -98,8 +81,7 @@ class cccl_async_resource_ref : public ResourceType {
   cccl_async_resource_ref(base const& other) : base(other) {}
   cccl_async_resource_ref(base&& other) : base(std::move(other)) {}
 
-  // BEGINNING OF LEGACY MR METHODS
-
+#ifdef RMM_ENABLE_LEGACY_MR_INTERFACE
   void* allocate(std::size_t bytes) { return this->allocate_sync(bytes); }
 
   void* allocate(std::size_t bytes, std::size_t alignment)
@@ -140,9 +122,8 @@ class cccl_async_resource_ref : public ResourceType {
     return this->deallocate(stream, ptr, bytes, alignment);
   }
 
-  // END OF LEGACY MR METHODS
+#endif  // RMM_ENABLE_LEGACY_MR_INTERFACE
 
-#if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 1)
   void* allocate_sync(std::size_t bytes) { return base::allocate_sync(bytes); }
 
   void* allocate_sync(std::size_t bytes, std::size_t alignment)
@@ -182,47 +163,6 @@ class cccl_async_resource_ref : public ResourceType {
   {
     return base::deallocate(stream, ptr, bytes, alignment);
   }
-#else
-  void* allocate_sync(std::size_t bytes) { return base::allocate(bytes); }
-
-  void* allocate_sync(std::size_t bytes, std::size_t alignment)
-  {
-    return base::allocate(bytes, alignment);
-  }
-
-  void deallocate_sync(void* ptr, std::size_t bytes) noexcept
-  {
-    return base::deallocate(ptr, bytes);
-  }
-
-  void deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment) noexcept
-  {
-    return base::deallocate(ptr, bytes, alignment);
-  }
-
-  void* allocate(cuda_stream_view stream, std::size_t bytes)
-  {
-    return base::allocate_async(bytes, stream);
-  }
-
-  void* allocate(cuda_stream_view stream, std::size_t bytes, std::size_t alignment)
-  {
-    return base::allocate_async(bytes, alignment, stream);
-  }
-
-  void deallocate(cuda_stream_view stream, void* ptr, std::size_t bytes) noexcept
-  {
-    return base::deallocate_async(ptr, bytes, stream);
-  }
-
-  void deallocate(cuda_stream_view stream,
-                  void* ptr,
-                  std::size_t bytes,
-                  std::size_t alignment) noexcept
-  {
-    return base::deallocate_async(ptr, bytes, alignment, stream);
-  }
-#endif
 };
 
 }  // namespace detail

cpp/include/rmm/detail/cuda_memory_resource.hpp

@@ -28,7 +28,6 @@ namespace RMM_NAMESPACE {
 namespace detail {
 namespace polyfill {
 
-#if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 1)
 template <class Resource>
 inline constexpr bool resource = cuda::mr::synchronous_resource<Resource>;
 template <class Resource, class... Properties>
@@ -37,16 +36,6 @@ template <class Resource>
 inline constexpr bool async_resource = cuda::mr::resource<Resource>;
 template <class Resource, class... Properties>
 inline constexpr bool async_resource_with = cuda::mr::resource_with<Resource, Properties...>;
-#else   // ^^^ CCCL >= 3.1 ^^^ / vvv CCCL < 3.1 vvv
-template <class Resource>
-inline constexpr bool resource = cuda::mr::resource<Resource>;
-template <class Resource, class... Properties>
-inline constexpr bool resource_with = cuda::mr::resource_with<Resource, Properties...>;
-template <class Resource>
-inline constexpr bool async_resource = cuda::mr::async_resource<Resource>;
-template <class Resource, class... Properties>
-inline constexpr bool async_resource_with = cuda::mr::async_resource_with<Resource, Properties...>;
-#endif  // CCCL < 3.1
 
 }  // namespace polyfill
 }  // namespace detail

cpp/include/rmm/mr/device/cuda_async_memory_resource.hpp

@@ -180,7 +180,7 @@ class cuda_async_memory_resource final : public device_memory_resource {
   void* do_allocate(std::size_t bytes, rmm::cuda_stream_view stream) override
   {
     void* ptr{nullptr};
-    ptr = pool_.allocate(bytes, stream);
+    ptr = pool_.allocate(stream, bytes);
     return ptr;
   }
 
@@ -194,7 +194,7 @@ class cuda_async_memory_resource final : public device_memory_resource {
    */
   void do_deallocate(void* ptr, std::size_t bytes, rmm::cuda_stream_view stream) noexcept override
   {
-    pool_.deallocate(ptr, bytes, stream);
+    pool_.deallocate(stream, ptr, bytes);
   }
 
   /**

cpp/include/rmm/mr/device/detail/arena.hpp

@@ -525,7 +525,7 @@ class global_arena final {
   ~global_arena()
   {
     std::lock_guard lock(mtx_);
-    upstream_mr_.deallocate(upstream_block_.pointer(), upstream_block_.size());
+    upstream_mr_.deallocate_sync(upstream_block_.pointer(), upstream_block_.size());
   }
 
   /**
@@ -701,7 +701,7 @@ class global_arena final {
    */
   void initialize(std::size_t size)
   {
-    upstream_block_ = {upstream_mr_.allocate(size), size};
+    upstream_block_ = {upstream_mr_.allocate_sync(size), size};
     superblocks_.emplace(upstream_block_.pointer(), size);
   }
 

cpp/include/rmm/mr/device/device_memory_resource.hpp

@@ -103,6 +103,7 @@ class device_memory_resource {
   device_memory_resource& operator=(device_memory_resource&&) noexcept =
     default;  ///< @default_move_assignment{device_memory_resource}
 
+#ifdef RMM_ENABLE_LEGACY_MR_INTERFACE
   /**
    * @brief Allocates memory of size at least \p bytes.
    *
@@ -151,24 +152,6 @@ class device_memory_resource {
     do_deallocate(ptr, bytes, stream);
   }
 
-  /**
-   * @brief Compare this resource to another.
-   *
-   * Two device_memory_resources compare equal if and only if memory allocated
-   * from one device_memory_resource can be deallocated from the other and vice
-   * versa.
-   *
-   * By default, simply checks if \p *this and \p other refer to the same
-   * object, i.e., does not check if they are two objects of the same class.
-   *
-   * @param other The other resource to compare to
-   * @returns If the two resources are equivalent
-   */
-  [[nodiscard]] bool is_equal(device_memory_resource const& other) const noexcept
-  {
-    return do_is_equal(other);
-  }
-
   /**
    * @brief Allocates memory of size at least \p bytes.
    *
@@ -293,9 +276,7 @@ class device_memory_resource {
     RMM_FUNC_RANGE();
     do_deallocate(ptr, bytes, stream);
   }
-
-#if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 1)
-  // CCCL >= 3.1 needs a different set of methods to satisfy the memory resource concepts
+#endif  // RMM_ENABLE_LEGACY_MR_INTERFACE
 
   /**
    * @brief Allocates memory of size at least \p bytes.
@@ -369,7 +350,24 @@ class device_memory_resource {
   {
     do_deallocate(ptr, rmm::align_up(bytes, alignment), stream);
   }
-#endif  // CCCL >= 3.1
+
+  /**
+   * @brief Compare this resource to another.
+   *
+   * Two device_memory_resources compare equal if and only if memory allocated
+   * from one device_memory_resource can be deallocated from the other and vice
+   * versa.
+   *
+   * By default, simply checks if \p *this and \p other refer to the same
+   * object, i.e., does not check if they are two objects of the same class.
+   *
+   * @param other The other resource to compare to
+   * @returns If the two resources are equivalent
+   */
+  [[nodiscard]] bool is_equal(device_memory_resource const& other) const noexcept
+  {
+    return do_is_equal(other);
+  }
 
   /**
    * @brief Comparison operator with another device_memory_resource

cpp/include/rmm/mr/device/owning_wrapper.hpp

@@ -170,7 +170,7 @@ class owning_wrapper : public device_memory_resource {
    */
   void* do_allocate(std::size_t bytes, cuda_stream_view stream) override
   {
-    return wrapped().allocate(bytes, stream);
+    return wrapped().allocate(stream, bytes);
   }
 
   /**
@@ -184,7 +184,7 @@ class owning_wrapper : public device_memory_resource {
    */
   void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) noexcept override
   {
-    wrapped().deallocate(ptr, bytes, stream);
+    wrapped().deallocate(stream, ptr, bytes);
   }
 
   /**

cpp/include/rmm/mr/host/host_memory_resource.hpp

@@ -68,7 +68,7 @@ class host_memory_resource {
   /**
    * @brief Allocates memory on the host of size at least `bytes` bytes.
    *
-   * The returned storage is aligned to the specified `alignment` if supported, and to
+   * The returned storage is aligned to the specified `alignment` if provided, and to
    * `alignof(std::max_align_t)` otherwise.
    *
    * @throws std::bad_alloc When the requested `bytes` and `alignment` cannot be allocated.
@@ -86,7 +86,7 @@ class host_memory_resource {
   /**
    * @brief Deallocate memory pointed to by `ptr`.
    *
-   * `ptr` must have been returned by a prior call to `allocate(bytes,alignment)` on a
+   * `ptr` must have been returned by a prior call to `allocate(bytes, alignment)` on a
    * `host_memory_resource` that compares equal to `*this`, and the storage it points to must not
    * yet have been deallocated, otherwise behavior is undefined.
    *
@@ -104,6 +104,45 @@ class host_memory_resource {
     do_deallocate(ptr, bytes, alignment);
   }
 
+  /**
+   * @brief Allocates memory on the host of size at least `bytes` bytes.
+   *
+   * The returned storage is aligned to the specified `alignment` if provided, and to
+   * `alignof(std::max_align_t)` otherwise.
+   *
+   * @throws std::bad_alloc When the requested `bytes` and `alignment` cannot be allocated.
+   *
+   * @param bytes The size of the allocation
+   * @param alignment Alignment of the allocation
+   * @return void* Pointer to the newly allocated memory
+   */
+  void* allocate_sync(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t))
+  {
+    RMM_FUNC_RANGE();
+    return do_allocate(bytes, alignment);
+  }
+
+  /**
+   * @brief Deallocate memory pointed to by `ptr`.
+   *
+   * `ptr` must have been returned by a prior call to `allocate(bytes, alignment)` on a
+   * `host_memory_resource` that compares equal to `*this`, and the storage it points to must not
+   * yet have been deallocated, otherwise behavior is undefined.
+   *
+   * @param ptr Pointer to be deallocated
+   * @param bytes The size in bytes of the allocation. This must be equal to the value of `bytes`
+   *              that was passed to the `allocate` call that returned `ptr`.
+   * @param alignment Alignment of the allocation. This must be equal to the value of `alignment`
+   *                  that was passed to the `allocate` call that returned `ptr`.
+   */
+  void deallocate_sync(void* ptr,
+                       std::size_t bytes,
+                       std::size_t alignment = alignof(std::max_align_t)) noexcept
+  {
+    RMM_FUNC_RANGE();
+    do_deallocate(ptr, bytes, alignment);
+  }
+
   /**
    * @brief Compare this resource to another.
    *
@@ -201,46 +240,6 @@ class host_memory_resource {
   {
     return this == &other;
   }
-
-#if CCCL_MAJOR_VERSION > 3 || (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION >= 1)
-
- public:
-  /**
-   * @brief Allocates memory on the host of size at least `bytes` bytes.
-   *
-   * The returned storage is aligned to the specified `alignment` if supported, and to
-   * `alignof(std::max_align_t)` otherwise.
-   *
-   * @throws std::bad_alloc When the requested `bytes` and `alignment` cannot be allocated.
-   *
-   * @param bytes The size of the allocation
-   * @param alignment Alignment of the allocation
-   * @return void* Pointer to the newly allocated memory
-   */
-  void* allocate_sync(std::size_t bytes, std::size_t alignment)
-  {
-    return allocate(bytes, alignment);
-  }
-
-  /**
-   * @brief Deallocate memory pointed to by `ptr`.
-   *
-   * `ptr` must have been returned by a prior call to `allocate(bytes,alignment)` on a
-   * `host_memory_resource` that compares equal to `*this`, and the storage it points to must not
-   * yet have been deallocated, otherwise behavior is undefined.
-   *
-   * @param ptr Pointer to be deallocated
-   * @param bytes The size in bytes of the allocation. This must be equal to the value of `bytes`
-   *              that was passed to the `allocate` call that returned `ptr`.
-   * @param alignment Alignment of the allocation. This must be equal to the value of `alignment`
-   *                  that was passed to the `allocate` call that returned `ptr`.
-   */
-  void deallocate_sync(void* ptr, std::size_t bytes, std::size_t alignment)
-  {
-    return deallocate(ptr, bytes, alignment);
-  }
-
-#endif
 };
 
 // static property checks