Limiting workspace memory resource

cpp/include/raft/core/device_resources.hpp

@@ -21,6 +21,7 @@
 
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <string>
 #include <unordered_map>
 #include <utility>
@@ -60,12 +61,12 @@ namespace raft {
 class device_resources : public resources {
  public:
   device_resources(const device_resources& handle,
-                   rmm::mr::device_memory_resource* workspace_resource)
+                   rmm::mr::device_memory_resource* workspace_resource,
+                   std::optional<std::size_t> allocation_limit = std::nullopt)
     : resources{handle}
   {
     // replace the resource factory for the workspace_resources
-    resources::add_resource_factory(
-      std::make_shared<resource::workspace_resource_factory>(workspace_resource));
+    resource::set_workspace_resource(*this, workspace_resource, allocation_limit);
   }
 
   device_resources(const device_resources& handle) : resources{handle} {}
@@ -80,19 +81,21 @@ class device_resources : public resources {
    * @param[in] stream_pool the stream pool used (which has default of nullptr if unspecified)
    * @param[in] workspace_resource an optional resource used by some functions for allocating
    *            temporary workspaces.
+   * @param[in] allocation_limit the total amount of memory in bytes available to the temporary
+   *            workspace resources.
    */
   device_resources(rmm::cuda_stream_view stream_view                  = rmm::cuda_stream_per_thread,
                    std::shared_ptr<rmm::cuda_stream_pool> stream_pool = {nullptr},
-                   rmm::mr::device_memory_resource* workspace_resource = nullptr)
+                   rmm::mr::device_memory_resource* workspace_resource = nullptr,
+                   std::optional<std::size_t> allocation_limit         = std::nullopt)
     : resources{}
   {
     resources::add_resource_factory(std::make_shared<resource::device_id_resource_factory>());
     resources::add_resource_factory(
       std::make_shared<resource::cuda_stream_resource_factory>(stream_view));
     resources::add_resource_factory(
       std::make_shared<resource::cuda_stream_pool_resource_factory>(stream_pool));
-    resources::add_resource_factory(
-      std::make_shared<resource::workspace_resource_factory>(workspace_resource));
+    resource::set_workspace_resource(*this, workspace_resource, allocation_limit);
   }
 
   /** Destroys all held-up resources */
@@ -255,4 +258,4 @@ class stream_syncer {
 
 }  // namespace raft
 
-#endif
+#endif

cpp/include/raft/core/resource/device_memory_resource.hpp

@@ -15,23 +15,48 @@
  */
 #pragma once
 
+#include <raft/core/operators.hpp>
 #include <raft/core/resource/resource_types.hpp>
 #include <raft/core/resources.hpp>
+#include <raft/util/cudart_utils.hpp>
+
 #include <rmm/mr/device/device_memory_resource.hpp>
+#include <rmm/mr/device/limiting_resource_adaptor.hpp>
+#include <rmm/mr/device/pool_memory_resource.hpp>
+
+#include <cstddef>
+#include <optional>
 
 namespace raft::resource {
-class device_memory_resource : public resource {
+class limited_memory_resource : public resource {
  public:
-  device_memory_resource(rmm::mr::device_memory_resource* mr_ = nullptr) : mr(mr_)
+  limited_memory_resource(std::shared_ptr<rmm::mr::device_memory_resource> mr,
+                          std::size_t allocation_limit,
+                          std::optional<std::size_t> alignment)
+    : upstream_(mr), mr_(make_adaptor(mr, allocation_limit, alignment))
   {
-    if (mr_ == nullptr) { mr = rmm::mr::get_current_device_resource(); }
   }
-  void* get_resource() override { return mr; }
 
-  ~device_memory_resource() override {}
+  auto get_resource() -> void* override { return &mr_; }
+
+  ~limited_memory_resource() override = default;
 
  private:
-  rmm::mr::device_memory_resource* mr;
+  std::shared_ptr<rmm::mr::device_memory_resource> upstream_;
+  rmm::mr::limiting_resource_adaptor<rmm::mr::device_memory_resource> mr_;
+
+  static inline auto make_adaptor(std::shared_ptr<rmm::mr::device_memory_resource> upstream,
+                                  std::size_t limit,
+                                  std::optional<std::size_t> alignment)
+    -> rmm::mr::limiting_resource_adaptor<rmm::mr::device_memory_resource>
+  {
+    auto p = upstream.get();
+    if (alignment.has_value()) {
+      return rmm::mr::limiting_resource_adaptor(p, limit, alignment.value());
+    } else {
+      return rmm::mr::limiting_resource_adaptor(p, limit);
+    }
+  }
 };
 
 /**
@@ -40,36 +65,109 @@ class device_memory_resource : public resource {
  */
 class workspace_resource_factory : public resource_factory {
  public:
-  workspace_resource_factory(rmm::mr::device_memory_resource* mr_ = nullptr) : mr(mr_) {}
-  resource_type get_resource_type() override { return resource_type::WORKSPACE_RESOURCE; }
-  resource* make_resource() override { return new device_memory_resource(mr); }
+  explicit workspace_resource_factory(
+    std::shared_ptr<rmm::mr::device_memory_resource> mr = {nullptr},
+    std::optional<std::size_t> allocation_limit         = std::nullopt,
+    std::optional<std::size_t> alignment                = std::nullopt)
+    : allocation_limit_(allocation_limit.value_or(default_allocation_limit())),
+      alignment_(alignment),
+      mr_(mr ? mr : default_memory_resource(allocation_limit_))
+  {
+  }
+
+  auto get_resource_type() -> resource_type override { return resource_type::WORKSPACE_RESOURCE; }
+  auto make_resource() -> resource* override
+  {
+    return new limited_memory_resource(mr_, allocation_limit_, alignment_);
+  }
 
  private:
-  rmm::mr::device_memory_resource* mr;
+  std::size_t allocation_limit_;
+  std::optional<std::size_t> alignment_;
+  std::shared_ptr<rmm::mr::device_memory_resource> mr_;
+
+  // Create a pool memory resource by default
+  static inline auto default_memory_resource(std::size_t limit)
+    -> std::shared_ptr<rmm::mr::device_memory_resource>
+  {
+    constexpr std::size_t kOneGb = 1024lu * 1024lu * 1024lu;
+    auto min_size                = std::min<std::size_t>(kOneGb, limit / 2);
+    auto max_size                = limit * 3lu / 2lu;
+    auto upstream                = rmm::mr::get_current_device_resource();
+    return std::make_shared<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>>(
+      upstream, min_size, max_size);
+  }
+
+  // Allow a fraction of available memory by default.
+  static inline auto default_allocation_limit() -> std::size_t
+  {
+    std::size_t free_size{};
+    std::size_t total_size{};
+    RAFT_CUDA_TRY(cudaMemGetInfo(&free_size, &total_size));
+    return free_size / 2;
+  }
 };
 
 /**
  * Load a temp workspace resource from a resources instance (and populate it on the res
  * if needed).
+ *
  * @param res raft resources object for managing resources
  * @return device memory resource object
  */
-inline rmm::mr::device_memory_resource* get_workspace_resource(resources const& res)
+inline auto get_workspace_resource(resources const& res)
+  -> rmm::mr::limiting_resource_adaptor<rmm::mr::device_memory_resource>*
 {
   if (!res.has_resource_factory(resource_type::WORKSPACE_RESOURCE)) {
     res.add_resource_factory(std::make_shared<workspace_resource_factory>());
   }
-  return res.get_resource<rmm::mr::device_memory_resource>(resource_type::WORKSPACE_RESOURCE);
+  return res.get_resource<rmm::mr::limiting_resource_adaptor<rmm::mr::device_memory_resource>>(
+    resource_type::WORKSPACE_RESOURCE);
+};
+
+/**
+ * Set a temporary workspace resource on a resources instance.
+ *
+ * @param res raft resources object for managing resources
+ * @param mr an optional RMM device_memory_resource
+ * @param allocation_limit
+ *   the total amount of memory in bytes available to the temporary workspace resources.
+ * @param alignment optional alignment requirements passed to RMM allocations
+ *
+ */
+inline void set_workspace_resource(resources const& res,
+                                   std::shared_ptr<rmm::mr::device_memory_resource> mr = {nullptr},
+                                   std::optional<std::size_t> allocation_limit = std::nullopt,
+                                   std::optional<std::size_t> alignment        = std::nullopt)
+{
+  res.add_resource_factory(
+    std::make_shared<workspace_resource_factory>(mr, allocation_limit, alignment));
 };
 
 /**
- * Set a temp workspace resource on a resources instance.
+ * Set a temporary workspace resource on a resources instance.
  *
  * @param res raft resources object for managing resources
- * @param mr a valid rmm device_memory_resource
+ * @param mr an optional RMM device_memory_resource;
+ *   note, the ownership of the object is not transferred with this raw pointer interface.
+ * @param allocation_limit
+ *   the total amount of memory in bytes available to the temporary workspace resources.
+ * @param alignment optional alignment requirements passed to RMM allocations
+ *
  */
-inline void set_workspace_resource(resources const& res, rmm::mr::device_memory_resource* mr)
+inline void set_workspace_resource(resources const& res,
+                                   rmm::mr::device_memory_resource* mr,
+                                   std::optional<std::size_t> allocation_limit = std::nullopt,
+                                   std::optional<std::size_t> alignment        = std::nullopt)
 {
-  res.add_resource_factory(std::make_shared<workspace_resource_factory>(mr));
+  // NB: to preserve the semantics of passing memory resource without transferring the ownership,
+  //     we create a shared pointer with a dummy deleter (void_op).
+  set_workspace_resource(res,
+                         mr != nullptr
+                           ? std::shared_ptr<rmm::mr::device_memory_resource>{mr, void_op{}}
+                           : std::shared_ptr<rmm::mr::device_memory_resource>{nullptr},
+                         allocation_limit,
+                         alignment);
 };
-}  // namespace raft::resource
+
+}  // namespace raft::resource

cpp/include/raft/core/resources.hpp

@@ -94,6 +94,11 @@ class resources {
     RAFT_EXPECTS(rtype != resource::resource_type::LAST_KEY,
                  "LAST_KEY is a placeholder and not a valid resource factory type.");
     factories_.at(rtype) = std::make_pair(rtype, factory);
+    // Clear the corresponding resource, so that on next `get_resource` the new factory is used
+    if (resources_.at(rtype).first != resource::resource_type::LAST_KEY) {
+      resources_.at(rtype) = std::make_pair(resource::resource_type::LAST_KEY,
+                                            std::make_shared<resource::empty_resource>());
+    }
   }
 
   /**
@@ -128,4 +133,4 @@ class resources {
   mutable std::vector<pair_res_factory> factories_;
   mutable std::vector<pair_resource> resources_;
 };
-}  // namespace raft
+}  // namespace raft

cpp/include/raft/neighbors/detail/ivf_pq_build.cuh

@@ -442,9 +442,6 @@ void train_per_subset(raft::device_resources const& handle,
                  index.pq_len(),
                  stream);
 
-    // clone the handle and attached the device memory resource to it
-    const device_resources new_handle(handle, device_memory);
-
     // train PQ codebook for this subspace
     auto sub_trainset_view =
       raft::make_device_matrix_view<const float, IdxT>(sub_trainset.data(), n_rows, index.pq_len());
@@ -458,7 +455,7 @@ void train_per_subset(raft::device_resources const& handle,
     raft::cluster::kmeans_balanced_params kmeans_params;
     kmeans_params.n_iters = kmeans_n_iters;
     kmeans_params.metric  = raft::distance::DistanceType::L2Expanded;
-    raft::cluster::kmeans_balanced::helpers::build_clusters(new_handle,
+    raft::cluster::kmeans_balanced::helpers::build_clusters(handle,
                                                             kmeans_params,
                                                             sub_trainset_view,
                                                             centers_tmp_view,
@@ -523,9 +520,6 @@ void train_per_cluster(raft::device_resources const& handle,
                      indices + cluster_offsets[l],
                      device_memory);
 
-    // clone the handle and attached the device memory resource to it
-    const device_resources new_handle(handle, device_memory);
-
     // limit the cluster size to bound the training time.
     // [sic] we interpret the data as pq_len-dimensional
     size_t big_enough     = 256ul * std::max<size_t>(index.pq_book_size(), index.pq_dim());
@@ -546,7 +540,7 @@ void train_per_cluster(raft::device_resources const& handle,
     raft::cluster::kmeans_balanced_params kmeans_params;
     kmeans_params.n_iters = kmeans_n_iters;
     kmeans_params.metric  = raft::distance::DistanceType::L2Expanded;
-    raft::cluster::kmeans_balanced::helpers::build_clusters(new_handle,
+    raft::cluster::kmeans_balanced::helpers::build_clusters(handle,
                                                             kmeans_params,
                                                             rot_vectors_view,
                                                             centers_tmp_view,

cpp/include/raft/neighbors/detail/ivf_pq_search.cuh

@@ -1262,10 +1262,10 @@ void ivfpq_search_worker(raft::device_resources const& handle,
                          IdxT* neighbors,                    // [n_queries, topK]
                          float* distances,                   // [n_queries, topK]
                          float scaling_factor,
-                         double preferred_shmem_carveout,
-                         rmm::mr::device_memory_resource* mr)
+                         double preferred_shmem_carveout)
 {
   auto stream = handle.get_stream();
+  auto mr     = handle.get_workspace_resource();
 
   bool manage_local_topk = is_local_topk_feasible(topK, n_probes, n_queries);
   auto topk_len          = manage_local_topk ? n_probes * topK : max_samples;
@@ -1554,8 +1554,7 @@ inline void search(raft::device_resources const& handle,
                    uint32_t n_queries,
                    uint32_t k,
                    IdxT* neighbors,
-                   float* distances,
-                   rmm::mr::device_memory_resource* mr = nullptr)
+                   float* distances)
 {
   static_assert(std::is_same_v<T, float> || std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t>,
                 "Unsupported element type.");
@@ -1601,11 +1600,7 @@ inline void search(raft::device_resources const& handle,
     max_samples = ms;
   }
 
-  auto pool_guard = raft::get_pool_memory_resource(mr, n_queries * n_probes * k * 16);
-  if (pool_guard) {
-    RAFT_LOG_DEBUG("ivf_pq::search: using pool memory resource with initial size %zu bytes",
-                   pool_guard->pool_size());
-  }
+  auto mr = handle.get_workspace_resource();
 
   // Maximum number of query vectors to search at the same time.
   const auto max_queries = std::min<uint32_t>(std::max<uint32_t>(n_queries, 1), 4096);
@@ -1669,8 +1664,7 @@ inline void search(raft::device_resources const& handle,
                       neighbors + uint64_t(k) * (offset_q + offset_b),
                       distances + uint64_t(k) * (offset_q + offset_b),
                       utils::config<T>::kDivisor / utils::config<float>::kDivisor,
-                      params.preferred_shmem_carveout,
-                      mr);
+                      params.preferred_shmem_carveout);
     }
   }
 }

cpp/include/raft/neighbors/ivf_pq.cuh

@@ -23,10 +23,13 @@
 
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/device_resources.hpp>
+#include <raft/core/operators.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 
+#include <memory>
+
 namespace raft::neighbors::ivf_pq {
 
 /**
@@ -182,8 +185,7 @@ void search(raft::device_resources const& handle,
                         static_cast<std::uint32_t>(queries.extent(0)),
                         k,
                         neighbors.data_handle(),
-                        distances.data_handle(),
-                        handle.get_workspace_resource());
+                        distances.data_handle());
 }
 
 /** @} */  // end group ivf_pq
@@ -349,7 +351,8 @@ void search(raft::device_resources const& handle,
             float* distances,
             rmm::mr::device_memory_resource* mr = nullptr)
 {
-  return detail::search(handle, params, idx, queries, n_queries, k, neighbors, distances, mr);
+  if (mr != nullptr) { resource::set_workspace_resource(handle, mr); }
+  return detail::search(handle, params, idx, queries, n_queries, k, neighbors, distances);
 }
 
 }  // namespace raft::neighbors::ivf_pq