Merge pull request #361 from rapidsai/branch-24.10

Forward-merge branch-24.10 into branch-24.12

cpp/include/cuvs/core/c_api.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <cuda_runtime.h>
+#include <stdbool.h>
 #include <stdint.h>
 
 #ifdef __cplusplus
@@ -138,10 +139,12 @@ cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes);
  * available memory
  * @param[in] max_pool_size_percent The maximum pool size as a percentage of the total
  * available memory
+ * @param[in] managed Whether to use a managed memory resource as upstream resource or not
  * @return cuvsError_t
  */
 cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent,
-                                            int max_pool_size_percent);
+                                            int max_pool_size_percent,
+                                            bool managed);
 /**
  * @brief Resets the memory resource to use the default memory resource (cuda_memory_resource)
  * @return cuvsError_t

cpp/src/core/c_api.cpp

@@ -21,6 +21,9 @@
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/mr/device/device_memory_resource.hpp>
+#include <rmm/mr/device/managed_memory_resource.hpp>
+#include <rmm/mr/device/owning_wrapper.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 #include <rmm/mr/device/pool_memory_resource.hpp>
 #include <thread>
@@ -83,10 +86,14 @@ extern "C" cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes)
   });
 }
 
-thread_local std::unique_ptr<rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>> pool_mr;
+thread_local std::shared_ptr<
+  rmm::mr::owning_wrapper<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>,
+                          rmm::mr::device_memory_resource>>
+  pool_mr;
 
 extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent,
-                                                       int max_pool_size_percent)
+                                                       int max_pool_size_percent,
+                                                       bool managed)
 {
   return cuvs::core::translate_exceptions([=] {
     // Upstream memory resource needs to be a cuda_memory_resource
@@ -95,10 +102,22 @@ extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_per
     if (cuda_mr_casted == nullptr) {
       throw std::runtime_error("Current memory resource is not a cuda_memory_resource");
     }
+
     auto initial_size = rmm::percent_of_free_device_memory(initial_pool_size_percent);
     auto max_size     = rmm::percent_of_free_device_memory(max_pool_size_percent);
-    pool_mr = std::make_unique<rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>>(
-      cuda_mr_casted, initial_size, max_size);
+
+    auto mr = std::shared_ptr<rmm::mr::device_memory_resource>();
+    if (managed) {
+      mr = std::static_pointer_cast<rmm::mr::device_memory_resource>(
+        std::make_shared<rmm::mr::managed_memory_resource>());
+    } else {
+      mr = std::static_pointer_cast<rmm::mr::device_memory_resource>(
+        std::make_shared<rmm::mr::cuda_memory_resource>());
+    }
+
+    pool_mr =
+      rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(mr, initial_size, max_size);
+
     rmm::mr::set_current_device_resource(pool_mr.get());
   });
 }

cpp/test/core/c_api.c

@@ -33,34 +33,49 @@ int main()
 
   // Allocate memory
   void* ptr;
-  size_t bytes            = 1024;
-  cuvsError_t alloc_error = cuvsRMMAlloc(res, &ptr, bytes);
-  if (alloc_error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+  size_t bytes      = 1024;
+  cuvsError_t error = cuvsRMMAlloc(res, &ptr, bytes);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
 
   // Free memory
-  cuvsError_t free_error = cuvsRMMFree(res, ptr, bytes);
-  if (free_error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+  error = cuvsRMMFree(res, ptr, bytes);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
 
   // Enable pool memory resource
-  cuvsError_t pool_error = cuvsRMMPoolMemoryResourceEnable(10, 100);
-  if (pool_error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+  error = cuvsRMMPoolMemoryResourceEnable(10, 100, false);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
 
   // Allocate memory again
-  void* ptr2;
-  cuvsError_t alloc_error_pool = cuvsRMMAlloc(res, &ptr2, 1024);
-  if (alloc_error_pool == CUVS_ERROR) { exit(EXIT_FAILURE); }
+  error = cuvsRMMAlloc(res, &ptr, 1024);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
 
   // Free memory
-  cuvsError_t free_error_pool = cuvsRMMFree(res, ptr2, 1024);
-  if (free_error_pool == CUVS_ERROR) { exit(EXIT_FAILURE); }
+  error = cuvsRMMFree(res, ptr, 1024);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
 
   // Reset pool memory resource
-  cuvsError_t reset_error = cuvsRMMMemoryResourceReset();
-  if (reset_error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+  error = cuvsRMMMemoryResourceReset();
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+
+  // Enable pool memory resource (managed)
+  error = cuvsRMMPoolMemoryResourceEnable(10, 100, true);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+
+  // Allocate memory again
+  error = cuvsRMMAlloc(res, &ptr, 1024);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+
+  // Free memory
+  error = cuvsRMMFree(res, ptr, 1024);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+
+  // Reset pool memory resource
+  error = cuvsRMMMemoryResourceReset();
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
 
   // Destroy resources
-  cuvsError_t destroy_error = cuvsResourcesDestroy(res);
-  if (destroy_error == CUVS_ERROR) { exit(EXIT_FAILURE); }
+  error = cuvsResourcesDestroy(res);
+  if (error == CUVS_ERROR) { exit(EXIT_FAILURE); }
 
   return 0;
 }