faiss 0.3.0 → 0.3.1

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

@@ -35,6 +35,22 @@ void pickEncoding(int& codes, int& dim) {
     }
 }
 
+void pickRaftEncoding(int& codes, int& dim, int bitsPerCode) {
+    // Above 32 doesn't work with no precomputed codes
+    std::vector<int> dimSizes{4, 8, 10, 12, 16, 20, 24, 28, 32};
+
+    while (true) {
+        codes = faiss::gpu::randVal(0, 96);
+        dim = codes * dimSizes[faiss::gpu::randVal(0, dimSizes.size() - 1)];
+
+        // for such a small test, super-low or high dim is more likely to
+        // generate comparison errors
+        if (dim < 256 && dim >= 64 && (codes * bitsPerCode) % 8 == 0) {
+            return;
+        }
+    }
+}
+
 struct Options {
     Options() {
         numAdd = faiss::gpu::randVal(2000, 5000);
@@ -43,9 +59,10 @@ struct Options {
 
         pickEncoding(codes, dim);
 
-        // TODO: Change back to `faiss::gpu::randVal(3, 7)` when we officially
-        //   support non-multiple of 8 subcodes for IVFPQ.
+        // TODO: Change back to `faiss::gpu::randVal(3, 7)` when we
+        // officially support non-multiple of 8 subcodes for IVFPQ.
         bitsPerCode = 8;
+
         nprobe = std::min(faiss::gpu::randVal(40, 1000), numCentroids);
         numQuery = faiss::gpu::randVal(4, 8);
 
@@ -66,6 +83,9 @@ struct Options {
         }
 
         device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
+
+        interleavedLayout = false;
+        useRaft = false;
     }
 
     std::string toString() const {
@@ -105,50 +125,66 @@ struct Options {
     faiss::gpu::IndicesOptions indicesOpt;
     bool useFloat16;
     int device;
+    bool interleavedLayout;
+    bool useRaft;
 };
 
-TEST(TestGpuIndexIVFPQ, Query_L2) {
-    for (int tries = 0; tries < 2; ++tries) {
-        Options opt;
+void queryTest(Options opt, faiss::MetricType metricType) {
+    std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
+    std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
 
-        std::vector<float> trainVecs =
-                faiss::gpu::randVecs(opt.numTrain, opt.dim);
-        std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
+    faiss::IndexFlatL2 coarseQuantizerL2(opt.dim);
+    faiss::IndexFlatIP coarseQuantizerIP(opt.dim);
+    faiss::Index* quantizer = metricType == faiss::METRIC_L2
+            ? (faiss::Index*)&coarseQuantizerL2
+            : (faiss::Index*)&coarseQuantizerIP;
 
-        faiss::IndexFlatL2 coarseQuantizer(opt.dim);
-        faiss::IndexIVFPQ cpuIndex(
-                &coarseQuantizer,
-                opt.dim,
-                opt.numCentroids,
-                opt.codes,
-                opt.bitsPerCode);
-        cpuIndex.nprobe = opt.nprobe;
-        cpuIndex.train(opt.numTrain, trainVecs.data());
-        cpuIndex.add(opt.numAdd, addVecs.data());
+    faiss::IndexIVFPQ cpuIndex(
+            quantizer, opt.dim, opt.numCentroids, opt.codes, opt.bitsPerCode);
+    cpuIndex.metric_type = metricType;
+    cpuIndex.nprobe = opt.nprobe;
+    cpuIndex.train(opt.numTrain, trainVecs.data());
+    cpuIndex.add(opt.numAdd, addVecs.data());
 
-        // Use the default temporary memory management to test the memory
-        // manager
-        faiss::gpu::StandardGpuResources res;
+    // Use the default temporary memory management to test the memory
+    // manager
+    faiss::gpu::StandardGpuResources res;
 
-        faiss::gpu::GpuIndexIVFPQConfig config;
-        config.device = opt.device;
-        config.usePrecomputedTables = (tries % 2 == 0);
-        config.indicesOptions = opt.indicesOpt;
-        config.useFloat16LookupTables = opt.useFloat16;
+    faiss::gpu::GpuIndexIVFPQConfig config;
+    config.device = opt.device;
+    config.usePrecomputedTables = opt.usePrecomputed;
+    config.indicesOptions = opt.indicesOpt;
+    config.useFloat16LookupTables = opt.useFloat16;
+    config.interleavedLayout = opt.interleavedLayout;
+    config.use_raft = opt.useRaft;
 
-        faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
-        gpuIndex.nprobe = opt.nprobe;
+    faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
+    gpuIndex.nprobe = opt.nprobe;
 
-        faiss::gpu::compareIndices(
-                cpuIndex,
-                gpuIndex,
-                opt.numQuery,
-                opt.dim,
-                opt.k,
-                opt.toString(),
-                opt.getCompareEpsilon(),
-                opt.getPctMaxDiff1(),
-                opt.getPctMaxDiffN());
+    faiss::gpu::compareIndices(
+            cpuIndex,
+            gpuIndex,
+            opt.numQuery,
+            opt.dim,
+            opt.k,
+            opt.toString(),
+            opt.getCompareEpsilon(),
+            opt.getPctMaxDiff1(),
+            opt.getPctMaxDiffN());
+}
+
+TEST(TestGpuIndexIVFPQ, Query_L2) {
+    for (int tries = 0; tries < 2; ++tries) {
+        Options opt;
+        opt.usePrecomputed = (tries % 2 == 0);
+        queryTest(opt, faiss::MetricType::METRIC_L2);
+    }
+}
+
+TEST(TestGpuIndexIVFPQ, Query_IP) {
+    for (int tries = 0; tries < 2; ++tries) {
+        Options opt;
+        queryTest(opt, faiss::MetricType::METRIC_INNER_PRODUCT);
     }
 }
 
@@ -161,45 +197,10 @@ TEST(TestGpuIndexIVFPQ, LargeBatch) {
         opt.dim = 4;
         opt.numQuery = 100000;
         opt.codes = 2;
+        opt.usePrecomputed = usePrecomputed;
+        opt.useFloat16 = false;
 
-        std::vector<float> trainVecs =
-                faiss::gpu::randVecs(opt.numTrain, opt.dim);
-        std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
-
-        faiss::IndexFlatL2 coarseQuantizer(opt.dim);
-        faiss::IndexIVFPQ cpuIndex(
-                &coarseQuantizer,
-                opt.dim,
-                opt.numCentroids,
-                opt.codes,
-                opt.bitsPerCode);
-        cpuIndex.nprobe = opt.nprobe;
-        cpuIndex.train(opt.numTrain, trainVecs.data());
-        cpuIndex.add(opt.numAdd, addVecs.data());
-
-        // Use the default temporary memory management to test the memory
-        // manager
-        faiss::gpu::StandardGpuResources res;
-
-        faiss::gpu::GpuIndexIVFPQConfig config;
-        config.device = opt.device;
-        config.usePrecomputedTables = usePrecomputed;
-        config.indicesOptions = opt.indicesOpt;
-        config.useFloat16LookupTables = false;
-
-        faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
-        gpuIndex.nprobe = opt.nprobe;
-
-        faiss::gpu::compareIndices(
-                cpuIndex,
-                gpuIndex,
-                opt.numQuery,
-                opt.dim,
-                opt.k,
-                opt.toString(),
-                opt.getCompareEpsilon(),
-                opt.getPctMaxDiff1(),
-                opt.getPctMaxDiffN());
+        queryTest(opt, faiss::MetricType::METRIC_L2);
     }
 }
 
@@ -234,6 +235,7 @@ void testMMCodeDistance(faiss::MetricType mt) {
         config.usePrecomputedTables = false;
         config.useMMCodeDistance = true;
         config.indicesOptions = opt.indicesOpt;
+        config.use_raft = false;
 
         // Make sure that the float16 version works as well
         config.useFloat16LookupTables = (tries % 2 == 0);
@@ -284,6 +286,7 @@ void testMMCodeDistance(faiss::MetricType mt) {
         config.device = opt.device;
         config.usePrecomputedTables = false;
         config.indicesOptions = opt.indicesOpt;
+        config.use_raft = false;
 
         // Make sure that the float16 version works as well
         config.useFloat16LookupTables = (dimPerSubQ == 7);
@@ -312,53 +315,6 @@ TEST(TestGpuIndexIVFPQ, Query_IP_MMCodeDistance) {
     testMMCodeDistance(faiss::MetricType::METRIC_INNER_PRODUCT);
 }
 
-TEST(TestGpuIndexIVFPQ, Query_IP) {
-    for (int tries = 0; tries < 2; ++tries) {
-        Options opt;
-
-        std::vector<float> trainVecs =
-                faiss::gpu::randVecs(opt.numTrain, opt.dim);
-        std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
-
-        faiss::IndexFlatIP coarseQuantizer(opt.dim);
-        faiss::IndexIVFPQ cpuIndex(
-                &coarseQuantizer,
-                opt.dim,
-                opt.numCentroids,
-                opt.codes,
-                opt.bitsPerCode);
-        cpuIndex.metric_type = faiss::MetricType::METRIC_INNER_PRODUCT;
-
-        cpuIndex.nprobe = opt.nprobe;
-        cpuIndex.train(opt.numTrain, trainVecs.data());
-        cpuIndex.add(opt.numAdd, addVecs.data());
-
-        // Use the default temporary memory management to test the memory
-        // manager
-        faiss::gpu::StandardGpuResources res;
-
-        faiss::gpu::GpuIndexIVFPQConfig config;
-        config.device = opt.device;
-        config.usePrecomputedTables = false; // not supported/required for IP
-        config.indicesOptions = opt.indicesOpt;
-        config.useFloat16LookupTables = opt.useFloat16;
-
-        faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
-        gpuIndex.nprobe = opt.nprobe;
-
-        faiss::gpu::compareIndices(
-                cpuIndex,
-                gpuIndex,
-                opt.numQuery,
-                opt.dim,
-                opt.k,
-                opt.toString(),
-                opt.getCompareEpsilon(),
-                opt.getPctMaxDiff1(),
-                opt.getPctMaxDiffN());
-    }
-}
-
 TEST(TestGpuIndexIVFPQ, Float16Coarse) {
     Options opt;
 
@@ -384,6 +340,7 @@ TEST(TestGpuIndexIVFPQ, Float16Coarse) {
     config.usePrecomputedTables = opt.usePrecomputed;
     config.indicesOptions = opt.indicesOpt;
     config.useFloat16LookupTables = opt.useFloat16;
+    config.use_raft = false;
 
     faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
     gpuIndex.nprobe = opt.nprobe;
@@ -403,104 +360,68 @@ TEST(TestGpuIndexIVFPQ, Float16Coarse) {
             opt.getPctMaxDiffN());
 }
 
-TEST(TestGpuIndexIVFPQ, Add_L2) {
-    for (int tries = 0; tries < 2; ++tries) {
-        Options opt;
+void addTest(Options opt, faiss::MetricType metricType) {
+    std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
+    std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
 
-        std::vector<float> trainVecs =
-                faiss::gpu::randVecs(opt.numTrain, opt.dim);
-        std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
+    faiss::IndexFlatL2 coarseQuantizerL2(opt.dim);
+    faiss::IndexFlatIP coarseQuantizerIP(opt.dim);
+    faiss::Index* quantizer = metricType == faiss::METRIC_L2
+            ? (faiss::Index*)&coarseQuantizerL2
+            : (faiss::Index*)&coarseQuantizerIP;
 
-        faiss::IndexFlatL2 coarseQuantizer(opt.dim);
-        faiss::IndexIVFPQ cpuIndex(
-                &coarseQuantizer,
-                opt.dim,
-                opt.numCentroids,
-                opt.codes,
-                opt.bitsPerCode);
-        cpuIndex.nprobe = opt.nprobe;
-        cpuIndex.train(opt.numTrain, trainVecs.data());
+    faiss::IndexIVFPQ cpuIndex(
+            quantizer, opt.dim, opt.numCentroids, opt.codes, opt.bitsPerCode);
+    cpuIndex.nprobe = opt.nprobe;
+    cpuIndex.metric_type = metricType;
+    cpuIndex.train(opt.numTrain, trainVecs.data());
 
-        // Use the default temporary memory management to test the memory
-        // manager
-        faiss::gpu::StandardGpuResources res;
+    // Use the default temporary memory management to test the memory
+    // manager
+    faiss::gpu::StandardGpuResources res;
 
-        faiss::gpu::GpuIndexIVFPQConfig config;
-        config.device = opt.device;
-        config.usePrecomputedTables = opt.usePrecomputed;
-        config.indicesOptions = opt.indicesOpt;
-        config.useFloat16LookupTables = opt.useFloat16;
+    faiss::gpu::GpuIndexIVFPQConfig config;
+    config.device = opt.device;
+    config.usePrecomputedTables = opt.usePrecomputed;
+    config.indicesOptions = opt.indicesOpt;
+    config.useFloat16LookupTables = opt.useFloat16;
+    config.interleavedLayout = opt.interleavedLayout;
+    config.use_raft = opt.useRaft;
 
-        faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
-        gpuIndex.nprobe = opt.nprobe;
+    faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
+    gpuIndex.nprobe = opt.nprobe;
 
-        gpuIndex.add(opt.numAdd, addVecs.data());
-        cpuIndex.add(opt.numAdd, addVecs.data());
+    gpuIndex.add(opt.numAdd, addVecs.data());
+    cpuIndex.add(opt.numAdd, addVecs.data());
 
-        faiss::gpu::compareIndices(
-                cpuIndex,
-                gpuIndex,
-                opt.numQuery,
-                opt.dim,
-                opt.k,
-                opt.toString(),
-                opt.getCompareEpsilon(),
-                opt.getPctMaxDiff1(),
-                opt.getPctMaxDiffN());
+    faiss::gpu::compareIndices(
+            cpuIndex,
+            gpuIndex,
+            opt.numQuery,
+            opt.dim,
+            opt.k,
+            opt.toString(),
+            opt.getCompareEpsilon(),
+            opt.getPctMaxDiff1(),
+            opt.getPctMaxDiffN());
+}
+
+TEST(TestGpuIndexIVFPQ, Add_L2) {
+    for (int tries = 0; tries < 2; ++tries) {
+        Options opt;
+        addTest(opt, faiss::METRIC_L2);
     }
 }
 
 TEST(TestGpuIndexIVFPQ, Add_IP) {
     for (int tries = 0; tries < 2; ++tries) {
         Options opt;
-
-        std::vector<float> trainVecs =
-                faiss::gpu::randVecs(opt.numTrain, opt.dim);
-        std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
-
-        faiss::IndexFlatIP coarseQuantizer(opt.dim);
-        faiss::IndexIVFPQ cpuIndex(
-                &coarseQuantizer,
-                opt.dim,
-                opt.numCentroids,
-                opt.codes,
-                opt.bitsPerCode);
-        cpuIndex.metric_type = faiss::MetricType::METRIC_INNER_PRODUCT;
-        cpuIndex.nprobe = opt.nprobe;
-        cpuIndex.train(opt.numTrain, trainVecs.data());
-
-        // Use the default temporary memory management to test the memory
-        // manager
-        faiss::gpu::StandardGpuResources res;
-
-        faiss::gpu::GpuIndexIVFPQConfig config;
-        config.device = opt.device;
-        config.usePrecomputedTables = opt.usePrecomputed;
-        config.indicesOptions = opt.indicesOpt;
-        config.useFloat16LookupTables = opt.useFloat16;
-
-        faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
-        gpuIndex.nprobe = opt.nprobe;
-
-        gpuIndex.add(opt.numAdd, addVecs.data());
-        cpuIndex.add(opt.numAdd, addVecs.data());
-
-        faiss::gpu::compareIndices(
-                cpuIndex,
-                gpuIndex,
-                opt.numQuery,
-                opt.dim,
-                opt.k,
-                opt.toString(),
-                opt.getCompareEpsilon(),
-                opt.getPctMaxDiff1(),
-                opt.getPctMaxDiffN());
+        addTest(opt, faiss::METRIC_INNER_PRODUCT);
     }
 }
 
-TEST(TestGpuIndexIVFPQ, CopyTo) {
+void copyToTest(Options opt) {
     for (int tries = 0; tries < 2; ++tries) {
-        Options opt;
         std::vector<float> trainVecs =
                 faiss::gpu::randVecs(opt.numTrain, opt.dim);
         std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
@@ -511,9 +432,11 @@ TEST(TestGpuIndexIVFPQ, CopyTo) {
 
         faiss::gpu::GpuIndexIVFPQConfig config;
         config.device = opt.device;
-        config.usePrecomputedTables = (tries % 2 == 0);
+        config.usePrecomputedTables = false;
         config.indicesOptions = opt.indicesOpt;
         config.useFloat16LookupTables = opt.useFloat16;
+        config.interleavedLayout = opt.interleavedLayout;
+        config.use_raft = opt.useRaft;
 
         faiss::gpu::GpuIndexIVFPQ gpuIndex(
                 &res,
@@ -561,8 +484,12 @@ TEST(TestGpuIndexIVFPQ, CopyTo) {
     }
 }
 
-TEST(TestGpuIndexIVFPQ, CopyFrom) {
+TEST(TestGpuIndexIVFPQ, CopyTo) {
     Options opt;
+    copyToTest(opt);
+}
+
+void copyFromTest(Options opt) {
     std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
     std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
 
@@ -585,6 +512,8 @@ TEST(TestGpuIndexIVFPQ, CopyFrom) {
     config.usePrecomputedTables = opt.usePrecomputed;
     config.indicesOptions = opt.indicesOpt;
     config.useFloat16LookupTables = opt.useFloat16;
+    config.interleavedLayout = opt.interleavedLayout;
+    config.use_raft = opt.useRaft;
 
     // Use garbage values to see if we overwrite them
     faiss::gpu::GpuIndexIVFPQ gpuIndex(
@@ -621,9 +550,12 @@ TEST(TestGpuIndexIVFPQ, CopyFrom) {
             opt.getPctMaxDiffN());
 }
 
-TEST(TestGpuIndexIVFPQ, QueryNaN) {
+TEST(TestGpuIndexIVFPQ, CopyFrom) {
     Options opt;
+    copyFromTest(opt);
+}
 
+void queryNaNTest(Options opt) {
     std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
     std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
 
@@ -635,6 +567,8 @@ TEST(TestGpuIndexIVFPQ, QueryNaN) {
     config.usePrecomputedTables = opt.usePrecomputed;
     config.indicesOptions = opt.indicesOpt;
     config.useFloat16LookupTables = opt.useFloat16;
+    config.use_raft = opt.useRaft;
+    config.interleavedLayout = opt.useRaft ? true : opt.interleavedLayout;
 
     faiss::gpu::GpuIndexIVFPQ gpuIndex(
             &res,
@@ -670,9 +604,13 @@ TEST(TestGpuIndexIVFPQ, QueryNaN) {
     }
 }
 
-TEST(TestGpuIndexIVFPQ, AddNaN) {
+TEST(TestGpuIndexIVFPQ, QueryNaN) {
     Options opt;
+    opt.useRaft = false;
+    queryNaNTest(opt);
+}
 
+void addNaNTest(Options opt) {
     // Use the default temporary memory management to test the memory manager
     faiss::gpu::StandardGpuResources res;
 
@@ -681,6 +619,8 @@ TEST(TestGpuIndexIVFPQ, AddNaN) {
     config.usePrecomputedTables = opt.usePrecomputed;
     config.indicesOptions = opt.indicesOpt;
     config.useFloat16LookupTables = opt.useFloat16;
+    config.interleavedLayout = opt.interleavedLayout;
+    config.use_raft = opt.useRaft;
 
     faiss::gpu::GpuIndexIVFPQ gpuIndex(
             &res,
@@ -722,6 +662,128 @@ TEST(TestGpuIndexIVFPQ, AddNaN) {
             indices.data());
 }
 
+TEST(TestGpuIndexIVFPQ, AddNaN) {
+    Options opt;
+    opt.useRaft = false;
+    addNaNTest(opt);
+}
+
+#if defined USE_NVIDIA_RAFT
+TEST(TestGpuIndexIVFPQ, Query_L2_Raft) {
+    for (int tries = 0; tries < 2; ++tries) {
+        Options opt;
+        opt.bitsPerCode = faiss::gpu::randVal(4, 8);
+        opt.useRaft = true;
+        opt.interleavedLayout = true;
+        opt.usePrecomputed = false;
+        opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+        pickRaftEncoding(opt.codes, opt.dim, opt.bitsPerCode);
+        queryTest(opt, faiss::MetricType::METRIC_L2);
+    }
+}
+
+TEST(TestGpuIndexIVFPQ, Query_IP_Raft) {
+    for (int tries = 0; tries < 2; ++tries) {
+        Options opt;
+        opt.bitsPerCode = faiss::gpu::randVal(4, 8);
+        opt.useRaft = true;
+        opt.interleavedLayout = true;
+        opt.usePrecomputed = false;
+        opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+        pickRaftEncoding(opt.codes, opt.dim, opt.bitsPerCode);
+        queryTest(opt, faiss::MetricType::METRIC_INNER_PRODUCT);
+    }
+}
+
+// Large batch sizes (>= 65536) should also work
+TEST(TestGpuIndexIVFPQ, LargeBatch_Raft) {
+    Options opt;
+
+    // override for large sizes
+    opt.dim = 4;
+    opt.numQuery = 100000;
+    opt.codes = 2;
+    opt.useRaft = true;
+    opt.interleavedLayout = true;
+    opt.usePrecomputed = false;
+    opt.useFloat16 = false;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    opt.bitsPerCode = 8;
+
+    queryTest(opt, faiss::MetricType::METRIC_L2);
+}
+
+TEST(TestGpuIndexIVFPQ, CopyFrom_Raft) {
+    Options opt;
+    opt.useRaft = true;
+    opt.interleavedLayout = true;
+    opt.bitsPerCode = faiss::gpu::randVal(4, 8);
+    opt.usePrecomputed = false;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    pickRaftEncoding(opt.codes, opt.dim, opt.bitsPerCode);
+    copyFromTest(opt);
+}
+
+TEST(TestGpuIndexIVFPQ, Add_L2_Raft) {
+    for (int tries = 0; tries < 2; ++tries) {
+        Options opt;
+        opt.useRaft = true;
+        opt.interleavedLayout = true;
+        opt.bitsPerCode = faiss::gpu::randVal(4, 8);
+        opt.usePrecomputed = false;
+        opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+        pickRaftEncoding(opt.codes, opt.dim, opt.bitsPerCode);
+        addTest(opt, faiss::METRIC_L2);
+    }
+}
+
+TEST(TestGpuIndexIVFPQ, Add_IP_Raft) {
+    for (int tries = 0; tries < 2; ++tries) {
+        Options opt;
+        opt.useRaft = true;
+        opt.interleavedLayout = true;
+        opt.bitsPerCode = faiss::gpu::randVal(4, 8);
+        opt.usePrecomputed = false;
+        opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+        pickRaftEncoding(opt.codes, opt.dim, opt.bitsPerCode);
+        addTest(opt, faiss::METRIC_INNER_PRODUCT);
+    }
+}
+
+TEST(TestGpuIndexIVFPQ, QueryNaN_Raft) {
+    Options opt;
+    opt.useRaft = true;
+    opt.interleavedLayout = true;
+    opt.bitsPerCode = faiss::gpu::randVal(4, 8);
+    opt.usePrecomputed = false;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    pickRaftEncoding(opt.codes, opt.dim, opt.bitsPerCode);
+    queryNaNTest(opt);
+}
+
+TEST(TestGpuIndexIVFPQ, AddNaN_Raft) {
+    Options opt;
+    opt.useRaft = true;
+    opt.interleavedLayout = true;
+    opt.bitsPerCode = faiss::gpu::randVal(4, 8);
+    opt.usePrecomputed = false;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    pickRaftEncoding(opt.codes, opt.dim, opt.bitsPerCode);
+    addNaNTest(opt);
+}
+
+TEST(TestGpuIndexIVFPQ, CopyTo_Raft) {
+    Options opt;
+    opt.useRaft = true;
+    opt.interleavedLayout = true;
+    opt.bitsPerCode = faiss::gpu::randVal(4, 8);
+    opt.usePrecomputed = false;
+    opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
+    pickRaftEncoding(opt.codes, opt.dim, opt.bitsPerCode);
+    copyToTest(opt);
+}
+#endif
+
 TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
     // Construct on a random device to test multi-device, if we have
     // multiple devices
@@ -762,6 +824,7 @@ TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
     faiss::gpu::GpuIndexIVFPQConfig config;
     config.device = device;
     config.memorySpace = faiss::gpu::MemorySpace::Unified;
+    config.use_raft = false;
 
     faiss::gpu::GpuIndexIVFPQ gpuIndex(
             &res,
@@ -784,6 +847,34 @@ TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
             0.015f,
             0.1f,
             0.015f);
+
+#if defined USE_NVIDIA_RAFT
+    config.interleavedLayout = true;
+    config.use_raft = true;
+    config.indicesOptions = faiss::gpu::INDICES_64_BIT;
+
+    faiss::gpu::GpuIndexIVFPQ raftGpuIndex(
+            &res,
+            dim,
+            numCentroids,
+            codes,
+            bitsPerCode,
+            faiss::METRIC_L2,
+            config);
+    raftGpuIndex.copyFrom(&cpuIndex);
+    raftGpuIndex.nprobe = nprobe;
+
+    faiss::gpu::compareIndices(
+            cpuIndex,
+            raftGpuIndex,
+            numQuery,
+            dim,
+            k,
+            "Unified Memory",
+            0.015f,
+            0.1f,
+            0.015f);
+#endif
 }
 
 int main(int argc, char** argv) {

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

@@ -31,6 +31,7 @@ TEST(TestGpuMemoryException, AddException) {
 
     faiss::gpu::GpuIndexFlatConfig config;
     config.device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
+    config.use_raft = false;
 
     faiss::gpu::GpuIndexFlatL2 gpuIndexL2Broken(
             &res, (int)brokenAddDims, config);