Add GPU-accelerated IVFPQ search to the Metal backend

[Evandabest]

Summary

This is a follow-up to #5288 which added MetalIndexIVFPQ with GPU-resident storage while delegating search to the CPU.

Adds GPU-accelerated IVFPQ search to the Metal backend.

• Add Metal kernels for 8-bit PQ lookup-table construction, inverted-list scanning, and top-k merging
• Add a precomputed-table path equivalent to the CPU IVFPQ optimization
• Precompute the query-independent centroid/PQ term once per trained index
• Compute the M x 256 query term once per query batch instead of once per (query, probe)
• Support both L2 and inner product metrics
• Add exact segmented top-k selection for inverted lists of any length
• Add grouped multi-round merging without a fixed nprobe * k candidate limit
• Reuse dynamically sized Metal scratch buffers across searches
• Retain the legacy per-probe LUT path with optional FP16 lookup tables
• Fall back to CPU search for unsupported configurations or Metal failures

Changes

• Modified: MetalDistance.metal - IVFPQ LUT construction, precomputed-table, list scan, segmented top-k, and grouped merge kernels
• Modified: MetalDistance.h/.mm - IVFPQ search orchestration and precomputed-table entry points
• Modified: MetalKernels.h/.mm - IVFPQ Metal kernel dispatch methods
• Modified: MetalIndexIVFPQ.h/.mm - GPU search path, reusable buffers, centroid uploads, precomputed terms, and CPU fallback
• Modified: MetalIndex.h - add the useFloat16 LUT configuration

Differences from CUDA IVFPQ

Training: MetalIndexIVFPQ::train delegates training to its CPU index, as introduced in #5288. CUDA can train the coarse and product quantizers on GPU. Training is a one-time cost.

Add path: Coarse assignment, residual computation, and PQ encoding remain on the CPU. The encoded PQ codes are stored in GPU-resident Metal buffers. CUDA performs these operations on GPU.

Coarse quantization: The coarse quantizer search runs on the CPU. The selected lists and coarse distances are passed to the Metal scan.

Lookup tables: The preferred Metal path mirrors the CPU IVFPQ precomputed-table decomposition. The query independent centroid/PQ term is computed once per trained index, while the query term is computed once per batch. This avoids materializing a separate lookup table for every (query, probe) pair.

List scanning: Metal performs asymmetric distance computation directly over the GPU-resident 8-bit PQ codes. Segmented selection maintains an exact running top-k for inverted lists of any length.

Top-k merge: Per-list results are merged in groups over multiple rounds. This avoids the fixed candidate limit of the original single-pass merge.

Fallback: The preferred precomputed-table path supports M <= 16, d / M <= 256, and k <= 512. Unsupported configurations fall back to the legacy Metal path or CPU search.

Build and test

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

[Evandabest]

Following up on #5288, here are the isolated PQ lookup-table construction results from the GPU-scan work (same as shown in first PR). All measurements were collected on an M3 Pro with d=128, 8-bit PQ, and FP32 lookup tables.

PQ LUT computation
Metal ivfpq_build_lut_l2 kernel vs faiss CPU ProductQuantizer::compute_distance_tables (same M×256 tables, numerically identical, max rel diff ~1e-7):

Case	CPU	Metal	Speedup
M=8, nq=10k	12.3 ms	4.4 ms	2.8x
M=16, nq=10k	17.1 ms	4.6 ms	3.7x
M=32, nq=10k	30.0 ms	8.1 ms	3.7x
M=16, nq=1k	1.78 ms	1.49 ms	1.2x
M=16, nq=50k	84.2 ms	21.9 ms	3.8x
M=16, nq=100k	160.2 ms	43.5 ms	3.7x

LUT speedup saturates at ~3.7-3.8x once the batch is large enough to hide kernel-launch overhead (nq ≥ 10k); ~1.2x for tiny batches.

[meta-codesync]

@mnorris11 has imported this pull request. If you are a Meta employee, you can view this in D113864779.