enable RAFT under the hood of GPU FAISS (facebookresearch#2840)

Summary: Pull Request resolved: facebookresearch#2840 Reviewed By: wickedfoo, mdouze Differential Revision: D45054274 fbshipit-source-id: 87889a30e209431908f488035bcf01a4b2bb2eee
blevesearch · May 9, 2023 · 81f2709 · 81f2709
1 parent 3219e3d
commit 81f2709
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/faiss/gpu/StandardGpuResources.cpp b/faiss/gpu/StandardGpuResources.cpp
@@ -312,7 +312,12 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
         // If this is the first device that we're initializing, create our
         // pinned memory allocation
         if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
-            pinnedMemAlloc_ = pmr->allocate(pinnedMemSize_);
+            try {
+                pinnedMemAlloc_ = pmr->allocate(pinnedMemSize_);
+            } catch (const std::bad_alloc& rmm_ex) {
+                FAISS_THROW_MSG("CUDA memory allocation error");
+            }
+
             pinnedMemAllocSize_ = pinnedMemSize_;
         }
 #else
@@ -490,7 +495,11 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
 
     } else if (adjReq.space == MemorySpace::Device) {
 #if defined USE_NVIDIA_RAFT
-        p = cmr->allocate(adjReq.size, adjReq.stream);
+        try {
+            p = cmr->allocate(adjReq.size, adjReq.stream);
+        } catch (const std::bad_alloc& rmm_ex) {
+            FAISS_THROW_MSG("CUDA memory allocation error");
+        }
 #else
         auto err = cudaMalloc(&p, adjReq.size);
 
@@ -516,7 +525,11 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
 #endif
     } else if (adjReq.space == MemorySpace::Unified) {
 #if defined USE_NVIDIA_RAFT
-        p = mmr->allocate(adjReq.size, adjReq.stream);
+        try {
+            p = mmr->allocate(adjReq.size, adjReq.stream);
+        } catch (const std::bad_alloc& rmm_ex) {
+            FAISS_THROW_MSG("CUDA memory allocation error");
+        }
 #else
         auto err = cudaMallocManaged(&p, adjReq.size);
 

diff --git a/faiss/gpu/test/test_gpu_index.py b/faiss/gpu/test/test_gpu_index.py
@@ -189,7 +189,7 @@ def test_ivfsq_cpu_coarse(self):
 
         self.assertGreaterEqual(knn_intersection_measure(i_c, i_g), 0.9)
 
-        self.assertTrue(np.allclose(d_g, d_c, rtol=5e-5, atol=5e-5))
+        self.assertTrue(np.allclose(d_g, d_c, rtol=2e-4, atol=2e-4))
 
     def test_ivfpq_cpu_coarse(self):
         res = faiss.StandardGpuResources()