Add MetalIndexIVFPQ with product quantization and residual encoding support

.github/actions/build_cmake/action.yml

@@ -247,7 +247,7 @@ runs:
               -DCMAKE_BUILD_TYPE=Release \
               -DCMAKE_PREFIX_PATH="$(brew --prefix libomp)" \
               .
-        cmake --build build --target faiss faiss_metal swigfaiss TestMetalIndexFlat TestMetalIndexIVFFlat -j$(sysctl -n hw.logicalcpu)
+        cmake --build build --target faiss faiss_metal swigfaiss TestMetalIndexFlat TestMetalIndexIVFFlat TestMetalIndexIVFPQ -j$(sysctl -n hw.logicalcpu)
     - name: C++ tests
       if: inputs.metal != 'ON'
       shell: bash

faiss/gpu_metal/CMakeLists.txt

@@ -16,7 +16,9 @@ set(FAISS_METAL_SRC
   MetalFlatKernels.mm
   MetalIndexFlat.mm
   MetalIndexIVFFlat.mm
+  MetalIndexIVFPQ.mm
   impl/MetalIVFFlat.mm
+  impl/MetalIVFPQ.mm
   StandardMetalResources.mm
   MetalCloner.mm
   MetalPythonBridge.mm

faiss/gpu_metal/MetalCloner.mm

@@ -9,10 +9,12 @@
 #import "MetalCloner.h"
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFFlat.h>
+#include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/FaissAssert.h>
 #include <cstring>
 #import "MetalIndexFlat.h"
 #import "MetalIndexIVFFlat.h"
+#import "MetalIndexIVFPQ.h"
 #import "StandardMetalResources.h"
 
 namespace faiss {
@@ -35,6 +37,14 @@ int get_num_gpus() {
     MetalIndexConfig config;
     config.device = 0;
 
+    const auto* ivfPQ = dynamic_cast<const faiss::IndexIVFPQ*>(index);
+    if (ivfPQ) {
+        FAISS_THROW_IF_NOT(
+                ivfPQ->metric_type == METRIC_L2 ||
+                ivfPQ->metric_type == METRIC_INNER_PRODUCT);
+        return new MetalIndexIVFPQ(res->getResources(), ivfPQ, config);
+    }
+
     const auto* ivfFlat = dynamic_cast<const faiss::IndexIVFFlat*>(index);
     if (ivfFlat) {
         FAISS_THROW_IF_NOT(
@@ -63,10 +73,26 @@ int get_num_gpus() {
 
     FAISS_THROW_MSG(
             "index_cpu_to_metal_gpu: unsupported index type "
-            "(supported: IndexFlat, IndexIVFFlat)");
+            "(supported: IndexFlat, IndexIVFFlat, IndexIVFPQ)");
 }
 
 faiss::Index* index_metal_gpu_to_cpu(const faiss::Index* index) {
+    const auto* metalIvfPQ = dynamic_cast<const MetalIndexIVFPQ*>(index);
+    if (metalIvfPQ) {
+        auto* quantizer = (metalIvfPQ->metric_type == METRIC_INNER_PRODUCT)
+                ? (faiss::IndexFlat*)new faiss::IndexFlatIP(metalIvfPQ->d)
+                : (faiss::IndexFlat*)new faiss::IndexFlatL2(metalIvfPQ->d);
+        auto* cpu = new faiss::IndexIVFPQ(
+                quantizer,
+                metalIvfPQ->d,
+                metalIvfPQ->nlist(),
+                metalIvfPQ->getNumSubQuantizers(),
+                8);
+        cpu->own_fields = true;
+        metalIvfPQ->copyTo(cpu);
+        return cpu;
+    }
+
     const auto* metalIvf = dynamic_cast<const MetalIndexIVFFlat*>(index);
     if (metalIvf) {
         auto* quantizer = (metalIvf->metric_type == METRIC_INNER_PRODUCT)
@@ -91,7 +117,7 @@ int get_num_gpus() {
 
     FAISS_THROW_MSG(
             "index_metal_gpu_to_cpu: unsupported index type "
-            "(supported: MetalIndexFlat, MetalIndexIVFFlat)");
+            "(supported: MetalIndexFlat, MetalIndexIVFFlat, MetalIndexIVFPQ)");
 }
 
 } // namespace gpu_metal

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

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