From 81f27099ec53e8ed4d5f2bd9d949cdba13b27ac7 Mon Sep 17 00:00:00 2001 From: Alexandr Guzhva Date: Tue, 9 May 2023 11:14:03 -0700 Subject: [PATCH] enable RAFT under the hood of GPU FAISS (#2840) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2840 Reviewed By: wickedfoo, mdouze Differential Revision: D45054274 fbshipit-source-id: 87889a30e209431908f488035bcf01a4b2bb2eee --- faiss/gpu/StandardGpuResources.cpp | 19 ++++++++++++++++--- faiss/gpu/test/test_gpu_index.py | 2 +- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/faiss/gpu/StandardGpuResources.cpp b/faiss/gpu/StandardGpuResources.cpp index c1be2fbf37..4b35aa4c0a 100644 --- 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 index fcc4f08130..36a1f8a64b 100755 --- 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()