Add GPU-accelerated IVFPQ search to the Metal backend

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

@@ -83,5 +83,87 @@ bool runMetalIVFFlatFullSearch(
         bool centroidsAreFP16 = false,
         bool waitForCompletion = true);
 
+/// Full IVF-PQ search on the GPU: builds per-(query, probe) PQ lookup tables,
+/// scans the assigned inverted lists with 8-bit ADC, and merges the per-list
+/// top-k into the final top-k. Coarse quantization is performed on the CPU and
+/// the assignments are passed in via @p coarseAssign.
+bool runMetalIVFPQFullSearch(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries,
+        id<MTLBuffer> coarseAssign,
+        id<MTLBuffer> coarseCentroids,
+        id<MTLBuffer> pqCentroids,
+        id<MTLBuffer> lookupTable,
+        id<MTLBuffer> codes,
+        id<MTLBuffer> ids,
+        id<MTLBuffer> listOffset,
+        id<MTLBuffer> listLength,
+        int nq,
+        int d,
+        int M,
+        int k,
+        int nprobe,
+        int nlist,
+        int avgListLen,
+        bool lookupFp16,
+        bool isL2,
+        id<MTLBuffer> outDistances,
+        id<MTLBuffer> outIndices,
+        id<MTLBuffer> perListDistBuf,
+        id<MTLBuffer> perListIdxBuf,
+        bool waitForCompletion = true);
+
+/// Build the query-independent precomputed term of the IVFPQ L2 distance
+/// decomposition on the GPU: term2[l][m][c] = ||pq[m][c]||^2 +
+/// 2 <coarse_l[m], pq[m][c]>, laid out as nlist * M * 256 floats. Run once
+/// per trained index; consumed by runMetalIVFPQPrecompSearch.
+bool runMetalIVFPQPrecomputeTerm2(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> coarseCentroids,
+        id<MTLBuffer> pqCentroids,
+        id<MTLBuffer> outTerm2,
+        int nlist,
+        int d,
+        int M);
+
+/// IVF-PQ search using the precomputed-table decomposition: a per-query
+/// M*256 term is built once per batch (instead of a full residual LUT per
+/// (query, probe)), combined with @p term2 in threadgroup memory by the scan
+/// kernel, and the coarse distance is added as the constant ||x - c||^2 term.
+/// The scan keeps an exact running top-k over arbitrarily long lists and the
+/// merge runs in rounds, so there is no list-length or nprobe*k cap; the
+/// remaining requirements are k <= 512, M <= 16 and d/M <= 256.
+/// @p term2 may be nil for inner product (qterm is the whole table);
+/// @p useDis0 adds the coarse distance (required for L2, by_residual for IP).
+bool runMetalIVFPQPrecompSearch(
+        id<MTLDevice> device,
+        id<MTLCommandQueue> queue,
+        id<MTLBuffer> queries,
+        id<MTLBuffer> coarseAssign,
+        id<MTLBuffer> coarseDist,
+        id<MTLBuffer> term2,
+        id<MTLBuffer> qtermScratch,
+        id<MTLBuffer> pqCentroids,
+        id<MTLBuffer> codes,
+        id<MTLBuffer> ids,
+        id<MTLBuffer> listOffset,
+        id<MTLBuffer> listLength,
+        int nq,
+        int d,
+        int M,
+        int k,
+        int nprobe,
+        bool isL2,
+        bool useDis0,
+        id<MTLBuffer> outDistances,
+        id<MTLBuffer> outIndices,
+        id<MTLBuffer> perListDistBuf,
+        id<MTLBuffer> perListIdxBuf,
+        id<MTLBuffer> mergeScratchDistBuf,
+        id<MTLBuffer> mergeScratchIdxBuf,
+        bool waitForCompletion = true);
+
 } // namespace gpu_metal
 } // namespace faiss