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

include/rmm/mr/device/cuda_async_memory_resource.hpp

@@ -22,13 +22,15 @@
 #include <rmm/mr/device/cuda_async_view_memory_resource.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
 
+#include <cuda/std/type_traits>
 #include <rmm/detail/thrust_namespace.h>
 #include <thrust/optional.h>
 
 #include <cuda_runtime_api.h>
 
 #include <cstddef>
 #include <limits>
+#include <optional>
 
 #if CUDART_VERSION >= 11020  // 11.2 introduced cudaMallocAsync
 #ifndef RMM_DISABLE_CUDA_MALLOC_ASYNC
@@ -66,6 +68,39 @@ class cuda_async_memory_resource final : public device_memory_resource {
     win32     = 0x2,              ///< Allows a Win32 NT handle to be used for exporting. (HANDLE)
     win32_kmt = 0x4  ///< Allows a Win32 KMT handle to be used for exporting. (D3DKMT_HANDLE)
   };
+  /**
+   * @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)
+  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
@@ -85,9 +120,9 @@ 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 = {})
+  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/cuda_async_view_memory_resource.hpp

@@ -22,7 +22,6 @@
 #include <rmm/mr/device/device_memory_resource.hpp>
 
 #include <rmm/detail/thrust_namespace.h>
-#include <thrust/optional.h>
 
 #include <cuda_runtime_api.h>
 

include/rmm/mr/device/pool_memory_resource.hpp

@@ -24,6 +24,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 <rmm/detail/thrust_namespace.h>
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/transform_iterator.h>
@@ -39,6 +40,7 @@
 #include <map>
 #include <mutex>
 #include <numeric>
+#include <optional>
 #include <set>
 #include <thread>
 #include <unordered_map>
@@ -108,6 +110,36 @@ class pool_memory_resource final
  public:
   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
@@ -128,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`.
@@ -155,8 +218,8 @@ 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)
   {
   }
@@ -176,9 +239,36 @@ class pool_memory_resource final
    * @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,
-                                thrust::optional<std::size_t> maximum_pool_size = thrust::nullopt)
+                                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`.
+   *
+   * @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.
+   */
+  explicit pool_memory_resource(Upstream* upstream_mr,
+                                std::size_t initial_pool_size,
+                                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;
@@ -192,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`.
@@ -211,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)
   {
   }
@@ -313,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;
@@ -377,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;
     }
   }
 
@@ -516,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

@@ -76,7 +76,7 @@ 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:
+cdef extern from "optional" namespace "std" nogil:
 
     struct nullopt_t:
         pass