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

Author: Alexandr Guzhva

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);
 

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()