mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
Evan Haque
76c0e01e0f
Summary: This is a follow-up to https://github.com/facebookresearch/faiss/issues/5288, which added `MetalIndexIVFPQ` with GPU-resident storage while delegating search to the CPU. Adds GPU-accelerated IVFPQ search to the Metal backend, including optimized scanning through `M <= 64`. - Add Metal kernels for 8-bit PQ distance computation, inverted-list scanning, and top-k merging - Support optimized GPU scanning through `M <= 64` - Compute distance contributions only for PQ codes encountered in each list when `d <= 256` - Add a 32-thread cooperative kernel for lists containing at most 128 vectors - Retain the cached precomputed-LUT kernel for suitable `M <= 16` workloads - Stream precomputed terms for larger dimensions without allocating an `M * 256` threadgroup LUT - 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 - Compact invalid per-list padding before merging when lists contain fewer than `k` results - 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, search parameters, or Metal failures ## Changes - **Modified:** `MetalDistance.metal` - on demand PQ scans, large-M precomputed scan, short-list specialization, segmented top-k, and compacting grouped merge - **Modified:** `MetalDistance.h/.mm` - IVFPQ search orchestration and support for `M <= 64` - **Modified:** `MetalKernels.h/.mm` - dispatch for cached, on-demand, short-list, and large-M scan kernels - **Modified:** `MetalIndexIVFPQ.h/.mm` - GPU scan selection, reusable buffers, centroid uploads, precomputed terms, and CPU fallback - **Modified:** `MetalIndex.h` - add the `useFloat16` LUT configuration - **Modified:** `TestMetalIndexIVFPQ.mm` - direct-GPU scan validation and search-parameter fallback coverage ## Differences from CUDA IVFPQ **Training:** `MetalIndexIVFPQ::train` delegates training to its CPU index, as introduced in https://github.com/facebookresearch/faiss/issues/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. **Distance computation:** For `d <= 256`, the preferred Metal path evaluates only the PQ codes encountered while scanning a list. This avoids constructing or loading all 256 lookup entries for every subquantizer, which is especially beneficial for short IVF lists. For other supported dimensions, Metal uses the CPU IVFPQ precomputed-table decomposition. The query-independent centroid/PQ term is computed once per trained index and the query term once per batch. **Threadgroup memory:** The existing cached lookup-table kernel remains available for suitable `M <= 16` workloads. Larger `M` values do not expand its `M * 256` threadgroup LUT, which would exceed Metal threadgroup-memory limits. Instead, they use the on-demand or streamed-precomputed scan paths. **List scanning:** Short lists use a 32-thread cooperative kernel with approximately 2 KiB of threadgroup scratch. Longer lists use segmented selection to maintain an exact running top-k. **Top-k merge:** Per-list results are merged in groups over multiple rounds. Invalid padding is compacted when lists contain fewer than `k` results. **Fallback:** The optimized scan supports `M <= 64`, `d / M <= 256`, and `k <= 512`. Unsupported configurations fall back to the legacy Metal path or CPU search. Search options that the GPU path cannot preserve, including selectors, scan budgets, and polysemous filtering, are forwarded to the CPU with their original parameters. ## 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 ``` All six IVFPQ tests pass. Direct-GPU regression coverage exercises 800 combinations across `M=8,16,17,32,48,64`, L2 and inner product, `k` through 512, short and long lists, both merge modes, and 64-bit IDs. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5449 Reviewed By: junjieqi Differential Revision: D113864779 Pulled By: mnorris11 fbshipit-source-id: 741c12b2b014ef28a82569c6b0344352351a7325
104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
"""Regression tests for the Metal IVF-PQ search backend.
|
|
|
|
The per-(query, probe) lookup-table path must add the coarse-quantizer score
|
|
for inner product: its tables hold <q_m, pq_m> only, while the CPU reference
|
|
adds dis0 = <q, centroid> (see QueryTables::precompute_list_tables_IP).
|
|
These tests force the lookup-table path with configurations the
|
|
precomputed-table path declines, and compare GPU results against the CPU
|
|
index.
|
|
"""
|
|
|
|
import functools
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import faiss
|
|
|
|
|
|
_HAS_METAL = "MAC_METAL" in faiss.get_compile_options()
|
|
needs_metal_build = unittest.skipUnless(
|
|
_HAS_METAL, "needs a faiss Metal build"
|
|
)
|
|
|
|
|
|
def needs_metal_device(test_func):
|
|
"""Skip the test when no Metal device is available."""
|
|
|
|
@functools.wraps(test_func)
|
|
def wrapper(*args, **kwargs):
|
|
if faiss.get_num_gpus() == 0:
|
|
raise unittest.SkipTest("no Metal device available")
|
|
return test_func(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
@needs_metal_build
|
|
class TestMetalIVFPQ(unittest.TestCase):
|
|
|
|
def check_matches_cpu(
|
|
self, metric, d, M, nlist, nb, nq, nprobe, k
|
|
):
|
|
rs = np.random.RandomState(1234)
|
|
xb = rs.rand(nb, d).astype(np.float32)
|
|
xq = rs.rand(nq, d).astype(np.float32)
|
|
|
|
quantizer = (
|
|
faiss.IndexFlatIP(d)
|
|
if metric == faiss.METRIC_INNER_PRODUCT
|
|
else faiss.IndexFlatL2(d)
|
|
)
|
|
cpu_index = faiss.IndexIVFPQ(quantizer, d, nlist, M, 8, metric)
|
|
cpu_index.nprobe = nprobe
|
|
cpu_index.train(xb)
|
|
cpu_index.add(xb)
|
|
ref_D, ref_I = cpu_index.search(xq, k)
|
|
|
|
res = faiss.StandardGpuResources()
|
|
gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)
|
|
got_D, got_I = gpu_index.search(xq, k)
|
|
|
|
for q in range(nq):
|
|
valid = ref_I[q] != -1
|
|
self.assertTrue(np.any(valid))
|
|
self.assertEqual(set(got_I[q][valid]), set(ref_I[q][valid]))
|
|
np.testing.assert_allclose(
|
|
got_D[q][valid], ref_D[q][valid], atol=1e-2
|
|
)
|
|
|
|
@needs_metal_device
|
|
def test_lookup_table_path_ip_large_k(self):
|
|
# k > 512 declines the precomputed-table path. With nprobe == 1 the
|
|
# missing coarse term keeps the ranking but shifts every score.
|
|
self.check_matches_cpu(
|
|
faiss.METRIC_INNER_PRODUCT,
|
|
d=64,
|
|
M=8,
|
|
nlist=8,
|
|
nb=500,
|
|
nq=5,
|
|
nprobe=1,
|
|
k=600,
|
|
)
|
|
|
|
@needs_metal_device
|
|
def test_lookup_table_path_l2_large_k(self):
|
|
# L2 tables hold the complete residual distance, so the coarse term
|
|
# must stay off: guard the flag plumbing in both directions.
|
|
self.check_matches_cpu(
|
|
faiss.METRIC_L2,
|
|
d=64,
|
|
M=8,
|
|
nlist=8,
|
|
nb=500,
|
|
nq=5,
|
|
nprobe=1,
|
|
k=600,
|
|
)
|