faiss 0.6.0 → 0.6.2

data/vendor/faiss/faiss/gpu_metal/MetalCloner.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.
+ *
+ * Clone CPU <-> Metal GPU. Mirrors GpuCloner roles for Metal backend.
+ */
+
+#pragma once
+
+#include <faiss/Index.h>
+
+namespace faiss {
+namespace gpu_metal {
+
+class StandardMetalResources;
+
+/// Returns the number of Metal "devices" (1 if Metal is available, else 0).
+int get_num_gpus();
+
+/// Clone a CPU index to Metal GPU. Supports IndexFlat, IndexFlatL2,
+/// IndexFlatIP. device must be 0. Caller owns the returned index.
+faiss::Index* index_cpu_to_metal_gpu(
+        StandardMetalResources* res,
+        int device,
+        const faiss::Index* index);
+
+/// Copy a Metal index back to CPU. Supports MetalIndexFlat -> IndexFlat.
+/// Caller owns the returned index.
+faiss::Index* index_metal_gpu_to_cpu(const faiss::Index* index);
+
+} // namespace gpu_metal
+} // namespace faiss

data/vendor/faiss/faiss/gpu_metal/MetalDistance.h

@@ -0,0 +1,87 @@
+// @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.
+ *
+ * IVF distance computation and scan dispatch for Metal backend.
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+
+class MetalResources;
+
+int getMetalDistanceMaxK();
+
+bool runMetalComputeNorms(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> vectors,
+        int nb,
+        int d,
+        id<MTLBuffer> normsBuf,
+        bool waitForCompletion = true);
+
+bool runMetalIVFFlatScan(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries,
+        id<MTLBuffer> codes,
+        id<MTLBuffer> ids,
+        id<MTLBuffer> listOffset,
+        id<MTLBuffer> listLength,
+        id<MTLBuffer> coarseAssign,
+        int nq,
+        int d,
+        int k,
+        int nprobe,
+        bool isL2,
+        id<MTLBuffer> outDistances,
+        id<MTLBuffer> outIndices,
+        id<MTLBuffer> perListDistBuf,
+        id<MTLBuffer> perListIdxBuf,
+        id<MTLBuffer> interleavedCodes = nil,
+        id<MTLBuffer> interleavedCodesOffset = nil,
+        bool waitForCompletion = true);
+
+bool runMetalIVFFlatFullSearch(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries,
+        int nq,
+        int d,
+        int k,
+        int nprobe,
+        bool isL2,
+        id<MTLBuffer> centroids,
+        int nlist,
+        id<MTLBuffer> codes,
+        id<MTLBuffer> ids,
+        id<MTLBuffer> listOffset,
+        id<MTLBuffer> listLength,
+        id<MTLBuffer> outDistances,
+        id<MTLBuffer> outIndices,
+        id<MTLBuffer> perListDistBuf,
+        id<MTLBuffer> perListIdxBuf,
+        id<MTLBuffer> coarseDistBuf,
+        id<MTLBuffer> coarseIdxBuf,
+        id<MTLBuffer> distMatrixBuf,
+        id<MTLBuffer> centroidNormsBuf = nil,
+        int avgListLen = 256,
+        id<MTLBuffer> interleavedCodes = nil,
+        id<MTLBuffer> interleavedCodesOffset = nil,
+        bool centroidsAreFP16 = false,
+        bool waitForCompletion = true);
+
+} // namespace gpu_metal
+} // namespace faiss

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,58 @@
+// @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/GpuIndicesOptions.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;
+
+    bool useFloat16CoarseQuantizer = false;
+
+    faiss::gpu::IndicesOptions indicesOptions = faiss::gpu::INDICES_64_BIT;
+
+    bool interleavedLayout = true;
+};
+
+/// 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/MetalIndexIVFFlat.h

@@ -0,0 +1,181 @@
+// @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.
+ *
+ * Minimal Metal IVFFlat wrapper.
+ *
+ */
+
+#pragma once
+
+#import <Metal/Metal.h>
+
+#include <faiss/IndexIVFFlat.h>
+#include <faiss/gpu/GpuIndicesOptions.h>
+#include <faiss/gpu_metal/MetalIndex.h>
+
+#include <memory>
+
+namespace faiss {
+namespace gpu_metal {
+class MetalIVFFlatImpl;
+} // namespace gpu_metal
+} // namespace faiss
+
+namespace faiss {
+namespace gpu_metal {
+
+/// IVFFlat index wrapper for Metal backend.
+/// Currently delegates to an internal CPU IndexIVFFlat; later phases
+/// may move list scanning to GPU.
+class MetalIndexIVFFlat : public MetalIndex {
+   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;
+    };
+
+    /// Construct empty IVFFlat index with its own CPU quantizer.
+    MetalIndexIVFFlat(
+            std::shared_ptr<MetalResources> resources,
+            int dims,
+            idx_t nlist,
+            faiss::MetricType metric,
+            float metricArg = 0.0f,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    /// Construct empty IVFFlat index with caller-provided coarse quantizer.
+    /// If ownFields is true, this index takes ownership of `coarseQuantizer`.
+    MetalIndexIVFFlat(
+            std::shared_ptr<MetalResources> resources,
+            faiss::Index* coarseQuantizer,
+            int dims,
+            idx_t nlist,
+            faiss::MetricType metric,
+            float metricArg = 0.0f,
+            MetalIndexConfig config = MetalIndexConfig(),
+            bool ownFields = false);
+
+    /// Construct from an existing CPU IndexIVFFlat (used by cloners later).
+    MetalIndexIVFFlat(
+            std::shared_ptr<MetalResources> resources,
+            const faiss::IndexIVFFlat* cpuIndex,
+            MetalIndexConfig config = MetalIndexConfig());
+
+    ~MetalIndexIVFFlat() 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;
+
+    /// Search with caller-provided coarse assignments (skips coarse quantizer).
+    /// @param assign      Coarse list assignments (n x nprobe), row-major idx_t
+    /// @param centroid_dis Distances to assigned centroids (n x nprobe); unused
+    ///                     by GPU scan but accepted for API compatibility
+    /// @param store_pairs  Ignored (always false for GPU path)
+    void search_preassigned(
+            idx_t n,
+            const float* x,
+            idx_t k,
+            const idx_t* assign,
+            const float* centroid_dis,
+            float* distances,
+            idx_t* labels,
+            bool store_pairs,
+            const IVFSearchParameters* params = nullptr,
+            IndexIVFStats* stats = nullptr) const;
+
+    /// Copy from a CPU IndexIVFFlat (helper for future cloner support).
+    void copyFrom(const faiss::IndexIVFFlat* index);
+
+    /// Copy to a CPU IndexIVFFlat.
+    void copyTo(faiss::IndexIVFFlat* index) const;
+
+    /// Reconstruct a single stored vector by internal key.
+    void reconstruct(idx_t key, float* recons) const override;
+
+    /// Reconstruct n contiguous stored vectors starting at i0.
+    void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
+
+    /// Re-upload coarse quantizer centroids to GPU after external changes.
+    void updateQuantizer();
+
+    /// Return the vector indices in inverted list `listId`.
+    std::vector<idx_t> getListIndices(idx_t listId) const;
+
+    /// Return raw vector data from inverted list `listId`.
+    std::vector<float> getListVectorData(idx_t listId) const;
+
+    /// Release unused GPU memory.
+    void reclaimMemory();
+
+    /// Pre-allocate GPU storage for the given total number of vectors.
+    void reserveMemory(idx_t numVecs);
+
+    /// Accessors (needed by cloner and tests).
+    idx_t nlist() const;
+    size_t nprobe() const;
+    bool interleavedLayout() const;
+    faiss::gpu::IndicesOptions indicesOptions() const;
+    AppendDebugStats appendDebugStats() const;
+    void resetAppendDebugStats();
+
+   private:
+    std::unique_ptr<faiss::IndexIVFFlat> cpuIndex_;
+    std::unique_ptr<MetalIVFFlatImpl> gpuIvf_;
+    faiss::gpu::IndicesOptions indicesOptions_;
+    bool interleavedLayout_;
+
+    // Persistent search buffers — allocated once, grown lazily.
+    // Declared mutable so search() (const) can resize them.
+    mutable id<MTLBuffer> searchQueriesBuf_ = nil;
+    mutable id<MTLBuffer> searchCoarseBuf_ = nil;
+    mutable id<MTLBuffer> searchOutDistBuf_ = nil;
+    mutable id<MTLBuffer> searchOutIdxBuf_ = nil;
+    mutable size_t searchQueriesCap_ = 0; // bytes
+    mutable size_t searchCoarseCap_ = 0;
+    mutable size_t searchOutDistCap_ = 0;
+    mutable size_t searchOutIdxCap_ = 0;
+    mutable id<MTLBuffer> searchPerListDistBuf_ = nil;
+    mutable id<MTLBuffer> searchPerListIdxBuf_ = nil;
+    mutable size_t searchPerListDistCap_ = 0;
+    mutable size_t searchPerListIdxCap_ = 0;
+
+    // GPU coarse quantizer buffers (cached, rebuilt on train)
+    mutable id<MTLBuffer> centroidBuf_ = nil;
+    mutable id<MTLBuffer> centroidNormsBuf_ = nil; // pre-computed ||c||²
+    mutable id<MTLBuffer> coarseOutDistBuf_ = nil;
+    mutable id<MTLBuffer> coarseOutIdxBuf_ = nil;
+    mutable size_t coarseOutDistCap_ = 0;
+    mutable size_t coarseOutIdxCap_ = 0;
+    mutable id<MTLBuffer> distMatrixBuf_ = nil;
+    mutable size_t distMatrixCap_ = 0;
+
+    /// Ensures buf is at least `needed` bytes, reallocating if necessary.
+    void ensureSearchBuf_(id<MTLBuffer>& buf, size_t& cap, size_t needed) const;
+
+    /// (Re)uploads quantizer centroids to centroidBuf_.
+    void uploadCentroids_() const;
+};
+
+} // namespace gpu_metal
+} // namespace faiss

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

@@ -0,0 +1,111 @@
+// @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 {
+
+enum class IVFScanVariant { Standard, Small, Interleaved };
+
+class MetalKernels {
+   public:
+    explicit MetalKernels(id<MTLDevice> device);
+    ~MetalKernels();
+
+    bool isValid() const;
+    static constexpr int kMaxK = 2048;
+
+    void encodeDistanceMatrix(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> vectors,
+            id<MTLBuffer> distances,
+            int nq,
+            int nb,
+            int d,
+            MetricType metric);
+
+    void encodeL2WithNorms(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> vectors,
+            id<MTLBuffer> distances,
+            id<MTLBuffer> vecNorms,
+            int nq,
+            int nb,
+            int d);
+
+    void encodeComputeNorms(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> vectors,
+            id<MTLBuffer> norms,
+            int nb,
+            int d);
+
+    void encodeTopKThreadgroup(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> distances,
+            id<MTLBuffer> outDist,
+            id<MTLBuffer> outIdx,
+            int nq,
+            int nb,
+            int k,
+            bool wantMin);
+
+    void encodeIVFScanList(
+            id<MTLComputeCommandEncoder> enc,
+            IVFScanVariant variant,
+            id<MTLBuffer> queries,
+            id<MTLBuffer> codes,
+            id<MTLBuffer> ids,
+            id<MTLBuffer> listOffset,
+            id<MTLBuffer> listLength,
+            id<MTLBuffer> coarseAssign,
+            id<MTLBuffer> perListDist,
+            id<MTLBuffer> perListIdx,
+            id<MTLBuffer> paramsBuf,
+            int nq,
+            int nprobe,
+            id<MTLBuffer> ilCodesOffset = nil);
+
+    void encodeIVFMergeLists(
+            id<MTLComputeCommandEncoder> enc,
+            id<MTLBuffer> perListDist,
+            id<MTLBuffer> perListIdx,
+            id<MTLBuffer> outDist,
+            id<MTLBuffer> outIdx,
+            id<MTLBuffer> paramsBuf,
+            int nq);
+
+    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, 512, 1024, 2048};
+    static constexpr int kNumTopKVariants = 7;
+};
+
+MetalKernels& getMetalKernels(id<MTLDevice> device);
+
+} // namespace gpu_metal
+} // namespace faiss

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