faiss 0.6.1 → 0.6.3

data/vendor/faiss/faiss/gpu_metal/MetalPythonBridge.h

@@ -0,0 +1,45 @@
+// @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.
+ *
+ * C++-only API for Python/SWIG. No Objective-C types so SWIG can
+ * parse it. Implemented in MetalPythonBridge.mm.
+ */
+
+#pragma once
+
+#include <faiss/Index.h>
+
+namespace faiss {
+namespace gpu_metal {
+
+/// Opaque holder for Metal resources.
+struct StandardMetalResourcesHolder {
+    void* impl = nullptr;
+    StandardMetalResourcesHolder();
+    ~StandardMetalResourcesHolder();
+    StandardMetalResourcesHolder(const StandardMetalResourcesHolder&) = delete;
+    StandardMetalResourcesHolder& operator=(
+            const StandardMetalResourcesHolder&) = delete;
+};
+
+/// Same names as GPU API for unified Python binding.
+int get_num_gpus();
+void gpu_profiler_start();
+void gpu_profiler_stop();
+void gpu_sync_all_devices();
+
+/// Clone CPU index to Metal GPU. Caller owns returned index.
+faiss::Index* index_cpu_to_gpu(
+        StandardMetalResourcesHolder* res,
+        int device,
+        const faiss::Index* index);
+
+/// Copy Metal index back to CPU. Caller owns returned index.
+faiss::Index* index_gpu_to_cpu(const faiss::Index* index);
+
+} // namespace gpu_metal
+} // namespace faiss

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

@@ -0,0 +1,193 @@
+// @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 Flat implementation: GPU-resident IVF list storage and helpers.
+ * Mirrors the roles of faiss/gpu/impl/IVFFlat.cuh (storage side only).
+ */
+
+#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/GpuIndicesOptions.h>
+#include <faiss/gpu_metal/MetalResources.h>
+
+namespace faiss {
+namespace gpu_metal {
+
+/// GPU-resident IVF list storage for flat (float32) codes.
+/// Layout: all lists are stored contiguously in a single codes/ids buffer;
+/// lists are described by (listOffset[list], listLength[list]).
+class MetalIVFFlatImpl {
+   public:
+    struct AppendDebugStats {
+        size_t relayoutEvents = 0;
+        size_t movedLists = 0;
+        size_t movedVectors = 0;
+        size_t reusedSegmentAllocs = 0;
+        size_t tailSegmentAllocs = 0;
+        size_t reusedCapacityVecs = 0;
+        size_t tailCapacityVecs = 0;
+        size_t tailShrinkEvents = 0;
+        size_t tailShrunkVecs = 0;
+    };
+
+    MetalIVFFlatImpl(
+            std::shared_ptr<MetalResources> resources,
+            int dim,
+            idx_t nlist,
+            faiss::MetricType metric,
+            float metricArg,
+            faiss::gpu::IndicesOptions indicesOptions,
+            bool interleavedLayout);
+
+    ~MetalIVFFlatImpl();
+
+    /// Reset all IVF lists and free GPU storage.
+    void reset();
+
+    /// Reserve host/GPU storage for at least totalVecs vectors.
+    void reserveMemory(idx_t totalVecs);
+
+    /// Append a batch of vectors to IVF lists.
+    /// - x: host pointer, size n * dim
+    /// - list_nos: host pointer, size n; -1 entries are skipped
+    /// - xids: host pointer, size n (may be null to use internal ids)
+    void appendVectors(
+            idx_t n,
+            const float* x,
+            const idx_t* list_nos,
+            const idx_t* xids);
+
+    /// Accessors for future GPU search path.
+    int dim() const {
+        return dim_;
+    }
+    idx_t nlist() const {
+        return nlist_;
+    }
+    faiss::MetricType metricType() const {
+        return metric_type_;
+    }
+    float metricArg() const {
+        return metric_arg_;
+    }
+
+    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_;
+    }
+    /// Pre-built GPU buffer of (nlist) uint32_t offsets (updated on every add).
+    id<MTLBuffer> listOffsetGpuBuffer() const {
+        return listOffsetBuf_;
+    }
+    /// Pre-built GPU buffer of (nlist) uint32_t lengths (updated on every add).
+    id<MTLBuffer> listLengthGpuBuffer() const {
+        return listLengthBuf_;
+    }
+
+    size_t totalVecs() const {
+        return totalVecs_;
+    }
+
+    /// Interleaved codes buffer (blocks of 32 vectors, dims interleaved).
+    id<MTLBuffer> interleavedCodesBuffer() const {
+        return interleavedCodesBuf_;
+    }
+    /// Per-list float offsets into the interleaved codes buffer.
+    id<MTLBuffer> interleavedCodesOffsetBuffer() const {
+        return interleavedCodesOffsetBuf_;
+    }
+    bool interleavedLayout() const {
+        return interleavedLayout_;
+    }
+    /// Rebuild interleaved buffers from host storage if they are stale.
+    void ensureInterleavedLayoutUpToDate();
+    const AppendDebugStats& appendDebugStats() const {
+        return appendStats_;
+    }
+    void resetAppendDebugStats() {
+        appendStats_ = AppendDebugStats{};
+    }
+
+   private:
+    struct FreeSegment {
+        size_t offset = 0;
+        size_t length = 0;
+    };
+
+    bool ensureCapacityForAppend_(
+            const std::vector<size_t>& addPerList,
+            std::vector<uint8_t>* movedLists);
+    void uploadToGpu_(
+            const std::vector<size_t>& oldLength,
+            const std::vector<size_t>& addPerList,
+            const std::vector<uint8_t>& movedLists,
+            bool forceFullUpload);
+    void rebuildInterleavedBuffers_();
+    size_t allocSegment_(size_t length);
+    void freeSegment_(
+            size_t offset,
+            size_t length,
+            bool allowTailShrink = true);
+    void coalesceFreeSegments_();
+    void tryShrinkTail_();
+
+    std::shared_ptr<MetalResources> resources_;
+
+    int dim_;
+    idx_t nlist_;
+    faiss::MetricType metric_type_;
+    float metric_arg_;
+    faiss::gpu::IndicesOptions indicesOptions_;
+    bool interleavedLayout_;
+
+    // Per-list metadata
+    std::vector<size_t> listLength_;
+    std::vector<size_t> listOffset_;
+    std::vector<size_t> listCapacity_;
+
+    // Host copies of IVF data (flat layout)
+    std::vector<float> hostCodes_; // size = totalVecs_ * dim_
+    std::vector<idx_t> hostIds_;   // size = totalVecs_
+    std::vector<FreeSegment> freeSegments_;
+    AppendDebugStats appendStats_;
+    size_t totalVecs_;
+    size_t totalCapacityVecs_;
+
+    // GPU storage
+    id<MTLBuffer> codesBuffer_;
+    id<MTLBuffer> idsBuffer_;
+    id<MTLBuffer> listOffsetBuf_; // (nlist) uint32_t, list element offsets
+    id<MTLBuffer> listLengthBuf_; // (nlist) uint32_t, list sizes
+
+    // Interleaved codes layout (blocks of 32 vectors, dims interleaved)
+    id<MTLBuffer> interleavedCodesBuf_;
+    id<MTLBuffer> interleavedCodesOffsetBuf_; // (nlist) uint32_t, float offsets
+    bool interleavedDirty_ = true;
+
+    static constexpr int kInterleavedGroupSize = 32;
+};
+
+} // 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