faiss 0.2.7 → 0.3.1

data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp

@@ -4,10 +4,26 @@
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  */
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/gpu/GpuIndexIVFFlat.h>
+#include <faiss/gpu/GpuIndicesOptions.h>
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/test/TestUtils.h>
 #include <faiss/gpu/utils/DeviceUtils.h>
@@ -40,6 +56,8 @@ struct Options {
                  faiss::gpu::INDICES_64_BIT});
 
         device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
+
+        useRaft = false;
     }
 
     std::string toString() const {
@@ -47,7 +65,7 @@ struct Options {
         str << "IVFFlat device " << device << " numVecs " << numAdd << " dim "
             << dim << " numCentroids " << numCentroids << " nprobe " << nprobe
             << " numQuery " << numQuery << " k " << k << " indicesOpt "
-            << indicesOpt;
+            << indicesOpt << " useRaft " << useRaft;
 
         return str.str();
     }
@@ -61,6 +79,7 @@ struct Options {
     int k;
     int device;
     faiss::gpu::IndicesOptions indicesOpt;
+    bool useRaft;
 };
 
 void queryTest(
@@ -91,6 +110,7 @@ void queryTest(
         config.device = opt.device;
         config.indicesOptions = opt.indicesOpt;
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
+        config.use_raft = opt.useRaft;
 
         faiss::gpu::GpuIndexIVFFlat gpuIndex(
                 &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
@@ -114,7 +134,10 @@ void queryTest(
     }
 }
 
-void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
+void addTest(
+        faiss::MetricType metricType,
+        bool useFloat16CoarseQuantizer,
+        bool useRaft) {
     for (int tries = 0; tries < 2; ++tries) {
         Options opt;
 
@@ -138,8 +161,10 @@ void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
 
         faiss::gpu::GpuIndexIVFFlatConfig config;
         config.device = opt.device;
-        config.indicesOptions = opt.indicesOpt;
+        config.indicesOptions =
+                useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
         config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
+        config.use_raft = useRaft;
 
         faiss::gpu::GpuIndexIVFFlat gpuIndex(
                 &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
@@ -163,7 +188,7 @@ void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
     }
 }
 
-void copyToTest(bool useFloat16CoarseQuantizer) {
+void copyToTest(bool useFloat16CoarseQuantizer, bool useRaft) {
     Options opt;
     std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
     std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
@@ -173,8 +198,10 @@ void copyToTest(bool useFloat16CoarseQuantizer) {
 
     faiss::gpu::GpuIndexIVFFlatConfig config;
     config.device = opt.device;
-    config.indicesOptions = opt.indicesOpt;
+    config.indicesOptions =
+            useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
     config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
+    config.use_raft = useRaft;
 
     faiss::gpu::GpuIndexIVFFlat gpuIndex(
             &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
@@ -214,7 +241,7 @@ void copyToTest(bool useFloat16CoarseQuantizer) {
             compFloat16 ? 0.30f : 0.015f);
 }
 
-void copyFromTest(bool useFloat16CoarseQuantizer) {
+void copyFromTest(bool useFloat16CoarseQuantizer, bool useRaft) {
     Options opt;
     std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
     std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
@@ -232,8 +259,10 @@ void copyFromTest(bool useFloat16CoarseQuantizer) {
 
     faiss::gpu::GpuIndexIVFFlatConfig config;
     config.device = opt.device;
-    config.indicesOptions = opt.indicesOpt;
+    config.indicesOptions =
+            useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
     config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
+    config.use_raft = useRaft;
 
     faiss::gpu::GpuIndexIVFFlat gpuIndex(&res, 1, 1, faiss::METRIC_L2, config);
     gpuIndex.nprobe = 1;
@@ -265,19 +294,35 @@ void copyFromTest(bool useFloat16CoarseQuantizer) {
 }
 
 TEST(TestGpuIndexIVFFlat, Float32_32_Add_L2) {
-    addTest(faiss::METRIC_L2, false);
+    addTest(faiss::METRIC_L2, false, false);
+
+#if defined USE_NVIDIA_RAFT
+    addTest(faiss::METRIC_L2, false, true);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float32_32_Add_IP) {
-    addTest(faiss::METRIC_INNER_PRODUCT, false);
+    addTest(faiss::METRIC_INNER_PRODUCT, false, false);
+
+#if defined USE_NVIDIA_RAFT
+    addTest(faiss::METRIC_INNER_PRODUCT, false, true);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float16_32_Add_L2) {
-    addTest(faiss::METRIC_L2, true);
+    addTest(faiss::METRIC_L2, true, false);
+
+#if defined USE_NVIDIA_RAFT
+    addTest(faiss::METRIC_L2, true, true);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
-    addTest(faiss::METRIC_INNER_PRODUCT, true);
+    addTest(faiss::METRIC_INNER_PRODUCT, true, false);
+
+#if defined USE_NVIDIA_RAFT
+    addTest(faiss::METRIC_INNER_PRODUCT, true, true);
+#endif
 }
 
 //
@@ -285,11 +330,25 @@ TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
 //
 
 TEST(TestGpuIndexIVFFlat, Float32_Query_L2) {
-    queryTest(Options(), faiss::METRIC_L2, false);
+    Options opt;
+    queryTest(opt, faiss::METRIC_L2, false);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_L2, false);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float32_Query_IP) {
-    queryTest(Options(), faiss::METRIC_INNER_PRODUCT, false);
+    Options opt;
+    queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, LargeBatch) {
@@ -297,16 +356,36 @@ TEST(TestGpuIndexIVFFlat, LargeBatch) {
     opt.dim = 3;
     opt.numQuery = 100000;
     queryTest(opt, faiss::METRIC_L2, false);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_L2, false);
+#endif
 }
 
 // float16 coarse quantizer
 
 TEST(TestGpuIndexIVFFlat, Float16_32_Query_L2) {
-    queryTest(Options(), faiss::METRIC_L2, true);
+    Options opt;
+    queryTest(opt, faiss::METRIC_L2, true);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_L2, true);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float16_32_Query_IP) {
-    queryTest(Options(), faiss::METRIC_INNER_PRODUCT, true);
+    Options opt;
+    queryTest(opt, faiss::METRIC_INNER_PRODUCT, true);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_INNER_PRODUCT, true);
+#endif
 }
 
 //
@@ -318,24 +397,48 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_L2_64) {
     Options opt;
     opt.dim = 64;
     queryTest(opt, faiss::METRIC_L2, false);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_L2, false);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float32_Query_IP_64) {
     Options opt;
     opt.dim = 64;
     queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float32_Query_L2_128) {
     Options opt;
     opt.dim = 128;
     queryTest(opt, faiss::METRIC_L2, false);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_L2, false);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
     Options opt;
     opt.dim = 128;
     queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
+
+#if defined USE_NVIDIA_RAFT
+    opt.useRaft = true;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
+#endif
 }
 
 //
@@ -343,11 +446,19 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
 //
 
 TEST(TestGpuIndexIVFFlat, Float32_32_CopyTo) {
-    copyToTest(false);
+    copyToTest(false, false);
+
+#if defined USE_NVIDIA_RAFT
+    copyToTest(false, true);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float32_32_CopyFrom) {
-    copyFromTest(false);
+    copyFromTest(false, false);
+
+#if defined USE_NVIDIA_RAFT
+    copyFromTest(false, true);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, Float32_negative) {
@@ -377,23 +488,24 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
     faiss::gpu::StandardGpuResources res;
     res.noTempMemory();
 
+    // Construct a positive test set
+    auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
+
+    // Put all vecs on positive size
+    for (auto& f : queryVecs) {
+        f = std::abs(f);
+    }
+
     faiss::gpu::GpuIndexIVFFlatConfig config;
     config.device = opt.device;
     config.indicesOptions = opt.indicesOpt;
+    config.use_raft = false;
 
     faiss::gpu::GpuIndexIVFFlat gpuIndex(
             &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
     gpuIndex.copyFrom(&cpuIndex);
     gpuIndex.nprobe = opt.nprobe;
 
-    // Construct a positive test set
-    auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
-
-    // Put all vecs on positive size
-    for (auto& f : queryVecs) {
-        f = std::abs(f);
-    }
-
     bool compFloat16 = false;
     faiss::gpu::compareIndices(
             queryVecs,
@@ -409,6 +521,31 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
             // in fp16. Figure out another way to test
             compFloat16 ? 0.99f : 0.1f,
             compFloat16 ? 0.65f : 0.015f);
+
+#if defined USE_NVIDIA_RAFT
+    config.use_raft = true;
+    config.indicesOptions = faiss::gpu::INDICES_64_BIT;
+
+    faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
+            &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
+    raftGpuIndex.copyFrom(&cpuIndex);
+    raftGpuIndex.nprobe = opt.nprobe;
+
+    faiss::gpu::compareIndices(
+            queryVecs,
+            cpuIndex,
+            raftGpuIndex,
+            opt.numQuery,
+            opt.dim,
+            opt.k,
+            opt.toString(),
+            compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
+            // FIXME: the fp16 bounds are
+            // useless when math (the accumulator) is
+            // in fp16. Figure out another way to test
+            compFloat16 ? 0.99f : 0.1f,
+            compFloat16 ? 0.65f : 0.015f);
+#endif
 }
 
 //
@@ -424,10 +561,18 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
     faiss::gpu::StandardGpuResources res;
     res.noTempMemory();
 
+    int numQuery = 10;
+    std::vector<float> nans(
+            numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
+
+    std::vector<float> distances(numQuery * opt.k, 0);
+    std::vector<faiss::idx_t> indices(numQuery * opt.k, 0);
+
     faiss::gpu::GpuIndexIVFFlatConfig config;
     config.device = opt.device;
     config.indicesOptions = opt.indicesOpt;
     config.flatConfig.useFloat16 = faiss::gpu::randBool();
+    config.use_raft = false;
 
     faiss::gpu::GpuIndexIVFFlat gpuIndex(
             &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
@@ -436,14 +581,31 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
     gpuIndex.train(opt.numTrain, trainVecs.data());
     gpuIndex.add(opt.numAdd, addVecs.data());
 
-    int numQuery = 10;
-    std::vector<float> nans(
-            numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
+    gpuIndex.search(
+            numQuery, nans.data(), opt.k, distances.data(), indices.data());
 
-    std::vector<float> distances(numQuery * opt.k, 0);
-    std::vector<faiss::idx_t> indices(numQuery * opt.k, 0);
+    for (int q = 0; q < numQuery; ++q) {
+        for (int k = 0; k < opt.k; ++k) {
+            EXPECT_EQ(indices[q * opt.k + k], -1);
+            EXPECT_EQ(
+                    distances[q * opt.k + k],
+                    std::numeric_limits<float>::max());
+        }
+    }
 
-    gpuIndex.search(
+#if defined USE_NVIDIA_RAFT
+    config.use_raft = true;
+    config.indicesOptions = faiss::gpu::INDICES_64_BIT;
+    std::fill(distances.begin(), distances.end(), 0);
+    std::fill(indices.begin(), indices.end(), 0);
+    faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
+            &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
+    raftGpuIndex.nprobe = opt.nprobe;
+
+    raftGpuIndex.train(opt.numTrain, trainVecs.data());
+    raftGpuIndex.add(opt.numAdd, addVecs.data());
+
+    raftGpuIndex.search(
             numQuery, nans.data(), opt.k, distances.data(), indices.data());
 
     for (int q = 0; q < numQuery; ++q) {
@@ -454,6 +616,7 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
                     std::numeric_limits<float>::max());
         }
     }
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, AddNaN) {
@@ -462,15 +625,6 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
     faiss::gpu::StandardGpuResources res;
     res.noTempMemory();
 
-    faiss::gpu::GpuIndexIVFFlatConfig config;
-    config.device = opt.device;
-    config.indicesOptions = opt.indicesOpt;
-    config.flatConfig.useFloat16 = faiss::gpu::randBool();
-
-    faiss::gpu::GpuIndexIVFFlat gpuIndex(
-            &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
-    gpuIndex.nprobe = opt.nprobe;
-
     int numNans = 10;
     std::vector<float> nans(
             numNans * opt.dim, std::numeric_limits<float>::quiet_NaN());
@@ -482,6 +636,15 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
     }
 
     std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
+
+    faiss::gpu::GpuIndexIVFFlatConfig config;
+    config.device = opt.device;
+    config.indicesOptions = opt.indicesOpt;
+    config.flatConfig.useFloat16 = faiss::gpu::randBool();
+    config.use_raft = false;
+    faiss::gpu::GpuIndexIVFFlat gpuIndex(
+            &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
+    gpuIndex.nprobe = opt.nprobe;
     gpuIndex.train(opt.numTrain, trainVecs.data());
 
     // should not crash
@@ -499,6 +662,27 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
             opt.k,
             distance.data(),
             indices.data());
+
+#if defined USE_NVIDIA_RAFT
+    config.use_raft = true;
+    config.indicesOptions = faiss::gpu::INDICES_64_BIT;
+    faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
+            &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
+    raftGpuIndex.nprobe = opt.nprobe;
+    raftGpuIndex.train(opt.numTrain, trainVecs.data());
+
+    // should not crash
+    EXPECT_EQ(raftGpuIndex.ntotal, 0);
+    raftGpuIndex.add(numNans, nans.data());
+
+    // should not crash
+    raftGpuIndex.search(
+            opt.numQuery,
+            queryVecs.data(),
+            opt.k,
+            distance.data(),
+            indices.data());
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
@@ -539,6 +723,7 @@ TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
     faiss::gpu::GpuIndexIVFFlatConfig config;
     config.device = device;
     config.memorySpace = faiss::gpu::MemorySpace::Unified;
+    config.use_raft = false;
 
     faiss::gpu::GpuIndexIVFFlat gpuIndex(
             &res, dim, numCentroids, faiss::METRIC_L2, config);
@@ -555,6 +740,26 @@ TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
             kF32MaxRelErr,
             0.1f,
             0.015f);
+
+#if defined USE_NVIDIA_RAFT
+    config.use_raft = true;
+    config.indicesOptions = faiss::gpu::INDICES_64_BIT;
+    faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
+            &res, dim, numCentroids, faiss::METRIC_L2, config);
+    raftGpuIndex.copyFrom(&cpuIndex);
+    raftGpuIndex.nprobe = nprobe;
+
+    faiss::gpu::compareIndices(
+            cpuIndex,
+            raftGpuIndex,
+            numQuery,
+            dim,
+            k,
+            "Unified Memory",
+            kF32MaxRelErr,
+            0.1f,
+            0.015f);
+#endif
 }
 
 TEST(TestGpuIndexIVFFlat, LongIVFList) {
@@ -596,6 +801,7 @@ TEST(TestGpuIndexIVFFlat, LongIVFList) {
 
     faiss::gpu::GpuIndexIVFFlatConfig config;
     config.device = device;
+    config.use_raft = false;
 
     faiss::gpu::GpuIndexIVFFlat gpuIndex(
             &res, dim, numCentroids, faiss::METRIC_L2, config);
@@ -613,6 +819,27 @@ TEST(TestGpuIndexIVFFlat, LongIVFList) {
             kF32MaxRelErr,
             0.1f,
             0.015f);
+
+#if defined USE_NVIDIA_RAFT
+    config.use_raft = true;
+    config.indicesOptions = faiss::gpu::INDICES_64_BIT;
+    faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
+            &res, dim, numCentroids, faiss::METRIC_L2, config);
+    raftGpuIndex.train(numTrain, trainVecs.data());
+    raftGpuIndex.add(numAdd, addVecs.data());
+    raftGpuIndex.nprobe = 1;
+
+    faiss::gpu::compareIndices(
+            cpuIndex,
+            raftGpuIndex,
+            numQuery,
+            dim,
+            k,
+            "Unified Memory",
+            kF32MaxRelErr,
+            0.1f,
+            0.015f);
+#endif
 }
 
 int main(int argc, char** argv) {