faiss 0.6.1 → 0.6.2

data/vendor/faiss/faiss/factory_tools.cpp

@@ -43,6 +43,10 @@ const std::map<faiss::ScalarQuantizer::QuantizerType, std::string> sq_types = {
         {faiss::ScalarQuantizer::QT_3bit_tqmse, "SQtqmse3"},
         {faiss::ScalarQuantizer::QT_4bit_tqmse, "SQtqmse4"},
         {faiss::ScalarQuantizer::QT_8bit_tqmse, "SQtqmse8"},
+        {faiss::ScalarQuantizer::QT_2bit_tq, "SQtq2"},
+        {faiss::ScalarQuantizer::QT_3bit_tq, "SQtq3"},
+        {faiss::ScalarQuantizer::QT_4bit_tq, "SQtq4"},
+        {faiss::ScalarQuantizer::QT_5bit_tq, "SQtq5"},
 };
 
 int get_hnsw_M(const faiss::IndexHNSW* index) {

data/vendor/faiss/faiss/gpu/GpuResources.h

@@ -33,7 +33,8 @@
 
 #if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <cuda/memory_resource>
+#include <optional>
 #endif
 
 namespace faiss {
@@ -163,7 +164,7 @@ struct AllocRequest : public AllocInfo {
     size_t size = 0;
 
 #if defined USE_NVIDIA_CUVS
-    rmm::mr::device_memory_resource* mr = nullptr;
+    std::optional<cuda::mr::any_resource<cuda::mr::device_accessible>> mr;
 #endif
 };
 

data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp

@@ -92,8 +92,8 @@ std::string allocsToString(const std::unordered_map<void*, AllocRequest>& map) {
 StandardGpuResourcesImpl::StandardGpuResourcesImpl()
         :
 #if defined USE_NVIDIA_CUVS
-          mmr_(new rmm::mr::managed_memory_resource),
-          pmr_(new rmm::mr::pinned_host_memory_resource),
+          mmr_{},
+          pmr_{},
 #endif
           pinnedMemAlloc_(nullptr),
           pinnedMemAllocSize_(0),
@@ -164,7 +164,7 @@ StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {
 
     if (pinnedMemAlloc_) {
 #if defined USE_NVIDIA_CUVS
-        pmr_->deallocate_sync(pinnedMemAlloc_, pinnedMemAllocSize_);
+        pmr_.deallocate_sync(pinnedMemAlloc_, pinnedMemAllocSize_);
 #else
         auto err = cudaFreeHost(pinnedMemAlloc_);
         FAISS_ASSERT_FMT(
@@ -350,7 +350,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
         // pinned memory allocation
         if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
             try {
-                pinnedMemAlloc_ = pmr_->allocate_sync(pinnedMemSize_);
+                pinnedMemAlloc_ = pmr_.allocate_sync(pinnedMemSize_);
             } catch (const std::bad_alloc& rmm_ex) {
                 FAISS_THROW_MSG("CUDA memory allocation error");
             }
@@ -546,10 +546,9 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
     } else if (adjReq.space == MemorySpace::Device) {
 #if defined USE_NVIDIA_CUVS
         try {
-            rmm::mr::device_memory_resource* current_mr =
-                    rmm::mr::get_per_device_resource(
-                            rmm::cuda_device_id{adjReq.device});
-            p = current_mr->allocate(adjReq.stream, adjReq.size);
+            auto current_mr = rmm::mr::get_per_device_resource_ref(
+                    rmm::cuda_device_id{adjReq.device});
+            p = current_mr.allocate(adjReq.stream, adjReq.size);
             adjReq.mr = current_mr;
         } catch (const std::bad_alloc& rmm_ex) {
             FAISS_THROW_MSG("CUDA memory allocation error");
@@ -562,7 +561,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
             // FIXME: as of CUDA 11, a memory allocation error appears to be
             // presented via cudaGetLastError as well, and needs to be
             // cleared. Just call the function to clear it
-            cudaGetLastError();
+            (void)cudaGetLastError();
 
             std::stringstream ss;
             ss << "StandardGpuResources: alloc fail " << adjReq.toString()
@@ -584,8 +583,8 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
             // TODO: change this to use the current device resource once RMM has
             // a way to retrieve a "guaranteed" managed memory resource for a
             // device.
-            p = mmr_->allocate(adjReq.stream, adjReq.size);
-            adjReq.mr = mmr_.get();
+            p = mmr_.allocate(adjReq.stream, adjReq.size);
+            adjReq.mr = mmr_;
         } catch (const std::bad_alloc& rmm_ex) {
             FAISS_THROW_MSG("CUDA memory allocation error");
         }
@@ -596,7 +595,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
             // FIXME: as of CUDA 11, a memory allocation error appears to be
             // presented via cudaGetLastError as well, and needs to be cleared.
             // Just call the function to clear it
-            cudaGetLastError();
+            (void)cudaGetLastError();
 
             std::stringstream ss;
             ss << "StandardGpuResources: alloc fail " << adjReq.toString()

data/vendor/faiss/faiss/gpu/StandardGpuResources.h

@@ -25,7 +25,7 @@
 
 #if defined USE_NVIDIA_CUVS
 #include <raft/core/device_resources.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/mr/managed_memory_resource.hpp>
 #include <rmm/mr/pinned_host_memory_resource.hpp>
 #endif
 
@@ -171,10 +171,10 @@ class StandardGpuResourcesImpl : public GpuResources {
      */
 
     // managed_memory_resource
-    std::unique_ptr<rmm::mr::device_memory_resource> mmr_;
+    rmm::mr::managed_memory_resource mmr_;
 
     // pinned_host_memory_resource
-    std::unique_ptr<rmm::mr::pinned_host_memory_resource> pmr_;
+    rmm::mr::pinned_host_memory_resource pmr_;
 #endif
 
     /// Pinned memory allocation for use with this GPU

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/MetalIndex.h

@@ -11,6 +11,7 @@
 #pragma once
 
 #include <faiss/Index.h>
+#include <faiss/gpu/GpuIndicesOptions.h>
 #include <faiss/gpu_metal/MetalResources.h>
 #include <memory>
 
@@ -20,6 +21,12 @@ 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.

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

@@ -19,13 +19,15 @@
 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 = 256;
+    static constexpr int kMaxK = 2048;
 
     void encodeDistanceMatrix(
             id<MTLComputeCommandEncoder> enc,
@@ -37,6 +39,23 @@ class MetalKernels {
             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,
@@ -47,6 +66,31 @@ class MetalKernels {
             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:
@@ -56,8 +100,9 @@ class MetalKernels {
     id<MTLLibrary> library_;
     std::unordered_map<std::string, id<MTLComputePipelineState>> cache_;
 
-    static constexpr int kTopKVariantSizes[] = {32, 64, 128, 256};
-    static constexpr int kNumTopKVariants = 4;
+    static constexpr int kTopKVariantSizes[] =
+            {32, 64, 128, 256, 512, 1024, 2048};
+    static constexpr int kNumTopKVariants = 7;
 };
 
 MetalKernels& getMetalKernels(id<MTLDevice> device);

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