faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/gpu_metal/MetalFlatKernels.h

@@ -0,0 +1,40 @@
+// @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.
+ *
+ * Objective-C++ header. Runs L2/IP distance + top-k via Metal compute.
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <cstddef>
+
+namespace faiss {
+namespace gpu_metal {
+
+/// Runs GPU search: distance matrix (L2 or IP) then top-k. Uses shared buffers
+/// (queries, vectors, outDistances, outIndices). outIndices are int32
+/// (0..nb-1). Maximum k supported by the GPU top-k kernel (256).
+int getMetalFlatSearchMaxK();
+
+/// Returns true on success; false if pipeline creation failed.
+bool runFlatSearchGPU(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries, // (nq * d) float, row-major
+        id<MTLBuffer> vectors, // (nb * d) float, row-major
+        int nq,
+        int nb,
+        int d,
+        int k,
+        bool isL2,                  // true = L2 squared, false = inner product
+        id<MTLBuffer> outDistances, // (nq * k) float
+        id<MTLBuffer> outIndices);  // (nq * k) int32
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalIndex.h

@@ -0,0 +1,51 @@
+// @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.
+ *
+ * Objective-C++ header (uses MetalResources).
+ */
+
+#pragma once
+
+#include <faiss/Index.h>
+#include <faiss/gpu_metal/MetalResources.h>
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+/// Configuration for Metal index (mirrors GpuIndexConfig roles).
+struct MetalIndexConfig {
+    int device = 0;
+};
+
+/// Base class for Metal-backed indexes. Mirrors faiss::gpu::GpuIndex.
+class MetalIndex : public faiss::Index {
+   public:
+    MetalIndex(
+            std::shared_ptr<MetalResources> resources,
+            int dims,
+            faiss::MetricType metric,
+            float metricArg,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    int getDevice() const {
+        return config_.device;
+    }
+    std::shared_ptr<MetalResources> getResources() {
+        return resources_;
+    }
+    std::shared_ptr<const MetalResources> getResources() const {
+        return resources_;
+    }
+
+   protected:
+    std::shared_ptr<MetalResources> resources_;
+    MetalIndexConfig config_;
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalIndexFlat.h

@@ -0,0 +1,65 @@
+// @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.
+ *
+ * Objective-C++ header (uses Metal types).
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <faiss/Index.h>
+#include <faiss/gpu_metal/MetalIndex.h>
+
+namespace faiss {
+struct IndexFlat;
+}
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+/// Flat index that stores vectors in an MTLBuffer. Supports L2 and inner
+/// product. Search runs on GPU via Metal compute (distance + top-k kernels).
+class MetalIndexFlat : public MetalIndex {
+   public:
+    MetalIndexFlat(
+            std::shared_ptr<MetalResources> resources,
+            int dims,
+            faiss::MetricType metric,
+            float metricArg = 0.0f,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    ~MetalIndexFlat() 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;
+
+    /// Copy vectors to a CPU IndexFlat (e.g. for index_metal_gpu_to_cpu).
+    void copyTo(::faiss::IndexFlat* index) const;
+
+   private:
+    /// Ensures vector buffer can hold at least \p newNtotal vectors; grows
+    /// buffer if necessary.
+    void ensureCapacity(idx_t newNtotal);
+
+    /// Vector storage (row-major, ntotal * d floats). Nil when empty.
+    id<MTLBuffer> vectorsBuffer_;
+    /// Capacity of vectorsBuffer_ in number of vectors (0 if buffer is nil).
+    size_t capacityVecs_;
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalKernels.h

@@ -0,0 +1,66 @@
+// @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.
+ *
+ * MetalKernels: typed wrapper around Metal compute kernels.
+ * Owns library compilation, pipeline caching, and dispatch encoding.
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+#include <faiss/MetricType.h>
+#include <string>
+#include <unordered_map>
+
+namespace faiss {
+namespace gpu_metal {
+
+class MetalKernels {
+   public:
+    explicit MetalKernels(id<MTLDevice> device);
+    ~MetalKernels();
+
+    bool isValid() const;
+    static constexpr int kMaxK = 256;
+
+    void encodeDistanceMatrix(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> vectors,
+            id<MTLBuffer> distances,
+            int nq,
+            int nb,
+            int d,
+            MetricType metric);
+
+    void encodeTopKThreadgroup(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> distances,
+            id<MTLBuffer> outDist,
+            id<MTLBuffer> outIdx,
+            int nq,
+            int nb,
+            int k,
+            bool wantMin);
+
+    static int selectTopKVariantIndex(int k);
+
+   private:
+    id<MTLComputePipelineState> pipeline(const char* name);
+
+    id<MTLDevice> device_;
+    id<MTLLibrary> library_;
+    std::unordered_map<std::string, id<MTLComputePipelineState>> cache_;
+
+    static constexpr int kTopKVariantSizes[] = {32, 64, 128, 256};
+    static constexpr int kNumTopKVariants = 4;
+};
+
+MetalKernels& getMetalKernels(id<MTLDevice> device);
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalResources.h

@@ -0,0 +1,79 @@
+// @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.
+ *
+ * This header uses Objective-C types (Metal framework: id, nil, MTLDevice,
+ * etc.). For correct IDE/linter behavior, associate this file with
+ * "Objective-C++":
+ *
+ */
+
+#pragma once
+
+#import <Foundation/Foundation.h>
+#import <Metal/Metal.h>
+
+#include <cstddef>
+
+namespace faiss {
+namespace gpu_metal {
+
+/// Allocation type for Metal buffers (mirrors faiss::gpu::AllocType roles).
+enum MetalAllocType {
+    Other = 0,
+    FlatData = 1,
+    IVFLists = 2,
+    Quantizer = 3,
+    QuantizerPrecomputedCodes = 4,
+    TemporaryMemoryBuffer = 10,
+    TemporaryMemoryOverflow = 11,
+};
+
+/// Owns Metal device, command queue, and provides buffer allocation.
+/// Mirrors the roles of faiss::gpu::GpuResources for the Metal backend.
+class MetalResources {
+   public:
+    MetalResources();
+    ~MetalResources();
+
+    MetalResources(const MetalResources&) = delete;
+    MetalResources& operator=(const MetalResources&) = delete;
+
+    /// Returns the Metal device (nil if no Metal-capable device is available).
+    id<MTLDevice> getDevice() const {
+        return device_;
+    }
+
+    /// Returns the command queue for the device (nil if device is nil).
+    id<MTLCommandQueue> getCommandQueue() const {
+        return commandQueue_;
+    }
+
+    /// Allocates a buffer of the given size (bytes). Caller owns the returned
+    /// buffer and must call deallocBuffer when done, or the buffer will leak.
+    /// Returns nil on failure (e.g. device nil or allocation failure).
+    id<MTLBuffer> allocBuffer(size_t size, MetalAllocType type);
+
+    /// Releases a buffer previously returned by allocBuffer. The caller must
+    /// not use the buffer after this call.
+    void deallocBuffer(id<MTLBuffer> buffer, MetalAllocType type);
+
+    /// Blocks until all work submitted to the default command queue has
+    /// completed.
+    void synchronize();
+
+    /// Returns true if the Metal device and queue are available.
+    bool isAvailable() const {
+        return device_ != nil && commandQueue_ != nil;
+    }
+
+   private:
+    id<MTLDevice> device_;
+    id<MTLCommandQueue> commandQueue_;
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/StandardMetalResources.h

@@ -0,0 +1,35 @@
+// @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.
+ *
+ * Mirrors the role of StandardGpuResources for the Metal backend.
+ */
+
+#pragma once
+
+#include <faiss/gpu_metal/MetalResources.h>
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+/// Default Metal resources (single device). Use with index_cpu_to_metal_gpu.
+class StandardMetalResources {
+   public:
+    StandardMetalResources();
+    std::shared_ptr<MetalResources> getResources() const {
+        return res_;
+    }
+    bool isAvailable() const {
+        return res_ && res_->isAvailable();
+    }
+
+   private:
+    std::shared_ptr<MetalResources> res_;
+};
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/impl/AdSampling.cpp

@@ -0,0 +1,103 @@
+/*
+ * 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.
+ */
+
+#include <faiss/impl/AdSampling.h>
+
+#include <cmath>
+
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss {
+namespace detail {
+
+double normal_quantile(double p) {
+    // Three-branch rational polynomial; branch breakpoint p_low = 0.02425.
+    static constexpr double a[] = {
+            -3.969683028665376e+01,
+            2.209460984245205e+02,
+            -2.759285104469687e+02,
+            1.383577518672690e+02,
+            -3.066479806614716e+01,
+            2.506628277459239e+00,
+    };
+    static constexpr double b[] = {
+            -5.447609879822406e+01,
+            1.615858368580409e+02,
+            -1.556989798598866e+02,
+            6.680131188771972e+01,
+            -1.328068155288572e+01,
+    };
+    static constexpr double c[] = {
+            -7.784894002430293e-03,
+            -3.223964580411365e-01,
+            -2.400758277161838e+00,
+            -2.549732539343734e+00,
+            4.374664141464968e+00,
+            2.938163982698783e+00,
+    };
+    static constexpr double d[] = {
+            7.784695709041462e-03,
+            3.224671290700398e-01,
+            2.445134137142996e+00,
+            3.754408661907416e+00,
+    };
+    constexpr double p_low = 0.02425;
+    constexpr double p_high = 1.0 - p_low;
+    if (p < p_low) {
+        const double q = std::sqrt(-2.0 * std::log(p));
+        return (((((c[0] * q + c[1]) * q + c[2]) * q + c[3]) * q + c[4]) * q +
+                c[5]) /
+                ((((d[0] * q + d[1]) * q + d[2]) * q + d[3]) * q + 1.0);
+    } else if (p <= p_high) {
+        const double q = p - 0.5;
+        const double r = q * q;
+        return (((((a[0] * r + a[1]) * r + a[2]) * r + a[3]) * r + a[4]) * r +
+                a[5]) *
+                q /
+                (((((b[0] * r + b[1]) * r + b[2]) * r + b[3]) * r + b[4]) * r +
+                 1.0);
+    } else {
+        const double q = std::sqrt(-2.0 * std::log(1.0 - p));
+        return -(((((c[0] * q + c[1]) * q + c[2]) * q + c[3]) * q + c[4]) * q +
+                 c[5]) /
+                ((((d[0] * q + d[1]) * q + d[2]) * q + d[3]) * q + 1.0);
+    }
+}
+
+double chi2_quantile_wh(int p, double alpha) {
+    FAISS_THROW_IF_NOT(p > 0);
+    // Wilson-Hilferty cube-root approximation:
+    //   ((X/p)^(1/3) - (1 - 2/(9p))) / sqrt(2/(9p)) ~ N(0,1)
+    // inverted into a quantile formula.
+    //
+    // Domain constraint: for very small alpha (< ~0.001) and small p
+    // (< 4), t can go negative, producing a negative chi-squared quantile
+    // (physically impossible). In practice this cannot happen here:
+    // precompute_ad_thresholds calls with alpha = 1 - epsilon where
+    // epsilon = ad_epsilon_factor / d, and d_prime_min >= 16, so
+    // p >= 16 and alpha >= 1 - 1/16 = 0.9375 — well inside the accurate
+    // region of the approximation.
+    const double z = normal_quantile(alpha);
+    const double t = 1.0 - 2.0 / (9.0 * p) + z * std::sqrt(2.0 / (9.0 * p));
+    return p * t * t * t;
+}
+
+std::vector<float> precompute_ad_thresholds(int d, double epsilon) {
+    FAISS_THROW_IF_NOT_MSG(
+            epsilon > 0.0 && epsilon < 1.0,
+            "precompute_ad_thresholds: epsilon must be in (0, 1)");
+    FAISS_THROW_IF_NOT_MSG(
+            d > 0, "precompute_ad_thresholds: d must be positive");
+    std::vector<float> coeff(d + 1);
+    for (int p = 1; p <= d; p++) {
+        coeff[p] = static_cast<float>(chi2_quantile_wh(p, 1.0 - epsilon) / d);
+    }
+    return coeff;
+}
+
+} // namespace detail
+} // namespace faiss

data/vendor/faiss/faiss/impl/AdSampling.h

@@ -0,0 +1,35 @@
+/*
+ * 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 <vector>
+
+namespace faiss {
+namespace detail {
+
+/** Inverse standard normal CDF. Three-branch rational polynomial,
+ * absolute error < 1.15e-9 over `p in (0, 1)`. Behavior at the boundaries
+ * (p <= 0 or p >= 1) is unspecified — returns NaN or +/-inf. */
+double normal_quantile(double p);
+
+/** Chi-squared quantile via cube-root approximation. Validated to within
+ * 2% of scipy for `p in [16, d]` and `alpha <= 1 - 1e-6`. Accuracy
+ * degrades for smaller `p` or for `alpha` near 1. */
+double chi2_quantile_wh(int p, double alpha);
+
+/** Build ADSampling threshold table of size `d + 1`:
+ *   coeff[p] = chi2_quantile_wh(p, 1 - epsilon) / d.
+ *
+ * Indexing: coeff[0] is reserved (left at 0.0f). coeff[1..15] are
+ * computed but NOT accuracy-bounded — callers requiring the 2% scipy
+ * tolerance must consume only coeff[16..d]. SuperKMeans enforces
+ * this via its `d_prime_min = 16` parameter. */
+std::vector<float> precompute_ad_thresholds(int d, double epsilon);
+
+} // namespace detail
+} // namespace faiss