Limiting workspace memory resource

cpp/include/raft/core/device_container_policy.hpp

@@ -164,10 +164,19 @@ class device_uvector_policy {
  public:
   auto create(raft::resources const& res, size_t n) -> container_type
   {
-    return container_type(n, resource::get_cuda_stream(res), resource::get_workspace_resource(res));
+    if (mr_ == nullptr) {
+      // NB: not using the workspace resource by default!
+      //     The workspace resource is for short-lived temporary allocations.
+      return container_type(n, resource::get_cuda_stream(res));
+    } else {
+      return container_type(n, resource::get_cuda_stream(res), mr_);
+    }
   }
 
-  device_uvector_policy() = default;
+  constexpr device_uvector_policy() = default;
+  constexpr explicit device_uvector_policy(rmm::mr::device_memory_resource* mr) noexcept : mr_(mr)
+  {
+  }
 
   [[nodiscard]] constexpr auto access(container_type& c, size_t n) const noexcept -> reference
   {
@@ -181,6 +190,9 @@ class device_uvector_policy {
 
   [[nodiscard]] auto make_accessor_policy() noexcept { return accessor_policy{}; }
   [[nodiscard]] auto make_accessor_policy() const noexcept { return const_accessor_policy{}; }
+
+ private:
+  rmm::mr::device_memory_resource* mr_{nullptr};
 };
 
 }  // namespace raft

cpp/include/raft/core/device_mdarray.hpp

@@ -112,7 +112,7 @@ auto make_device_mdarray(raft::resources const& handle,
   using mdarray_t = device_mdarray<ElementType, decltype(exts), LayoutPolicy>;
 
   typename mdarray_t::mapping_type layout{exts};
-  typename mdarray_t::container_policy_type policy{};
+  typename mdarray_t::container_policy_type policy{mr};
 
   return mdarray_t{handle, layout, policy};
 }

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)
+                   std::shared_ptr<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,23 @@ 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)
+                   std::shared_ptr<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));
+    if (workspace_resource) {
+      resource::set_workspace_resource(*this, workspace_resource, allocation_limit);
+    }
   }
 
   /** Destroys all held-up resources */
@@ -255,4 +260,4 @@ class stream_syncer {
 
 }  // namespace raft
 
-#endif
+#endif

cpp/include/raft/core/handle.hpp

@@ -32,7 +32,8 @@ namespace raft {
  */
 class handle_t : public raft::device_resources {
  public:
-  handle_t(const handle_t& handle, rmm::mr::device_memory_resource* workspace_resource)
+  handle_t(const handle_t& handle,
+           std::shared_ptr<rmm::mr::device_memory_resource> workspace_resource)
     : device_resources(handle, workspace_resource)
   {
   }
@@ -51,9 +52,9 @@ class handle_t : public raft::device_resources {
    * @param[in] workspace_resource an optional resource used by some functions for allocating
    *            temporary workspaces.
    */
-  handle_t(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)
+  handle_t(rmm::cuda_stream_view stream_view                  = rmm::cuda_stream_per_thread,
+           std::shared_ptr<rmm::cuda_stream_pool> stream_pool = {nullptr},
+           std::shared_ptr<rmm::mr::device_memory_resource> workspace_resource = {nullptr})
     : device_resources{stream_view, stream_pool, workspace_resource}
   {
   }

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

@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <raft/core/logger.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
+
+#include <rmm/mr/device/device_memory_resource.hpp>
+
+#include <mutex>
+#include <set>
+#include <string>
+
+namespace raft::resource::detail {
+
+/**
+ * Warn a user of the calling algorithm if they use the default non-pooled memory allocator,
+ * as it may hurt the performance.
+ *
+ * This helper function is designed to produce the warning once for a given `user_name`.
+ *
+ * @param[in] res
+ * @param[in] user_name the name of the algorithm or any other identification.
+ *
+ */
+inline void warn_non_pool_workspace(resources const& res, std::string user_name)
+{
+  // Detect if the plain cuda memory resource is used for the workspace
+  if (rmm::mr::cuda_memory_resource{}.is_equal(*get_workspace_resource(res)->get_upstream())) {
+    static std::set<std::string> notified_names{};
+    static std::mutex mutex{};
+    std::lock_guard<std::mutex> guard(mutex);
+    auto [it, inserted] = notified_names.insert(std::move(user_name));
+    if (inserted) {
+      RAFT_LOG_WARN(
+        "[%s] the default cuda resource is used for the raft workspace allocations. This may lead "
+        "to a significant slowdown for this algorithm. Consider using the default pool resource "
+        "(`raft::resource::set_workspace_to_pool_resource`) or set your own resource explicitly "
+        "(`raft::resource::set_workspace_resource`).",
+        it->c_str());
+    }
+  }
+}
+
+}  // namespace raft::resource::detail