faiss 0.6.2 → 0.6.3

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

@@ -11,8 +11,10 @@
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/test/TestUtils.h>
 #include <faiss/gpu/utils/DeviceUtils.h>
+#include <faiss/impl/IDSelector.h>
 #include <gtest/gtest.h>
 #include <cmath>
+#include <limits> // IWYU pragma: keep
 #include <sstream>
 #include <vector>
 
@@ -74,6 +76,7 @@ void runCopyToTest(faiss::ScalarQuantizer::QuantizerType qtype) {
 
     auto config = GpuIndexIVFScalarQuantizerConfig();
     config.device = opt.device;
+    config.use_cuvs = false;
 
     GpuIndexIVFScalarQuantizer gpuIndex(
             &res, opt.dim, opt.numCentroids, qtype, METRIC_L2, true, config);
@@ -163,6 +166,7 @@ void runCopyFromTest(faiss::ScalarQuantizer::QuantizerType qtype) {
 
     auto config = GpuIndexIVFScalarQuantizerConfig();
     config.device = opt.device;
+    config.use_cuvs = false;
 
     GpuIndexIVFScalarQuantizer gpuIndex(
             &res,
@@ -223,6 +227,182 @@ TEST(TestGpuIndexIVFScalarQuantizer, CopyFrom_4bit_uniform) {
     runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform);
 }
 
+#if defined USE_NVIDIA_CUVS
+void runCuvsSQ8Test(faiss::MetricType metricType) {
+    using namespace faiss;
+    using namespace faiss::gpu;
+
+    Options opt;
+    opt.indicesOpt = INDICES_64_BIT;
+
+    std::vector<float> trainVecs = randVecs(opt.numTrain, opt.dim);
+    std::vector<float> addVecs = randVecs(opt.numAdd, opt.dim);
+    std::vector<float> queryVecs = randVecs(opt.numQuery, opt.dim);
+
+    StandardGpuResources res;
+    res.noTempMemory();
+
+    auto config = GpuIndexIVFScalarQuantizerConfig();
+    config.device = opt.device;
+    config.indicesOptions = INDICES_64_BIT;
+    config.use_cuvs = true;
+    if (!should_use_cuvs(config)) {
+        GTEST_SKIP() << "cuVS is not enabled for this GPU";
+    }
+
+    GpuIndexIVFScalarQuantizer gpuIndex(
+            &res,
+            opt.dim,
+            opt.numCentroids,
+            ScalarQuantizer::QT_8bit,
+            metricType,
+            true,
+            config);
+    gpuIndex.train(opt.numTrain, trainVecs.data());
+    gpuIndex.add(opt.numAdd, addVecs.data());
+    gpuIndex.nprobe = opt.nprobe;
+
+    std::vector<float> distances(opt.numQuery * opt.k);
+    std::vector<idx_t> labels(opt.numQuery * opt.k);
+    gpuIndex.search(
+            opt.numQuery,
+            queryVecs.data(),
+            opt.k,
+            distances.data(),
+            labels.data());
+    EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
+    EXPECT_NE(labels[0], idx_t(-1));
+
+    IndexFlatL2 quantizerL2(opt.dim);
+    IndexFlatIP quantizerIP(opt.dim);
+    Index* quantizer = metricType == METRIC_L2 ? (Index*)&quantizerL2
+                                               : (Index*)&quantizerIP;
+
+    IndexIVFScalarQuantizer cpuCopy(
+            quantizer,
+            opt.dim,
+            opt.numCentroids,
+            ScalarQuantizer::QT_8bit,
+            metricType,
+            true);
+    gpuIndex.copyTo(&cpuCopy);
+    cpuCopy.nprobe = opt.nprobe;
+
+    EXPECT_EQ(cpuCopy.ntotal, gpuIndex.ntotal);
+    EXPECT_EQ(cpuCopy.d, gpuIndex.d);
+    EXPECT_EQ(cpuCopy.nlist, gpuIndex.getNumLists());
+    EXPECT_EQ(cpuCopy.sq.qtype, ScalarQuantizer::QT_8bit);
+    EXPECT_EQ(cpuCopy.sq.trained.size(), 2 * opt.dim);
+    EXPECT_TRUE(cpuCopy.by_residual);
+
+    compareIndices(
+            cpuCopy,
+            gpuIndex,
+            opt.numQuery,
+            opt.dim,
+            opt.k,
+            opt.toString(),
+            0.08f,
+            0.2f,
+            0.05f);
+
+    GpuIndexIVFScalarQuantizer copiedGpuIndex(
+            &res, 1, 1, ScalarQuantizer::QT_4bit, METRIC_L2, false, config);
+    copiedGpuIndex.copyFrom(&cpuCopy);
+    copiedGpuIndex.nprobe = opt.nprobe;
+    EXPECT_EQ(copiedGpuIndex.ntotal, cpuCopy.ntotal);
+    testIVFEquality(cpuCopy, copiedGpuIndex);
+
+    compareIndices(
+            cpuCopy,
+            copiedGpuIndex,
+            opt.numQuery,
+            opt.dim,
+            opt.k,
+            opt.toString(),
+            0.08f,
+            0.2f,
+            0.05f);
+
+    SearchParametersIVF searchParams;
+    searchParams.nprobe = opt.nprobe;
+    TestIDSelectorStruct selectorStruct(opt.numAdd);
+    for (auto& [selectorName, selector] : selectorStruct.selector_map) {
+        searchParams.sel = selector.get();
+        testIDSelectorSearch(
+                &gpuIndex,
+                &searchParams,
+                queryVecs,
+                opt.numQuery,
+                opt.k,
+                selectorName);
+    }
+
+    GpuIndexIVFScalarQuantizer shiftedGpuIndex(
+            &res,
+            opt.dim,
+            opt.numCentroids,
+            ScalarQuantizer::QT_8bit,
+            metricType,
+            true,
+            config);
+    shiftedGpuIndex.train(opt.numTrain, trainVecs.data());
+    std::vector<idx_t> shiftedIds(opt.numAdd);
+    idx_t shiftedIdBase = opt.numAdd * 4 + 123;
+    for (int i = 0; i < opt.numAdd; ++i) {
+        shiftedIds[i] = shiftedIdBase + i;
+    }
+    shiftedGpuIndex.add_with_ids(opt.numAdd, addVecs.data(), shiftedIds.data());
+    shiftedGpuIndex.nprobe = opt.nprobe;
+
+    SearchParametersIVF shiftedSearchParams;
+    shiftedSearchParams.nprobe = opt.nprobe;
+    IDSelectorRange shiftedSelector(shiftedIdBase, shiftedIdBase + opt.numAdd);
+    shiftedSearchParams.sel = &shiftedSelector;
+    std::vector<float> shiftedDistances(opt.numQuery * opt.k);
+    std::vector<idx_t> shiftedLabels(opt.numQuery * opt.k, -1);
+    shiftedGpuIndex.search(
+            opt.numQuery,
+            queryVecs.data(),
+            opt.k,
+            shiftedDistances.data(),
+            shiftedLabels.data(),
+            &shiftedSearchParams);
+
+    bool foundShiftedLabel = false;
+    for (idx_t label : shiftedLabels) {
+        if (label >= 0) {
+            foundShiftedLabel = true;
+            EXPECT_TRUE(shiftedSelector.is_member(label));
+        }
+    }
+    EXPECT_TRUE(foundShiftedLabel);
+
+    auto nanQueries = queryVecs;
+    nanQueries[opt.dim + 3] = std::numeric_limits<float>::quiet_NaN();
+    gpuIndex.search(
+            2, nanQueries.data(), opt.k, distances.data(), labels.data());
+    for (int j = 0; j < opt.k; ++j) {
+        EXPECT_EQ(labels[opt.k + j], idx_t(-1));
+        EXPECT_EQ(distances[opt.k + j], std::numeric_limits<float>::max());
+    }
+
+    gpuIndex.reset();
+    EXPECT_EQ(gpuIndex.ntotal, 0);
+    for (int i = 0; i < opt.numCentroids; ++i) {
+        EXPECT_EQ(gpuIndex.getListLength(i), 0);
+    }
+}
+
+TEST(TestCuvsGpuIndexIVFScalarQuantizer, SQ8_L2) {
+    runCuvsSQ8Test(faiss::METRIC_L2);
+}
+
+TEST(TestCuvsGpuIndexIVFScalarQuantizer, SQ8_IP) {
+    runCuvsSQ8Test(faiss::METRIC_INNER_PRODUCT);
+}
+#endif
+
 int main(int argc, char** argv) {
     testing::InitGoogleTest(&argc, argv);
 

data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFFlat.h

@@ -21,12 +21,8 @@
 
 namespace faiss {
 namespace gpu_metal {
-class MetalIVFFlatImpl;
-} // namespace gpu_metal
-} // namespace faiss
 
-namespace faiss {
-namespace gpu_metal {
+class MetalIVFFlatImpl;
 
 /// IVFFlat index wrapper for Metal backend.
 /// Currently delegates to an internal CPU IndexIVFFlat; later phases

data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFPQ.h

@@ -0,0 +1,88 @@
+// @lint-ignore-every LICENSELINT
+/**
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * Metal IVF-PQ index: 8-bit product quantization with precomputed
+ * per-query lookup tables and GPU IVF list scanning.
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <faiss/IndexIVFPQ.h>
+#include <faiss/gpu_metal/MetalIndex.h>
+
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+class MetalIVFPQImpl;
+
+class MetalIndexIVFPQ : public MetalIndex {
+   public:
+    MetalIndexIVFPQ(
+            std::shared_ptr<MetalResources> resources,
+            int dims,
+            idx_t nlist,
+            int M,
+            int nbitsPerIdx,
+            faiss::MetricType metric,
+            float metricArg = 0.0f,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    MetalIndexIVFPQ(
+            std::shared_ptr<MetalResources> resources,
+            const faiss::IndexIVFPQ* cpuIndex,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    ~MetalIndexIVFPQ() override;
+
+    void train(idx_t n, const float* x) override;
+    void add(idx_t n, const float* x) override;
+    void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
+    void reset() override;
+
+    void search(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            float* distances,
+            idx_t* labels,
+            const SearchParameters* params = nullptr) const override;
+
+    void copyFrom(const faiss::IndexIVFPQ* index);
+    void copyTo(faiss::IndexIVFPQ* index) const;
+
+    void updateQuantizer();
+    std::vector<idx_t> getListIndices(idx_t listId) const;
+    void reclaimMemory();
+
+    /// Pre-allocate GPU storage for the given total number of vectors.
+    void reserveMemory(idx_t numVecs);
+
+    idx_t nlist() const;
+    size_t nprobe() const;
+    int getNumSubQuantizers() const;
+    void setUsePrecomputedTables(bool enable);
+    bool getUsePrecomputedTables() const;
+
+   private:
+    std::unique_ptr<faiss::IndexIVFPQ> cpuIndex_;
+    std::unique_ptr<MetalIVFPQImpl> gpuIvf_;
+
+    void verifyPQSettings_() const;
+
+    void encodeResidualAndAppend_(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            const idx_t* xids);
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/impl/MetalIVFPQ.h

@@ -0,0 +1,134 @@
+// @lint-ignore-every LICENSELINT
+/**
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * Metal IVF PQ implementation: GPU-resident IVF list storage for
+ * 8-bit product-quantized codes. Each vector is M bytes (one byte
+ * per subquantizer, ksub=256).
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include <faiss/Index.h>
+#include <faiss/MetricType.h>
+#include <faiss/gpu_metal/MetalResources.h>
+
+namespace faiss {
+namespace gpu_metal {
+
+class MetalIVFPQImpl {
+   public:
+    MetalIVFPQImpl(
+            std::shared_ptr<MetalResources> resources,
+            int dim,
+            idx_t nlist,
+            int numSubQuantizers,
+            int bitsPerCode,
+            faiss::MetricType metric,
+            float metricArg);
+
+    ~MetalIVFPQImpl();
+
+    void reset();
+    void reserveMemory(idx_t totalVecs);
+
+    /// Append PQ-encoded vectors. codes: n * codeSize bytes.
+    void appendCodes(
+            idx_t n,
+            const uint8_t* codes,
+            const idx_t* list_nos,
+            const idx_t* xids);
+
+    /// Upload PQ centroids: M * ksub * dsub floats, row-major.
+    /// Layout: pqCentroids[m][c][dsub_dim]
+    void setPQCentroids(const float* centroids);
+
+    int dim() const {
+        return dim_;
+    }
+    idx_t nlist() const {
+        return nlist_;
+    }
+    int numSubQuantizers() const {
+        return M_;
+    }
+    int bitsPerCode() const {
+        return bitsPerCode_;
+    }
+    int ksub() const {
+        return ksub_;
+    }
+    int dsub() const {
+        return dsub_;
+    }
+    size_t codeSize() const {
+        return codeSize_;
+    }
+
+    const std::vector<size_t>& listLength() const {
+        return listLength_;
+    }
+    const std::vector<size_t>& listOffset() const {
+        return listOffset_;
+    }
+
+    id<MTLBuffer> codesBuffer() const {
+        return codesBuffer_;
+    }
+    id<MTLBuffer> idsBuffer() const {
+        return idsBuffer_;
+    }
+    id<MTLBuffer> listOffsetGpuBuffer() const {
+        return listOffsetBuf_;
+    }
+    id<MTLBuffer> listLengthGpuBuffer() const {
+        return listLengthBuf_;
+    }
+    id<MTLBuffer> pqCentroidsBuffer() const {
+        return pqCentroidsBuf_;
+    }
+    size_t totalVecs() const {
+        return totalVecs_;
+    }
+
+   private:
+    void uploadToGpu();
+
+    std::shared_ptr<MetalResources> resources_;
+
+    int dim_;
+    idx_t nlist_;
+    int M_;
+    int bitsPerCode_;
+    int ksub_;
+    int dsub_;
+    size_t codeSize_;
+    faiss::MetricType metric_type_;
+    float metric_arg_;
+
+    std::vector<size_t> listLength_;
+    std::vector<size_t> listOffset_;
+
+    std::vector<uint8_t> hostCodes_;
+    std::vector<idx_t> hostIds_;
+    size_t totalVecs_;
+
+    id<MTLBuffer> codesBuffer_;
+    id<MTLBuffer> idsBuffer_;
+    id<MTLBuffer> listOffsetBuf_;
+    id<MTLBuffer> listLengthBuf_;
+    id<MTLBuffer> pqCentroidsBuf_;
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp

@@ -157,8 +157,8 @@ void ClusteringInitialization::init_centroids(
             n,
             k);
     FAISS_THROW_IF_NOT(d > 0);
-    FAISS_THROW_IF_NOT(x != nullptr);
-    FAISS_THROW_IF_NOT(centroids != nullptr);
+    FAISS_THROW_IF_NOT(x);
+    FAISS_THROW_IF_NOT(centroids);
     FAISS_THROW_IF_NOT(
             n_existing_centroids == 0 || existing_centroids != nullptr);
 

data/vendor/faiss/faiss/impl/DistanceComputer.h

@@ -130,6 +130,26 @@ struct FlatCodesDistanceComputer : DistanceComputer {
         return distance_to_code(codes + i * code_size);
     }
 
+    void distances_batch_4(
+            const idx_t idx0,
+            const idx_t idx1,
+            const idx_t idx2,
+            const idx_t idx3,
+            float& dis0,
+            float& dis1,
+            float& dis2,
+            float& dis3) override {
+        distance_to_code_batch_4(
+                codes + idx0 * code_size,
+                codes + idx1 * code_size,
+                codes + idx2 * code_size,
+                codes + idx3 * code_size,
+                dis0,
+                dis1,
+                dis2,
+                dis3);
+    }
+
     /// Computes a partial dot product over a slice of the query vector.
     /// The slice is defined by the following parameters:
     ///   — `offset`: the starting index of the first component to include
@@ -167,6 +187,20 @@ struct FlatCodesDistanceComputer : DistanceComputer {
 
     /// compute distance of current query to an encoded vector
     virtual float distance_to_code(const uint8_t* code) = 0;
+    virtual void distance_to_code_batch_4(
+            const uint8_t* c1,
+            const uint8_t* c2,
+            const uint8_t* c3,
+            const uint8_t* c4,
+            float& d1,
+            float& d2,
+            float& d3,
+            float& d4) {
+        d1 = distance_to_code(c1);
+        d2 = distance_to_code(c2);
+        d3 = distance_to_code(c3);
+        d4 = distance_to_code(c4);
+    }
 
     /// Compute partial dot products of current query to 4 stored vectors.
     /// See `partial_dot_product` for more details.

data/vendor/faiss/faiss/impl/EDENQuantizer.h

@@ -0,0 +1,119 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+#include <faiss/MetricType.h>
+#include <faiss/impl/DistanceComputer.h>
+#include <faiss/impl/ScalarQuantizer.h>
+#include <faiss/impl/platform_macros.h>
+
+namespace faiss {
+
+enum EDENScaleType {
+    // Original EDEN unbiased scale.
+    EDENScaleType_UNBIASED = 1,
+    // Biased DRIVE scale (NeurIPS 2021, arXiv:2105.08339), generalized to
+    // multi-bit EDEN; see also https://arxiv.org/abs/2604.18555.
+    EDENScaleType_BIASED = 2,
+};
+
+FAISS_PACK_STRUCTS_BEGIN
+struct FAISS_PACKED EDENCodeFactors {
+    // L2 term used by the distance computer. For the unbiased scale this is
+    // ||r||^2. For the biased scale this is S^2 * ||q_code||^2, the
+    // reconstructed-code norm.
+    float l2_norm_term = 0;
+    // Per-vector reconstruction scale.
+    float scale = 0;
+};
+FAISS_PACK_STRUCTS_END
+
+struct EDENFlatCodesDistanceComputer : FlatCodesDistanceComputer {
+    using FlatCodesDistanceComputer::FlatCodesDistanceComputer;
+
+    virtual void consecutive_distances_batch_8(idx_t first, float* distances) {
+        distances_batch_4(
+                first,
+                first + 1,
+                first + 2,
+                first + 3,
+                distances[0],
+                distances[1],
+                distances[2],
+                distances[3]);
+        distances_batch_4(
+                first + 4,
+                first + 5,
+                first + 6,
+                first + 7,
+                distances[4],
+                distances[5],
+                distances[6],
+                distances[7]);
+    }
+
+    virtual void consecutive_distances_batch_16(idx_t first, float* distances) {
+        consecutive_distances_batch_8(first, distances);
+        consecutive_distances_batch_8(first + 8, distances + 8);
+    }
+};
+
+// EDEN Lloyd-Max quantizer from the EDEN ICML 2022 paper:
+// https://proceedings.mlr.press/v162/vargaftik22a.html.
+//
+// EDEN operates on the vectors supplied to it. Optional preprocessing such as
+// a random rotation can be applied externally with IndexPreTransform. The
+// scalar assignment itself is a ScalarQuantizer::QT_*_eden qtype. EDEN adds
+// per-vector scale factors after the packed scalar code and computes them
+// against an externally supplied centroid; nullptr means a zero centroid. The
+// default stored scale is EDEN's original unbiased scale. The biased scale
+// follows DRIVE's MSE-minimizing scale for the chosen Lloyd-Max codeword; see
+// also
+// https://arxiv.org/abs/2604.18555.
+
+namespace eden_utils {
+
+ScalarQuantizer::QuantizerType quantizer_type_for_bits(size_t nb_bits);
+
+bool is_eden_quantizer_type(ScalarQuantizer::QuantizerType qtype);
+
+size_t nb_bits_for_qtype(ScalarQuantizer::QuantizerType qtype);
+
+size_t packed_code_size(size_t d, size_t nb_bits);
+
+size_t code_size(size_t d, size_t nb_bits);
+
+uint8_t extract_code(const uint8_t* codes, size_t index, size_t nb_bits);
+
+void compute_codes(
+        const ScalarQuantizer& sq,
+        MetricType metric_type,
+        EDENScaleType scale_type,
+        const float* x,
+        uint8_t* codes,
+        size_t n,
+        const float* centroid = nullptr);
+
+void decode(
+        const ScalarQuantizer& sq,
+        const uint8_t* codes,
+        float* x,
+        size_t n,
+        const float* centroid = nullptr);
+
+EDENFlatCodesDistanceComputer* get_distance_computer(
+        const ScalarQuantizer& sq,
+        MetricType metric_type,
+        const float* centroid = nullptr);
+
+} // namespace eden_utils
+
+} // namespace faiss