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
This commit is contained in:
Evan Haque
2026-06-04 14:14:58 -07:00
committed by meta-codesync[bot]
parent 74d36198fa
commit 10b6b2a1e5
22 changed files with 5447 additions and 54 deletions
+15 -4
View File
@@ -179,7 +179,9 @@ runs:
- name: Setup macOS Metal environment
if: inputs.metal == 'ON'
shell: bash
run: brew install cmake libomp gflags
run: |
brew install cmake libomp gflags swig
pip3 install --break-system-packages numpy pytest
- name: Build all targets
if: inputs.metal != 'ON'
shell: bash
@@ -212,12 +214,12 @@ runs:
cmake -B build \
-DFAISS_ENABLE_METAL=ON \
-DFAISS_ENABLE_GPU=OFF \
-DFAISS_ENABLE_PYTHON=OFF \
-DFAISS_ENABLE_PYTHON=ON \
-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)
cmake --build build --target faiss faiss_metal swigfaiss TestMetalIndexFlat TestMetalIndexIVFFlat -j$(sysctl -n hw.logicalcpu)
- name: C++ tests
if: inputs.metal != 'ON'
shell: bash
@@ -228,7 +230,16 @@ runs:
- name: C++ tests (Metal)
if: inputs.metal == 'ON'
shell: bash
run: cd build && ctest -R TestMetalIndexFlat --output-on-failure
run: cd build && ctest -R TestMetalIndex --output-on-failure
- name: Install Python extension (Metal)
if: inputs.metal == 'ON'
shell: bash
working-directory: build/faiss/python
run: pip3 install --break-system-packages .
- name: Python tests (Metal)
if: inputs.metal == 'ON'
shell: bash
run: python3 -m pytest faiss/gpu_metal/test/test_metal_python.py -v
- name: C++ perf benchmarks
if: inputs.rocm == 'OFF' && inputs.metal != 'ON'
shell: bash