Add MetalIndexIVFPQ with product quantization and residual encoding support (#5288)

Summary:
Adds IVF-PQ (inverted file with product quantization) index support to the Metal GPU backend

- Add `MetalIndexIVFPQ` with full train/add/search/reset/copyFrom/copyTo support
- Add `MetalIVFPQImpl` GPU-resident IVF list storage for PQ codes (segment allocator, same pattern as IVFFlat)
- Support 8-bit product quantization with precomputed per-query lookup tables
- Support both L2 and inner product metrics
- Residual encoding when `by_residual=true` (default for L2)
- CPU-side PQ lookup table computation with precomputed tables optimization for L2
- GPU scan path via `runMetalIVFPQFullSearch` with CPU LUT fallback via `runMetalIVFPQScan`
- Update `MetalCloner` to support IVFPQ in `index_cpu_to_metal_gpu` / `index_metal_gpu_to_cpu`

## Changes
- **New:** `MetalIndexIVFPQ.h/.mm` - IVFPQ index class (train, add, search, reset, copyFrom/copyTo, cloner support)
- **New:** `impl/MetalIVFPQ.h/.mm` - GPU-resident IVF list storage with segment allocator for PQ codes
- **New:** `test/TestMetalIndexIVFPQ.mm` - 4 C++ tests (L2, IP, reset, CPU↔GPU round-trip)
- **Modified:** `test/CMakeLists.txt` - added TestMetalIndexIVFPQ build target

## Differences from CUDA IVFPQ

**Training:** Delegates to CPU (`IndexIVFPQ::train`). CUDA can train on GPU. Same rationale as IVFFlat - training is a one-time cost.

**Add path:** Coarse quantization and PQ encoding run on CPU, then codes are copied to GPU storage. CUDA does both on GPU. On Apple Silicon with unified memory, the copy cost is minimal.

**Residual encoding:** When `by_residual=true`, residuals (x - coarse_centroid) are computed on CPU before PQ encoding. CUDA computes residuals on GPU. Functionally equivalent.

**Lookup tables:** PQ distance lookup tables are computed on CPU and uploaded to GPU for the scan phase. CUDA computes LUTs on GPU. CPU LUT computation is fast relative to the scan and avoids a separate GPU kernel launch.

**IVF list storage:** Same segment allocator pattern as IVFFlat - single contiguous buffer rather than CUDA's per-list `DeviceVector` allocations.

## Note
FP16 coarse quantizer and GPU merge kernel are planned optimizations for a future PR. Both apply across all IVF index types (IVFFlat, IVFPQ, IVFSQ).

## Build and test
```bash
cmake -B build \
  -DFAISS_ENABLE_GPU=OFF \
  -DFAISS_ENABLE_METAL=ON \
  -DBUILD_TESTING=ON \
  -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_PREFIX_PATH="$(brew --prefix libomp)" \
  .
cmake --build build --target faiss faiss_metal TestMetalIndexIVFPQ -j$(sysctl -n hw.logicalcpu)
cd build && ctest -R TestMetalIndexIVFPQ --output-on-failure
```

Pull Request resolved: #5288

Reviewed By: alibeklfc

Differential Revision: D113037852

Pulled By: mnorris11

fbshipit-source-id: 5fdb729b10ad9884a673fc3483543b4c5fcdb621

Author: Evandabest

.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