Unify memory resources

cpp/include/raft/core/device_container_policy.hpp

@@ -1,6 +1,6 @@
 /*
  * SPDX-FileCopyrightText: Copyright (2019) Sandia Corporation
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
  */
 /*
@@ -134,7 +134,7 @@ class device_uvector {
  * @brief A container policy for device mdarray.
  */
 template <typename ElementType>
-class device_uvector_policy {
+class device_container_policy {
  public:
   using element_type   = ElementType;
   using container_type = device_uvector<element_type>;
@@ -153,8 +153,8 @@ class device_uvector_policy {
     return container_type(n, resource::get_cuda_stream(res), mr_);
   }
 
-  constexpr device_uvector_policy() = default;
-  explicit device_uvector_policy(rmm::device_async_resource_ref mr) noexcept : mr_(mr) {}
+  constexpr device_container_policy() = default;
+  explicit device_container_policy(rmm::device_async_resource_ref mr) noexcept : mr_(mr) {}
 
   [[nodiscard]] constexpr auto access(container_type& c, size_t n) const noexcept -> reference
   {
@@ -170,7 +170,7 @@ class device_uvector_policy {
   [[nodiscard]] auto make_accessor_policy() const noexcept { return const_accessor_policy{}; }
 
  private:
-  rmm::device_async_resource_ref mr_{rmm::mr::get_current_device_resource()};
+  rmm::device_async_resource_ref mr_{rmm::mr::get_current_device_resource_ref()};
 };
 
 }  // namespace raft
@@ -189,7 +189,7 @@ template <typename T>
 using device_uvector = detail::fail_container<T>;
 
 template <typename ElementType>
-using device_uvector_policy = detail::fail_container_policy<ElementType>;
+using device_container_policy = detail::fail_container_policy<ElementType>;
 
 }  // namespace raft
 #endif

cpp/include/raft/core/device_coo_matrix.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -29,7 +29,7 @@ using device_coordinate_structure_view = coordinate_structure_view<RowType, ColT
 template <typename RowType,
           typename ColType,
           typename NZType,
-          template <typename T> typename ContainerPolicy = device_uvector_policy>
+          template <typename T> typename ContainerPolicy = device_container_policy>
 using device_coordinate_structure =
   coordinate_structure<RowType, ColType, NZType, true, ContainerPolicy>;
 
@@ -43,7 +43,7 @@ template <typename ElementType,
           typename RowType,
           typename ColType,
           typename NZType,
-          template <typename T> typename ContainerPolicy = device_uvector_policy,
+          template <typename T> typename ContainerPolicy = device_container_policy,
           SparsityType sparsity_type                     = SparsityType::OWNING>
 using device_coo_matrix =
   coo_matrix<ElementType, RowType, ColType, NZType, true, ContainerPolicy, sparsity_type>;
@@ -55,15 +55,15 @@ template <typename ElementType,
           typename RowType,
           typename ColType,
           typename NZType,
-          template <typename T> typename ContainerPolicy = device_uvector_policy>
+          template <typename T> typename ContainerPolicy = device_container_policy>
 using device_sparsity_owning_coo_matrix =
   coo_matrix<ElementType, RowType, ColType, NZType, true, ContainerPolicy>;
 
 template <typename ElementType,
           typename RowType,
           typename ColType,
           typename NZType,
-          template <typename T> typename ContainerPolicy = device_uvector_policy>
+          template <typename T> typename ContainerPolicy = device_container_policy>
 using device_sparsity_preserving_coo_matrix = coo_matrix<ElementType,
                                                          RowType,
                                                          ColType,

cpp/include/raft/core/device_csr_matrix.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -33,7 +33,7 @@ using device_compressed_structure_view =
 template <typename IndptrType,
           typename IndicesType,
           typename NZType,
-          template <typename T> typename ContainerPolicy = device_uvector_policy>
+          template <typename T> typename ContainerPolicy = device_container_policy>
 using device_compressed_structure =
   compressed_structure<IndptrType, IndicesType, NZType, true, ContainerPolicy>;
 
@@ -47,7 +47,7 @@ template <typename ElementType,
           typename IndptrType,
           typename IndicesType,
           typename NZType,
-          template <typename T> typename ContainerPolicy = device_uvector_policy,
+          template <typename T> typename ContainerPolicy = device_container_policy,
           SparsityType sparsity_type                     = SparsityType::OWNING>
 using device_csr_matrix =
   csr_matrix<ElementType, IndptrType, IndicesType, NZType, true, ContainerPolicy, sparsity_type>;
@@ -59,7 +59,7 @@ template <typename ElementType,
           typename IndptrType,
           typename IndicesType,
           typename NZType,
-          template <typename T> typename ContainerPolicy = device_uvector_policy>
+          template <typename T> typename ContainerPolicy = device_container_policy>
 using device_sparsity_owning_csr_matrix =
   csr_matrix<ElementType, IndptrType, IndicesType, NZType, true, ContainerPolicy>;
 
@@ -70,7 +70,7 @@ template <typename ElementType,
           typename IndptrType,
           typename IndicesType,
           typename NZType,
-          template <typename T> typename ContainerPolicy = device_uvector_policy>
+          template <typename T> typename ContainerPolicy = device_container_policy>
 using device_sparsity_preserving_csr_matrix = csr_matrix<ElementType,
                                                          IndptrType,
                                                          IndicesType,

cpp/include/raft/core/device_mdarray.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -26,7 +26,7 @@ namespace raft {
 template <typename ElementType,
           typename Extents,
           typename LayoutPolicy    = layout_c_contiguous,
-          typename ContainerPolicy = device_uvector_policy<ElementType>>
+          typename ContainerPolicy = device_container_policy<ElementType>>
 using device_mdarray =
   mdarray<ElementType, Extents, LayoutPolicy, device_accessor<ContainerPolicy>>;
 
@@ -130,12 +130,30 @@ auto make_device_matrix(raft::resources const& handle, IndexType n_rows, IndexTy
 }
 
 /**
- * @brief Create a device scalar from v.
+ * @brief Create an uninitialized device scalar.
  *
  * @tparam ElementType the data type of the scalar element
  * @tparam IndexType the index type of the extents
  * @param[in] handle raft handle for managing expensive cuda resources
- * @param[in] v scalar to wrap on device
+ * @return raft::device_scalar
+ */
+template <typename ElementType, typename IndexType = std::uint32_t>
+auto make_device_scalar(raft::resources const& handle)
+{
+  scalar_extent<IndexType> extents;
+  using policy_t = typename device_scalar<ElementType, IndexType>::container_policy_type;
+  policy_t policy{};
+  return device_scalar<ElementType, IndexType>{handle, extents, policy};
+}
+
+/**
+ * @brief Create a device scalar from v
+ * (async copy in the resource-provided stream).
+ *
+ * @tparam ElementType the data type of the scalar element
+ * @tparam IndexType the index type of the extents
+ * @param[in] handle raft handle for managing expensive cuda resources
+ * @param[in] v scalar to copy to device
  * @return raft::device_scalar
  */
 template <typename ElementType, typename IndexType = std::uint32_t>

cpp/include/raft/core/host_container_policy.hpp

@@ -1,6 +1,6 @@
 /*
- * SPDX-FileCopyrightText: Copyright (2019) Sandia Corporation
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019 Sandia Corporation
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
  */
 /*
@@ -14,18 +14,23 @@
 
 #include <raft/core/mdspan_types.hpp>
 #include <raft/core/resources.hpp>
-#include <raft/util/integer_utils.hpp>
-
-#include <memory_resource>
+#include <raft/mr/host_memory_resource.hpp>
 
 namespace raft {
 
 /**
- * @brief A container using the std::pmr::memory_resource for allocations.
+ * @brief A container backed by a host-accessible cuda::mr::synchronous_resource.
  *
+ * @tparam T element type
+ * @tparam MR a type satisfying cuda::mr::synchronous_resource_with<cuda::mr::host_accessible>
  */
-template <typename T>
-struct host_container {
+template <typename T, typename MR>
+#ifdef __cpp_concepts
+requires cuda::mr::synchronous_resource_with<MR, cuda::mr::host_accessible>
+#endif
+  struct host_container {
+  static_assert(cuda::mr::synchronous_resource_with<MR, cuda::mr::host_accessible>,
+                "MR must be a host-accessible synchronous resource");
   using value_type = std::remove_cv_t<T>;
   using size_type  = std::size_t;
 
@@ -39,25 +44,26 @@ struct host_container {
   using const_iterator = const_pointer;
 
  private:
-  std::pmr::memory_resource* mr_;
+  MR mr_;
   size_type bytesize_ = 0;
   value_type* data_   = nullptr;
 
  public:
-  host_container(size_type count, std::pmr::memory_resource* mr = nullptr)
-    : mr_(mr == nullptr ? std::pmr::get_default_resource() : mr),
+  host_container(size_type count,
+                 MR mr)  // NB: pass by value, as we expect a resource_ref anyway
+    : mr_(std::move(mr)),
       bytesize_(sizeof(value_type) * count),
-      data_(bytesize_ > 0 ? static_cast<pointer>(mr_->allocate(bytesize_)) : nullptr)
+      data_(bytesize_ > 0 ? static_cast<pointer>(mr_.allocate_sync(bytesize_)) : nullptr)
   {
   }
 
   ~host_container() noexcept
   {
-    if (bytesize_ > 0 && data_ != nullptr) { mr_->deallocate(data_, bytesize_); }
+    if (bytesize_ > 0 && data_ != nullptr) { mr_.deallocate_sync(data_, bytesize_); }
   }
 
   host_container(host_container&& other) noexcept
-    : mr_{std::exchange(other.mr_, nullptr)},
+    : mr_{std::move(other.mr_)},
       bytesize_{std::exchange(other.bytesize_, 0)},
       data_{std::exchange(other.data_, nullptr)}
   {
@@ -103,25 +109,29 @@ struct host_container {
 };
 
 /**
- * @brief A container policy for host mdarray.
+ * @brief Container policy for host mdarray.
+ *
+ * Defaults to raft::mr::get_default_host_resource().
  */
 template <typename ElementType>
 class host_container_policy {
  public:
   using element_type          = ElementType;
-  using container_type        = host_container<element_type>;
+  using container_type        = host_container<element_type, raft::mr::host_resource_ref>;
   using pointer               = typename container_type::pointer;
   using const_pointer         = typename container_type::const_pointer;
   using reference             = typename container_type::reference;
   using const_reference       = typename container_type::const_reference;
   using accessor_policy       = cuda::std::default_accessor<element_type>;
   using const_accessor_policy = cuda::std::default_accessor<element_type const>;
 
- public:
-  auto create(raft::resources const&, size_t n) -> container_type { return container_type(n, mr_); }
+  host_container_policy() = default;
+  explicit host_container_policy(raft::mr::host_resource_ref ref) noexcept : ref_(ref) {}
 
-  constexpr host_container_policy() noexcept = default;
-  explicit host_container_policy(std::pmr::memory_resource* mr) noexcept : mr_(mr) {}
+  auto create(raft::resources const&, size_t n) -> container_type
+  {
+    return container_type(n, ref_);
+  }
 
   [[nodiscard]] constexpr auto access(container_type& c, size_t n) const noexcept -> reference
   {
@@ -137,7 +147,7 @@ class host_container_policy {
   [[nodiscard]] auto make_accessor_policy() const noexcept { return const_accessor_policy{}; }
 
  private:
-  std::pmr::memory_resource* mr_{std::pmr::get_default_resource()};
+  raft::mr::host_resource_ref ref_ = raft::mr::get_default_host_resource();
 };
 
 }  // namespace raft

cpp/include/raft/core/host_mdarray.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -74,7 +74,7 @@ template <typename ElementType,
           typename IndexType    = std::uint32_t,
           typename LayoutPolicy = layout_c_contiguous,
           size_t... Extents>
-auto make_host_mdarray(raft::resources& res, extents<IndexType, Extents...> exts)
+auto make_host_mdarray(raft::resources const& res, extents<IndexType, Extents...> exts)
 {
   using mdarray_t = host_mdarray<ElementType, decltype(exts), LayoutPolicy>;
 
@@ -85,21 +85,21 @@ auto make_host_mdarray(raft::resources& res, extents<IndexType, Extents...> exts
 }
 
 /**
- * @brief Create a host mdarray.
+ * @brief Create a host mdarray with a custom memory resource.
  * @tparam ElementType the data type of the matrix elements
  * @tparam IndexType the index type of the extents
  * @tparam LayoutPolicy policy for strides and layout ordering
  * @param res raft::resources
- * @param mr std::pmr::memory_resource used for allocating the memory for the array
+ * @param mr host memory resource reference used for allocating the memory for the array
  * @param exts dimensionality of the array (series of integers)
- * @return raft::device_mdarray
+ * @return raft::host_mdarray
  */
 template <typename ElementType,
           typename IndexType    = std::uint32_t,
           typename LayoutPolicy = layout_c_contiguous,
           size_t... Extents>
 auto make_host_mdarray(raft::resources const& res,
-                       std::pmr::memory_resource* mr,
+                       raft::mr::host_resource_ref mr,
                        extents<IndexType, Extents...> exts)
 {
   using mdarray_t = host_mdarray<ElementType, decltype(exts), LayoutPolicy>;
@@ -154,7 +154,7 @@ auto make_host_mdarray(extents<IndexType, Extents...> exts)
 template <typename ElementType,
           typename IndexType    = std::uint32_t,
           typename LayoutPolicy = layout_c_contiguous>
-auto make_host_matrix(raft::resources& res, IndexType n_rows, IndexType n_cols)
+auto make_host_matrix(raft::resources const& res, IndexType n_rows, IndexType n_cols)
 {
   return make_host_mdarray<ElementType, IndexType, LayoutPolicy>(
     res, make_extents<IndexType>(n_rows, n_cols));
@@ -181,6 +181,28 @@ auto make_host_matrix(IndexType n_rows, IndexType n_cols)
     make_extents<IndexType>(n_rows, n_cols));
 }
 
+/**
+ * @ingroup host_mdarray_factories
+ * @brief Create an uninitialized host scalar.
+ *
+ * @tparam ElementType the data type of the scalar element
+ * @tparam IndexType the index type of the extents
+ * @param[in] res raft handle for managing expensive resources
+ * @return raft::host_scalar
+ */
+template <typename ElementType, typename IndexType = std::uint32_t>
+auto make_host_scalar(raft::resources const& res)
+{
+  // FIXME(jiamingy): We can optimize this by using std::array as container policy, which
+  // requires some more compile time dispatching. This is enabled in the ref impl but
+  // hasn't been ported here yet.
+  scalar_extent<IndexType> extents;
+  using policy_t = typename host_scalar<ElementType, IndexType>::container_policy_type;
+  policy_t policy;
+  auto scalar = host_scalar<ElementType, IndexType>{res, extents, policy};
+  return scalar;
+}
+
 /**
  * @ingroup host_mdarray_factories
  * @brief Create a host scalar from v.
@@ -192,7 +214,7 @@ auto make_host_matrix(IndexType n_rows, IndexType n_cols)
  * @return raft::host_scalar
  */
 template <typename ElementType, typename IndexType = std::uint32_t>
-auto make_host_scalar(raft::resources& res, ElementType const& v)
+auto make_host_scalar(raft::resources const& res, ElementType const& v)
 {
   // FIXME(jiamingy): We can optimize this by using std::array as container policy, which
   // requires some more compile time dispatching. This is enabled in the ref impl but
@@ -244,7 +266,7 @@ auto make_host_scalar(ElementType const& v)
 template <typename ElementType,
           typename IndexType    = std::uint32_t,
           typename LayoutPolicy = layout_c_contiguous>
-auto make_host_vector(raft::resources& res, IndexType n)
+auto make_host_vector(raft::resources const& res, IndexType n)
 {
   return make_host_mdarray<ElementType, IndexType, LayoutPolicy>(res, make_extents<IndexType>(n));
 }