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

Reviewed By: junjieqi

Differential Revision: D55577576

Author: ramilbakhshyiev

faiss/gpu/GpuCloner.cpp

@@ -205,6 +205,7 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
         config.usePrecomputedTables = usePrecomputed;
         config.use_raft = use_raft;
         config.interleavedLayout = use_raft;
+        config.enableCpuFallback = enableCpuFallback;
 
         GpuIndexIVFPQ* res = new GpuIndexIVFPQ(provider, ipq, config);
 
@@ -214,8 +215,11 @@ Index* ToGpuCloner::clone_Index(const Index* index) {
 
         return res;
     } else {
-        // default: use CPU cloner
-        return Cloner::clone_Index(index);
+        // use CPU cloner if CPU fallback is enabled
+        if (enableCpuFallback) {
+            return Cloner::clone_Index(index);
+        }
+        FAISS_THROW_MSG("This index type is not implemented on GPU.");
     }
 }
 
@@ -224,8 +228,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/GpuClonerOptions.h

@@ -43,6 +43,11 @@ struct GpuClonerOptions {
 #else
     bool use_raft = false;
 #endif
+
+    /// enable CPU fallback; when set to true, the cloner will clone to CPU
+    /// the index components that are not impemented on GPU; when set to false,
+    /// the cloner will throw an exception if it cannot convert the index to GPU
+    bool enableCpuFallback = false;
 };
 
 struct GpuMultipleClonerOptions : public GpuClonerOptions {

faiss/gpu/GpuIndex.h

@@ -43,6 +43,11 @@ struct GpuIndexConfig {
 #else
     bool use_raft = false;
 #endif
+
+    /// enable CPU fallback; when set to true, the cloner will clone to CPU
+    /// the index components that are not impemented on GPU; when set to false,
+    /// the cloner will throw an exception if it cannot convert the index to GPU
+    bool enableCpuFallback = false;
 };
 
 /// A centralized function that determines whether RAFT should

faiss/gpu/GpuIndexIVF.cu

@@ -173,6 +173,7 @@ void GpuIndexIVF::copyFrom(const faiss::IndexIVF* index) {
         GpuResourcesProviderFromInstance pfi(getResources());
 
         GpuClonerOptions options;
+        options.enableCpuFallback = config_.enableCpuFallback;
         auto cloner = ToGpuCloner(&pfi, getDevice(), options);
 
         quantizer = cloner.clone_Index(index->quantizer);

faiss/gpu/test/test_gpu_index.py

@@ -589,7 +589,10 @@ class TestGpuAutoTune(unittest.TestCase):
 
     def test_params(self):
         index = faiss.index_factory(32, "IVF65536_HNSW,PQ16")
-        index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0, index)
+        res = faiss.StandardGpuResources()
+        options = faiss.GpuClonerOptions()
+        options.enableCpuFallback = True
+        index = faiss.index_cpu_to_gpu(res, 0, index, options)
         ps = faiss.GpuParameterSpace()
         ps.initialize(index)
         for i in range(ps.parameter_ranges.size()):

faiss/gpu/test/test_index_cpu_to_gpu.py

@@ -4,26 +4,80 @@
 
 
 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)
+        if factory_string.startswith("IDMap"):
+            index.add_with_ids(db_vectors, np.arange(n))
+        else:
+            index.add(db_vectors)
+        return index
+
+    def create_and_clone(self, factory_string,
+                         enableCpuFallback=None,
+                         use_raft=None):
+        idx = self.create_index(factory_string)
+        config = faiss.GpuClonerOptions()
+        if enableCpuFallback is not None:
+            config.enableCpuFallback = enableCpuFallback
+        if use_raft is not None:
+            config.use_raft = use_raft
+        faiss.index_cpu_to_gpu(self.res, 0, idx, config)
+
+    def verify_throws_on_unsupported_index(self, factory_string):
+        try:
+            self.create_and_clone(factory_string)
+        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, use_raft=None):
         try:
-            faiss.index_cpu_to_gpu(res, 0, index_flat)
-        except Exception:
-            self.fail("index_cpu_to_gpu() threw an unexpected exception.")
+            self.create_and_clone(factory_string, use_raft=use_raft)
+        except Exception as e:
+            self.fail("Unexpected exception thrown factory_string: "
+                      "%s; error message: %s." % (factory_string, str(e)))
+
+    def verify_succeeds_on_unsupported_index_with_fallback_enabled(
+            self, factory_string, use_raft=None):
+        try:
+            self.create_and_clone(factory_string, enableCpuFallback=True,
+                                  use_raft=use_raft)
+        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")
+        self.verify_throws_on_unsupported_index("HNSW,PQ16")
+        self.verify_throws_on_unsupported_index("IDMap,PQ16")
+        self.verify_throws_on_unsupported_index("IVF256,ITQ64,SH1.2")
+
+    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,PQ8")
+
+        # set use_raft to false, this index type is not supported on RAFT
+        self.verify_succeeds_on_supported_index("IVF32,SQ8", use_raft=False)
+
+    def test_index_cpu_to_gpu_unsupported_indices_with_fallback_enabled(self):
+        self.verify_succeeds_on_unsupported_index_with_fallback_enabled("IDMap,Flat")
+        self.verify_succeeds_on_unsupported_index_with_fallback_enabled("PCA12,IVF32,Flat")
+        self.verify_succeeds_on_unsupported_index_with_fallback_enabled("PCA32,IVF32,PQ8")
+        self.verify_succeeds_on_unsupported_index_with_fallback_enabled("PCA32,IVF32,PQ8np")
+
+        # set use_raft to false, this index type is not supported on RAFT
+        self.verify_succeeds_on_unsupported_index_with_fallback_enabled(
+            "PCA32,IVF32,SQ8", use_raft=False)

faiss/gpu/test/test_multi_gpu.py

@@ -244,5 +244,7 @@ def test_cpu_to_gpu_IVFFlat(self):
     def test_set_gpu_param(self):
         index = faiss.index_factory(12, "PCAR8,IVF10,PQ4")
         res = faiss.StandardGpuResources()
-        gpu_index = faiss.index_cpu_to_gpu(res, 0, index)
+        options = faiss.GpuClonerOptions()
+        options.enableCpuFallback = True
+        gpu_index = faiss.index_cpu_to_gpu(res, 0, index, options)
         faiss.GpuParameterSpace().set_index_parameter(gpu_index, "nprobe", 3)

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)) {                                          \