[FEA]: Use std::optional instead of thrust::optional

include/rmm/mr/device/cuda_async_memory_resource.hpp

@@ -23,11 +23,13 @@
 #include <rmm/mr/device/cuda_async_view_memory_resource.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
 
+#include <cuda/std/type_traits>
 #include <cuda_runtime_api.h>
 #include <thrust/optional.h>
 
 #include <cstddef>
 #include <limits>
+#include <optional>
 
 #if CUDART_VERSION >= 11020  // 11.2 introduced cudaMallocAsync
 #ifndef RMM_DISABLE_CUDA_MALLOC_ASYNC
@@ -84,9 +86,43 @@ class cuda_async_memory_resource final : public device_memory_resource {
    * `cudaMemHandleTypeNone` for no IPC support.
    */
   // NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
-  cuda_async_memory_resource(thrust::optional<std::size_t> initial_pool_size             = {},
-                             thrust::optional<std::size_t> release_threshold             = {},
-                             thrust::optional<allocation_handle_type> export_handle_type = {})
+  template <class Optional,
+            cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
+                                                        thrust::optional<std::size_t>>,
+                                   int> = 0>
+  [[deprecated("Use std::optional instead of thrust::optional.")]]  //
+  explicit cuda_async_memory_resource(
+    Optional initial_pool_size,
+    Optional release_threshold                                  = {},
+    thrust::optional<allocation_handle_type> export_handle_type = {})
+    : cuda_async_memory_resource(initial_pool_size.value_or(std::nullopt),
+                                 release_threshold.value_or(std::nullopt),
+                                 export_handle_type.value_or(std::nullopt))
+
+  {
+  }
+
+  /**
+   * @brief Constructs a cuda_async_memory_resource with the optionally specified initial pool size
+   * and release threshold.
+   *
+   * If the pool size grows beyond the release threshold, unused memory held by the pool will be
+   * released at the next synchronization event.
+   *
+   * @throws rmm::logic_error if the CUDA version does not support `cudaMallocAsync`
+   *
+   * @param initial_pool_size Optional initial size in bytes of the pool. If no value is provided,
+   * initial pool size is half of the available GPU memory.
+   * @param release_threshold Optional release threshold size in bytes of the pool. If no value is
+   * provided, the release threshold is set to the total amount of memory on the current device.
+   * @param export_handle_type Optional `cudaMemAllocationHandleType` that allocations from this
+   * resource should support interprocess communication (IPC). Default is
+   * `cudaMemHandleTypeNone` for no IPC support.
+   */
+  // NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
+  cuda_async_memory_resource(std::optional<std::size_t> initial_pool_size             = {},
+                             std::optional<std::size_t> release_threshold             = {},
+                             std::optional<allocation_handle_type> export_handle_type = {})
   {
 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
     // Check if cudaMallocAsync Memory pool supported

include/rmm/mr/device/pool_memory_resource.hpp

@@ -25,6 +25,7 @@
 #include <rmm/mr/device/detail/stream_ordered_memory_resource.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
 
+#include <cuda/std/type_traits>
 #include <cuda_runtime_api.h>
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/transform_iterator.h>
@@ -38,6 +39,7 @@
 #include <map>
 #include <mutex>
 #include <numeric>
+#include <optional>
 #include <set>
 #include <thread>
 #include <unordered_map>
@@ -108,6 +110,37 @@ class pool_memory_resource final
   friend class detail::stream_ordered_memory_resource<pool_memory_resource<Upstream>,
                                                       detail::coalescing_free_list>;
 
+  /**
+   * @brief Construct a `pool_memory_resource` and allocate the initial device memory
+   * pool using `upstream_mr`.
+   *
+   * @deprecated Use the constructor that takes an explicit initial pool size instead.
+   *
+   * @throws rmm::logic_error if `upstream_mr == nullptr`
+   * @throws rmm::logic_error if `initial_pool_size` is neither the default nor aligned to a
+   * multiple of pool_memory_resource::allocation_alignment bytes.
+   * @throws rmm::logic_error if `maximum_pool_size` is neither the default nor aligned to a
+   * multiple of pool_memory_resource::allocation_alignment bytes.
+   *
+   * @param upstream_mr The memory_resource from which to allocate blocks for the pool.
+   * @param initial_pool_size Minimum size, in bytes, of the initial pool. Defaults to zero.
+   * @param maximum_pool_size Maximum size, in bytes, that the pool can grow to. Defaults to all
+   * of the available memory from the upstream resource.
+   */
+  template <class Optional,
+            cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
+                                                        thrust::optional<std::size_t>>,
+                                   int> = 0>
+  [[deprecated(
+    "Must specify initial_pool_size and use std::optional instead of thrust::optional.")]]  //
+  explicit pool_memory_resource(Upstream* upstream_mr,
+                                Optional initial_pool_size,
+                                Optional maximum_pool_size = thrust::nullopt)
+    : pool_memory_resource(
+        upstream_mr, initial_pool_size.value_or(0), maximum_pool_size.value_or(std::nullopt))
+  {
+  }
+
   /**
    * @brief Construct a `pool_memory_resource` and allocate the initial device memory
    * pool using `upstream_mr`.
@@ -127,12 +160,43 @@ class pool_memory_resource final
    */
   [[deprecated("Must specify initial_pool_size")]]  //
   explicit pool_memory_resource(Upstream* upstream_mr,
-                                thrust::optional<std::size_t> initial_pool_size = thrust::nullopt,
-                                thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
+                                std::optional<std::size_t> initial_pool_size = std::nullopt,
+                                std::optional<std::size_t> maximum_pool_size = std::nullopt)
     : pool_memory_resource(upstream_mr, initial_pool_size.value_or(0), maximum_pool_size)
   {
   }
 
+  /**
+   * @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
+   * `upstream_mr`.
+   *
+   * @deprecated Use the constructor that takes an explicit initial pool size instead.
+   *
+   * @throws rmm::logic_error if `upstream_mr == nullptr`
+   * @throws rmm::logic_error if `initial_pool_size` is neither the default nor aligned to a
+   * multiple of pool_memory_resource::allocation_alignment bytes.
+   * @throws rmm::logic_error if `maximum_pool_size` is neither the default nor aligned to a
+   * multiple of pool_memory_resource::allocation_alignment bytes.
+   *
+   * @param upstream_mr The memory_resource from which to allocate blocks for the pool.
+   * @param initial_pool_size Minimum size, in bytes, of the initial pool. Defaults to zero.
+   * @param maximum_pool_size Maximum size, in bytes, that the pool can grow to. Defaults to all
+   * of the available memory from the upstream resource.
+   */
+  template <class Optional,
+            cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
+                                                        thrust::optional<std::size_t>>,
+                                   int> = 0>
+  [[deprecated(
+    "Must specify initial_pool_size and use std::optional instead of thrust::optional.")]]  //
+  explicit pool_memory_resource(Upstream& upstream_mr,
+                                Optional initial_pool_size,
+                                Optional maximum_pool_size = thrust::nullopt)
+    : pool_memory_resource(
+        upstream_mr, initial_pool_size.value_or(0), maximum_pool_size.value_or(std::nullopt))
+  {
+  }
+
   /**
    * @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
    * `upstream_mr`.
@@ -154,12 +218,39 @@ class pool_memory_resource final
             cuda::std::enable_if_t<cuda::mr::async_resource<Upstream2>, int> = 0>
   [[deprecated("Must specify initial_pool_size")]]  //
   explicit pool_memory_resource(Upstream2& upstream_mr,
-                                thrust::optional<std::size_t> initial_pool_size = thrust::nullopt,
-                                thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
+                                std::optional<std::size_t> initial_pool_size = std::nullopt,
+                                std::optional<std::size_t> maximum_pool_size = std::nullopt)
     : pool_memory_resource(upstream_mr, initial_pool_size.value_or(0), maximum_pool_size)
   {
   }
 
+  /**
+   * @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
+   * `upstream_mr`.
+   *
+   * @throws rmm::logic_error if `upstream_mr == nullptr`
+   * @throws rmm::logic_error if `initial_pool_size` is not aligned to a multiple of
+   * pool_memory_resource::allocation_alignment bytes.
+   * @throws rmm::logic_error if `maximum_pool_size` is neither the default nor aligned to a
+   * multiple of pool_memory_resource::allocation_alignment bytes.
+   *
+   * @param upstream_mr The memory_resource from which to allocate blocks for the pool.
+   * @param initial_pool_size Minimum size, in bytes, of the initial pool.
+   * @param maximum_pool_size Maximum size, in bytes, that the pool can grow to. Defaults to all
+   * of the available from the upstream resource.
+   */
+  template <class Optional,
+            cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
+                                                        thrust::optional<std::size_t>>,
+                                   int> = 0>
+  [[deprecated("Use std::optional instead of thrust::optional.")]]  //
+  explicit pool_memory_resource(Upstream* upstream_mr,
+                                std::size_t initial_pool_size,
+                                Optional maximum_pool_size)
+    : pool_memory_resource(upstream_mr, initial_pool_size, maximum_pool_size.value_or(std::nullopt))
+  {
+  }
+
   /**
    * @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
    * `upstream_mr`.
@@ -177,7 +268,7 @@ class pool_memory_resource final
    */
   explicit pool_memory_resource(Upstream* upstream_mr,
                                 std::size_t initial_pool_size,
-                                thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
+                                std::optional<std::size_t> maximum_pool_size = std::nullopt)
     : upstream_mr_{[upstream_mr]() {
         RMM_EXPECTS(nullptr != upstream_mr, "Unexpected null upstream pointer.");
         return upstream_mr;
@@ -191,6 +282,35 @@ class pool_memory_resource final
     initialize_pool(initial_pool_size, maximum_pool_size);
   }
 
+  /**
+   * @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
+   * `upstream_mr`.
+   *
+   * @throws rmm::logic_error if `upstream_mr == nullptr`
+   * @throws rmm::logic_error if `initial_pool_size` is not aligned to a multiple of
+   * pool_memory_resource::allocation_alignment bytes.
+   * @throws rmm::logic_error if `maximum_pool_size` is neither the default nor aligned to a
+   * multiple of pool_memory_resource::allocation_alignment bytes.
+   *
+   * @param upstream_mr The memory_resource from which to allocate blocks for the pool.
+   * @param initial_pool_size Minimum size, in bytes, of the initial pool.
+   * @param maximum_pool_size Maximum size, in bytes, that the pool can grow to. Defaults to all
+   * of the available memory from the upstream resource.
+   */
+  template <class Optional,
+            cuda::std::enable_if_t<cuda::std::is_same_v<cuda::std::remove_cvref_t<Optional>,
+                                                        thrust::optional<std::size_t>>,
+                                   int> = 0>
+  [[deprecated("Use std::optional instead of thrust::optional.")]]  //
+  explicit pool_memory_resource(Upstream& upstream_mr,
+                                std::size_t initial_pool_size,
+                                Optional maximum_pool_size)
+    : pool_memory_resource(cuda::std::addressof(upstream_mr),
+                           initial_pool_size,
+                           maximum_pool_size.value_or(std::nullopt))
+  {
+  }
+
   /**
    * @brief Construct a `pool_memory_resource` and allocate the initial device memory pool using
    * `upstream_mr`.
@@ -210,7 +330,7 @@ class pool_memory_resource final
             cuda::std::enable_if_t<cuda::mr::async_resource<Upstream2>, int> = 0>
   explicit pool_memory_resource(Upstream2& upstream_mr,
                                 std::size_t initial_pool_size,
-                                thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
+                                std::optional<std::size_t> maximum_pool_size = std::nullopt)
     : pool_memory_resource(cuda::std::addressof(upstream_mr), initial_pool_size, maximum_pool_size)
   {
   }
@@ -312,7 +432,7 @@ class pool_memory_resource final
    *
    * @throws logic_error if @p initial_size is larger than @p maximum_size (if set).
    */
-  void initialize_pool(std::size_t initial_size, thrust::optional<std::size_t> maximum_size)
+  void initialize_pool(std::size_t initial_size, std::optional<std::size_t> maximum_size)
   {
     current_pool_size_ = 0;  // try_to_expand will set this if it succeeds
     maximum_pool_size_ = maximum_size;
@@ -376,18 +496,18 @@ class pool_memory_resource final
    * @param stream The stream on which the memory is to be used.
    * @return block_type The allocated block
    */
-  thrust::optional<block_type> block_from_upstream(std::size_t size, cuda_stream_view stream)
+  std::optional<block_type> block_from_upstream(std::size_t size, cuda_stream_view stream)
   {
     RMM_LOG_DEBUG("[A][Stream {}][Upstream {}B]", fmt::ptr(stream.value()), size);
 
     if (size == 0) { return {}; }
 
     try {
       void* ptr = get_upstream()->allocate_async(size, stream);
-      return thrust::optional<block_type>{
+      return std::optional<block_type>{
         *upstream_blocks_.emplace(static_cast<char*>(ptr), size, true).first};
     } catch (std::exception const& e) {
-      return thrust::nullopt;
+      return std::nullopt;
     }
   }
 
@@ -515,7 +635,7 @@ class pool_memory_resource final
  private:
   Upstream* upstream_mr_;  // The "heap" to allocate the pool from
   std::size_t current_pool_size_{};
-  thrust::optional<std::size_t> maximum_pool_size_{};
+  std::optional<std::size_t> maximum_pool_size_{};
 
 #ifdef RMM_POOL_TRACK_ALLOCATIONS
   std::set<block_type, rmm::mr::detail::compare_blocks<block_type>> allocated_blocks_;

python/rmm/_lib/memory_resource.pyx

@@ -25,6 +25,7 @@ from libc.stddef cimport size_t
 from libc.stdint cimport int8_t, int64_t, uintptr_t
 from libcpp cimport bool
 from libcpp.memory cimport make_unique, unique_ptr
+from libcpp.optional cimport optional
 from libcpp.pair cimport pair
 from libcpp.string cimport string
 
@@ -76,19 +77,6 @@ cdef extern from *:
 
 # NOTE: Keep extern declarations in .pyx file as much as possible to avoid
 # leaking dependencies when importing RMM Cython .pxd files
-cdef extern from "thrust/optional.h" namespace "thrust" nogil:
-
-    struct nullopt_t:
-        pass
-
-    cdef nullopt_t nullopt
-
-    cdef cppclass optional[T]:
-        optional()
-        optional(T v)
-
-    cdef optional[T] make_optional[T](T v)
-
 cdef extern from "rmm/mr/device/cuda_memory_resource.hpp" \
         namespace "rmm::mr" nogil:
     cdef cppclass cuda_memory_resource(device_memory_resource):
@@ -322,13 +310,13 @@ cdef class CudaAsyncMemoryResource(DeviceMemoryResource):
         cdef optional[size_t] c_initial_pool_size = (
             optional[size_t]()
             if initial_pool_size is None
-            else optional[size_t](initial_pool_size)
+            else optional[size_t](<size_t> initial_pool_size)
         )
 
         cdef optional[size_t] c_release_threshold = (
             optional[size_t]()
             if release_threshold is None
-            else optional[size_t](release_threshold)
+            else optional[size_t](<size_t> release_threshold)
         )
 
         # If IPC memory handles are not supported, the constructor below will
@@ -382,7 +370,7 @@ cdef class PoolMemoryResource(UpstreamResourceAdaptor):
         c_maximum_pool_size = (
             optional[size_t]() if
             maximum_pool_size is None
-            else make_optional[size_t](maximum_pool_size)
+            else optional[size_t](<size_t> maximum_pool_size)
         )
         self.c_obj.reset(
             new pool_memory_resource[device_memory_resource](