Change index_cpu_to_gpu to throw for indices not implemented on GPU (#…

…3336)

Summary:

Issue: #3269

List of implemented GPU indices: https://github.com/facebookresearch/faiss/wiki/Faiss-on-the-GPU#implemented-indexes

Differential Revision: D55577576

faiss/gpu/GpuCloner.cpp

@@ -214,8 +214,7 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
 
         return res;
     } else {
-        // default: use CPU cloner
-        return Cloner::clone_Index(index);
+        FAISS_THROW_MSG("This index type is not implemented on GPU.");
     }
 }
 
@@ -224,8 +223,6 @@ faiss::Index* index_cpu_to_gpu(
         int device,
         const faiss::Index* index,
         const GpuClonerOptions* options) {
-    auto index_pq = dynamic_cast<const faiss::IndexPQ*>(index);
-    FAISS_THROW_IF_MSG(index_pq, "This index type is not implemented on GPU.");
     GpuClonerOptions defaults;
     ToGpuCloner cl(provider, device, options ? *options : defaults);
     return cl.clone_Index(index);

faiss/gpu/test/test_index_cpu_to_gpu.py

@@ -4,26 +4,44 @@
 
 
 class TestMoveToGpu(unittest.TestCase):
-    def test_index_cpu_to_gpu(self):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.res = faiss.StandardGpuResources()
+
+    def create_index(self, factory_string):
         dimension = 128
         n = 2500
         db_vectors = np.random.random((n, dimension)).astype('float32')
-        code_size = 16
-        res = faiss.StandardGpuResources()
-        index_pq = faiss.IndexPQ(dimension, code_size, 6)
-        index_pq.train(db_vectors)
-        index_pq.add(db_vectors)
-        self.assertRaisesRegex(Exception, ".*not implemented.*",
-                               faiss.index_cpu_to_gpu, res, 0, index_pq)
-
-    def test_index_cpu_to_gpu_does_not_throw_with_index_flat(self):
-        dimension = 128
-        n = 100
-        db_vectors = np.random.random((n, dimension)).astype('float32')
-        res = faiss.StandardGpuResources()
-        index_flat = faiss.IndexFlatL2(dimension)
-        index_flat.add(db_vectors)
+        index = faiss.index_factory(dimension, factory_string)
+        index.train(db_vectors)
+        index.add(db_vectors)
+        return index
+
+    def verify_throws_on_unsupported_index(self, factory_string):
+        idx = self.create_index(factory_string)
         try:
-            faiss.index_cpu_to_gpu(res, 0, index_flat)
-        except Exception:
-            self.fail("index_cpu_to_gpu() threw an unexpected exception.")
+            faiss.index_cpu_to_gpu(self.res, 0, idx)
+        except Exception as e:
+            if "not implemented" not in str(e):
+                self.fail("Expected an exception but no exception was "
+                          "thrown for factory_string: %s." % factory_string)
+
+    def verify_succeeds_on_supported_index(self, factory_string):
+        idx = self.create_index(factory_string)
+        try:
+            faiss.index_cpu_to_gpu(self.res, 0, idx)
+        except Exception as e:
+            self.fail("Unexpected exception thrown factory_string: "
+                      "%s; error message: %s." % (factory_string, str(e)))
+
+    def test_index_cpu_to_gpu_unsupported_indices(self):
+        self.verify_throws_on_unsupported_index("PQ16")
+        self.verify_throws_on_unsupported_index("LSHrt")
+        self.verify_throws_on_unsupported_index("HNSW,PQ16")
+
+    def test_index_cpu_to_gpu_supported_indices(self):
+        self.verify_succeeds_on_supported_index("Flat")
+        self.verify_succeeds_on_supported_index("IVF1,Flat")
+        self.verify_succeeds_on_supported_index("IVF32,SQ8")
+        self.verify_succeeds_on_supported_index("IVF32,PQ8")

faiss/impl/FaissAssert.h

@@ -94,13 +94,6 @@
         }                                              \
     } while (false)
 
-#define FAISS_THROW_IF_MSG(X, MSG)                           \
-    do {                                                     \
-        if (X) {                                             \
-            FAISS_THROW_FMT("Error: '%s' failed: " MSG, #X); \
-        }                                                    \
-    } while (false)
-
 #define FAISS_THROW_IF_NOT_MSG(X, MSG)                       \
     do {                                                     \
         if (!(X)) {                                          \