From 09914a1255d499b482231c2d6d27e3fd18a8ebea Mon Sep 17 00:00:00 2001 From: Gergely Szilvasy Date: Wed, 12 Jun 2024 21:25:44 +0000 Subject: [PATCH 1/2] add use_raft to knn_gpu (torch) --- contrib/torch_utils.py | 5 +++-- faiss/gpu/test/torch_test_contrib_gpu.py | 17 ++++++++++------- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/contrib/torch_utils.py b/contrib/torch_utils.py index e371932c9f..20bc15ba8b 100644 --- a/contrib/torch_utils.py +++ b/contrib/torch_utils.py @@ -493,7 +493,7 @@ def torch_replacement_sa_decode(self, codes, x=None): handle_torch_Index(the_class) # allows torch tensor usage with bfKnn -def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRIC_L2, device=-1): +def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRIC_L2, device=-1, use_raft=False): if type(xb) is np.ndarray: # Forward to faiss __init__.py base method return faiss.knn_gpu_numpy(res, xq, xb, k, D, I, metric, device) @@ -574,9 +574,10 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRI args.outIndices = I_ptr args.outIndicesType = I_type args.device = device + args.use_raft = use_raft with using_stream(res): - faiss.bfKnn(res, args) + faiss.bfKnn(res, args) return D, I diff --git a/faiss/gpu/test/torch_test_contrib_gpu.py b/faiss/gpu/test/torch_test_contrib_gpu.py index 0c949c29f2..85be8d51ee 100644 --- a/faiss/gpu/test/torch_test_contrib_gpu.py +++ b/faiss/gpu/test/torch_test_contrib_gpu.py @@ -249,7 +249,7 @@ def test_sa_encode_decode(self): return class TestTorchUtilsKnnGpu(unittest.TestCase): - def test_knn_gpu(self): + def test_knn_gpu(self, use_raft=False): torch.manual_seed(10) d = 32 nb = 1024 @@ -286,7 +286,7 @@ def test_knn_gpu(self): else: xb_c = xb_np - D, I = faiss.knn_gpu(res, xq_c, xb_c, k) + D, I = faiss.knn_gpu(res, xq_c, xb_c, k, use_raft=use_raft) self.assertTrue(torch.equal(torch.from_numpy(I), gt_I)) self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1e-4) @@ -312,7 +312,7 @@ def test_knn_gpu(self): xb_c = to_column_major_torch(xb) assert not xb_c.is_contiguous() - D, I = faiss.knn_gpu(res, xq_c, xb_c, k) + D, I = faiss.knn_gpu(res, xq_c, xb_c, k, use_raft=use_raft) self.assertTrue(torch.equal(I.cpu(), gt_I)) self.assertLess((D.cpu() - gt_D).abs().max(), 1e-4) @@ -320,7 +320,7 @@ def test_knn_gpu(self): # test on subset try: # This internally uses the current pytorch stream - D, I = faiss.knn_gpu(res, xq_c[6:8], xb_c, k) + D, I = faiss.knn_gpu(res, xq_c[6:8], xb_c, k, use_raft=use_raft) except TypeError: if not xq_row_major: # then it is expected @@ -331,7 +331,10 @@ def test_knn_gpu(self): self.assertTrue(torch.equal(I.cpu(), gt_I[6:8])) self.assertLess((D.cpu() - gt_D[6:8]).abs().max(), 1e-4) - def test_knn_gpu_datatypes(self): + def test_knn_gpu_raft(self): + self.test_knn_gpu(use_raft=True) + + def test_knn_gpu_datatypes(self, use_raft=False): torch.manual_seed(10) d = 10 nb = 1024 @@ -354,7 +357,7 @@ def test_knn_gpu_datatypes(self): D = torch.zeros(nq, k, device=xb_c.device, dtype=torch.float32) I = torch.zeros(nq, k, device=xb_c.device, dtype=torch.int32) - faiss.knn_gpu(res, xq_c, xb_c, k, D, I) + faiss.knn_gpu(res, xq_c, xb_c, k, D, I, use_raft=use_raft) self.assertTrue(torch.equal(I.long().cpu(), gt_I)) self.assertLess((D.float().cpu() - gt_D).abs().max(), 1.5e-3) @@ -366,7 +369,7 @@ def test_knn_gpu_datatypes(self): xb_c = xb.half().numpy() xq_c = xq.half().numpy() - faiss.knn_gpu(res, xq_c, xb_c, k, D, I) + faiss.knn_gpu(res, xq_c, xb_c, k, D, I, use_raft=use_raft) self.assertTrue(torch.equal(torch.from_numpy(I).long(), gt_I)) self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1.5e-3) From 449ec4c128fb0a086c85ca889b2454b507abe69a Mon Sep 17 00:00:00 2001 From: Gergely Szilvasy Date: Thu, 13 Jun 2024 00:28:18 +0000 Subject: [PATCH 2/2] skip test unless RAFT --- contrib/torch_utils.py | 3 ++- faiss/gpu/test/torch_test_contrib_gpu.py | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/contrib/torch_utils.py b/contrib/torch_utils.py index 20bc15ba8b..18f136e914 100644 --- a/contrib/torch_utils.py +++ b/contrib/torch_utils.py @@ -492,6 +492,7 @@ def torch_replacement_sa_decode(self, codes, x=None): if issubclass(the_class, faiss.Index): handle_torch_Index(the_class) + # allows torch tensor usage with bfKnn def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRIC_L2, device=-1, use_raft=False): if type(xb) is np.ndarray: @@ -577,7 +578,7 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRI args.use_raft = use_raft with using_stream(res): - faiss.bfKnn(res, args) + faiss.bfKnn(res, args) return D, I diff --git a/faiss/gpu/test/torch_test_contrib_gpu.py b/faiss/gpu/test/torch_test_contrib_gpu.py index 85be8d51ee..f7444337f1 100644 --- a/faiss/gpu/test/torch_test_contrib_gpu.py +++ b/faiss/gpu/test/torch_test_contrib_gpu.py @@ -331,6 +331,9 @@ def test_knn_gpu(self, use_raft=False): self.assertTrue(torch.equal(I.cpu(), gt_I[6:8])) self.assertLess((D.cpu() - gt_D[6:8]).abs().max(), 1e-4) + @unittest.skipUnless( + "RAFT" in faiss.get_compile_options(), + "only if RAFT is compiled in") def test_knn_gpu_raft(self): self.test_knn_gpu(use_raft=True)