mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
main
4
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Evan Haque
|
76c0e01e0f |
Add GPU-accelerated IVFPQ search to the Metal backend (#5449)
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 |
||
|
Evan Haque
|
3110fd2424 |
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: https://github.com/facebookresearch/faiss/pull/5288 Reviewed By: alibeklfc Differential Revision: D113037852 Pulled By: mnorris11 fbshipit-source-id: 5fdb729b10ad9884a673fc3483543b4c5fcdb621 |
||
|
Evan Haque
|
10b6b2a1e5 |
Add MetalIndexIVFFlat with IVF scan/merge kernels and expanded top-k support (#5202)
Summary:
## Summary
Adds IVFFlat index support to the Metal GPU backend
- Add `MetalIndexIVFFlat` with full train/add/search/reset/copyFrom/copyTo support
- Add IVF scan/merge MSL kernels (`ivf_scan_list`, `ivf_scan_list_small`, `ivf_scan_list_interleaved`, `ivf_merge_lists`) to `MetalDistance.metal`
- Add `MetalDistance.h/.mm` orchestration layer for IVF search dispatch (coarse quantization + scan + merge in a single command buffer)
- Add `compute_norms` and `l2_with_norms` kernels for efficient L2 coarse quantization
- Expand top-k support from k <= 256 to k <= 2048 (`topk_threadgroup_{512,1024,2048}`)
- Raise dimension limit from d <= 512 to d <= 2048 (supports 768, 1024, 1536 embeddings)
- Update `MetalCloner` to support IVFFlat in `index_cpu_to_metal_gpu` / `index_metal_gpu_to_cpu`
- Rename `nil` parameter to `n_il` in `InvertedLists.h/.cpp` to fix Objective-C++ keyword conflict
## Changes
- **New:** `MetalIndexIVFFlat.h/.mm` - IVFFlat index class (train, add, search, reset, cloner support)
- **New:** `impl/MetalIVFFlat.h/.mm` - GPU-resident IVF list storage with segment allocator
- **New:** `MetalDistance.h/.mm` - IVF search dispatch (scan + merge orchestration)
- **New:** `test/TestMetalIndexIVFFlat.mm` - 4 tests (L2, IP, reset, cpu to gpu round trip)
- **Modified:** `MetalDistance.metal` - IVF kernels, compute_norms, l2_with_norms, topk 512/1024/2048
- **Modified:** `MetalKernels.h/.mm` - IVF encode methods, expanded topk variants, kMaxK=2048
- **Modified:** `MetalIndex.h` - added IVF config fields (indicesOptions, interleavedLayout, useFloat16CoarseQuantizer)
- **Modified:** `MetalCloner.mm` - IVFFlat support in cpu <-> gpu cloning
- **Modified:** `CMakeLists.txt` - added new source files
- **Modified:** `.github/actions/build_cmake/action.yml` - added TestMetalIndexIVFFlat to CI
- **Modified:** `faiss/invlists/InvertedLists.h/.cpp` - renamed `nil` parameter to `n_il` (Objective-C++ keyword conflict)
## Differences from CUDA IVFFlat
**Training:** Delegates to CPU (`IndexIVFFlat::train`). CUDA can train on GPU via k-means or cuVS. GPU training is possible with Metal but low priority since training is a one-time cost.
**Add path:** Coarse quantization runs on CPU, then vectors are copied to GPU storage. CUDA does both on GPU. Moving coarse assignment to Metal is doable using the existing flat search kernels; planned for a future optimization.
**Reconstruct:** Reads from CPU host mirror rather than GPU buffers. On Apple Silicon with unified memory, CPU and GPU see the same physical memory, so this is functionally equivalent.
**IVF list storage:** Uses a single contiguous buffer with a segment allocator, rather than CUDA's per-list `DeviceVector` allocations. This avoids managing thousands of small `MTLBuffer` objects and reduces fragmentation; a better fit for Apple Silicon's unified memory pool.
## Notes on other changes
**`nil` → `n_il` rename in `InvertedLists.h/.cpp`:** ObjC++ treats `nil` as a keyword (equivalent to `nullptr` for Objective-C objects). Any `.mm` file that transitively includes `InvertedLists.h` fails to compile because the parameter name `int nil` conflicts with the keyword. This is a declaration only rename, no behavior, ABI, or API change. Does not affect CUDA, CPU, or any other build.
**`centroidsAreFP16` parameter in `runMetalIVFFlatFullSearch`:** The Metal-backend branch passes a `useFloat16CoarseQuantizer` config flag through the search dispatch chain. CUDA handles this differently, FP16 centroid storage is a property of the coarse quantizer object (`GpuIndexFlatConfig::useFloat16`), not a function parameter. The parameter is currently a no-op but is included to maintain the correct function signature and avoid a parameter alignment bug where a trailing `bool` would be misinterpreted as `waitForCompletion`.
## Build and test
```bash
cmake -B build \
-DFAISS_ENABLE_METAL=ON \
-DFAISS_ENABLE_GPU=OFF \
-DFAISS_ENABLE_PYTHON=OFF \
-DBUILD_TESTING=ON \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_PREFIX_PATH="$(brew --prefix libomp)" \
.
cmake --build build --target faiss_metal TestMetalIndexFlat TestMetalIndexIVFFlat -j$(sysctl -n hw.logicalcpu)
cd build && ctest -R TestMetalIndex --output-on-failure
```
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5202
Reviewed By: mdouze
Differential Revision: D107120463
Pulled By: mnorris11
fbshipit-source-id: 75a33d233d4844bec046818418ba75c85abb0402
|
||
|
Evandabest
|
f37041d073 |
Extract Metal shaders to standalone .metal file with kernel wrapper class (#5167)
Summary:
## Summary
Before submitting PRs for IVFFlat, IVFPQ, IVFSQ, and BinaryFlat implementations, I wanted to run this by you guys to see if this is a good approach for organizing the Metal shader code. Currently the MSL source is embedded as an inline string in the Objective C++ file, this extracts it into a standalone `.metal` file with a wrapper class, which will be the foundation for all the upcoming index type PRs. This also upgrades the top-k kernel from a single-thread insertion sort to a 256-thread parallel bitonic sort for better GPU utilization.
- Extract inline MSL shader source from `MetalFlatKernels.mm` into a standalone`MetalDistance.metal` file for proper syntax highlighting, easier editing, and reviewability
- Add `MetalKernels` wrapper class that loads a precompiled `.metallib`, caches pipeline states, and provides typed `encode*()` dispatch methods (previously the library was recompiled from source on every search call)
- Precompile `.metal` → `.metallib` via `xcrun metal`/`xcrun metallib` in CMake, catching shader errors at build time instead of runtime
- Upgrade distance kernels from naive per-element to tiled GEMM-style (32x32 tiles with shared memory) and upgrades top-k to a parallel threadgroup-based selection kernel (256-thread parallel bitonic sort)
- Foundation for future index types (IVFFlat, IVFPQ, etc.) to add kernels to the same `.metal` file and methods to `MetalKernels`
## Changes
- **New:** `MetalDistance.metal`, `l2_squared_matrix`, `ip_matrix`, `topk_threadgroup_{32,64,128,256}`
- **New:** `MetalKernels.h/.mm` runtime shader loading, pipeline caching, typed `encode*()` dispatch methods, per-device singleton
- **Modified:** `MetalFlatKernels.mm` replaced inline MSL + manual dispatch with `MetalKernels` calls
- **Modified:** `CMakeLists.txt` — added `MetalKernels.mm` to build, added custom commands to precompile `.metal` → `.air` → `.metallib`
## Build and test
```bash
cmake -B build \
-DFAISS_ENABLE_METAL=ON \
-DFAISS_ENABLE_GPU=OFF \
-DFAISS_ENABLE_PYTHON=OFF \
-DBUILD_TESTING=ON \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_PREFIX_PATH="$(brew --prefix libomp)" \
.
cmake --build build --target faiss_metal TestMetalIndexFlat -j$(sysctl -n hw.logicalcpu)
cd build && ctest -R TestMetalIndexFlat --output-on-failure
```
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5167
Reviewed By: alibeklfc
Differential Revision: D103979808
Pulled By: mnorris11
fbshipit-source-id: 9583d6767e589ad74f7aa58141cd2e9034e5698b
|