faiss-1.14-cuvs-26.06.diff

diff --git a/faiss/gpu/GpuResources.h b/faiss/gpu/GpuResources.h
index 61d9d4dbe..477d27cb2 100644
--- a/faiss/gpu/GpuResources.h
+++ b/faiss/gpu/GpuResources.h
@@ -33,7 +33,8 @@

 #if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <cuda/memory_resource>
+#include <optional>
 #endif

 namespace faiss {
@@ -163,7 +164,7 @@ struct AllocRequest : public AllocInfo {
     size_t size = 0;

 #if defined USE_NVIDIA_CUVS
-    rmm::mr::device_memory_resource* mr = nullptr;
+    std::optional<cuda::mr::any_resource<cuda::mr::device_accessible>> mr;
 #endif
 };

diff --git a/faiss/gpu/StandardGpuResources.cpp b/faiss/gpu/StandardGpuResources.cpp
index 548618262..3be071550 100644
--- a/faiss/gpu/StandardGpuResources.cpp
+++ b/faiss/gpu/StandardGpuResources.cpp
@@ -92,8 +92,8 @@ std::string allocsToString(const std::unordered_map<void*, AllocRequest>& map) {
 StandardGpuResourcesImpl::StandardGpuResourcesImpl()
         :
 #if defined USE_NVIDIA_CUVS
-          mmr_(new rmm::mr::managed_memory_resource),
-          pmr_(new rmm::mr::pinned_host_memory_resource),
+          mmr_{},
+          pmr_{},
 #endif
           pinnedMemAlloc_(nullptr),
           pinnedMemAllocSize_(0),
@@ -164,7 +164,7 @@ StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {

     if (pinnedMemAlloc_) {
 #if defined USE_NVIDIA_CUVS
-        pmr_->deallocate_sync(pinnedMemAlloc_, pinnedMemAllocSize_);
+        pmr_.deallocate_sync(pinnedMemAlloc_, pinnedMemAllocSize_);
 #else
         auto err = cudaFreeHost(pinnedMemAlloc_);
         FAISS_ASSERT_FMT(
@@ -350,7 +350,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
         // pinned memory allocation
         if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
             try {
-                pinnedMemAlloc_ = pmr_->allocate_sync(pinnedMemSize_);
+                pinnedMemAlloc_ = pmr_.allocate_sync(pinnedMemSize_);
             } catch (const std::bad_alloc& rmm_ex) {
                 FAISS_THROW_MSG("CUDA memory allocation error");
             }
@@ -546,10 +546,10 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
     } else if (adjReq.space == MemorySpace::Device) {
 #if defined USE_NVIDIA_CUVS
         try {
-            rmm::mr::device_memory_resource* current_mr =
-                    rmm::mr::get_per_device_resource(
+            auto current_mr =
+                    rmm::mr::get_per_device_resource(
                             rmm::cuda_device_id{adjReq.device});
-            p = current_mr->allocate(adjReq.stream, adjReq.size);
+            p = current_mr.allocate(adjReq.stream, adjReq.size);
             adjReq.mr = current_mr;
         } catch (const std::bad_alloc& rmm_ex) {
             FAISS_THROW_MSG("CUDA memory allocation error");
@@ -584,8 +584,8 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
             // TODO: change this to use the current device resource once RMM has
             // a way to retrieve a "guaranteed" managed memory resource for a
             // device.
-            p = mmr_->allocate(adjReq.stream, adjReq.size);
-            adjReq.mr = mmr_.get();
+            p = mmr_.allocate(adjReq.stream, adjReq.size);
+            adjReq.mr = mmr_;
         } catch (const std::bad_alloc& rmm_ex) {
             FAISS_THROW_MSG("CUDA memory allocation error");
         }
diff --git a/faiss/gpu/StandardGpuResources.h b/faiss/gpu/StandardGpuResources.h
index 3ba606606..4c1df7212 100644
--- a/faiss/gpu/StandardGpuResources.h
+++ b/faiss/gpu/StandardGpuResources.h
@@ -25,7 +25,7 @@

 #if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/mr/managed_memory_resource.hpp>
 #include <rmm/mr/pinned_host_memory_resource.hpp>
 #endif

@@ -171,10 +171,10 @@ class StandardGpuResourcesImpl : public GpuResources {
      */

     // managed_memory_resource
-    std::unique_ptr<rmm::mr::device_memory_resource> mmr_;
+    rmm::mr::managed_memory_resource mmr_;

     // pinned_host_memory_resource
-    std::unique_ptr<rmm::mr::pinned_host_memory_resource> pmr_;
+    rmm::mr::pinned_host_memory_resource pmr_;
 #endif

     /// Pinned memory allocation for use with this GPU