From 10b6b2a1e554afc4cbd66fe73d864728ff9e310c Mon Sep 17 00:00:00 2001 From: Evan Haque Date: Thu, 4 Jun 2026 14:14:58 -0700 Subject: [PATCH] Add MetalIndexIVFFlat with IVF scan/merge kernels and expanded top-k support (#5202) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .github/actions/build_cmake/action.yml | 19 +- faiss/gpu_metal/CMakeLists.txt | 4 + faiss/gpu_metal/MetalCloner.mm | 84 +- faiss/gpu_metal/MetalDistance.h | 87 + faiss/gpu_metal/MetalDistance.metal | 838 ++++++ faiss/gpu_metal/MetalDistance.mm | 359 +++ faiss/gpu_metal/MetalIndex.h | 7 + faiss/gpu_metal/MetalIndexIVFFlat.h | 181 ++ faiss/gpu_metal/MetalIndexIVFFlat.mm | 2308 +++++++++++++++++ faiss/gpu_metal/MetalKernels.h | 51 +- faiss/gpu_metal/MetalKernels.mm | 125 +- faiss/gpu_metal/MetalPythonBridge.h | 45 + faiss/gpu_metal/MetalPythonBridge.mm | 79 + faiss/gpu_metal/impl/MetalIVFFlat.h | 193 ++ faiss/gpu_metal/impl/MetalIVFFlat.mm | 632 +++++ faiss/gpu_metal/test/CMakeLists.txt | 16 + faiss/gpu_metal/test/TestMetalIndexIVFFlat.mm | 200 ++ faiss/gpu_metal/test/test_metal_python.py | 186 ++ faiss/invlists/InvertedLists.cpp | 26 +- faiss/invlists/InvertedLists.h | 4 +- faiss/python/CMakeLists.txt | 15 + faiss/python/swigfaiss.swig | 42 +- 22 files changed, 5447 insertions(+), 54 deletions(-) create mode 100644 faiss/gpu_metal/MetalDistance.h create mode 100644 faiss/gpu_metal/MetalDistance.mm create mode 100644 faiss/gpu_metal/MetalIndexIVFFlat.h create mode 100644 faiss/gpu_metal/MetalIndexIVFFlat.mm create mode 100644 faiss/gpu_metal/MetalPythonBridge.h create mode 100644 faiss/gpu_metal/MetalPythonBridge.mm create mode 100644 faiss/gpu_metal/impl/MetalIVFFlat.h create mode 100644 faiss/gpu_metal/impl/MetalIVFFlat.mm create mode 100644 faiss/gpu_metal/test/TestMetalIndexIVFFlat.mm create mode 100644 faiss/gpu_metal/test/test_metal_python.py diff --git a/.github/actions/build_cmake/action.yml b/.github/actions/build_cmake/action.yml index a66c75b10..81042eac5 100644 --- a/.github/actions/build_cmake/action.yml +++ b/.github/actions/build_cmake/action.yml @@ -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 diff --git a/faiss/gpu_metal/CMakeLists.txt b/faiss/gpu_metal/CMakeLists.txt index 656bfe80d..b6b216ccc 100644 --- a/faiss/gpu_metal/CMakeLists.txt +++ b/faiss/gpu_metal/CMakeLists.txt @@ -12,10 +12,14 @@ set(FAISS_METAL_SRC MetalResources.mm MetalIndex.mm MetalKernels.mm + MetalDistance.mm MetalFlatKernels.mm MetalIndexFlat.mm + MetalIndexIVFFlat.mm + impl/MetalIVFFlat.mm StandardMetalResources.mm MetalCloner.mm + MetalPythonBridge.mm ) add_library(faiss_metal STATIC ${FAISS_METAL_SRC}) diff --git a/faiss/gpu_metal/MetalCloner.mm b/faiss/gpu_metal/MetalCloner.mm index 07061807e..aa02de9f1 100644 --- a/faiss/gpu_metal/MetalCloner.mm +++ b/faiss/gpu_metal/MetalCloner.mm @@ -8,9 +8,11 @@ #import "MetalCloner.h" #include +#include #include #include #import "MetalIndexFlat.h" +#import "MetalIndexIVFFlat.h" #import "StandardMetalResources.h" namespace faiss { @@ -30,42 +32,66 @@ faiss::Index* index_cpu_to_metal_gpu( FAISS_THROW_IF_NOT(res->getResources()->isAvailable()); FAISS_THROW_IF_NOT_MSG(device == 0, "Metal backend supports only device 0"); - const auto* flat = dynamic_cast(index); - if (!flat) { - FAISS_THROW_MSG( - "index_cpu_to_metal_gpu: only IndexFlat (and L2/IP) supported"); - } - FAISS_THROW_IF_NOT( - flat->metric_type == METRIC_L2 || - flat->metric_type == METRIC_INNER_PRODUCT); - MetalIndexConfig config; config.device = 0; - auto* metal = new MetalIndexFlat( - res->getResources(), - flat->d, - flat->metric_type, - flat->metric_arg, - config); - if (flat->ntotal > 0) { - const float* xb = flat->get_xb(); - metal->add(flat->ntotal, xb); + + const auto* ivfFlat = dynamic_cast(index); + if (ivfFlat) { + FAISS_THROW_IF_NOT( + ivfFlat->metric_type == METRIC_L2 || + ivfFlat->metric_type == METRIC_INNER_PRODUCT); + return new MetalIndexIVFFlat(res->getResources(), ivfFlat, config); } - return metal; + + const auto* flat = dynamic_cast(index); + if (flat) { + FAISS_THROW_IF_NOT( + flat->metric_type == METRIC_L2 || + flat->metric_type == METRIC_INNER_PRODUCT); + auto* metal = new MetalIndexFlat( + res->getResources(), + flat->d, + flat->metric_type, + flat->metric_arg, + config); + if (flat->ntotal > 0) { + const float* xb = flat->get_xb(); + metal->add(flat->ntotal, xb); + } + return metal; + } + + FAISS_THROW_MSG( + "index_cpu_to_metal_gpu: unsupported index type " + "(supported: IndexFlat, IndexIVFFlat)"); } faiss::Index* index_metal_gpu_to_cpu(const faiss::Index* index) { - const auto* metal = dynamic_cast(index); - if (!metal) { - FAISS_THROW_MSG( - "index_metal_gpu_to_cpu: only MetalIndexFlat supported"); + const auto* metalIvf = dynamic_cast(index); + if (metalIvf) { + auto* quantizer = (metalIvf->metric_type == METRIC_INNER_PRODUCT) + ? (faiss::IndexFlat*)new faiss::IndexFlatIP(metalIvf->d) + : (faiss::IndexFlat*)new faiss::IndexFlatL2(metalIvf->d); + auto* cpu = new faiss::IndexIVFFlat( + quantizer, metalIvf->d, metalIvf->nlist()); + cpu->own_fields = true; + metalIvf->copyTo(cpu); + return cpu; } - faiss::IndexFlat* cpu = (metal->metric_type == METRIC_INNER_PRODUCT) - ? (faiss::IndexFlat*)new faiss::IndexFlatIP(metal->d) - : (faiss::IndexFlat*)new faiss::IndexFlatL2(metal->d); - cpu->metric_arg = metal->metric_arg; - metal->copyTo(cpu); - return cpu; + + const auto* metal = dynamic_cast(index); + if (metal) { + faiss::IndexFlat* cpu = (metal->metric_type == METRIC_INNER_PRODUCT) + ? (faiss::IndexFlat*)new faiss::IndexFlatIP(metal->d) + : (faiss::IndexFlat*)new faiss::IndexFlatL2(metal->d); + cpu->metric_arg = metal->metric_arg; + metal->copyTo(cpu); + return cpu; + } + + FAISS_THROW_MSG( + "index_metal_gpu_to_cpu: unsupported index type " + "(supported: MetalIndexFlat, MetalIndexIVFFlat)"); } } // namespace gpu_metal diff --git a/faiss/gpu_metal/MetalDistance.h b/faiss/gpu_metal/MetalDistance.h new file mode 100644 index 000000000..2785f5d47 --- /dev/null +++ b/faiss/gpu_metal/MetalDistance.h @@ -0,0 +1,87 @@ +// @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. + * + * IVF distance computation and scan dispatch for Metal backend. + */ + +#pragma once + +#import + +#include +#include +#include + +namespace faiss { +namespace gpu_metal { + +class MetalResources; + +int getMetalDistanceMaxK(); + +bool runMetalComputeNorms( + id device, + id queue, + id vectors, + int nb, + int d, + id normsBuf, + bool waitForCompletion = true); + +bool runMetalIVFFlatScan( + id device, + id queue, + id queries, + id codes, + id ids, + id listOffset, + id listLength, + id coarseAssign, + int nq, + int d, + int k, + int nprobe, + bool isL2, + id outDistances, + id outIndices, + id perListDistBuf, + id perListIdxBuf, + id interleavedCodes = nil, + id interleavedCodesOffset = nil, + bool waitForCompletion = true); + +bool runMetalIVFFlatFullSearch( + id device, + id queue, + id queries, + int nq, + int d, + int k, + int nprobe, + bool isL2, + id centroids, + int nlist, + id codes, + id ids, + id listOffset, + id listLength, + id outDistances, + id outIndices, + id perListDistBuf, + id perListIdxBuf, + id coarseDistBuf, + id coarseIdxBuf, + id distMatrixBuf, + id centroidNormsBuf = nil, + int avgListLen = 256, + id interleavedCodes = nil, + id interleavedCodesOffset = nil, + bool centroidsAreFP16 = false, + bool waitForCompletion = true); + +} // namespace gpu_metal +} // namespace faiss diff --git a/faiss/gpu_metal/MetalDistance.metal b/faiss/gpu_metal/MetalDistance.metal index cf9e829e7..d2f96de47 100644 --- a/faiss/gpu_metal/MetalDistance.metal +++ b/faiss/gpu_metal/MetalDistance.metal @@ -238,4 +238,842 @@ TOPK_THREADGROUP_VARIANT(32) TOPK_THREADGROUP_VARIANT(64) TOPK_THREADGROUP_VARIANT(128) TOPK_THREADGROUP_VARIANT(256) +TOPK_THREADGROUP_VARIANT(512) +TOPK_THREADGROUP_VARIANT(1024) #undef TOPK_THREADGROUP_VARIANT +kernel void topk_threadgroup_2048( + device const float* distances [[buffer(0)]], + device float* outDistances [[buffer(1)]], + device int* outIndices [[buffer(2)]], + device const uint* params [[buffer(3)]], + uint qi [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + constexpr uint TG_SIZE = 256; + constexpr uint R = 16; + constexpr uint CANDIDATES = TG_SIZE * R; // 4096 + threadgroup float tgDist[CANDIDATES]; + threadgroup int tgIdx[CANDIDATES]; + uint nq = params[0], nb = params[1], k = params[2], want_min = params[3]; + if (qi >= nq || k == 0) return; + const device float* row = distances + qi * nb; + uint kk = min(k, nb); + uint K_out = min((uint)2048, kk); + + if (nb <= CANDIDATES) { + for (uint i = tid; i < CANDIDATES; i += TG_SIZE) { + if (i < nb) { + tgDist[i] = row[i]; + tgIdx[i] = (int)i; + } else { + tgDist[i] = want_min ? 1e38f : -1e38f; + tgIdx[i] = -1; + } + } + } else { + float localDist[R]; + int localIdx[R]; + uint localCount = 0; + + for (uint j = tid; j < nb; j += TG_SIZE) { + float v = row[j]; + if (localCount < R) { + uint pos = localCount; + while (pos > 0 && ((want_min && v < localDist[pos-1]) || (!want_min && v > localDist[pos-1]))) { + localDist[pos] = localDist[pos-1]; + localIdx[pos] = localIdx[pos-1]; + pos--; + } + localDist[pos] = v; + localIdx[pos] = (int)j; + localCount++; + } else { + bool better = want_min ? (v < localDist[R-1]) : (v > localDist[R-1]); + if (better) { + uint pos = R - 1; + while (pos > 0 && ((want_min && v < localDist[pos-1]) || (!want_min && v > localDist[pos-1]))) { + localDist[pos] = localDist[pos-1]; + localIdx[pos] = localIdx[pos-1]; + pos--; + } + localDist[pos] = v; + localIdx[pos] = (int)j; + } + } + } + + for (uint i = 0; i < R; i++) { + uint idx = tid * R + i; + if (i < localCount) { + tgDist[idx] = localDist[i]; + tgIdx[idx] = localIdx[i]; + } else { + tgDist[idx] = want_min ? 1e38f : -1e38f; + tgIdx[idx] = -1; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k2 = 2; k2 <= CANDIDATES; k2 *= 2) { + for (uint j = k2 >> 1; j > 0; j >>= 1) { + for (uint idx = tid; idx < CANDIDATES; idx += TG_SIZE) { + uint partner = idx ^ j; + if (partner < CANDIDATES && partner > idx) { + bool ascending = ((idx & k2) == 0); + bool partnerBetter = want_min ? (tgDist[partner] < tgDist[idx] || (tgDist[partner] == tgDist[idx] && tgIdx[partner] < tgIdx[idx])) + : (tgDist[partner] > tgDist[idx] || (tgDist[partner] == tgDist[idx] && tgIdx[partner] < tgIdx[idx])); + bool idxBetter = want_min ? (tgDist[idx] < tgDist[partner] || (tgDist[idx] == tgDist[partner] && tgIdx[idx] < tgIdx[partner])) + : (tgDist[idx] > tgDist[partner] || (tgDist[idx] == tgDist[partner] && tgIdx[idx] < tgIdx[partner])); + bool swap = ascending ? partnerBetter : idxBetter; + if (swap) { + float td = tgDist[idx]; tgDist[idx] = tgDist[partner]; tgDist[partner] = td; + int ti = tgIdx[idx]; tgIdx[idx] = tgIdx[partner]; tgIdx[partner] = ti; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + for (uint i = tid; i < K_out; i += TG_SIZE) { + outDistances[qi * k + i] = tgDist[i]; + outIndices[qi * k + i] = tgIdx[i]; + } + for (uint i = tid; i < k - K_out; i += TG_SIZE) { + outDistances[qi * k + K_out + i] = want_min ? 1e38f : -1e38f; + outIndices[qi * k + K_out + i] = -1; + } +} + +// Bitonic merge kernel: merges two sorted lists of length K′ into one sorted list of length K′. +// Input: two buffers A and B, each (nq, K′) sorted (ascending for L2, descending for IP). +// Output: one buffer C (nq, K′) sorted, containing best K′ from A and B combined. +// Uses one threadgroup per query with bitonic merge network (compare-exchange pattern). +#define BITONIC_MERGE_TWO_SORTED_VARIANT(K) \ + +kernel void compute_norms( + device const float* vectors [[buffer(0)]], + device float* norms [[buffer(1)]], + device const uint* params [[buffer(2)]], // nb, d + uint vid [[thread_position_in_grid]] +) { + uint nb = params[0], d = params[1]; + if (vid >= nb) return; + const device float* v = vectors + vid * d; + float sum = 0.0f; + uint d4 = d / 4; + const device float4* v4 = (const device float4*)v; + for (uint t = 0; t < d4; t++) { + sum += dot(v4[t], v4[t]); + } + for (uint t = d4 * 4; t < d; t++) { + sum += v[t] * v[t]; + } + norms[vid] = sum; +} + +// L2 distance using pre-computed vector (centroid) norms: +// dist[i][j] = queryNorm[i] + vecNorm[j] - 2 * dot(query[i], vec[j]) +// We compute queryNorm on-the-fly (cheap, nq rows) but reuse vecNorms. +kernel void l2_with_norms( + device const float* queries [[buffer(0)]], + device const float* vectors [[buffer(1)]], + device float* distances [[buffer(2)]], + device const uint* params [[buffer(3)]], // nq, nb, d + device const float* vecNorms [[buffer(4)]], + uint2 gid [[thread_position_in_grid]] +) { + uint nq = params[0], nb = params[1], d = params[2]; + uint i = gid.y; + uint j = gid.x; + if (i >= nq || j >= nb) return; + + const device float* q = queries + i * d; + const device float* v = vectors + j * d; + float dot_val = 0.0f; + uint d4 = d / 4; + const device float4* q4 = (const device float4*)q; + const device float4* v4 = (const device float4*)v; + for (uint t = 0; t < d4; t++) { + dot_val += dot(q4[t], v4[t]); + } + for (uint t = d4 * 4; t < d; t++) { + dot_val += q[t] * v[t]; + } + + // Compute query norm inline (avoids separate query norm buffer). + float qNorm = 0.0f; + for (uint t = 0; t < d4; t++) { + qNorm += dot(q4[t], q4[t]); + } + for (uint t = d4 * 4; t < d; t++) { + qNorm += q[t] * q[t]; + } + + distances[i * nb + j] = qNorm + vecNorms[j] - 2.0f * dot_val; +} + +// ============================================================ + +// ============================================================ +// IVF Flat scan — two-pass design (mirrors CUDA IVF): +// +// Pass 1 ivf_scan_list +// Grid: (nq * nprobe) threadgroups, 256 threads each. +// Each threadgroup scans ONE inverted list for one query +// and writes a per-list top-k. +// Output: perListDist/perListIdx — (nq * nprobe * k). +// +// Pass 2 ivf_merge_lists +// Grid: (nq) threadgroups, 256 threads each. +// Merges nprobe per-list top-k into final top-k per query. +// +// Both use float4 vectorised loads for memory throughput. +// ============================================================ +// Shared params layout (device const uint*): +// [0] nq [1] d [2] k [3] nprobe [4] want_min + +inline bool ivf_better_int(float aDist, int aIdx, float bDist, int bIdx, uint want_min) { + constexpr float tie_eps = 1e-6f; + float delta = aDist - bDist; + bool strictly = want_min ? (delta < -tie_eps) : (delta > tie_eps); + if (strictly) return true; + if (fabs(delta) <= tie_eps) return aIdx < bIdx; + return false; +} + +inline bool ivf_better_long(float aDist, long aIdx, float bDist, long bIdx, uint want_min) { + constexpr float tie_eps = 1e-6f; + float delta = aDist - bDist; + bool strictly = want_min ? (delta < -tie_eps) : (delta > tie_eps); + if (strictly) return true; + if (fabs(delta) <= tie_eps) return aIdx < bIdx; + return false; +} + +kernel void ivf_scan_list( + device const float* queries [[buffer(0)]], + device const float* codes [[buffer(1)]], + device const long* ids [[buffer(2)]], + device const uint* listOffset [[buffer(3)]], + device const uint* listLength [[buffer(4)]], + device const int* coarseAssign [[buffer(5)]], + device float* perListDist [[buffer(6)]], + device long* perListIdx [[buffer(7)]], + device const uint* params [[buffer(8)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + constexpr uint TG_SIZE = 256; + constexpr uint LOCAL_K = 4; + + uint nq = params[0]; + uint d = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + + uint qi = tgid / nprobe; + uint pi = tgid % nprobe; + if (qi >= nq || k == 0) return; + + float sentinel = want_min ? 1e38f : -1e38f; + uint outBase = (qi * nprobe + pi) * k; + + int list_no = coarseAssign[qi * nprobe + pi]; + if (list_no < 0) { + for (uint i = tid; i < k; i += TG_SIZE) { + perListDist[outBase + i] = sentinel; + perListIdx [outBase + i] = -1; + } + return; + } + + uint lOff = listOffset[(uint)list_no]; + uint lLen = listLength[(uint)list_no]; + if (lLen == 0) { + for (uint i = tid; i < k; i += TG_SIZE) { + perListDist[outBase + i] = sentinel; + perListIdx [outBase + i] = -1; + } + return; + } + + // Cache query vector in threadgroup memory (read once from device, reused + // by all threads for every vector in this list). + threadgroup float tgQuery[2048]; // max d supported; only first d floats used + const device float* qvecDev = queries + qi * d; + for (uint i = tid; i < d; i += TG_SIZE) { + tgQuery[i] = qvecDev[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Sort internally by (dist, vecIdx) using 32-bit vecIdx for speed; + // resolve to 64-bit ID only when writing output. + float localDist[LOCAL_K]; + int localIdx [LOCAL_K]; + uint localCount = 0; + for (uint i = 0; i < LOCAL_K; i++) { + localDist[i] = sentinel; + localIdx [i] = -1; + } + + uint d4 = d / 4; + + for (uint li = tid; li < lLen; li += TG_SIZE) { + uint vecIdx = lOff + li; + const device float* vvec = codes + vecIdx * d; + + float dist = 0.0f; + if (want_min) { + const threadgroup float4* q4 = (const threadgroup float4*)tgQuery; + const device float4* v4 = (const device float4*)vvec; + for (uint t = 0; t < d4; t++) { + float4 diff = q4[t] - v4[t]; + dist += dot(diff, diff); + } + for (uint t = d4 * 4; t < d; t++) { + float diff = tgQuery[t] - vvec[t]; + dist += diff * diff; + } + } else { + const threadgroup float4* q4 = (const threadgroup float4*)tgQuery; + const device float4* v4 = (const device float4*)vvec; + for (uint t = 0; t < d4; t++) { + dist += dot(q4[t], v4[t]); + } + for (uint t = d4 * 4; t < d; t++) { + dist += tgQuery[t] * vvec[t]; + } + } + + int vi = (int)vecIdx; + + bool better = want_min ? (dist < localDist[LOCAL_K-1]) + : (dist > localDist[LOCAL_K-1]); + if (localCount < LOCAL_K || better) { + uint pos = (localCount < LOCAL_K) ? localCount : LOCAL_K - 1; + localDist[pos] = dist; + localIdx [pos] = vi; + while (pos > 0) { + bool sw = ivf_better_int( + localDist[pos], localIdx[pos], + localDist[pos - 1], localIdx[pos - 1], + want_min); + if (!sw) break; + float td = localDist[pos]; localDist[pos] = localDist[pos-1]; localDist[pos-1] = td; + int ti = localIdx [pos]; localIdx [pos] = localIdx [pos-1]; localIdx [pos-1] = ti; + pos--; + } + if (localCount < LOCAL_K) localCount++; + } + } + + constexpr uint CAND = TG_SIZE * LOCAL_K; // 1024 + threadgroup float tgDist[CAND]; + threadgroup int tgIdx [CAND]; + + for (uint i = 0; i < LOCAL_K; i++) { + tgDist[tid * LOCAL_K + i] = (i < localCount) ? localDist[i] : sentinel; + tgIdx [tid * LOCAL_K + i] = (i < localCount) ? localIdx [i] : -1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k2 = 2; k2 <= CAND; k2 *= 2) { + for (uint j = k2 >> 1; j > 0; j >>= 1) { + for (uint idx = tid; idx < CAND; idx += TG_SIZE) { + uint partner = idx ^ j; + if (partner < CAND && partner > idx) { + bool ascending = ((idx & k2) == 0); + bool pB = ivf_better_int( + tgDist[partner], tgIdx[partner], + tgDist[idx], tgIdx[idx], + want_min); + bool iB = ivf_better_int( + tgDist[idx], tgIdx[idx], + tgDist[partner], tgIdx[partner], + want_min); + if (ascending ? pB : iB) { + float td = tgDist[idx]; tgDist[idx] = tgDist[partner]; tgDist[partner] = td; + int ti = tgIdx [idx]; tgIdx [idx] = tgIdx [partner]; tgIdx [partner] = ti; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + // Write output: resolve 32-bit vecIdx → 64-bit ID from ids buffer. + uint kk = min(k, CAND); + for (uint i = tid; i < kk; i += TG_SIZE) { + int vi = tgIdx[i]; + perListDist[outBase + i] = tgDist[i]; + perListIdx [outBase + i] = (vi < 0) ? -1L : ids[vi]; + } + for (uint i = tid; i < k - kk; i += TG_SIZE) { + perListDist[outBase + kk + i] = sentinel; + perListIdx [outBase + kk + i] = -1L; + } +} + +// ---- Small-list variant: 32 threads, 32-element bitonic sort ---- +// Used when avg list size ≤ 32 (most threads in 256-thread version idle). +// Saves ~90% of bitonic sort barriers and threadgroup memory. +kernel void ivf_scan_list_small( + device const float* queries [[buffer(0)]], + device const float* codes [[buffer(1)]], + device const long* ids [[buffer(2)]], + device const uint* listOffset [[buffer(3)]], + device const uint* listLength [[buffer(4)]], + device const int* coarseAssign [[buffer(5)]], + device float* perListDist [[buffer(6)]], + device long* perListIdx [[buffer(7)]], + device const uint* params [[buffer(8)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + constexpr uint TG_SIZE = 32; + constexpr uint LOCAL_K = 1; + + uint nq = params[0]; + uint d = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + + uint qi = tgid / nprobe; + uint pi = tgid % nprobe; + if (qi >= nq || k == 0) return; + + float sentinel = want_min ? 1e38f : -1e38f; + uint outBase = (qi * nprobe + pi) * k; + + int list_no = coarseAssign[qi * nprobe + pi]; + if (list_no < 0) { + for (uint i = tid; i < k; i += TG_SIZE) { + perListDist[outBase + i] = sentinel; + perListIdx [outBase + i] = -1; + } + return; + } + + uint lOff = listOffset[(uint)list_no]; + uint lLen = listLength[(uint)list_no]; + if (lLen == 0) { + for (uint i = tid; i < k; i += TG_SIZE) { + perListDist[outBase + i] = sentinel; + perListIdx [outBase + i] = -1; + } + return; + } + + // Query cache in threadgroup memory. + threadgroup float tgQuery[2048]; + const device float* qvecDev = queries + qi * d; + for (uint i = tid; i < d; i += TG_SIZE) { + tgQuery[i] = qvecDev[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float bestDist = sentinel; + int bestIdx = -1; + uint d4 = d / 4; + + for (uint li = tid; li < lLen; li += TG_SIZE) { + uint vecIdx = lOff + li; + const device float* vvec = codes + vecIdx * d; + + float dist = 0.0f; + if (want_min) { + const threadgroup float4* q4 = (const threadgroup float4*)tgQuery; + const device float4* v4 = (const device float4*)vvec; + for (uint t = 0; t < d4; t++) { + float4 diff = q4[t] - v4[t]; + dist += dot(diff, diff); + } + for (uint t = d4 * 4; t < d; t++) { + float diff = tgQuery[t] - vvec[t]; + dist += diff * diff; + } + } else { + const threadgroup float4* q4 = (const threadgroup float4*)tgQuery; + const device float4* v4 = (const device float4*)vvec; + for (uint t = 0; t < d4; t++) { + dist += dot(q4[t], v4[t]); + } + for (uint t = d4 * 4; t < d; t++) { + dist += tgQuery[t] * vvec[t]; + } + } + + int vi = (int)vecIdx; + + bool better = ivf_better_int(dist, vi, bestDist, bestIdx, want_min); + if (better) { + bestDist = dist; + bestIdx = vi; + } + } + + constexpr uint CAND = TG_SIZE; // 32 + threadgroup float tgDist[CAND]; + threadgroup int tgIdx [CAND]; + tgDist[tid] = bestDist; + tgIdx [tid] = bestIdx; + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k2 = 2; k2 <= CAND; k2 *= 2) { + for (uint j = k2 >> 1; j > 0; j >>= 1) { + uint partner = tid ^ j; + if (partner < CAND && partner > tid) { + bool ascending = ((tid & k2) == 0); + bool pB = ivf_better_int( + tgDist[partner], tgIdx[partner], + tgDist[tid], tgIdx[tid], + want_min); + bool iB = ivf_better_int( + tgDist[tid], tgIdx[tid], + tgDist[partner], tgIdx[partner], + want_min); + if (ascending ? pB : iB) { + float td = tgDist[tid]; tgDist[tid] = tgDist[partner]; tgDist[partner] = td; + int ti = tgIdx [tid]; tgIdx [tid] = tgIdx [partner]; tgIdx [partner] = ti; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + uint kk = min(k, CAND); + for (uint i = tid; i < kk; i += TG_SIZE) { + int vi = tgIdx[i]; + perListDist[outBase + i] = tgDist[i]; + perListIdx [outBase + i] = (vi < 0) ? -1L : ids[vi]; + } + for (uint i = tid; i < k - kk; i += TG_SIZE) { + perListDist[outBase + kk + i] = sentinel; + perListIdx [outBase + kk + i] = -1L; + } +} + +// ---- Interleaved layout scan: 32-vector blocks, dimensions interleaved ---- +// Memory layout per block: [v0d0 v1d0 .. v31d0] [v0d1 v1d1 .. v31d1] ... +// Each simdgroup (32 threads) processes one block cooperatively, +// achieving coalesced reads: all threads read the same dimension simultaneously. +kernel void ivf_scan_list_interleaved( + device const float* queries [[buffer(0)]], + device const float* codes [[buffer(1)]], + device const long* ids [[buffer(2)]], + device const uint* listOffset [[buffer(3)]], + device const uint* listLength [[buffer(4)]], + device const int* coarseAssign [[buffer(5)]], + device float* perListDist [[buffer(6)]], + device long* perListIdx [[buffer(7)]], + device const uint* params [[buffer(8)]], + device const uint* ilCodesOffset [[buffer(9)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + constexpr uint TG_SIZE = 256; + constexpr uint SIMD_W = 32; + constexpr uint NUM_SIMDS = TG_SIZE / SIMD_W; // 8 + constexpr uint LOCAL_K = 4; + + uint nq = params[0]; + uint d = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + + uint qi = tgid / nprobe; + uint pi = tgid % nprobe; + if (qi >= nq || k == 0) return; + + float sentinel = want_min ? 1e38f : -1e38f; + uint outBase = (qi * nprobe + pi) * k; + + int list_no = coarseAssign[qi * nprobe + pi]; + if (list_no < 0) { + for (uint i = tid; i < k; i += TG_SIZE) { + perListDist[outBase + i] = sentinel; + perListIdx [outBase + i] = -1L; + } + return; + } + + uint idOff = listOffset[(uint)list_no]; + uint lLen = listLength[(uint)list_no]; + uint cOff = ilCodesOffset[(uint)list_no]; + if (lLen == 0) { + for (uint i = tid; i < k; i += TG_SIZE) { + perListDist[outBase + i] = sentinel; + perListIdx [outBase + i] = -1L; + } + return; + } + + uint numBlocks = (lLen + SIMD_W - 1) / SIMD_W; + uint laneId = tid % SIMD_W; + uint simdId = tid / SIMD_W; + + threadgroup float tgQuery[2048]; + const device float* qvecDev = queries + qi * d; + for (uint i = tid; i < d; i += TG_SIZE) { + tgQuery[i] = qvecDev[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float localDist[LOCAL_K]; + int localIdx [LOCAL_K]; + uint localCount = 0; + for (uint i = 0; i < LOCAL_K; i++) { + localDist[i] = sentinel; + localIdx [i] = -1; + } + + for (uint blk = simdId; blk < numBlocks; blk += NUM_SIMDS) { + uint vecInList = blk * SIMD_W + laneId; + bool valid = vecInList < lLen; + + const device float* blockPtr = codes + cOff + blk * SIMD_W * d; + + float dist = 0.0f; + if (want_min) { + for (uint dd = 0; dd < d; dd += 4) { + float v0 = valid ? blockPtr[(dd + 0) * SIMD_W + laneId] : 0.0f; + float v1 = valid ? blockPtr[(dd + 1) * SIMD_W + laneId] : 0.0f; + float v2 = valid ? blockPtr[(dd + 2) * SIMD_W + laneId] : 0.0f; + float v3 = valid ? blockPtr[(dd + 3) * SIMD_W + laneId] : 0.0f; + float4 diff = float4( + tgQuery[dd] - v0, + tgQuery[dd + 1] - v1, + tgQuery[dd + 2] - v2, + tgQuery[dd + 3] - v3); + dist += dot(diff, diff); + } + } else { + for (uint dd = 0; dd < d; dd += 4) { + float v0 = valid ? blockPtr[(dd + 0) * SIMD_W + laneId] : 0.0f; + float v1 = valid ? blockPtr[(dd + 1) * SIMD_W + laneId] : 0.0f; + float v2 = valid ? blockPtr[(dd + 2) * SIMD_W + laneId] : 0.0f; + float v3 = valid ? blockPtr[(dd + 3) * SIMD_W + laneId] : 0.0f; + dist += tgQuery[dd] * v0 + + tgQuery[dd + 1] * v1 + + tgQuery[dd + 2] * v2 + + tgQuery[dd + 3] * v3; + } + } + + if (!valid) dist = sentinel; + int vi = valid ? (int)(idOff + vecInList) : -1; + + bool better = want_min ? (dist < localDist[LOCAL_K-1]) + : (dist > localDist[LOCAL_K-1]); + if (localCount < LOCAL_K || better) { + uint pos = (localCount < LOCAL_K) ? localCount : LOCAL_K - 1; + localDist[pos] = dist; + localIdx [pos] = vi; + while (pos > 0) { + bool sw = ivf_better_int( + localDist[pos], localIdx[pos], + localDist[pos - 1], localIdx[pos - 1], + want_min); + if (!sw) break; + float td = localDist[pos]; localDist[pos] = localDist[pos-1]; localDist[pos-1] = td; + int ti = localIdx [pos]; localIdx [pos] = localIdx [pos-1]; localIdx [pos-1] = ti; + pos--; + } + if (localCount < LOCAL_K) localCount++; + } + } + + constexpr uint CAND = TG_SIZE * LOCAL_K; // 1024 + threadgroup float tgDist[CAND]; + threadgroup int tgIdx [CAND]; + + for (uint i = 0; i < LOCAL_K; i++) { + tgDist[tid * LOCAL_K + i] = (i < localCount) ? localDist[i] : sentinel; + tgIdx [tid * LOCAL_K + i] = (i < localCount) ? localIdx [i] : -1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k2 = 2; k2 <= CAND; k2 *= 2) { + for (uint j = k2 >> 1; j > 0; j >>= 1) { + for (uint idx = tid; idx < CAND; idx += TG_SIZE) { + uint partner = idx ^ j; + if (partner < CAND && partner > idx) { + bool ascending = ((idx & k2) == 0); + bool pB = ivf_better_int( + tgDist[partner], tgIdx[partner], + tgDist[idx], tgIdx[idx], + want_min); + bool iB = ivf_better_int( + tgDist[idx], tgIdx[idx], + tgDist[partner], tgIdx[partner], + want_min); + if (ascending ? pB : iB) { + float td = tgDist[idx]; tgDist[idx] = tgDist[partner]; tgDist[partner] = td; + int ti = tgIdx [idx]; tgIdx [idx] = tgIdx [partner]; tgIdx [partner] = ti; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + uint kk = min(k, CAND); + for (uint i = tid; i < kk; i += TG_SIZE) { + int vi = tgIdx[i]; + perListDist[outBase + i] = tgDist[i]; + perListIdx [outBase + i] = (vi < 0) ? -1L : ids[vi]; + } + for (uint i = tid; i < k - kk; i += TG_SIZE) { + perListDist[outBase + kk + i] = sentinel; + perListIdx [outBase + kk + i] = -1L; + } +} + +// ---- Pass 2: merge nprobe per-list results per query ---- +// Strided-scan pattern: each thread scans every TG_SIZE-th candidate across +// all nprobe×k entries, keeping its best LOCAL_K in registers. Then dump to +// threadgroup memory and bitonic-sort. Handles any nprobe×k without overflow. +kernel void ivf_merge_lists( + device const float* perListDist [[buffer(0)]], + device const long* perListIdx [[buffer(1)]], + device float* outDistances [[buffer(2)]], + device long* outIndices [[buffer(3)]], + device const uint* params [[buffer(4)]], + uint qi [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + constexpr uint TG_SIZE = 256; + constexpr uint LOCAL_K = 4; + constexpr uint CAND = TG_SIZE * LOCAL_K; // 1024 + + uint nq = params[0]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + + if (qi >= nq || k == 0) return; + + float sentinel = want_min ? 1e38f : -1e38f; + uint totalCand = nprobe * k; + uint inputBase = qi * totalCand; + + float localDist[LOCAL_K]; + long localIdx [LOCAL_K]; + uint localCount = 0; + for (uint i = 0; i < LOCAL_K; i++) { + localDist[i] = sentinel; + localIdx [i] = -1L; + } + + for (uint i = tid; i < totalCand; i += TG_SIZE) { + float d = perListDist[inputBase + i]; + long v = perListIdx [inputBase + i]; + + bool better = want_min ? (d < localDist[LOCAL_K-1]) + : (d > localDist[LOCAL_K-1]); + if (localCount < LOCAL_K || better) { + if (v < 0 && localCount >= LOCAL_K) continue; + uint pos = (localCount < LOCAL_K) ? localCount : LOCAL_K - 1; + localDist[pos] = d; + localIdx [pos] = v; + while (pos > 0) { + bool sw = ivf_better_long( + localDist[pos], localIdx[pos], + localDist[pos - 1], localIdx[pos - 1], + want_min); + if (!sw) break; + float td = localDist[pos]; localDist[pos] = localDist[pos-1]; localDist[pos-1] = td; + long ti = localIdx [pos]; localIdx [pos] = localIdx [pos-1]; localIdx [pos-1] = ti; + pos--; + } + if (localCount < LOCAL_K) localCount++; + } + } + + threadgroup float tgDist[CAND]; + threadgroup long tgIdx [CAND]; + + for (uint i = 0; i < LOCAL_K; i++) { + tgDist[tid * LOCAL_K + i] = (i < localCount) ? localDist[i] : sentinel; + tgIdx [tid * LOCAL_K + i] = (i < localCount) ? localIdx [i] : -1L; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (uint k2 = 2; k2 <= CAND; k2 *= 2) { + for (uint j = k2 >> 1; j > 0; j >>= 1) { + for (uint idx = tid; idx < CAND; idx += TG_SIZE) { + uint partner = idx ^ j; + if (partner < CAND && partner > idx) { + bool ascending = ((idx & k2) == 0); + bool pB = ivf_better_long( + tgDist[partner], tgIdx[partner], + tgDist[idx], tgIdx[idx], + want_min); + bool iB = ivf_better_long( + tgDist[idx], tgIdx[idx], + tgDist[partner], tgIdx[partner], + want_min); + if (ascending ? pB : iB) { + float td = tgDist[idx]; tgDist[idx] = tgDist[partner]; tgDist[partner] = td; + long ti = tgIdx [idx]; tgIdx [idx] = tgIdx [partner]; tgIdx [partner] = ti; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } + + uint kk = min(k, CAND); + for (uint i = tid; i < kk; i += TG_SIZE) { + outDistances[qi * k + i] = tgDist[i]; + outIndices [qi * k + i] = tgIdx [i]; + } + for (uint i = tid; i < k - kk; i += TG_SIZE) { + outDistances[qi * k + kk + i] = sentinel; + outIndices [qi * k + kk + i] = -1L; + } +} + +// ============================================================ +// IVF Scalar Quantizer scan kernels +// ============================================================ +// +// QT_8bit: codes are uchar[d] per vector; decode: vmin[dim] + code * vdiff[dim] +// SQ tables buffer layout: vmin[0..d-1], vdiff[0..d-1] (2*d floats) +// QT_4bit: packed 2 values per byte, low/high nibble +// QT_6bit: packed 4 values per 3 bytes +// QT_fp16: codes are half[d] per vector; no SQ tables needed +// +// Both share the same top-k selection as ivf_scan_list. + +inline float sq4_decode_component(device const uchar* code, uint i) { + uchar b = code[i >> 1]; + uchar bits = (i & 1u) ? (b >> 4) : (b & 0x0Fu); + return (float(bits) + 0.5f) / 15.0f; +} + +inline float sq6_decode_component(device const uchar* code, uint i) { + const device uchar* p = code + (i >> 2) * 3; + uchar bits = 0; + switch (i & 3u) { + case 0: + bits = p[0] & 0x3Fu; + break; + case 1: + bits = (p[0] >> 6) | ((p[1] & 0x0Fu) << 2); + break; + case 2: + bits = (p[1] >> 4) | ((p[2] & 0x03u) << 4); + break; + default: + bits = p[2] >> 2; + break; + } + return (float(bits) + 0.5f) / 63.0f; +} + diff --git a/faiss/gpu_metal/MetalDistance.mm b/faiss/gpu_metal/MetalDistance.mm new file mode 100644 index 000000000..a17e9d9ae --- /dev/null +++ b/faiss/gpu_metal/MetalDistance.mm @@ -0,0 +1,359 @@ +// @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. + * + * IVF distance computation and scan dispatch for Metal backend. + */ + +#import "MetalDistance.h" +#include +#include +#include +#import "MetalKernels.h" + +namespace faiss { +namespace gpu_metal { + +constexpr uint32_t kIvfReduceThreadgroupSize = 256; +constexpr uint32_t kIvfReduceLocalK = 4; +constexpr uint32_t kIvfReduceExactCandidates = + kIvfReduceThreadgroupSize * kIvfReduceLocalK; // 1024 +constexpr uint32_t kIvfSmallExactCandidates = 32; + +inline bool fitsU32(int v) { + return v >= 0 && + static_cast(v) <= + static_cast(std::numeric_limits::max()); +} + +inline bool ivfMergeExactnessHolds(int nprobe, int k) { + if (nprobe <= 0 || k <= 0) + return false; + const uint64_t totalCandidates = (uint64_t)nprobe * (uint64_t)k; + return totalCandidates <= (uint64_t)kIvfReduceExactCandidates; +} + +inline bool ivfScanExactnessHoldsForAllLists( + id listLength, + int nlist, + uint32_t maxCandidatesPerList) { + if (!listLength || nlist <= 0) + return false; + const uint32_t* listLens = + reinterpret_cast([listLength contents]); + if (!listLens) + return false; + for (int i = 0; i < nlist; ++i) { + if (listLens[i] > maxCandidatesPerList) + return false; + } + return true; +} + +inline bool ivfScanExactnessHoldsForAssignments( + id listLength, + id coarseAssign, + int nq, + int nprobe) { + if (!listLength || !coarseAssign || nq <= 0 || nprobe <= 0) + return false; + const uint32_t* listLens = + reinterpret_cast([listLength contents]); + const int32_t* assign = + reinterpret_cast([coarseAssign contents]); + if (!listLens || !assign) + return false; + const size_t nlist = [listLength length] / sizeof(uint32_t); + if (nlist == 0) + return false; + const size_t total = (size_t)nq * (size_t)nprobe; + for (size_t i = 0; i < total; ++i) { + const int32_t listNo = assign[i]; + if (listNo < 0) + continue; + if ((size_t)listNo >= nlist) + return false; + if (listLens[listNo] > kIvfReduceExactCandidates) + return false; + } + return true; +} + +int getMetalDistanceMaxK() { + return MetalKernels::kMaxK; +} + +bool runMetalComputeNorms( + id device, + id queue, + id vectors, + int nb, + int d, + id normsBuf, + bool waitForCompletion) { + if (!device || !queue || !vectors || !normsBuf || nb <= 0 || d <= 0) + return false; + MetalKernels& K = getMetalKernels(device); + if (!K.isValid()) + return false; + + id cmdBuf = [queue commandBuffer]; + id enc = [cmdBuf computeCommandEncoder]; + K.encodeComputeNorms(enc, vectors, normsBuf, nb, d); + [enc endEncoding]; + [cmdBuf commit]; + if (!waitForCompletion) + return true; + [cmdBuf waitUntilCompleted]; + return cmdBuf.status == MTLCommandBufferStatusCompleted; +} + +bool runMetalIVFFlatScan( + id device, + id queue, + id queries, + id codes, + id ids, + id listOffset, + id listLength, + id coarseAssign, + int nq, + int d, + int k, + int nprobe, + bool isL2, + id outDistances, + id outIndices, + id perListDistBuf, + id perListIdxBuf, + id interleavedCodes, + id interleavedCodesOffset, + bool waitForCompletion) { + bool useIL = (interleavedCodes != nil && interleavedCodesOffset != nil); + if (!device || !queue || !queries || (!codes && !useIL) || !ids || + !listOffset || !listLength || !coarseAssign || !outDistances || + !outIndices || !perListDistBuf || !perListIdxBuf) + return false; + if (k <= 0 || nq <= 0 || nprobe <= 0) + return false; + if (d <= 0 || d > 2048) + return false; + if (useIL && (d % 4 != 0)) + return false; + if (!ivfMergeExactnessHolds(nprobe, k)) + return false; + if (!ivfScanExactnessHoldsForAssignments( + listLength, coarseAssign, nq, nprobe)) + return false; + if (!fitsU32(nq) || !fitsU32(d) || !fitsU32(k) || !fitsU32(nprobe)) + return false; + + MetalKernels& K = getMetalKernels(device); + if (!K.isValid()) + return false; + + IVFScanVariant variant = + useIL ? IVFScanVariant::Interleaved : IVFScanVariant::Standard; + + uint32_t sp[5] = { + (uint32_t)nq, + (uint32_t)d, + (uint32_t)k, + (uint32_t)nprobe, + isL2 ? 1u : 0u}; + id paramsBuf = + [device newBufferWithBytes:sp + length:sizeof(sp) + options:MTLResourceStorageModeShared]; + if (!paramsBuf) + return false; + + id cmdBuf = [queue commandBuffer]; + id enc = [cmdBuf computeCommandEncoder]; + + K.encodeIVFScanList( + enc, + variant, + queries, + useIL ? interleavedCodes : codes, + ids, + listOffset, + listLength, + coarseAssign, + perListDistBuf, + perListIdxBuf, + paramsBuf, + nq, + nprobe, + useIL ? interleavedCodesOffset : nil); + K.encodeIVFMergeLists( + enc, + perListDistBuf, + perListIdxBuf, + outDistances, + outIndices, + paramsBuf, + nq); + + [enc endEncoding]; + [cmdBuf commit]; + if (!waitForCompletion) + return true; + [cmdBuf waitUntilCompleted]; + return cmdBuf.status == MTLCommandBufferStatusCompleted; +} + +bool runMetalIVFFlatFullSearch( + id device, + id queue, + id queries, + int nq, + int d, + int k, + int nprobe, + bool isL2, + id centroids, + int nlist, + id codes, + id ids, + id listOffset, + id listLength, + id outDistances, + id outIndices, + id perListDistBuf, + id perListIdxBuf, + id coarseDistBuf, + id coarseIdxBuf, + id distMatrixBuf, + id centroidNormsBuf, + int avgListLen, + id interleavedCodes, + id interleavedCodesOffset, + bool centroidsAreFP16, + bool waitForCompletion) { + (void)centroidsAreFP16; + bool useIL = (interleavedCodes != nil && interleavedCodesOffset != nil); + if (!device || !queue || !queries || !centroids || (!codes && !useIL) || + !ids || !listOffset || !listLength || !outDistances || !outIndices || + !perListDistBuf || !perListIdxBuf || !coarseDistBuf || !coarseIdxBuf || + !distMatrixBuf) + return false; + if (k <= 0 || nq <= 0 || nprobe <= 0 || nlist <= 0) + return false; + if (d <= 0 || d > 2048) + return false; + if (useIL && (d % 4 != 0)) + return false; + (void)avgListLen; + if (!ivfMergeExactnessHolds(nprobe, k)) + return false; + if (!ivfScanExactnessHoldsForAllLists( + listLength, nlist, kIvfReduceExactCandidates)) + return false; + if (!fitsU32(nq) || !fitsU32(d) || !fitsU32(k) || !fitsU32(nprobe)) + return false; + + MetalKernels& K = getMetalKernels(device); + if (!K.isValid()) + return false; + + const bool useSmall = !useIL && k <= (int)kIvfSmallExactCandidates && + ivfScanExactnessHoldsForAllLists( + listLength, nlist, kIvfSmallExactCandidates); + IVFScanVariant scanV = useIL ? IVFScanVariant::Interleaved + : useSmall ? IVFScanVariant::Small + : IVFScanVariant::Standard; + + id cmdBuf = [queue commandBuffer]; + id enc = [cmdBuf computeCommandEncoder]; + + // Step 1: coarse distance matrix + bool fusedL2 = isL2 && centroidNormsBuf != nil; + if (fusedL2) { + K.encodeL2WithNorms( + enc, + queries, + centroids, + distMatrixBuf, + centroidNormsBuf, + nq, + nlist, + d); + } else { + K.encodeDistanceMatrix( + enc, + queries, + centroids, + distMatrixBuf, + nq, + nlist, + d, + isL2 ? METRIC_L2 : METRIC_INNER_PRODUCT); + } + + // Step 2: coarse top-nprobe + K.encodeTopKThreadgroup( + enc, + distMatrixBuf, + coarseDistBuf, + coarseIdxBuf, + nq, + nlist, + nprobe, + isL2); + + // Step 3: IVF scan + uint32_t sp[5] = { + (uint32_t)nq, + (uint32_t)d, + (uint32_t)k, + (uint32_t)nprobe, + isL2 ? 1u : 0u}; + id paramsBuf = + [device newBufferWithBytes:sp + length:sizeof(sp) + options:MTLResourceStorageModeShared]; + if (!paramsBuf) { + [enc endEncoding]; + return false; + } + + K.encodeIVFScanList( + enc, + scanV, + queries, + useIL ? interleavedCodes : codes, + ids, + listOffset, + listLength, + coarseIdxBuf, + perListDistBuf, + perListIdxBuf, + paramsBuf, + nq, + nprobe, + useIL ? interleavedCodesOffset : nil); + + // Step 4: merge + K.encodeIVFMergeLists( + enc, + perListDistBuf, + perListIdxBuf, + outDistances, + outIndices, + paramsBuf, + nq); + + [enc endEncoding]; + [cmdBuf commit]; + if (!waitForCompletion) + return true; + [cmdBuf waitUntilCompleted]; + return cmdBuf.status == MTLCommandBufferStatusCompleted; +} + +} // namespace gpu_metal +} // namespace faiss diff --git a/faiss/gpu_metal/MetalIndex.h b/faiss/gpu_metal/MetalIndex.h index 0031bcccb..f40a4ca49 100644 --- a/faiss/gpu_metal/MetalIndex.h +++ b/faiss/gpu_metal/MetalIndex.h @@ -11,6 +11,7 @@ #pragma once #include +#include #include #include @@ -20,6 +21,12 @@ namespace gpu_metal { /// Configuration for Metal index (mirrors GpuIndexConfig roles). struct MetalIndexConfig { int device = 0; + + bool useFloat16CoarseQuantizer = false; + + faiss::gpu::IndicesOptions indicesOptions = faiss::gpu::INDICES_64_BIT; + + bool interleavedLayout = true; }; /// Base class for Metal-backed indexes. Mirrors faiss::gpu::GpuIndex. diff --git a/faiss/gpu_metal/MetalIndexIVFFlat.h b/faiss/gpu_metal/MetalIndexIVFFlat.h new file mode 100644 index 000000000..d7faccfdd --- /dev/null +++ b/faiss/gpu_metal/MetalIndexIVFFlat.h @@ -0,0 +1,181 @@ +// @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. + * + * Minimal Metal IVFFlat wrapper. + * + */ + +#pragma once + +#import + +#include +#include +#include + +#include + +namespace faiss { +namespace gpu_metal { +class MetalIVFFlatImpl; +} // namespace gpu_metal +} // namespace faiss + +namespace faiss { +namespace gpu_metal { + +/// IVFFlat index wrapper for Metal backend. +/// Currently delegates to an internal CPU IndexIVFFlat; later phases +/// may move list scanning to GPU. +class MetalIndexIVFFlat : public MetalIndex { + public: + struct AppendDebugStats { + size_t relayoutEvents = 0; + size_t movedLists = 0; + size_t movedVectors = 0; + size_t reusedSegmentAllocs = 0; + size_t tailSegmentAllocs = 0; + size_t reusedCapacityVecs = 0; + size_t tailCapacityVecs = 0; + size_t tailShrinkEvents = 0; + size_t tailShrunkVecs = 0; + }; + + /// Construct empty IVFFlat index with its own CPU quantizer. + MetalIndexIVFFlat( + std::shared_ptr resources, + int dims, + idx_t nlist, + faiss::MetricType metric, + float metricArg = 0.0f, + MetalIndexConfig config = MetalIndexConfig()); + + /// Construct empty IVFFlat index with caller-provided coarse quantizer. + /// If ownFields is true, this index takes ownership of `coarseQuantizer`. + MetalIndexIVFFlat( + std::shared_ptr resources, + faiss::Index* coarseQuantizer, + int dims, + idx_t nlist, + faiss::MetricType metric, + float metricArg = 0.0f, + MetalIndexConfig config = MetalIndexConfig(), + bool ownFields = false); + + /// Construct from an existing CPU IndexIVFFlat (used by cloners later). + MetalIndexIVFFlat( + std::shared_ptr resources, + const faiss::IndexIVFFlat* cpuIndex, + MetalIndexConfig config = MetalIndexConfig()); + + ~MetalIndexIVFFlat() 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; + + /// Search with caller-provided coarse assignments (skips coarse quantizer). + /// @param assign Coarse list assignments (n x nprobe), row-major idx_t + /// @param centroid_dis Distances to assigned centroids (n x nprobe); unused + /// by GPU scan but accepted for API compatibility + /// @param store_pairs Ignored (always false for GPU path) + void search_preassigned( + idx_t n, + const float* x, + idx_t k, + const idx_t* assign, + const float* centroid_dis, + float* distances, + idx_t* labels, + bool store_pairs, + const IVFSearchParameters* params = nullptr, + IndexIVFStats* stats = nullptr) const; + + /// Copy from a CPU IndexIVFFlat (helper for future cloner support). + void copyFrom(const faiss::IndexIVFFlat* index); + + /// Copy to a CPU IndexIVFFlat. + void copyTo(faiss::IndexIVFFlat* index) const; + + /// Reconstruct a single stored vector by internal key. + void reconstruct(idx_t key, float* recons) const override; + + /// Reconstruct n contiguous stored vectors starting at i0. + void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override; + + /// Re-upload coarse quantizer centroids to GPU after external changes. + void updateQuantizer(); + + /// Return the vector indices in inverted list `listId`. + std::vector getListIndices(idx_t listId) const; + + /// Return raw vector data from inverted list `listId`. + std::vector getListVectorData(idx_t listId) const; + + /// Release unused GPU memory. + void reclaimMemory(); + + /// Pre-allocate GPU storage for the given total number of vectors. + void reserveMemory(idx_t numVecs); + + /// Accessors (needed by cloner and tests). + idx_t nlist() const; + size_t nprobe() const; + bool interleavedLayout() const; + faiss::gpu::IndicesOptions indicesOptions() const; + AppendDebugStats appendDebugStats() const; + void resetAppendDebugStats(); + + private: + std::unique_ptr cpuIndex_; + std::unique_ptr gpuIvf_; + faiss::gpu::IndicesOptions indicesOptions_; + bool interleavedLayout_; + + // Persistent search buffers — allocated once, grown lazily. + // Declared mutable so search() (const) can resize them. + mutable id searchQueriesBuf_ = nil; + mutable id searchCoarseBuf_ = nil; + mutable id searchOutDistBuf_ = nil; + mutable id searchOutIdxBuf_ = nil; + mutable size_t searchQueriesCap_ = 0; // bytes + mutable size_t searchCoarseCap_ = 0; + mutable size_t searchOutDistCap_ = 0; + mutable size_t searchOutIdxCap_ = 0; + mutable id searchPerListDistBuf_ = nil; + mutable id searchPerListIdxBuf_ = nil; + mutable size_t searchPerListDistCap_ = 0; + mutable size_t searchPerListIdxCap_ = 0; + + // GPU coarse quantizer buffers (cached, rebuilt on train) + mutable id centroidBuf_ = nil; + mutable id centroidNormsBuf_ = nil; // pre-computed ||c||² + mutable id coarseOutDistBuf_ = nil; + mutable id coarseOutIdxBuf_ = nil; + mutable size_t coarseOutDistCap_ = 0; + mutable size_t coarseOutIdxCap_ = 0; + mutable id distMatrixBuf_ = nil; + mutable size_t distMatrixCap_ = 0; + + /// Ensures buf is at least `needed` bytes, reallocating if necessary. + void ensureSearchBuf_(id& buf, size_t& cap, size_t needed) const; + + /// (Re)uploads quantizer centroids to centroidBuf_. + void uploadCentroids_() const; +}; + +} // namespace gpu_metal +} // namespace faiss diff --git a/faiss/gpu_metal/MetalIndexIVFFlat.mm b/faiss/gpu_metal/MetalIndexIVFFlat.mm new file mode 100644 index 000000000..1a15c16b2 --- /dev/null +++ b/faiss/gpu_metal/MetalIndexIVFFlat.mm @@ -0,0 +1,2308 @@ +// @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. + */ + +#import "MetalIndexIVFFlat.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace { +constexpr size_t kDefaultIvfQueryTileBudgetBytes = 256ULL * 1024 * 1024; +constexpr size_t kMinIvfQueryTileBudgetBytes = 16ULL * 1024 * 1024; +constexpr size_t kMaxIvfQueryTileBudgetBytes = 4ULL * 1024 * 1024 * 1024; +constexpr size_t kDefaultIvfFullCoarseMaxBytes = 16ULL * 1024 * 1024; +constexpr faiss::idx_t kIVFFlatSupportedMaxK = 1024; +constexpr faiss::idx_t kAutoReserveMinBatch = 0; +constexpr size_t kIvfExactCandidates = 1024; +constexpr size_t kMinIvfExactCandidates = 64; + +size_t getIvfQueryTileBudgetBytes() { + const char* envBytes = std::getenv("FAISS_METAL_IVF_QUERY_TILE_BYTES"); + if (envBytes && envBytes[0] != '\0') { + char* end = nullptr; + unsigned long long val = std::strtoull(envBytes, &end, 10); + if (end != envBytes && val != 0) { + size_t out = static_cast(val); + out = std::max(out, kMinIvfQueryTileBudgetBytes); + out = std::min(out, kMaxIvfQueryTileBudgetBytes); + return out; + } + } + return kDefaultIvfQueryTileBudgetBytes; +} + +faiss::idx_t getIvfAutoReserveMinBatch() { + const char* env = std::getenv("FAISS_METAL_IVF_AUTO_RESERVE_MIN_BATCH"); + if (!env || env[0] == '\0') { + return kAutoReserveMinBatch; + } + char* end = nullptr; + unsigned long long v = std::strtoull(env, &end, 10); + if (end == env) { + return kAutoReserveMinBatch; + } + if (v > (unsigned long long)std::numeric_limits::max()) { + return std::numeric_limits::max(); + } + return (faiss::idx_t)v; +} + +size_t chooseIvfSearchTileRows( + size_t nq, + int d, + faiss::idx_t k, + size_t nprobe, + int nlist) { + size_t perQuery = 0; + perQuery += (size_t)d * sizeof(float); + perQuery += (size_t)k * (sizeof(float) + sizeof(int64_t)); + perQuery += nprobe * (size_t)k * (sizeof(float) + sizeof(int64_t)); + perQuery += nprobe * (sizeof(float) + sizeof(int32_t)); + perQuery += (size_t)nlist * sizeof(float); + if (perQuery == 0) { + return nq; + } + + size_t tile = getIvfQueryTileBudgetBytes() / perQuery; + tile = std::max(tile, 1); + tile = std::min(tile, nq); + return tile; +} + +size_t chooseIvfPreassignedTileRows( + size_t nq, + int d, + faiss::idx_t k, + size_t nprobe) { + size_t perQuery = 0; + perQuery += (size_t)d * sizeof(float); + perQuery += (size_t)k * (sizeof(float) + sizeof(int64_t)); + perQuery += nprobe * (size_t)k * (sizeof(float) + sizeof(int64_t)); + perQuery += nprobe * sizeof(int32_t); + if (perQuery == 0) { + return nq; + } + + size_t tile = getIvfQueryTileBudgetBytes() / perQuery; + tile = std::max(tile, 1); + tile = std::min(tile, nq); + return tile; +} + +bool allowCpuFallbackForIvf() { + const char* env = std::getenv("FAISS_METAL_IVF_ALLOW_CPU_FALLBACK"); + if (!env || env[0] == '\0') { + return true; + } + if (env[0] == '0' || env[0] == 'n' || env[0] == 'N' || env[0] == 'f' || + env[0] == 'F') { + return false; + } + return true; +} + +bool logCpuFallbackForIvf() { + const char* env = std::getenv("FAISS_METAL_IVF_LOG_CPU_FALLBACK"); + if (!env || env[0] == '\0') { + return true; + } + if (env[0] == '0' || env[0] == 'n' || env[0] == 'N' || env[0] == 'f' || + env[0] == 'F') { + return false; + } + return true; +} + +bool logSyncProfileForIvf() { + const char* env = std::getenv("FAISS_METAL_IVF_LOG_SYNC_PROFILE"); + if (!env || env[0] == '\0') { + return false; + } + if (env[0] == '0' || env[0] == 'n' || env[0] == 'N' || env[0] == 'f' || + env[0] == 'F') { + return false; + } + return true; +} + +size_t getIvfFullCoarseMaxBytes() { + const char* env = std::getenv("FAISS_METAL_IVF_FULL_COARSE_MAX_BYTES"); + if (!env || env[0] == '\0') { + return kDefaultIvfFullCoarseMaxBytes; + } + char* end = nullptr; + unsigned long long v = std::strtoull(env, &end, 10); + if (end == env) { + return kDefaultIvfFullCoarseMaxBytes; + } + return static_cast(v); +} + +bool useFullCoarseGpuForIvf() { + const char* env = std::getenv("FAISS_METAL_IVF_USE_FULL_COARSE"); + if (!env || env[0] == '\0') { + return true; + } + if (env[0] == '0' || env[0] == 'n' || env[0] == 'N' || env[0] == 'f' || + env[0] == 'F') { + return false; + } + return true; +} + +bool forceChunkedIvfSelectionPath() { + const char* env = std::getenv("FAISS_METAL_IVF_FORCE_CHUNKED_SELECTION"); + if (!env || env[0] == '\0') { + return false; + } + if (env[0] == '0' || env[0] == 'n' || env[0] == 'N' || env[0] == 'f' || + env[0] == 'F') { + return false; + } + return true; +} + +size_t getIvfExactCandidateBudget() { + const char* env = std::getenv("FAISS_METAL_IVF_EXACT_CANDIDATES"); + if (!env || env[0] == '\0') { + return kIvfExactCandidates; + } + char* end = nullptr; + unsigned long long v = std::strtoull(env, &end, 10); + if (end == env || v == 0) { + return kIvfExactCandidates; + } + size_t out = static_cast(v); + out = std::max(out, kMinIvfExactCandidates); + out = std::min(out, kIvfExactCandidates); + return out; +} + +struct ScopedIds { + const faiss::InvertedLists* inv = nullptr; + size_t listNo = 0; + const faiss::idx_t* ptr = nullptr; + + ScopedIds() = default; + ScopedIds(const faiss::InvertedLists* inv_, size_t listNo_) + : inv(inv_), listNo(listNo_) { + if (inv) { + ptr = inv->get_ids(listNo); + } + } + ~ScopedIds() { + if (inv && ptr) { + inv->release_ids(listNo, ptr); + } + } +}; + +struct ScopedCodes { + const faiss::InvertedLists* inv = nullptr; + size_t listNo = 0; + const uint8_t* ptr = nullptr; + + ScopedCodes() = default; + ScopedCodes(const faiss::InvertedLists* inv_, size_t listNo_) + : inv(inv_), listNo(listNo_) { + if (inv) { + ptr = inv->get_codes(listNo); + } + } + ~ScopedCodes() { + if (inv && ptr) { + inv->release_codes(listNo, ptr); + } + } +}; + +void floatToHalf(const float* src, uint16_t* dst, size_t n) { + for (size_t i = 0; i < n; ++i) { + __fp16 h = (__fp16)src[i]; + std::memcpy(&dst[i], &h, sizeof(uint16_t)); + } +} + +inline faiss::idx_t decodeCpuLabelFromPair( + const faiss::IndexIVFFlat* cpuIndex, + int64_t pairLabel) { + const uint64_t pair = static_cast(pairLabel); + const uint64_t listNo = faiss::lo_listno(pair); + const uint64_t offset = faiss::lo_offset(pair); + if (!cpuIndex || listNo >= (uint64_t)cpuIndex->nlist) { + return -1; + } + const size_t sz = cpuIndex->invlists->list_size((size_t)listNo); + if (offset >= sz) { + return -1; + } + ScopedIds ids(cpuIndex->invlists, (size_t)listNo); + return ids.ptr ? ids.ptr[offset] : -1; +} + +inline int checkedIdxToInt(faiss::idx_t v, const char* what) { + if (!(v >= 0 && v <= (faiss::idx_t)std::numeric_limits::max())) { + FAISS_THROW_FMT("%s", what); + } + return (int)v; +} + +inline int checkedSizeToInt(size_t v, const char* what) { + if (!(v <= (size_t)std::numeric_limits::max())) { + FAISS_THROW_FMT("%s", what); + } + return (int)v; +} + +bool assignedListsWithinLimit( + const faiss::IndexIVFFlat* cpuIndex, + const faiss::idx_t* assign, + size_t count, + size_t maxListSize, + faiss::idx_t& badListId, + size_t& badListSize) { + if (!cpuIndex || !cpuIndex->invlists || !assign) { + return false; + } + for (size_t i = 0; i < count; ++i) { + const faiss::idx_t listId = assign[i]; + if (listId < 0 || listId >= cpuIndex->nlist) { + badListId = listId; + badListSize = 0; + return false; + } + const size_t sz = cpuIndex->invlists->list_size((size_t)listId); + if (sz > maxListSize) { + badListId = listId; + badListSize = sz; + return false; + } + } + return true; +} + +std::string explainIvfScanFailureEnvelope( + int d, + faiss::idx_t k, + size_t exactCandidateBudget, + size_t nprobe, + bool useInterleaved, + const faiss::IndexIVFFlat* cpuIndex, + const faiss::idx_t* assign, + size_t assignCount) { + if (!(d > 0 && d <= 512)) { + return "scan envelope: d must be in [1,512]"; + } + if (useInterleaved && (d % 4 != 0)) { + return "scan envelope: interleaved layout requires d % 4 == 0"; + } + if (nprobe > 0 && (size_t)k > (exactCandidateBudget / nprobe)) { + return std::string("scan envelope: nprobe*k exceeds ") + + std::to_string(exactCandidateBudget) + " exact-candidate bound"; + } + faiss::idx_t badListId = -1; + size_t badListSize = 0; + if (!assignedListsWithinLimit( + cpuIndex, + assign, + assignCount, + exactCandidateBudget, + badListId, + badListSize)) { + if (badListSize > 0) { + return std::string("scan envelope: selected list size exceeds ") + + std::to_string(exactCandidateBudget) + + " (list=" + std::to_string((long long)badListId) + + ", size=" + std::to_string(badListSize) + ")"; + } + return "scan envelope: invalid coarse assignment list id"; + } + return {}; +} + +bool selectedListStats( + const faiss::IndexIVFFlat* cpuIndex, + const faiss::idx_t* assign, + size_t count, + std::vector& usedLists, + size_t& maxSelectedListSize, + faiss::idx_t& badListId) { + if (!cpuIndex || !cpuIndex->invlists || !assign || cpuIndex->nlist < 0) { + return false; + } + usedLists.assign((size_t)cpuIndex->nlist, 0); + maxSelectedListSize = 0; + badListId = -1; + for (size_t i = 0; i < count; ++i) { + faiss::idx_t li = assign[i]; + if (li < 0 || li >= cpuIndex->nlist) { + badListId = li; + return false; + } + if (!usedLists[(size_t)li]) { + usedLists[(size_t)li] = 1; + const size_t sz = cpuIndex->invlists->list_size((size_t)li); + maxSelectedListSize = std::max(maxSelectedListSize, sz); + } + } + return true; +} + +inline void heapSiftDown( + float* dist, + int64_t* idx, + int root, + int n, + bool isL2) { + while (true) { + int worst = root; + int left = 2 * root + 1; + int right = 2 * root + 2; + if (isL2) { + if (left < n && dist[left] > dist[worst]) + worst = left; + if (right < n && dist[right] > dist[worst]) + worst = right; + } else { + if (left < n && dist[left] < dist[worst]) + worst = left; + if (right < n && dist[right] < dist[worst]) + worst = right; + } + if (worst == root) + break; + std::swap(dist[root], dist[worst]); + std::swap(idx[root], idx[worst]); + root = worst; + } +} + +inline void insertTopKCandidate( + float candDist, + int64_t candIdx, + bool isL2, + float* bestDist, + int64_t* bestIdx, + int k) { + if (candIdx < 0 || k <= 0) + return; + bool dominated = + isL2 ? (candDist >= bestDist[0]) : (candDist <= bestDist[0]); + if (dominated) + return; + bestDist[0] = candDist; + bestIdx[0] = candIdx; + heapSiftDown(bestDist, bestIdx, 0, k, isL2); +} + +inline void heapSort(float* dist, int64_t* idx, int k, bool isL2) { + for (int i = k - 1; i > 0; --i) { + std::swap(dist[0], dist[i]); + std::swap(idx[0], idx[i]); + heapSiftDown(dist, idx, 0, i, isL2); + } +} +} // namespace + +namespace faiss { +namespace gpu_metal { + +MetalIndexIVFFlat::MetalIndexIVFFlat( + std::shared_ptr resources, + int dims, + idx_t nlist, + faiss::MetricType metric, + float metricArg, + MetalIndexConfig config) + : MetalIndex(resources, dims, metric, metricArg, config), + indicesOptions_(config.indicesOptions), + interleavedLayout_(config.interleavedLayout) { + // Simple CPU quantizer: IndexFlatL2 or IndexFlatIP + faiss::IndexFlat* quantizer = (metric == METRIC_INNER_PRODUCT) + ? (faiss::IndexFlat*)new faiss::IndexFlatIP(dims) + : (faiss::IndexFlat*)new faiss::IndexFlatL2(dims); + cpuIndex_ = std::make_unique( + quantizer, (size_t)d, (size_t)nlist, metric); + cpuIndex_->own_fields = true; + gpuIvf_ = std::make_unique( + resources, + (int)d, + nlist, + metric, + metricArg, + indicesOptions_, + interleavedLayout_); +} + +MetalIndexIVFFlat::MetalIndexIVFFlat( + std::shared_ptr resources, + faiss::Index* coarseQuantizer, + int dims, + idx_t nlist, + faiss::MetricType metric, + float metricArg, + MetalIndexConfig config, + bool ownFields) + : MetalIndex(resources, dims, metric, metricArg, config), + indicesOptions_(config.indicesOptions), + interleavedLayout_(config.interleavedLayout) { + FAISS_THROW_IF_NOT_MSG( + coarseQuantizer != nullptr, + "MetalIndexIVFFlat: coarseQuantizer must be non-null"); + cpuIndex_ = std::make_unique( + coarseQuantizer, (size_t)d, (size_t)nlist, metric); + cpuIndex_->own_fields = ownFields; + gpuIvf_ = std::make_unique( + resources, + (int)d, + nlist, + metric, + metricArg, + indicesOptions_, + interleavedLayout_); +} + +MetalIndexIVFFlat::MetalIndexIVFFlat( + std::shared_ptr resources, + const faiss::IndexIVFFlat* cpuIndex, + MetalIndexConfig config) + : MetalIndex( + resources, + (int)cpuIndex->d, + cpuIndex->metric_type, + cpuIndex->metric_arg, + config), + indicesOptions_(config.indicesOptions), + interleavedLayout_(config.interleavedLayout) { + faiss::IndexFlat* quantizer = + (cpuIndex->metric_type == METRIC_INNER_PRODUCT) + ? (faiss::IndexFlat*)new faiss::IndexFlatIP((int)cpuIndex->d) + : (faiss::IndexFlat*)new faiss::IndexFlatL2((int)cpuIndex->d); + cpuIndex_ = std::make_unique( + quantizer, cpuIndex->d, cpuIndex->nlist, cpuIndex->metric_type); + cpuIndex_->own_fields = true; + gpuIvf_ = std::make_unique( + resources, + (int)cpuIndex->d, + cpuIndex->nlist, + cpuIndex->metric_type, + cpuIndex->metric_arg, + indicesOptions_, + interleavedLayout_); + copyFrom(cpuIndex); +} + +MetalIndexIVFFlat::~MetalIndexIVFFlat() = default; + +void MetalIndexIVFFlat::uploadCentroids_() const { + if (!cpuIndex_ || !cpuIndex_->quantizer || !resources_) { + return; + } + auto* q = cpuIndex_->quantizer; + if (!q || q->ntotal == 0) { + centroidBuf_ = nil; + centroidNormsBuf_ = nil; + return; + } + size_t nCentroids = (size_t)q->ntotal; + std::vector centroids(nCentroids * (size_t)d); + q->reconstruct_n(0, (idx_t)nCentroids, centroids.data()); + + const bool fp16 = config_.useFloat16CoarseQuantizer; + size_t elemSize = fp16 ? sizeof(uint16_t) : sizeof(float); + size_t bytes = nCentroids * (size_t)d * elemSize; + id device = resources_->getDevice(); + if (!device) { + return; + } + centroidBuf_ = [device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (centroidBuf_) { + const float* src = centroids.data(); + if (fp16) { + floatToHalf( + src, + reinterpret_cast([centroidBuf_ contents]), + nCentroids * (size_t)d); + } else { + std::memcpy([centroidBuf_ contents], src, bytes); + } + } + + // Pre-compute centroid L2 norms on GPU (float32 centroids only). + if (centroidBuf_ && metric_type == METRIC_L2 && !fp16) { + size_t normBytes = nCentroids * sizeof(float); + centroidNormsBuf_ = + [device newBufferWithLength:normBytes + options:MTLResourceStorageModeShared]; + if (centroidNormsBuf_) { + id queue = resources_->getCommandQueue(); + const int nCentroidsI = checkedSizeToInt( + nCentroids, + "MetalIndexIVFFlat: centroid count exceeds int range"); + if (!runMetalComputeNorms( + device, + queue, + centroidBuf_, + nCentroidsI, + d, + centroidNormsBuf_, + false)) { + centroidNormsBuf_ = nil; + } + } + } else { + centroidNormsBuf_ = nil; + } +} + +void MetalIndexIVFFlat::ensureSearchBuf_( + id& buf, + size_t& cap, + size_t needed) const { + if (buf != nil && cap >= needed) { + return; // already large enough + } + // Grow by 2× to amortise future reallocations. + size_t newCap = std::max(needed, cap * 2); + id device = resources_->getDevice(); + buf = [device newBufferWithLength:newCap + options:MTLResourceStorageModeShared]; + cap = buf ? newCap : 0; +} + +void MetalIndexIVFFlat::train(idx_t n, const float* x) { + FAISS_THROW_IF_NOT(cpuIndex_); + cpuIndex_->train(n, x); + is_trained = cpuIndex_->is_trained; + if (is_trained) { + uploadCentroids_(); + } +} + +void MetalIndexIVFFlat::add(idx_t n, const float* x) { + FAISS_THROW_IF_NOT(cpuIndex_); + if (n == 0) { + return; + } + std::vector list_nos(n); + cpuIndex_->quantizer->assign(n, x, list_nos.data()); + + idx_t oldNt = cpuIndex_->ntotal; + const idx_t autoReserveMinBatch = getIvfAutoReserveMinBatch(); + if (gpuIvf_ && autoReserveMinBatch > 0 && n >= autoReserveMinBatch) { + gpuIvf_->reserveMemory(oldNt + n); + } + cpuIndex_->add_core(n, x, nullptr, list_nos.data()); + ntotal = cpuIndex_->ntotal; + + if (gpuIvf_) { + std::vector ids(n); + for (idx_t i = 0; i < n; ++i) { + ids[i] = oldNt + i; + } + gpuIvf_->appendVectors(n, x, list_nos.data(), ids.data()); + } +} + +void MetalIndexIVFFlat::add_with_ids( + idx_t n, + const float* x, + const idx_t* xids) { + FAISS_THROW_IF_NOT(cpuIndex_); + if (n == 0) { + return; + } + FAISS_THROW_IF_NOT(xids != nullptr); + + idx_t oldNt = cpuIndex_->ntotal; + const idx_t autoReserveMinBatch = getIvfAutoReserveMinBatch(); + if (gpuIvf_ && autoReserveMinBatch > 0 && n >= autoReserveMinBatch) { + // Pre-reserve before append to reduce relayout spikes on large batches. + gpuIvf_->reserveMemory(oldNt + n); + } + + std::vector list_nos(n); + cpuIndex_->quantizer->assign(n, x, list_nos.data()); + + cpuIndex_->add_core(n, x, xids, list_nos.data()); + ntotal = cpuIndex_->ntotal; + + if (gpuIvf_) { + gpuIvf_->appendVectors(n, x, list_nos.data(), xids); + } +} + +void MetalIndexIVFFlat::reset() { + FAISS_THROW_IF_NOT(cpuIndex_); + cpuIndex_->reset(); + ntotal = 0; + if (gpuIvf_) { + gpuIvf_->reset(); + } +} + +void MetalIndexIVFFlat::search( + idx_t n, + const float* x, + idx_t k, + float* distances, + idx_t* labels, + const SearchParameters* params) const { + FAISS_THROW_IF_NOT(cpuIndex_); + FAISS_THROW_IF_NOT(k > 0); + FAISS_THROW_IF_NOT_MSG( + k <= kIVFFlatSupportedMaxK, + "MetalIndexIVFFlat supports k <= 1024; larger k is not yet supported"); + + const float inf = std::numeric_limits::infinity(); + const float negInf = -std::numeric_limits::infinity(); + + // Empty index: mirror CPU IndexIVF behavior (labels = -1). + if (cpuIndex_->ntotal == 0 || n == 0) { + for (idx_t i = 0; i < n * k; ++i) { + labels[i] = -1; + distances[i] = (metric_type == METRIC_L2) ? inf : negInf; + } + return; + } + + // Only L2 and IP are supported. + FAISS_THROW_IF_NOT( + metric_type == METRIC_L2 || metric_type == METRIC_INNER_PRODUCT); + const bool isL2 = (metric_type == METRIC_L2); + + // Determine nprobe from params or index. + size_t nprobe = cpuIndex_->nprobe; + if (auto* ivfParams = dynamic_cast(params)) { + if (ivfParams->nprobe > 0) { + nprobe = ivfParams->nprobe; + } + } + nprobe = std::min(nprobe, cpuIndex_->nlist); + if (nprobe == 0) { + for (idx_t i = 0; i < n * k; ++i) { + labels[i] = -1; + distances[i] = isL2 ? inf : negInf; + } + return; + } + + auto cpuFallbackSearch = [&](idx_t qBase, idx_t qCount) { + const float* xTile = x + (size_t)qBase * (size_t)d; + float* distancesTile = distances + (size_t)qBase * (size_t)k; + idx_t* labelsTile = labels + (size_t)qBase * (size_t)k; + if (indicesOptions_ == faiss::gpu::INDICES_IVF) { + std::vector coarseDist((size_t)qCount * nprobe); + std::vector coarseAssign((size_t)qCount * nprobe); + cpuIndex_->quantizer->search( + qCount, + xTile, + (idx_t)nprobe, + coarseDist.data(), + coarseAssign.data()); + cpuIndex_->search_preassigned( + qCount, + xTile, + k, + coarseAssign.data(), + coarseDist.data(), + distancesTile, + labelsTile, + true, + dynamic_cast(params), + nullptr); + } else { + cpuIndex_->search(qCount, xTile, k, distancesTile, labelsTile); + } + }; + const bool allowCpuFallback = allowCpuFallbackForIvf(); + const bool logCpuFallback = logCpuFallbackForIvf(); + const bool logSyncProfile = logSyncProfileForIvf(); + const bool forceChunkedSelection = forceChunkedIvfSelectionPath(); + const size_t exactCandidateBudget = getIvfExactCandidateBudget(); + idx_t fallbackCount = 0; + std::string firstFallbackReason; + idx_t fullSearchCalls = 0; + idx_t scanCalls = 0; + idx_t syncScanCalls = 0; + idx_t asyncScanCalls = 0; + idx_t asyncBatchSyncs = 0; + auto fallbackOrThrow = [&](idx_t qBase, idx_t qCount, const char* reason) { + if (!allowCpuFallback) { + FAISS_THROW_FMT( + "MetalIndexIVFFlat::search requires GPU execution but hit fallback (%s) " + "at qBase=%lld qCount=%lld", + reason, + (long long)qBase, + (long long)qCount); + } + ++fallbackCount; + if (firstFallbackReason.empty()) { + firstFallbackReason = reason ? reason : "unknown"; + } + cpuFallbackSearch(qBase, qCount); + }; + + id device = resources_->getDevice(); + id queue = resources_->getCommandQueue(); + if (gpuIvf_) { + gpuIvf_->ensureInterleavedLayoutUpToDate(); + } + + const bool hasFlatCodes = gpuIvf_ && gpuIvf_->codesBuffer(); + const bool hasInterleavedCodes = gpuIvf_ && + gpuIvf_->interleavedCodesBuffer() && + gpuIvf_->interleavedCodesOffsetBuffer(); + const bool hasScanCodes = hasFlatCodes || hasInterleavedCodes; + + // Fall back to CPU if Metal is not available or GPU IVF storage not ready. + if (!device || !queue || !gpuIvf_ || !hasScanCodes || + !gpuIvf_->idsBuffer() || !gpuIvf_->listOffsetGpuBuffer() || + !gpuIvf_->listLengthGpuBuffer()) { + fallbackOrThrow(0, n, "missing Metal device/queue or IVF buffers"); + return; + } + + // ---- Upload centroid buffer if needed ----------------------------------- + if (!centroidBuf_) { + uploadCentroids_(); + } + + int nlist = (int)cpuIndex_->nlist; + + const size_t tileRows = + chooseIvfSearchTileRows((size_t)n, d, k, nprobe, nlist); + const int avgListLen = + (ntotal > 0 && nlist > 0) ? (int)(ntotal / nlist) : 256; + + for (idx_t qBase = 0; qBase < n; qBase += (idx_t)tileRows) { + idx_t qCount = std::min((idx_t)tileRows, n - qBase); + const float* xTile = x + (size_t)qBase * (size_t)d; + + size_t queriesBytes = (size_t)qCount * (size_t)d * sizeof(float); + size_t outDistBytes = (size_t)qCount * (size_t)k * sizeof(float); + size_t outIdxBytes = (size_t)qCount * (size_t)k * sizeof(int64_t); + size_t perListBytes = + (size_t)qCount * nprobe * (size_t)k * sizeof(float); + size_t perListIdxB = + (size_t)qCount * nprobe * (size_t)k * sizeof(int64_t); + size_t coarseDistBytes = (size_t)qCount * nprobe * sizeof(float); + size_t coarseIdxBytes = (size_t)qCount * nprobe * sizeof(int32_t); + size_t distMatBytes = (size_t)qCount * (size_t)nlist * sizeof(float); + + ensureSearchBuf_(searchQueriesBuf_, searchQueriesCap_, queriesBytes); + ensureSearchBuf_(searchOutDistBuf_, searchOutDistCap_, outDistBytes); + ensureSearchBuf_(searchOutIdxBuf_, searchOutIdxCap_, outIdxBytes); + const bool wantsFullCoarse = useFullCoarseGpuForIvf(); + const size_t fullCoarseMaxBytes = getIvfFullCoarseMaxBytes(); + bool tryFullCoarse = wantsFullCoarse && centroidBuf_ && + nprobe <= (size_t)getMetalDistanceMaxK() && + distMatBytes <= fullCoarseMaxBytes; + if (forceChunkedSelection) { + // Experimental path for selection-scalability work: skip monolithic + // full-search path and route through coarse+chunked scan/merge. + tryFullCoarse = false; + } + + ensureSearchBuf_( + searchPerListDistBuf_, searchPerListDistCap_, perListBytes); + ensureSearchBuf_( + searchPerListIdxBuf_, searchPerListIdxCap_, perListIdxB); + if (tryFullCoarse) { + ensureSearchBuf_( + coarseOutDistBuf_, coarseOutDistCap_, coarseDistBytes); + ensureSearchBuf_( + coarseOutIdxBuf_, coarseOutIdxCap_, coarseIdxBytes); + ensureSearchBuf_(distMatrixBuf_, distMatrixCap_, distMatBytes); + if (!coarseOutDistBuf_ || !coarseOutIdxBuf_ || !distMatrixBuf_) { + // Fall back to CPU coarse assignment path if matrix scratch + // cannot be allocated for this tile. + tryFullCoarse = false; + } + } + + if (!searchQueriesBuf_ || !searchOutDistBuf_ || !searchOutIdxBuf_ || + !searchPerListDistBuf_ || !searchPerListIdxBuf_) { + fallbackOrThrow( + qBase, qCount, "failed to allocate tiled search buffers"); + continue; + } + + std::memcpy([searchQueriesBuf_ contents], xTile, queriesBytes); + const int qCountI = checkedIdxToInt( + qCount, "MetalIndexIVFFlat: qCount exceeds int range"); + const int kI = + checkedIdxToInt(k, "MetalIndexIVFFlat: k exceeds int range"); + const int nprobeI = checkedSizeToInt( + nprobe, "MetalIndexIVFFlat: nprobe exceeds int range"); + + bool ok = false; + if (tryFullCoarse) { + ++fullSearchCalls; + ok = runMetalIVFFlatFullSearch( + device, + queue, + searchQueriesBuf_, + qCountI, + d, + kI, + nprobeI, + isL2, + centroidBuf_, + nlist, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + gpuIvf_->listOffsetGpuBuffer(), + gpuIvf_->listLengthGpuBuffer(), + searchOutDistBuf_, + searchOutIdxBuf_, + searchPerListDistBuf_, + searchPerListIdxBuf_, + coarseOutDistBuf_, + coarseOutIdxBuf_, + distMatrixBuf_, + centroidNormsBuf_, + avgListLen, + gpuIvf_->interleavedCodesBuffer(), + gpuIvf_->interleavedCodesOffsetBuffer(), + config_.useFloat16CoarseQuantizer); + } + + if (!ok) { + std::vector coarseDistVec((size_t)qCount * nprobe); + std::vector coarseAssignVec((size_t)qCount * nprobe); + cpuIndex_->quantizer->search( + qCount, + xTile, + (idx_t)nprobe, + coarseDistVec.data(), + coarseAssignVec.data()); + size_t coarseBytes = (size_t)qCount * nprobe * sizeof(int32_t); + ensureSearchBuf_(searchCoarseBuf_, searchCoarseCap_, coarseBytes); + if (!searchCoarseBuf_) { + fallbackOrThrow( + qBase, + qCount, + "failed to allocate coarse assign buffer"); + continue; + } + auto* dst = reinterpret_cast([searchCoarseBuf_ contents]); + for (size_t i = 0; i < (size_t)qCount * nprobe; ++i) { + FAISS_THROW_IF_NOT_MSG( + coarseAssignVec[i] >= (idx_t)std::numeric_limits< + int32_t>::min() && + coarseAssignVec[i] <= + (idx_t)std::numeric_limits< + int32_t>::max(), + "MetalIndexIVFFlat: coarse assignment exceeds int32 range"); + dst[i] = (int32_t)coarseAssignVec[i]; + } + const bool useInterleaved = + gpuIvf_->interleavedCodesBuffer() != nil && + gpuIvf_->interleavedCodesOffsetBuffer() != nil; + const std::string envelopeReason = explainIvfScanFailureEnvelope( + d, + k, + exactCandidateBudget, + nprobe, + useInterleaved, + cpuIndex_.get(), + coarseAssignVec.data(), + (size_t)qCount * nprobe); + bool usedChunkedScan = false; + std::vector chunkMergedDist; + std::vector chunkMergedIdx; + + std::vector usedLists; + size_t maxSelectedListSize = 0; + idx_t badListId = -1; + bool haveListStats = selectedListStats( + cpuIndex_.get(), + coarseAssignVec.data(), + (size_t)qCount * nprobe, + usedLists, + maxSelectedListSize, + badListId); + const bool needProbeChunk = + nprobe > 0 && (size_t)k > (exactCandidateBudget / nprobe); + const bool needListChunk = + haveListStats && maxSelectedListSize > exactCandidateBudget; + + if (needProbeChunk || needListChunk) { + usedChunkedScan = true; + chunkMergedDist.resize( + (size_t)qCount * (size_t)k, isL2 ? inf : negInf); + chunkMergedIdx.resize((size_t)qCount * (size_t)k, -1); + ok = true; + + const size_t maxProbePerChunk = needProbeChunk + ? std::max(1, exactCandidateBudget / (size_t)k) + : nprobe; + const size_t listChunk = exactCandidateBudget; + const size_t numListPass = needListChunk + ? ((maxSelectedListSize + listChunk - 1) / listChunk) + : 1; + // List-chunk remapping only updates list offset/length + // metadata. Interleaved offset metadata is list-base-relative, + // so force flat code path for list-chunked passes to keep + // ids/codes aligned. + id chunkInterleavedCodes = + needListChunk ? nil : gpuIvf_->interleavedCodesBuffer(); + id chunkInterleavedOffsets = needListChunk + ? nil + : gpuIvf_->interleavedCodesOffsetBuffer(); + + const uint32_t* baseListOffset = + reinterpret_cast( + [gpuIvf_->listOffsetGpuBuffer() contents]); + const uint32_t* baseListLength = + reinterpret_cast( + [gpuIvf_->listLengthGpuBuffer() contents]); + id listOffsetPassBuf = nil; + id listLengthPassBuf = nil; + std::vector passOffset; + std::vector passLength; + if (needListChunk) { + size_t metaBytes = (size_t)nlist * sizeof(uint32_t); + listOffsetPassBuf = [device + newBufferWithLength:metaBytes + options:MTLResourceStorageModeShared]; + listLengthPassBuf = [device + newBufferWithLength:metaBytes + options:MTLResourceStorageModeShared]; + if (!listOffsetPassBuf || !listLengthPassBuf) { + ok = false; + } else { + passOffset.resize((size_t)nlist); + passLength.resize((size_t)nlist); + } + } + + const bool canAsyncProbeBatch = + needProbeChunk && !needListChunk && numListPass == 1; + struct ProbePass { + size_t p0 = 0; + size_t chunkProbe = 0; + id coarseBuf = nil; + id outDistBuf = nil; + id outIdxBuf = nil; + }; + std::vector probePasses; + std::vector coarseChunk; + for (size_t p0 = 0; ok && p0 < nprobe; p0 += maxProbePerChunk) { + size_t chunkProbe = std::min(maxProbePerChunk, nprobe - p0); + coarseChunk.resize((size_t)qCount * chunkProbe); + for (idx_t qi = 0; qi < qCount; ++qi) { + const size_t srcBase = (size_t)qi * nprobe + p0; + const size_t dstBase = (size_t)qi * chunkProbe; + std::memcpy( + coarseChunk.data() + dstBase, + dst + srcBase, + chunkProbe * sizeof(int32_t)); + } + + size_t coarseChunkBytes = + coarseChunk.size() * sizeof(int32_t); + id coarseChunkBuf = searchCoarseBuf_; + if (canAsyncProbeBatch) { + coarseChunkBuf = [device + newBufferWithLength:coarseChunkBytes + options:MTLResourceStorageModeShared]; + if (!coarseChunkBuf) { + ok = false; + break; + } + } else { + ensureSearchBuf_( + searchCoarseBuf_, + searchCoarseCap_, + coarseChunkBytes); + coarseChunkBuf = searchCoarseBuf_; + if (!coarseChunkBuf) { + ok = false; + break; + } + } + std::memcpy( + [coarseChunkBuf contents], + coarseChunk.data(), + coarseChunkBytes); + + size_t perListBytesChunk = (size_t)qCount * chunkProbe * + (size_t)k * sizeof(float); + size_t perListIdxBytesChunk = (size_t)qCount * chunkProbe * + (size_t)k * sizeof(int64_t); + ensureSearchBuf_( + searchPerListDistBuf_, + searchPerListDistCap_, + perListBytesChunk); + ensureSearchBuf_( + searchPerListIdxBuf_, + searchPerListIdxCap_, + perListIdxBytesChunk); + if (!searchPerListDistBuf_ || !searchPerListIdxBuf_) { + ok = false; + break; + } + + const int chunkProbeI = checkedSizeToInt( + chunkProbe, + "MetalIndexIVFFlat: chunk nprobe exceeds int range"); + if (canAsyncProbeBatch) { + id passOutDistBuf = [device + newBufferWithLength:outDistBytes + options:MTLResourceStorageModeShared]; + id passOutIdxBuf = [device + newBufferWithLength:outIdxBytes + options:MTLResourceStorageModeShared]; + if (!passOutDistBuf || !passOutIdxBuf) { + ok = false; + break; + } + + bool chunkOk = runMetalIVFFlatScan( + device, + queue, + searchQueriesBuf_, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + gpuIvf_->listOffsetGpuBuffer(), + gpuIvf_->listLengthGpuBuffer(), + coarseChunkBuf, + qCountI, + d, + kI, + chunkProbeI, + isL2, + passOutDistBuf, + passOutIdxBuf, + searchPerListDistBuf_, + searchPerListIdxBuf_, + chunkInterleavedCodes, + chunkInterleavedOffsets, + false /* waitForCompletion */); + ++scanCalls; + ++asyncScanCalls; + if (!chunkOk) { + ok = false; + break; + } + probePasses.push_back( + ProbePass{ + p0, + chunkProbe, + coarseChunkBuf, + passOutDistBuf, + passOutIdxBuf}); + continue; + } + + if (needListChunk && numListPass > 1) { + std::vector> passOutDistBufs( + numListPass, nil); + std::vector> passOutIdxBufs( + numListPass, nil); + std::vector> passListOffsetBufs( + numListPass, nil); + std::vector> passListLengthBufs( + numListPass, nil); + + for (size_t lp = 0; ok && lp < numListPass; ++lp) { + const uint64_t shift = + (uint64_t)lp * (uint64_t)listChunk; + for (int li = 0; li < nlist; ++li) { + uint32_t off = baseListOffset[li]; + uint32_t len = baseListLength[li]; + if (!usedLists[(size_t)li]) { + passOffset[(size_t)li] = off; + passLength[(size_t)li] = 0; + continue; + } + if (shift >= len) { + passOffset[(size_t)li] = off + len; + passLength[(size_t)li] = 0; + continue; + } + const uint32_t delta = (uint32_t)shift; + const uint32_t rem = len - delta; + passOffset[(size_t)li] = off + delta; + passLength[(size_t)li] = std::min( + (uint32_t)listChunk, rem); + } + + passListOffsetBufs[lp] = [device + newBufferWithLength:(size_t)nlist * + sizeof(uint32_t) + options:MTLResourceStorageModeShared]; + passListLengthBufs[lp] = [device + newBufferWithLength:(size_t)nlist * + sizeof(uint32_t) + options:MTLResourceStorageModeShared]; + passOutDistBufs[lp] = [device + newBufferWithLength:outDistBytes + options:MTLResourceStorageModeShared]; + passOutIdxBufs[lp] = [device + newBufferWithLength:outIdxBytes + options:MTLResourceStorageModeShared]; + if (!passListOffsetBufs[lp] || + !passListLengthBufs[lp] || + !passOutDistBufs[lp] || !passOutIdxBufs[lp]) { + ok = false; + break; + } + std::memcpy( + [passListOffsetBufs[lp] contents], + passOffset.data(), + passOffset.size() * sizeof(uint32_t)); + std::memcpy( + [passListLengthBufs[lp] contents], + passLength.data(), + passLength.size() * sizeof(uint32_t)); + + bool chunkOk = runMetalIVFFlatScan( + device, + queue, + searchQueriesBuf_, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + passListOffsetBufs[lp], + passListLengthBufs[lp], + searchCoarseBuf_, + qCountI, + d, + kI, + chunkProbeI, + isL2, + passOutDistBufs[lp], + passOutIdxBufs[lp], + searchPerListDistBuf_, + searchPerListIdxBuf_, + chunkInterleavedCodes, + chunkInterleavedOffsets, + false /* waitForCompletion */); + ++scanCalls; + ++asyncScanCalls; + if (!chunkOk) { + ok = false; + break; + } + } + + if (ok) { + resources_->synchronize(); + ++asyncBatchSyncs; + for (size_t lp = 0; lp < numListPass; ++lp) { + const float* chunkDist = + reinterpret_cast( + [passOutDistBufs[lp] contents]); + const int64_t* chunkIdx = + reinterpret_cast( + [passOutIdxBufs[lp] contents]); + for (idx_t qi = 0; qi < qCount; ++qi) { + float* bestDist = chunkMergedDist.data() + + (size_t)qi * (size_t)k; + int64_t* bestIdx = chunkMergedIdx.data() + + (size_t)qi * (size_t)k; + for (idx_t j = 0; j < k; ++j) { + const size_t pos = + (size_t)qi * (size_t)k + + (size_t)j; + insertTopKCandidate( + chunkDist[pos], + chunkIdx[pos], + isL2, + bestDist, + bestIdx, + kI); + } + } + } + } + } else { + for (size_t lp = 0; ok && lp < numListPass; ++lp) { + id passListOffsetBuf = + gpuIvf_->listOffsetGpuBuffer(); + id passListLengthBuf = + gpuIvf_->listLengthGpuBuffer(); + if (needListChunk) { + const uint64_t shift = + (uint64_t)lp * (uint64_t)listChunk; + for (int li = 0; li < nlist; ++li) { + uint32_t off = baseListOffset[li]; + uint32_t len = baseListLength[li]; + if (!usedLists[(size_t)li]) { + passOffset[(size_t)li] = off; + passLength[(size_t)li] = 0; + continue; + } + if (shift >= len) { + passOffset[(size_t)li] = off + len; + passLength[(size_t)li] = 0; + continue; + } + const uint32_t delta = (uint32_t)shift; + const uint32_t rem = len - delta; + passOffset[(size_t)li] = off + delta; + passLength[(size_t)li] = std::min( + (uint32_t)listChunk, rem); + } + std::memcpy( + [listOffsetPassBuf contents], + passOffset.data(), + passOffset.size() * sizeof(uint32_t)); + std::memcpy( + [listLengthPassBuf contents], + passLength.data(), + passLength.size() * sizeof(uint32_t)); + passListOffsetBuf = listOffsetPassBuf; + passListLengthBuf = listLengthPassBuf; + } + + bool chunkOk = runMetalIVFFlatScan( + device, + queue, + searchQueriesBuf_, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + passListOffsetBuf, + passListLengthBuf, + searchCoarseBuf_, + qCountI, + d, + kI, + chunkProbeI, + isL2, + searchOutDistBuf_, + searchOutIdxBuf_, + searchPerListDistBuf_, + searchPerListIdxBuf_, + chunkInterleavedCodes, + chunkInterleavedOffsets); + ++scanCalls; + ++syncScanCalls; + if (!chunkOk) { + ok = false; + break; + } + + const float* chunkDist = + reinterpret_cast( + [searchOutDistBuf_ contents]); + const int64_t* chunkIdx = + reinterpret_cast( + [searchOutIdxBuf_ contents]); + for (idx_t qi = 0; qi < qCount; ++qi) { + float* bestDist = chunkMergedDist.data() + + (size_t)qi * (size_t)k; + int64_t* bestIdx = chunkMergedIdx.data() + + (size_t)qi * (size_t)k; + for (idx_t j = 0; j < k; ++j) { + const size_t pos = + (size_t)qi * (size_t)k + (size_t)j; + insertTopKCandidate( + chunkDist[pos], + chunkIdx[pos], + isL2, + bestDist, + bestIdx, + kI); + } + } + } + } + } + if (ok && canAsyncProbeBatch && !probePasses.empty()) { + resources_->synchronize(); + ++asyncBatchSyncs; + for (const auto& pass : probePasses) { + (void)pass.p0; + (void)pass.chunkProbe; + (void)pass.coarseBuf; + const float* chunkDist = reinterpret_cast( + [pass.outDistBuf contents]); + const int64_t* chunkIdx = + reinterpret_cast( + [pass.outIdxBuf contents]); + for (idx_t qi = 0; qi < qCount; ++qi) { + float* bestDist = chunkMergedDist.data() + + (size_t)qi * (size_t)k; + int64_t* bestIdx = chunkMergedIdx.data() + + (size_t)qi * (size_t)k; + for (idx_t j = 0; j < k; ++j) { + const size_t pos = + (size_t)qi * (size_t)k + (size_t)j; + insertTopKCandidate( + chunkDist[pos], + chunkIdx[pos], + isL2, + bestDist, + bestIdx, + kI); + } + } + } + } + } else { + ok = runMetalIVFFlatScan( + device, + queue, + searchQueriesBuf_, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + gpuIvf_->listOffsetGpuBuffer(), + gpuIvf_->listLengthGpuBuffer(), + searchCoarseBuf_, + qCountI, + d, + kI, + nprobeI, + isL2, + searchOutDistBuf_, + searchOutIdxBuf_, + searchPerListDistBuf_, + searchPerListIdxBuf_, + gpuIvf_->interleavedCodesBuffer(), + gpuIvf_->interleavedCodesOffsetBuffer()); + ++scanCalls; + ++syncScanCalls; + } + if (!ok && !envelopeReason.empty() && !usedChunkedScan) { + fallbackOrThrow(qBase, qCount, envelopeReason.c_str()); + continue; + } + if (ok && usedChunkedScan) { + for (idx_t qi = 0; qi < qCount; ++qi) { + heapSort( + chunkMergedDist.data() + (size_t)qi * (size_t)k, + chunkMergedIdx.data() + (size_t)qi * (size_t)k, + kI, + isL2); + } + for (idx_t qi = 0; qi < qCount; ++qi) { + for (idx_t j = 0; j < k; ++j) { + const size_t pos = (size_t)qi * (size_t)k + (size_t)j; + const size_t globalPos = + (size_t)(qBase + qi) * (size_t)k + (size_t)j; + const int64_t globalId = chunkMergedIdx[pos]; + if (globalId < 0) { + labels[globalPos] = -1; + } else if (indicesOptions_ == faiss::gpu::INDICES_CPU) { + labels[globalPos] = decodeCpuLabelFromPair( + cpuIndex_.get(), globalId); + } else if ( + indicesOptions_ == faiss::gpu::INDICES_32_BIT) { + labels[globalPos] = (idx_t)(int32_t)globalId; + } else { + labels[globalPos] = (idx_t)globalId; + } + distances[globalPos] = chunkMergedDist[pos]; + } + } + continue; + } + } + + if (!ok) { + fallbackOrThrow( + qBase, qCount, "GPU IVF scan failed (runtime/kernel)"); + continue; + } + + const float* outDistPtr = + reinterpret_cast([searchOutDistBuf_ contents]); + const int64_t* outIdxPtr = + reinterpret_cast([searchOutIdxBuf_ contents]); + + for (idx_t qi = 0; qi < qCount; ++qi) { + for (idx_t j = 0; j < k; ++j) { + size_t localPos = (size_t)qi * (size_t)k + (size_t)j; + size_t globalPos = (size_t)(qBase + qi) * (size_t)k + (size_t)j; + int64_t globalId = outIdxPtr[localPos]; + if (globalId < 0) { + labels[globalPos] = -1; + } else if (indicesOptions_ == faiss::gpu::INDICES_CPU) { + labels[globalPos] = + decodeCpuLabelFromPair(cpuIndex_.get(), globalId); + } else if (indicesOptions_ == faiss::gpu::INDICES_32_BIT) { + labels[globalPos] = (idx_t)(int32_t)globalId; + } else { + labels[globalPos] = (idx_t)globalId; + } + distances[globalPos] = outDistPtr[localPos]; + } + } + } + + if (allowCpuFallback && logCpuFallback && fallbackCount > 0) { + std::fprintf( + stderr, + "IVF_CPU_FALLBACK,api=search,count=%lld,first_reason=%s\n", + (long long)fallbackCount, + firstFallbackReason.c_str()); + } + if (logSyncProfile) { + const idx_t estimatedWaits = + fullSearchCalls + syncScanCalls + asyncBatchSyncs; + std::fprintf( + stderr, + "IVF_SYNC_PROFILE,api=search,full_calls=%lld,scan_calls=%lld,sync_scan_calls=%lld,async_scan_calls=%lld,async_batch_syncs=%lld,estimated_waits=%lld\n", + (long long)fullSearchCalls, + (long long)scanCalls, + (long long)syncScanCalls, + (long long)asyncScanCalls, + (long long)asyncBatchSyncs, + (long long)estimatedWaits); + } +} + +void MetalIndexIVFFlat::search_preassigned( + idx_t n, + const float* x, + idx_t k, + const idx_t* assign, + const float* centroid_dis, + float* distances, + idx_t* labels, + bool store_pairs, + const IVFSearchParameters* params, + IndexIVFStats* stats) const { + FAISS_THROW_IF_NOT(cpuIndex_); + FAISS_THROW_IF_NOT(k > 0); + FAISS_THROW_IF_NOT(assign); + FAISS_THROW_IF_NOT_MSG(stats == nullptr, "IVF stats not supported"); + FAISS_THROW_IF_NOT_MSG( + !store_pairs, + "MetalIndexIVFFlat::search_preassigned does not currently support store_pairs"); + if (params) { + FAISS_THROW_IF_NOT_FMT( + params->max_codes == 0, + "Metal IVF index does not currently support " + "SearchParametersIVF::max_codes (passed %zu, must be 0)", + params->max_codes); + } + FAISS_THROW_IF_NOT_MSG( + k <= kIVFFlatSupportedMaxK, + "MetalIndexIVFFlat supports k <= 1024; larger k is not yet supported"); + + const float inf = std::numeric_limits::infinity(); + const float negInf = -std::numeric_limits::infinity(); + const bool isL2 = (metric_type == METRIC_L2); + + if (cpuIndex_->ntotal == 0 || n == 0) { + for (idx_t i = 0; i < n * k; ++i) { + labels[i] = -1; + distances[i] = isL2 ? inf : negInf; + } + return; + } + + size_t nprobe = cpuIndex_->nprobe; + if (params && params->nprobe > 0) { + nprobe = params->nprobe; + } + nprobe = std::min(nprobe, cpuIndex_->nlist); + + id device = resources_->getDevice(); + id queue = resources_->getCommandQueue(); + if (gpuIvf_) { + gpuIvf_->ensureInterleavedLayoutUpToDate(); + } + + const bool hasFlatCodes = gpuIvf_ && gpuIvf_->codesBuffer(); + const bool hasInterleavedCodes = gpuIvf_ && + gpuIvf_->interleavedCodesBuffer() && + gpuIvf_->interleavedCodesOffsetBuffer(); + const bool hasScanCodes = hasFlatCodes || hasInterleavedCodes; + + auto cpuFallbackSearch = [&](idx_t qBase, idx_t qCount) { + const float* xTile = x + (size_t)qBase * (size_t)d; + const idx_t* assignTile = assign + (size_t)qBase * nprobe; + const float* centroidTile = centroid_dis + ? (centroid_dis + (size_t)qBase * nprobe) + : nullptr; + float* distTile = distances + (size_t)qBase * (size_t)k; + idx_t* labelsTile = labels + (size_t)qBase * (size_t)k; + cpuIndex_->search_preassigned( + qCount, + xTile, + k, + assignTile, + centroidTile, + distTile, + labelsTile, + indicesOptions_ == faiss::gpu::INDICES_IVF ? true : store_pairs, + params, + stats); + }; + const bool allowCpuFallback = allowCpuFallbackForIvf(); + const bool logCpuFallback = logCpuFallbackForIvf(); + const bool logSyncProfile = logSyncProfileForIvf(); + const size_t exactCandidateBudget = getIvfExactCandidateBudget(); + idx_t fallbackCount = 0; + std::string firstFallbackReason; + idx_t scanCalls = 0; + idx_t syncScanCalls = 0; + idx_t asyncScanCalls = 0; + idx_t asyncBatchSyncs = 0; + auto fallbackOrThrow = [&](idx_t qBase, idx_t qCount, const char* reason) { + if (!allowCpuFallback) { + FAISS_THROW_FMT( + "MetalIndexIVFFlat::search_preassigned requires GPU execution but hit " + "fallback (%s) at qBase=%lld qCount=%lld", + reason, + (long long)qBase, + (long long)qCount); + } + ++fallbackCount; + if (firstFallbackReason.empty()) { + firstFallbackReason = reason ? reason : "unknown"; + } + cpuFallbackSearch(qBase, qCount); + }; + + if (!device || !queue || !gpuIvf_ || !hasScanCodes || + !gpuIvf_->idsBuffer() || !gpuIvf_->listOffsetGpuBuffer() || + !gpuIvf_->listLengthGpuBuffer()) { + fallbackOrThrow(0, n, "missing Metal device/queue or IVF buffers"); + return; + } + + const int nlist = (int)cpuIndex_->nlist; + const size_t tileRows = + chooseIvfPreassignedTileRows((size_t)n, d, k, nprobe); + for (idx_t qBase = 0; qBase < n; qBase += (idx_t)tileRows) { + idx_t qCount = std::min((idx_t)tileRows, n - qBase); + const float* xTile = x + (size_t)qBase * (size_t)d; + const idx_t* assignTile = assign + (size_t)qBase * nprobe; + + size_t queriesBytes = (size_t)qCount * (size_t)d * sizeof(float); + size_t outDistBytes = (size_t)qCount * (size_t)k * sizeof(float); + size_t outIdxBytes = (size_t)qCount * (size_t)k * sizeof(int64_t); + size_t perListBytes = + (size_t)qCount * nprobe * (size_t)k * sizeof(float); + size_t perListIdxB = + (size_t)qCount * nprobe * (size_t)k * sizeof(int64_t); + size_t coarseBytes = (size_t)qCount * nprobe * sizeof(int32_t); + + ensureSearchBuf_(searchQueriesBuf_, searchQueriesCap_, queriesBytes); + ensureSearchBuf_(searchOutDistBuf_, searchOutDistCap_, outDistBytes); + ensureSearchBuf_(searchOutIdxBuf_, searchOutIdxCap_, outIdxBytes); + ensureSearchBuf_( + searchPerListDistBuf_, searchPerListDistCap_, perListBytes); + ensureSearchBuf_( + searchPerListIdxBuf_, searchPerListIdxCap_, perListIdxB); + ensureSearchBuf_(searchCoarseBuf_, searchCoarseCap_, coarseBytes); + + if (!searchQueriesBuf_ || !searchOutDistBuf_ || !searchOutIdxBuf_ || + !searchPerListDistBuf_ || !searchPerListIdxBuf_ || + !searchCoarseBuf_) { + fallbackOrThrow( + qBase, + qCount, + "failed to allocate tiled preassigned buffers"); + continue; + } + + std::memcpy([searchQueriesBuf_ contents], xTile, queriesBytes); + const int qCountI = checkedIdxToInt( + qCount, + "MetalIndexIVFFlat::search_preassigned: qCount exceeds int range"); + const int kI = checkedIdxToInt( + k, + "MetalIndexIVFFlat::search_preassigned: k exceeds int range"); + const int nprobeI = checkedSizeToInt( + nprobe, + "MetalIndexIVFFlat::search_preassigned: nprobe exceeds int range"); + + auto* coarseDst = + reinterpret_cast([searchCoarseBuf_ contents]); + for (size_t i = 0; i < (size_t)qCount * nprobe; ++i) { + FAISS_THROW_IF_NOT_MSG( + assignTile[i] >= (idx_t)std::numeric_limits< + int32_t>::min() && + assignTile[i] <= + (idx_t)std::numeric_limits::max(), + "MetalIndexIVFFlat: preassigned list id exceeds int32 range"); + coarseDst[i] = (int32_t)assignTile[i]; + } + + const bool useInterleaved = gpuIvf_->interleavedCodesBuffer() != nil && + gpuIvf_->interleavedCodesOffsetBuffer() != nil; + const std::string envelopeReason = explainIvfScanFailureEnvelope( + d, + k, + exactCandidateBudget, + nprobe, + useInterleaved, + cpuIndex_.get(), + assignTile, + (size_t)qCount * nprobe); + bool ok = false; + bool usedChunkedScan = false; + std::vector chunkMergedDist; + std::vector chunkMergedIdx; + std::vector usedLists; + size_t maxSelectedListSize = 0; + idx_t badListId = -1; + bool haveListStats = selectedListStats( + cpuIndex_.get(), + assignTile, + (size_t)qCount * nprobe, + usedLists, + maxSelectedListSize, + badListId); + const bool needProbeChunk = + nprobe > 0 && (size_t)k > (exactCandidateBudget / nprobe); + const bool needListChunk = + haveListStats && maxSelectedListSize > exactCandidateBudget; + + if (needProbeChunk || needListChunk) { + const size_t maxProbePerChunk = needProbeChunk + ? std::max(1, exactCandidateBudget / (size_t)k) + : nprobe; + usedChunkedScan = true; + chunkMergedDist.resize( + (size_t)qCount * (size_t)k, isL2 ? inf : negInf); + chunkMergedIdx.resize((size_t)qCount * (size_t)k, -1); + ok = true; + + const size_t listChunk = exactCandidateBudget; + const size_t numListPass = needListChunk + ? ((maxSelectedListSize + listChunk - 1) / listChunk) + : 1; + // List-chunk remapping only updates list offset/length metadata. + // Interleaved offset metadata is list-base-relative, so force flat + // code path for list-chunked passes to keep ids/codes aligned. + id chunkInterleavedCodes = + needListChunk ? nil : gpuIvf_->interleavedCodesBuffer(); + id chunkInterleavedOffsets = needListChunk + ? nil + : gpuIvf_->interleavedCodesOffsetBuffer(); + const uint32_t* baseListOffset = reinterpret_cast( + [gpuIvf_->listOffsetGpuBuffer() contents]); + const uint32_t* baseListLength = reinterpret_cast( + [gpuIvf_->listLengthGpuBuffer() contents]); + id listOffsetPassBuf = nil; + id listLengthPassBuf = nil; + std::vector passOffset; + std::vector passLength; + if (needListChunk) { + size_t metaBytes = (size_t)nlist * sizeof(uint32_t); + listOffsetPassBuf = [device + newBufferWithLength:metaBytes + options:MTLResourceStorageModeShared]; + listLengthPassBuf = [device + newBufferWithLength:metaBytes + options:MTLResourceStorageModeShared]; + if (!listOffsetPassBuf || !listLengthPassBuf) { + ok = false; + } else { + passOffset.resize((size_t)nlist); + passLength.resize((size_t)nlist); + } + } + + std::vector fullCoarseVec( + reinterpret_cast( + [searchCoarseBuf_ contents]), + reinterpret_cast( + [searchCoarseBuf_ contents]) + + (size_t)qCount * nprobe); + std::vector coarseChunk; + const int32_t* fullCoarse = fullCoarseVec.data(); + + for (size_t p0 = 0; ok && p0 < nprobe; p0 += maxProbePerChunk) { + size_t chunkProbe = std::min(maxProbePerChunk, nprobe - p0); + coarseChunk.resize((size_t)qCount * chunkProbe); + for (idx_t qi = 0; qi < qCount; ++qi) { + const size_t srcBase = (size_t)qi * nprobe + p0; + const size_t dstBase = (size_t)qi * chunkProbe; + std::memcpy( + coarseChunk.data() + dstBase, + fullCoarse + srcBase, + chunkProbe * sizeof(int32_t)); + } + + size_t coarseChunkBytes = coarseChunk.size() * sizeof(int32_t); + ensureSearchBuf_( + searchCoarseBuf_, searchCoarseCap_, coarseChunkBytes); + if (!searchCoarseBuf_) { + ok = false; + break; + } + std::memcpy( + [searchCoarseBuf_ contents], + coarseChunk.data(), + coarseChunkBytes); + + size_t perListBytesChunk = + (size_t)qCount * chunkProbe * (size_t)k * sizeof(float); + size_t perListIdxBytesChunk = (size_t)qCount * chunkProbe * + (size_t)k * sizeof(int64_t); + ensureSearchBuf_( + searchPerListDistBuf_, + searchPerListDistCap_, + perListBytesChunk); + ensureSearchBuf_( + searchPerListIdxBuf_, + searchPerListIdxCap_, + perListIdxBytesChunk); + if (!searchPerListDistBuf_ || !searchPerListIdxBuf_) { + ok = false; + break; + } + + const int chunkProbeI = checkedSizeToInt( + chunkProbe, + "MetalIndexIVFFlat::search_preassigned: chunk nprobe exceeds int range"); + if (needListChunk && numListPass > 1) { + std::vector> passOutDistBufs( + numListPass, nil); + std::vector> passOutIdxBufs(numListPass, nil); + std::vector> passListOffsetBufs( + numListPass, nil); + std::vector> passListLengthBufs( + numListPass, nil); + + for (size_t lp = 0; ok && lp < numListPass; ++lp) { + const uint64_t shift = + (uint64_t)lp * (uint64_t)listChunk; + for (int li = 0; li < nlist; ++li) { + uint32_t off = baseListOffset[li]; + uint32_t len = baseListLength[li]; + if (!usedLists[(size_t)li]) { + passOffset[(size_t)li] = off; + passLength[(size_t)li] = 0; + continue; + } + if (shift >= len) { + passOffset[(size_t)li] = off + len; + passLength[(size_t)li] = 0; + continue; + } + const uint32_t delta = (uint32_t)shift; + const uint32_t rem = len - delta; + passOffset[(size_t)li] = off + delta; + passLength[(size_t)li] = std::min( + (uint32_t)listChunk, rem); + } + + passListOffsetBufs[lp] = [device + newBufferWithLength:(size_t)nlist * + sizeof(uint32_t) + options:MTLResourceStorageModeShared]; + passListLengthBufs[lp] = [device + newBufferWithLength:(size_t)nlist * + sizeof(uint32_t) + options:MTLResourceStorageModeShared]; + passOutDistBufs[lp] = [device + newBufferWithLength:outDistBytes + options:MTLResourceStorageModeShared]; + passOutIdxBufs[lp] = [device + newBufferWithLength:outIdxBytes + options:MTLResourceStorageModeShared]; + if (!passListOffsetBufs[lp] || + !passListLengthBufs[lp] || !passOutDistBufs[lp] || + !passOutIdxBufs[lp]) { + ok = false; + break; + } + std::memcpy( + [passListOffsetBufs[lp] contents], + passOffset.data(), + passOffset.size() * sizeof(uint32_t)); + std::memcpy( + [passListLengthBufs[lp] contents], + passLength.data(), + passLength.size() * sizeof(uint32_t)); + + bool chunkOk = runMetalIVFFlatScan( + device, + queue, + searchQueriesBuf_, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + passListOffsetBufs[lp], + passListLengthBufs[lp], + searchCoarseBuf_, + qCountI, + d, + kI, + chunkProbeI, + isL2, + passOutDistBufs[lp], + passOutIdxBufs[lp], + searchPerListDistBuf_, + searchPerListIdxBuf_, + chunkInterleavedCodes, + chunkInterleavedOffsets, + false /* waitForCompletion */); + ++scanCalls; + ++asyncScanCalls; + if (!chunkOk) { + ok = false; + break; + } + } + + if (ok) { + resources_->synchronize(); + ++asyncBatchSyncs; + for (size_t lp = 0; lp < numListPass; ++lp) { + const float* chunkDist = + reinterpret_cast( + [passOutDistBufs[lp] contents]); + const int64_t* chunkIdx = + reinterpret_cast( + [passOutIdxBufs[lp] contents]); + for (idx_t qi = 0; qi < qCount; ++qi) { + float* bestDist = chunkMergedDist.data() + + (size_t)qi * (size_t)k; + int64_t* bestIdx = chunkMergedIdx.data() + + (size_t)qi * (size_t)k; + for (idx_t j = 0; j < k; ++j) { + const size_t pos = + (size_t)qi * (size_t)k + (size_t)j; + insertTopKCandidate( + chunkDist[pos], + chunkIdx[pos], + isL2, + bestDist, + bestIdx, + kI); + } + } + } + } + } else { + for (size_t lp = 0; ok && lp < numListPass; ++lp) { + id passListOffsetBuf = + gpuIvf_->listOffsetGpuBuffer(); + id passListLengthBuf = + gpuIvf_->listLengthGpuBuffer(); + if (needListChunk) { + const uint64_t shift = + (uint64_t)lp * (uint64_t)listChunk; + for (int li = 0; li < nlist; ++li) { + uint32_t off = baseListOffset[li]; + uint32_t len = baseListLength[li]; + if (!usedLists[(size_t)li]) { + passOffset[(size_t)li] = off; + passLength[(size_t)li] = 0; + continue; + } + if (shift >= len) { + passOffset[(size_t)li] = off + len; + passLength[(size_t)li] = 0; + continue; + } + const uint32_t delta = (uint32_t)shift; + const uint32_t rem = len - delta; + passOffset[(size_t)li] = off + delta; + passLength[(size_t)li] = std::min( + (uint32_t)listChunk, rem); + } + std::memcpy( + [listOffsetPassBuf contents], + passOffset.data(), + passOffset.size() * sizeof(uint32_t)); + std::memcpy( + [listLengthPassBuf contents], + passLength.data(), + passLength.size() * sizeof(uint32_t)); + passListOffsetBuf = listOffsetPassBuf; + passListLengthBuf = listLengthPassBuf; + } + + bool chunkOk = runMetalIVFFlatScan( + device, + queue, + searchQueriesBuf_, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + passListOffsetBuf, + passListLengthBuf, + searchCoarseBuf_, + qCountI, + d, + kI, + chunkProbeI, + isL2, + searchOutDistBuf_, + searchOutIdxBuf_, + searchPerListDistBuf_, + searchPerListIdxBuf_, + chunkInterleavedCodes, + chunkInterleavedOffsets); + ++scanCalls; + ++syncScanCalls; + if (!chunkOk) { + ok = false; + break; + } + + const float* chunkDist = reinterpret_cast( + [searchOutDistBuf_ contents]); + const int64_t* chunkIdx = + reinterpret_cast( + [searchOutIdxBuf_ contents]); + for (idx_t qi = 0; qi < qCount; ++qi) { + float* bestDist = chunkMergedDist.data() + + (size_t)qi * (size_t)k; + int64_t* bestIdx = chunkMergedIdx.data() + + (size_t)qi * (size_t)k; + for (idx_t j = 0; j < k; ++j) { + const size_t pos = + (size_t)qi * (size_t)k + (size_t)j; + insertTopKCandidate( + chunkDist[pos], + chunkIdx[pos], + isL2, + bestDist, + bestIdx, + kI); + } + } + } + } + } + } else { + ok = runMetalIVFFlatScan( + device, + queue, + searchQueriesBuf_, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + gpuIvf_->listOffsetGpuBuffer(), + gpuIvf_->listLengthGpuBuffer(), + searchCoarseBuf_, + qCountI, + d, + kI, + nprobeI, + isL2, + searchOutDistBuf_, + searchOutIdxBuf_, + searchPerListDistBuf_, + searchPerListIdxBuf_, + gpuIvf_->interleavedCodesBuffer(), + gpuIvf_->interleavedCodesOffsetBuffer()); + ++scanCalls; + ++syncScanCalls; + } + if (!ok && !envelopeReason.empty() && !usedChunkedScan) { + fallbackOrThrow(qBase, qCount, envelopeReason.c_str()); + continue; + } + + if (!ok) { + fallbackOrThrow( + qBase, qCount, "GPU IVF scan failed (runtime/kernel)"); + continue; + } + + if (usedChunkedScan) { + for (idx_t qi = 0; qi < qCount; ++qi) { + heapSort( + chunkMergedDist.data() + (size_t)qi * (size_t)k, + chunkMergedIdx.data() + (size_t)qi * (size_t)k, + kI, + isL2); + } + for (idx_t qi = 0; qi < qCount; ++qi) { + for (idx_t j = 0; j < k; ++j) { + const size_t pos = (size_t)qi * (size_t)k + (size_t)j; + const size_t globalPos = + (size_t)(qBase + qi) * (size_t)k + (size_t)j; + const int64_t globalId = chunkMergedIdx[pos]; + if (globalId < 0) { + labels[globalPos] = -1; + } else if (indicesOptions_ == faiss::gpu::INDICES_CPU) { + labels[globalPos] = decodeCpuLabelFromPair( + cpuIndex_.get(), globalId); + } else if (indicesOptions_ == faiss::gpu::INDICES_32_BIT) { + labels[globalPos] = (idx_t)(int32_t)globalId; + } else { + labels[globalPos] = (idx_t)globalId; + } + distances[globalPos] = chunkMergedDist[pos]; + } + } + continue; + } + + const float* outDistPtr = + reinterpret_cast([searchOutDistBuf_ contents]); + const int64_t* outIdxPtr = + reinterpret_cast([searchOutIdxBuf_ contents]); + for (idx_t qi = 0; qi < qCount; ++qi) { + for (idx_t j = 0; j < k; ++j) { + size_t localPos = (size_t)qi * (size_t)k + (size_t)j; + size_t globalPos = (size_t)(qBase + qi) * (size_t)k + (size_t)j; + int64_t globalId = outIdxPtr[localPos]; + if (globalId < 0) { + labels[globalPos] = -1; + } else if (indicesOptions_ == faiss::gpu::INDICES_CPU) { + labels[globalPos] = + decodeCpuLabelFromPair(cpuIndex_.get(), globalId); + } else if (indicesOptions_ == faiss::gpu::INDICES_32_BIT) { + labels[globalPos] = (idx_t)(int32_t)globalId; + } else { + labels[globalPos] = (idx_t)globalId; + } + distances[globalPos] = outDistPtr[localPos]; + } + } + } + + if (allowCpuFallback && logCpuFallback && fallbackCount > 0) { + std::fprintf( + stderr, + "IVF_CPU_FALLBACK,api=search_preassigned,count=%lld,first_reason=%s\n", + (long long)fallbackCount, + firstFallbackReason.c_str()); + } + if (logSyncProfile) { + const idx_t estimatedWaits = syncScanCalls + asyncBatchSyncs; + std::fprintf( + stderr, + "IVF_SYNC_PROFILE,api=search_preassigned,scan_calls=%lld,sync_scan_calls=%lld,async_scan_calls=%lld,async_batch_syncs=%lld,estimated_waits=%lld\n", + (long long)scanCalls, + (long long)syncScanCalls, + (long long)asyncScanCalls, + (long long)asyncBatchSyncs, + (long long)estimatedWaits); + } +} + +void MetalIndexIVFFlat::reconstruct(idx_t key, float* recons) const { + FAISS_THROW_IF_NOT_MSG(cpuIndex_, "MetalIndexIVFFlat: no internal index"); + cpuIndex_->reconstruct(key, recons); +} + +void MetalIndexIVFFlat::reconstruct_n(idx_t i0, idx_t ni, float* recons) const { + FAISS_THROW_IF_NOT_MSG(cpuIndex_, "MetalIndexIVFFlat: no internal index"); + cpuIndex_->reconstruct_n(i0, ni, recons); +} + +void MetalIndexIVFFlat::updateQuantizer() { + uploadCentroids_(); +} + +std::vector MetalIndexIVFFlat::getListIndices(idx_t listId) const { + FAISS_THROW_IF_NOT(cpuIndex_); + FAISS_THROW_IF_NOT(listId >= 0 && listId < cpuIndex_->nlist); + size_t ls = cpuIndex_->invlists->list_size(listId); + if (ls == 0) + return {}; + ScopedIds ids(cpuIndex_->invlists, (size_t)listId); + return ids.ptr ? std::vector(ids.ptr, ids.ptr + ls) + : std::vector{}; +} + +std::vector MetalIndexIVFFlat::getListVectorData(idx_t listId) const { + FAISS_THROW_IF_NOT(cpuIndex_); + FAISS_THROW_IF_NOT(listId >= 0 && listId < cpuIndex_->nlist); + size_t ls = cpuIndex_->invlists->list_size(listId); + if (ls == 0) + return {}; + ScopedCodes codes(cpuIndex_->invlists, (size_t)listId); + if (!codes.ptr) { + return {}; + } + size_t floatCount = ls * (size_t)d; + const float* fptr = reinterpret_cast(codes.ptr); + return std::vector(fptr, fptr + floatCount); +} + +void MetalIndexIVFFlat::reclaimMemory() { + // No-op for now: Metal unified memory doesn't benefit from explicit + // reclaim in the same way as discrete CUDA GPU memory. +} + +void MetalIndexIVFFlat::reserveMemory(idx_t numVecs) { + if (gpuIvf_) { + gpuIvf_->reserveMemory(numVecs); + } +} + +idx_t MetalIndexIVFFlat::nlist() const { + return cpuIndex_ ? cpuIndex_->nlist : 0; +} + +size_t MetalIndexIVFFlat::nprobe() const { + return cpuIndex_ ? cpuIndex_->nprobe : 1; +} + +bool MetalIndexIVFFlat::interleavedLayout() const { + return interleavedLayout_; +} + +faiss::gpu::IndicesOptions MetalIndexIVFFlat::indicesOptions() const { + return indicesOptions_; +} + +MetalIndexIVFFlat::AppendDebugStats MetalIndexIVFFlat::appendDebugStats() + const { + AppendDebugStats out{}; + if (!gpuIvf_) { + return out; + } + const auto& s = gpuIvf_->appendDebugStats(); + out.relayoutEvents = s.relayoutEvents; + out.movedLists = s.movedLists; + out.movedVectors = s.movedVectors; + out.reusedSegmentAllocs = s.reusedSegmentAllocs; + out.tailSegmentAllocs = s.tailSegmentAllocs; + out.reusedCapacityVecs = s.reusedCapacityVecs; + out.tailCapacityVecs = s.tailCapacityVecs; + out.tailShrinkEvents = s.tailShrinkEvents; + out.tailShrunkVecs = s.tailShrunkVecs; + return out; +} + +void MetalIndexIVFFlat::resetAppendDebugStats() { + if (gpuIvf_) { + gpuIvf_->resetAppendDebugStats(); + } +} + +void MetalIndexIVFFlat::copyFrom(const faiss::IndexIVFFlat* src) { + FAISS_THROW_IF_NOT(cpuIndex_); + FAISS_THROW_IF_NOT(src); + FAISS_THROW_IF_NOT_FMT( + src->nlist == cpuIndex_->nlist, + "copyFrom: nlist mismatch (%zd vs %zd)", + (size_t)src->nlist, + (size_t)cpuIndex_->nlist); + reset(); + + if (!src->is_trained) { + is_trained = false; + return; + } + + // Copy quantizer centroids (allow non-IndexFlat CPU coarse quantizers by + // reconstructing centroid vectors). + FAISS_THROW_IF_NOT_MSG( + src->quantizer, "copyFrom: source quantizer is null"); + auto* ourQ = cpuIndex_->quantizer; + FAISS_THROW_IF_NOT_MSG(ourQ, "copyFrom: internal quantizer is null"); + ourQ->reset(); + if (src->nlist > 0) { + std::vector coarse((size_t)src->nlist * d); + src->quantizer->reconstruct_n(0, src->nlist, coarse.data()); + if (!ourQ->is_trained) { + ourQ->train(src->nlist, coarse.data()); + } + ourQ->add(src->nlist, coarse.data()); + } + cpuIndex_->is_trained = true; + cpuIndex_->nprobe = src->nprobe; + is_trained = true; + + // Gather all vectors from inverted lists for a single GPU upload. + size_t totalN = 0; + for (size_t l = 0; l < (size_t)src->nlist; ++l) { + totalN += src->invlists->list_size(l); + } + + if (totalN > 0) { + std::vector allCodes(totalN * (size_t)d); + std::vector allListNos(totalN); + std::vector allIds(totalN); + size_t pos = 0; + + for (size_t l = 0; l < (size_t)src->nlist; ++l) { + size_t ls = src->invlists->list_size(l); + if (ls == 0) { + continue; + } + ScopedCodes codes(src->invlists, l); + ScopedIds ids(src->invlists, l); + + cpuIndex_->invlists->add_entries(l, ls, ids.ptr, codes.ptr); + + std::memcpy( + allCodes.data() + pos * (size_t)d, + codes.ptr, + ls * (size_t)d * sizeof(float)); + std::memcpy(allIds.data() + pos, ids.ptr, ls * sizeof(idx_t)); + for (size_t i = 0; i < ls; ++i) { + allListNos[pos + i] = (idx_t)l; + } + pos += ls; + } + + cpuIndex_->ntotal = (idx_t)totalN; + ntotal = (idx_t)totalN; + + if (gpuIvf_) { + gpuIvf_->appendVectors( + (idx_t)totalN, + allCodes.data(), + allListNos.data(), + allIds.data()); + } + } + + uploadCentroids_(); +} + +void MetalIndexIVFFlat::copyTo(faiss::IndexIVFFlat* dst) const { + FAISS_THROW_IF_NOT(cpuIndex_); + FAISS_THROW_IF_NOT(dst); + + auto* srcQ = cpuIndex_->quantizer; + auto* dstQ = dst->quantizer; + FAISS_THROW_IF_NOT_MSG(srcQ, "copyTo: internal quantizer is null"); + FAISS_THROW_IF_NOT_MSG(dstQ, "copyTo: destination quantizer is null"); + + dstQ->reset(); + if (srcQ->ntotal > 0) { + std::vector coarse((size_t)srcQ->ntotal * d); + srcQ->reconstruct_n(0, srcQ->ntotal, coarse.data()); + if (!dstQ->is_trained) { + dstQ->train(srcQ->ntotal, coarse.data()); + } + dstQ->add(srcQ->ntotal, coarse.data()); + } + + dst->metric_type = cpuIndex_->metric_type; + dst->metric_arg = cpuIndex_->metric_arg; + dst->d = cpuIndex_->d; + dst->nlist = cpuIndex_->nlist; + dst->nprobe = cpuIndex_->nprobe; + dst->is_trained = cpuIndex_->is_trained; + + for (size_t l = 0; l < (size_t)cpuIndex_->nlist; ++l) { + size_t ls = cpuIndex_->invlists->list_size(l); + if (ls == 0) { + continue; + } + ScopedCodes codes(cpuIndex_->invlists, l); + ScopedIds ids(cpuIndex_->invlists, l); + dst->invlists->add_entries(l, ls, ids.ptr, codes.ptr); + } + dst->ntotal = cpuIndex_->ntotal; +} + +} // namespace gpu_metal +} // namespace faiss diff --git a/faiss/gpu_metal/MetalKernels.h b/faiss/gpu_metal/MetalKernels.h index 586103b26..92fb79df5 100644 --- a/faiss/gpu_metal/MetalKernels.h +++ b/faiss/gpu_metal/MetalKernels.h @@ -19,13 +19,15 @@ namespace faiss { namespace gpu_metal { +enum class IVFScanVariant { Standard, Small, Interleaved }; + class MetalKernels { public: explicit MetalKernels(id device); ~MetalKernels(); bool isValid() const; - static constexpr int kMaxK = 256; + static constexpr int kMaxK = 2048; void encodeDistanceMatrix( id enc, @@ -37,6 +39,23 @@ class MetalKernels { int d, MetricType metric); + void encodeL2WithNorms( + id enc, + id queries, + id vectors, + id distances, + id vecNorms, + int nq, + int nb, + int d); + + void encodeComputeNorms( + id enc, + id vectors, + id norms, + int nb, + int d); + void encodeTopKThreadgroup( id enc, id distances, @@ -47,6 +66,31 @@ class MetalKernels { int k, bool wantMin); + void encodeIVFScanList( + id enc, + IVFScanVariant variant, + id queries, + id codes, + id ids, + id listOffset, + id listLength, + id coarseAssign, + id perListDist, + id perListIdx, + id paramsBuf, + int nq, + int nprobe, + id ilCodesOffset = nil); + + void encodeIVFMergeLists( + id enc, + id perListDist, + id perListIdx, + id outDist, + id outIdx, + id paramsBuf, + int nq); + static int selectTopKVariantIndex(int k); private: @@ -56,8 +100,9 @@ class MetalKernels { id library_; std::unordered_map> cache_; - static constexpr int kTopKVariantSizes[] = {32, 64, 128, 256}; - static constexpr int kNumTopKVariants = 4; + static constexpr int kTopKVariantSizes[] = + {32, 64, 128, 256, 512, 1024, 2048}; + static constexpr int kNumTopKVariants = 7; }; MetalKernels& getMetalKernels(id device); diff --git a/faiss/gpu_metal/MetalKernels.mm b/faiss/gpu_metal/MetalKernels.mm index 42270d51b..d0356af59 100644 --- a/faiss/gpu_metal/MetalKernels.mm +++ b/faiss/gpu_metal/MetalKernels.mm @@ -10,6 +10,7 @@ */ #import "MetalKernels.h" +#include #include #ifndef FAISS_METALLIB_PATH @@ -25,7 +26,10 @@ static const char* kThreadgroupNames[] = { "topk_threadgroup_32", "topk_threadgroup_64", "topk_threadgroup_128", - "topk_threadgroup_256"}; + "topk_threadgroup_256", + "topk_threadgroup_512", + "topk_threadgroup_1024", + "topk_threadgroup_2048"}; } // namespace @@ -100,6 +104,47 @@ void MetalKernels::encodeDistanceMatrix( threadsPerThreadgroup:MTLSizeMake(16, 16, 1)]; } +void MetalKernels::encodeL2WithNorms( + id enc, + id queries, + id vectors, + id distances, + id vecNorms, + int nq, + int nb, + int d) { + [enc setComputePipelineState:pipeline("l2_with_norms")]; + [enc setBuffer:queries offset:0 atIndex:0]; + [enc setBuffer:vectors offset:0 atIndex:1]; + [enc setBuffer:distances offset:0 atIndex:2]; + uint32_t args[3] = {(uint32_t)nq, (uint32_t)nb, (uint32_t)d}; + [enc setBytes:args length:sizeof(args) atIndex:3]; + [enc setBuffer:vecNorms offset:0 atIndex:4]; + const NSUInteger w = 16, h = 16; + MTLSize grid = MTLSizeMake( + ((NSUInteger)nb + w - 1) / w, ((NSUInteger)nq + h - 1) / h, 1); + [enc dispatchThreadgroups:grid threadsPerThreadgroup:MTLSizeMake(w, h, 1)]; +} + +void MetalKernels::encodeComputeNorms( + id enc, + id vectors, + id norms, + int nb, + int d) { + auto ps = pipeline("compute_norms"); + [enc setComputePipelineState:ps]; + [enc setBuffer:vectors offset:0 atIndex:0]; + [enc setBuffer:norms offset:0 atIndex:1]; + uint32_t args[2] = {(uint32_t)nb, (uint32_t)d}; + [enc setBytes:args length:sizeof(args) atIndex:2]; + NSUInteger tgSize = + std::min((NSUInteger)256, ps.maxTotalThreadsPerThreadgroup); + NSUInteger groups = ((NSUInteger)nb + tgSize - 1) / tgSize; + [enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(tgSize, 1, 1)]; +} + void MetalKernels::encodeTopKThreadgroup( id enc, id distances, @@ -121,6 +166,73 @@ void MetalKernels::encodeTopKThreadgroup( threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; } +void MetalKernels::encodeIVFScanList( + id enc, + IVFScanVariant variant, + id queries, + id codes, + id ids, + id listOffset, + id listLength, + id coarseAssign, + id perListDist, + id perListIdx, + id paramsBuf, + int nq, + int nprobe, + id ilCodesOffset) { + const char* name; + NSUInteger tgSize; + switch (variant) { + case IVFScanVariant::Small: + name = "ivf_scan_list_small"; + tgSize = 32; + break; + case IVFScanVariant::Interleaved: + name = "ivf_scan_list_interleaved"; + tgSize = 256; + break; + default: + name = "ivf_scan_list"; + tgSize = 256; + break; + } + [enc setComputePipelineState:pipeline(name)]; + [enc setBuffer:queries offset:0 atIndex:0]; + [enc setBuffer:codes offset:0 atIndex:1]; + [enc setBuffer:ids offset:0 atIndex:2]; + [enc setBuffer:listOffset offset:0 atIndex:3]; + [enc setBuffer:listLength offset:0 atIndex:4]; + [enc setBuffer:coarseAssign offset:0 atIndex:5]; + [enc setBuffer:perListDist offset:0 atIndex:6]; + [enc setBuffer:perListIdx offset:0 atIndex:7]; + [enc setBuffer:paramsBuf offset:0 atIndex:8]; + if (variant == IVFScanVariant::Interleaved && ilCodesOffset) { + [enc setBuffer:ilCodesOffset offset:0 atIndex:9]; + } + NSUInteger totalTGs = (NSUInteger)nq * (NSUInteger)nprobe; + [enc dispatchThreadgroups:MTLSizeMake(totalTGs, 1, 1) + threadsPerThreadgroup:MTLSizeMake(tgSize, 1, 1)]; +} + +void MetalKernels::encodeIVFMergeLists( + id enc, + id perListDist, + id perListIdx, + id outDist, + id outIdx, + id paramsBuf, + int nq) { + [enc setComputePipelineState:pipeline("ivf_merge_lists")]; + [enc setBuffer:perListDist offset:0 atIndex:0]; + [enc setBuffer:perListIdx offset:0 atIndex:1]; + [enc setBuffer:outDist offset:0 atIndex:2]; + [enc setBuffer:outIdx offset:0 atIndex:3]; + [enc setBuffer:paramsBuf offset:0 atIndex:4]; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nq, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; +} + MetalKernels& getMetalKernels(id device) { static std::mutex mu; static std::unordered_map> map; @@ -133,4 +245,15 @@ MetalKernels& getMetalKernels(id device) { } } // namespace gpu_metal + +extern std::string& ref_gpu_compile_options(); + +struct InitMetalCompileOptions { + InitMetalCompileOptions() { + ref_gpu_compile_options() = "MAC_METAL "; + } +}; + +InitMetalCompileOptions InitMetalCompileOptions_instance; + } // namespace faiss diff --git a/faiss/gpu_metal/MetalPythonBridge.h b/faiss/gpu_metal/MetalPythonBridge.h new file mode 100644 index 000000000..592ac11ad --- /dev/null +++ b/faiss/gpu_metal/MetalPythonBridge.h @@ -0,0 +1,45 @@ +// @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. + * + * C++-only API for Python/SWIG. No Objective-C types so SWIG can + * parse it. Implemented in MetalPythonBridge.mm. + */ + +#pragma once + +#include + +namespace faiss { +namespace gpu_metal { + +/// Opaque holder for Metal resources. +struct StandardMetalResourcesHolder { + void* impl = nullptr; + StandardMetalResourcesHolder(); + ~StandardMetalResourcesHolder(); + StandardMetalResourcesHolder(const StandardMetalResourcesHolder&) = delete; + StandardMetalResourcesHolder& operator=( + const StandardMetalResourcesHolder&) = delete; +}; + +/// Same names as GPU API for unified Python binding. +int get_num_gpus(); +void gpu_profiler_start(); +void gpu_profiler_stop(); +void gpu_sync_all_devices(); + +/// Clone CPU index to Metal GPU. Caller owns returned index. +faiss::Index* index_cpu_to_gpu( + StandardMetalResourcesHolder* res, + int device, + const faiss::Index* index); + +/// Copy Metal index back to CPU. Caller owns returned index. +faiss::Index* index_gpu_to_cpu(const faiss::Index* index); + +} // namespace gpu_metal +} // namespace faiss diff --git a/faiss/gpu_metal/MetalPythonBridge.mm b/faiss/gpu_metal/MetalPythonBridge.mm new file mode 100644 index 000000000..d019c248e --- /dev/null +++ b/faiss/gpu_metal/MetalPythonBridge.mm @@ -0,0 +1,79 @@ +// @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. + */ + +#import "MetalPythonBridge.h" +#include +#include +#include +#include +#include +#import "MetalCloner.h" +#import "MetalResources.h" +#import "StandardMetalResources.h" + +namespace faiss { +namespace gpu_metal { + +namespace { +std::atomic gMetalProfilerDepth{0}; +} // namespace + +StandardMetalResourcesHolder::StandardMetalResourcesHolder() { + impl = new StandardMetalResources(); +} + +StandardMetalResourcesHolder::~StandardMetalResourcesHolder() { + delete static_cast(impl); + impl = nullptr; +} + +void gpu_profiler_start() { + const int depth = gMetalProfilerDepth.fetch_add(1) + 1; + if (std::getenv("FAISS_METAL_PROFILE_LOG")) { + std::fprintf( + stderr, "[faiss_metal] gpu_profiler_start depth=%d\n", depth); + } +} + +void gpu_profiler_stop() { + int depth = gMetalProfilerDepth.load(); + while (depth > 0 && + !gMetalProfilerDepth.compare_exchange_weak(depth, depth - 1)) { + } + if (std::getenv("FAISS_METAL_PROFILE_LOG")) { + std::fprintf( + stderr, + "[faiss_metal] gpu_profiler_stop depth=%d\n", + gMetalProfilerDepth.load()); + } +} + +void gpu_sync_all_devices() { + auto res = std::make_shared(); + if (res && res->isAvailable()) { + res->synchronize(); + } +} + +faiss::Index* index_cpu_to_gpu( + StandardMetalResourcesHolder* res, + int device, + const faiss::Index* index) { + if (!res || !res->impl) { + return nullptr; + } + return index_cpu_to_metal_gpu( + static_cast(res->impl), device, index); +} + +faiss::Index* index_gpu_to_cpu(const faiss::Index* index) { + return index_metal_gpu_to_cpu(index); +} + +} // namespace gpu_metal +} // namespace faiss diff --git a/faiss/gpu_metal/impl/MetalIVFFlat.h b/faiss/gpu_metal/impl/MetalIVFFlat.h new file mode 100644 index 000000000..93c8f896d --- /dev/null +++ b/faiss/gpu_metal/impl/MetalIVFFlat.h @@ -0,0 +1,193 @@ +// @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 Flat implementation: GPU-resident IVF list storage and helpers. + * Mirrors the roles of faiss/gpu/impl/IVFFlat.cuh (storage side only). + */ + +#pragma once + +#import + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace faiss { +namespace gpu_metal { + +/// GPU-resident IVF list storage for flat (float32) codes. +/// Layout: all lists are stored contiguously in a single codes/ids buffer; +/// lists are described by (listOffset[list], listLength[list]). +class MetalIVFFlatImpl { + public: + struct AppendDebugStats { + size_t relayoutEvents = 0; + size_t movedLists = 0; + size_t movedVectors = 0; + size_t reusedSegmentAllocs = 0; + size_t tailSegmentAllocs = 0; + size_t reusedCapacityVecs = 0; + size_t tailCapacityVecs = 0; + size_t tailShrinkEvents = 0; + size_t tailShrunkVecs = 0; + }; + + MetalIVFFlatImpl( + std::shared_ptr resources, + int dim, + idx_t nlist, + faiss::MetricType metric, + float metricArg, + faiss::gpu::IndicesOptions indicesOptions, + bool interleavedLayout); + + ~MetalIVFFlatImpl(); + + /// Reset all IVF lists and free GPU storage. + void reset(); + + /// Reserve host/GPU storage for at least totalVecs vectors. + void reserveMemory(idx_t totalVecs); + + /// Append a batch of vectors to IVF lists. + /// - x: host pointer, size n * dim + /// - list_nos: host pointer, size n; -1 entries are skipped + /// - xids: host pointer, size n (may be null to use internal ids) + void appendVectors( + idx_t n, + const float* x, + const idx_t* list_nos, + const idx_t* xids); + + /// Accessors for future GPU search path. + int dim() const { + return dim_; + } + idx_t nlist() const { + return nlist_; + } + faiss::MetricType metricType() const { + return metric_type_; + } + float metricArg() const { + return metric_arg_; + } + + const std::vector& listLength() const { + return listLength_; + } + const std::vector& listOffset() const { + return listOffset_; + } + + id codesBuffer() const { + return codesBuffer_; + } + id idsBuffer() const { + return idsBuffer_; + } + /// Pre-built GPU buffer of (nlist) uint32_t offsets (updated on every add). + id listOffsetGpuBuffer() const { + return listOffsetBuf_; + } + /// Pre-built GPU buffer of (nlist) uint32_t lengths (updated on every add). + id listLengthGpuBuffer() const { + return listLengthBuf_; + } + + size_t totalVecs() const { + return totalVecs_; + } + + /// Interleaved codes buffer (blocks of 32 vectors, dims interleaved). + id interleavedCodesBuffer() const { + return interleavedCodesBuf_; + } + /// Per-list float offsets into the interleaved codes buffer. + id interleavedCodesOffsetBuffer() const { + return interleavedCodesOffsetBuf_; + } + bool interleavedLayout() const { + return interleavedLayout_; + } + /// Rebuild interleaved buffers from host storage if they are stale. + void ensureInterleavedLayoutUpToDate(); + const AppendDebugStats& appendDebugStats() const { + return appendStats_; + } + void resetAppendDebugStats() { + appendStats_ = AppendDebugStats{}; + } + + private: + struct FreeSegment { + size_t offset = 0; + size_t length = 0; + }; + + bool ensureCapacityForAppend_( + const std::vector& addPerList, + std::vector* movedLists); + void uploadToGpu_( + const std::vector& oldLength, + const std::vector& addPerList, + const std::vector& movedLists, + bool forceFullUpload); + void rebuildInterleavedBuffers_(); + size_t allocSegment_(size_t length); + void freeSegment_( + size_t offset, + size_t length, + bool allowTailShrink = true); + void coalesceFreeSegments_(); + void tryShrinkTail_(); + + std::shared_ptr resources_; + + int dim_; + idx_t nlist_; + faiss::MetricType metric_type_; + float metric_arg_; + faiss::gpu::IndicesOptions indicesOptions_; + bool interleavedLayout_; + + // Per-list metadata + std::vector listLength_; + std::vector listOffset_; + std::vector listCapacity_; + + // Host copies of IVF data (flat layout) + std::vector hostCodes_; // size = totalVecs_ * dim_ + std::vector hostIds_; // size = totalVecs_ + std::vector freeSegments_; + AppendDebugStats appendStats_; + size_t totalVecs_; + size_t totalCapacityVecs_; + + // GPU storage + id codesBuffer_; + id idsBuffer_; + id listOffsetBuf_; // (nlist) uint32_t, list element offsets + id listLengthBuf_; // (nlist) uint32_t, list sizes + + // Interleaved codes layout (blocks of 32 vectors, dims interleaved) + id interleavedCodesBuf_; + id interleavedCodesOffsetBuf_; // (nlist) uint32_t, float offsets + bool interleavedDirty_ = true; + + static constexpr int kInterleavedGroupSize = 32; +}; + +} // namespace gpu_metal +} // namespace faiss diff --git a/faiss/gpu_metal/impl/MetalIVFFlat.mm b/faiss/gpu_metal/impl/MetalIVFFlat.mm new file mode 100644 index 000000000..be41a771a --- /dev/null +++ b/faiss/gpu_metal/impl/MetalIVFFlat.mm @@ -0,0 +1,632 @@ +// @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. + */ + +#import "MetalIVFFlat.h" + +#include +#include +#include +#include +#include + +#include +#include + +namespace faiss { +namespace gpu_metal { + +namespace { +inline uint32_t checkedToU32(size_t v, const char* what) { + if (v > (size_t)std::numeric_limits::max()) { + FAISS_THROW_MSG(what); + } + return static_cast(v); +} + +constexpr size_t kDefaultMinTailShrinkVecs = 8192; +constexpr size_t kMaxU32SizeT = (size_t)std::numeric_limits::max(); + +size_t getTailShrinkMinVecs() { + const char* env = std::getenv("FAISS_METAL_IVF_TAIL_SHRINK_MIN_VECS"); + if (!env || env[0] == '\0') { + return kDefaultMinTailShrinkVecs; + } + char* end = nullptr; + unsigned long long v = std::strtoull(env, &end, 10); + if (end == env) { + return kDefaultMinTailShrinkVecs; + } + return (size_t)v; +} +} // namespace + +MetalIVFFlatImpl::MetalIVFFlatImpl( + std::shared_ptr resources, + int dim, + idx_t nlist, + faiss::MetricType metric, + float metricArg, + faiss::gpu::IndicesOptions indicesOptions, + bool interleavedLayout) + : resources_(std::move(resources)), + dim_(dim), + nlist_(nlist), + metric_type_(metric), + metric_arg_(metricArg), + indicesOptions_(indicesOptions), + interleavedLayout_(interleavedLayout), + listLength_(nlist_, 0), + listOffset_(nlist_, 0), + listCapacity_(nlist_, 0), + totalVecs_(0), + totalCapacityVecs_(0), + codesBuffer_(nil), + idsBuffer_(nil), + listOffsetBuf_(nil), + listLengthBuf_(nil), + interleavedCodesBuf_(nil), + interleavedCodesOffsetBuf_(nil), + interleavedDirty_(true) { + FAISS_THROW_IF_NOT(dim_ > 0); + FAISS_THROW_IF_NOT(nlist_ >= 0); +} + +MetalIVFFlatImpl::~MetalIVFFlatImpl() { + reset(); +} + +void MetalIVFFlatImpl::reset() { + hostCodes_.clear(); + hostIds_.clear(); + freeSegments_.clear(); + appendStats_ = AppendDebugStats{}; + totalVecs_ = 0; + totalCapacityVecs_ = 0; + + std::fill(listLength_.begin(), listLength_.end(), 0); + std::fill(listOffset_.begin(), listOffset_.end(), 0); + std::fill(listCapacity_.begin(), listCapacity_.end(), 0); + + if (codesBuffer_ != nil) { + resources_->deallocBuffer(codesBuffer_, MetalAllocType::IVFLists); + codesBuffer_ = nil; + } + if (idsBuffer_ != nil) { + resources_->deallocBuffer(idsBuffer_, MetalAllocType::IVFLists); + idsBuffer_ = nil; + } + if (listOffsetBuf_ != nil) { + resources_->deallocBuffer(listOffsetBuf_, MetalAllocType::IVFLists); + listOffsetBuf_ = nil; + } + if (listLengthBuf_ != nil) { + resources_->deallocBuffer(listLengthBuf_, MetalAllocType::IVFLists); + listLengthBuf_ = nil; + } + if (interleavedCodesBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesBuf_, MetalAllocType::IVFLists); + interleavedCodesBuf_ = nil; + } + if (interleavedCodesOffsetBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesOffsetBuf_, MetalAllocType::IVFLists); + interleavedCodesOffsetBuf_ = nil; + } + interleavedDirty_ = true; +} + +void MetalIVFFlatImpl::reserveMemory(idx_t totalVecs) { + if (totalVecs <= 0) { + return; + } + size_t t = (size_t)totalVecs; + FAISS_THROW_IF_NOT_MSG( + t <= kMaxU32SizeT, + "MetalIVFFlatImpl::reserveMemory exceeds uint32 list-offset range"); + if (t <= totalCapacityVecs_) { + return; + } + const size_t oldCap = totalCapacityVecs_; + const size_t extra = t - oldCap; + hostCodes_.resize(t * (size_t)dim_, 0.0f); + hostIds_.resize(t, (idx_t)-1); + totalCapacityVecs_ = t; + // Add reserved tail as a reusable free segment for future list growth. + freeSegment_(oldCap, extra, false); + + // Materialize GPU allocations up-front so future appends can avoid + // reallocating until reserved capacity is exceeded. + std::vector oldLength((size_t)nlist_, 0); + std::vector addPerList((size_t)nlist_, 0); + std::vector movedLists((size_t)nlist_, 0); + uploadToGpu_(oldLength, addPerList, movedLists, true); +} + +void MetalIVFFlatImpl::appendVectors( + idx_t n, + const float* x, + const idx_t* list_nos, + const idx_t* xids) { + if (n == 0) { + return; + } + FAISS_THROW_IF_NOT(list_nos != nullptr); + + // Count how many vectors go to each list. + std::vector addPerList(nlist_, 0); + for (idx_t i = 0; i < n; ++i) { + idx_t list = list_nos[i]; + if (list < 0 || list >= nlist_) { + continue; + } + addPerList[(size_t)list]++; + } + + // Early-out: nothing to append. + size_t batchNew = 0; + for (size_t l = 0; l < (size_t)nlist_; ++l) { + batchNew += addPerList[l]; + } + if (batchNew == 0) { + return; + } + + const std::vector oldLength = listLength_; + std::vector movedLists((size_t)nlist_, 0); + bool forceFullUpload = ensureCapacityForAppend_(addPerList, &movedLists); + + for (idx_t i = 0; i < n; ++i) { + idx_t list = list_nos[i]; + if (list < 0 || list >= nlist_) { + continue; + } + size_t l = (size_t)list; + size_t dstIndex = listOffset_[l] + listLength_[l]; + size_t listOffset = listLength_[l]; + listLength_[l]++; + + // Copy vector + const float* xi = x + (size_t)i * (size_t)dim_; + std::memcpy( + hostCodes_.data() + dstIndex * (size_t)dim_, + xi, + (size_t)dim_ * sizeof(float)); + + // Copy id + idx_t id = -1; + if (indicesOptions_ == faiss::gpu::INDICES_CPU || + indicesOptions_ == faiss::gpu::INDICES_IVF) { + id = (idx_t)faiss::lo_build((uint64_t)l, (uint64_t)listOffset); + } else { + id = xids ? xids[i] : (idx_t)(totalVecs_ + (size_t)(i)); + if (indicesOptions_ == faiss::gpu::INDICES_32_BIT) { + FAISS_THROW_IF_NOT_MSG( + id >= (idx_t)std::numeric_limits::min() && + id <= (idx_t)std::numeric_limits< + int32_t>::max(), + "MetalIVFFlatImpl: id out of int32 range"); + id = (idx_t)(int32_t)id; + } + } + hostIds_[dstIndex] = id; + } + + totalVecs_ += batchNew; + uploadToGpu_(oldLength, addPerList, movedLists, forceFullUpload); +} + +bool MetalIVFFlatImpl::ensureCapacityForAppend_( + const std::vector& addPerList, + std::vector* movedLists) { + if (movedLists) { + movedLists->assign((size_t)nlist_, 0); + } + + std::vector newCap((size_t)nlist_, 0); + bool anyMoved = false; + + for (size_t l = 0; l < (size_t)nlist_; ++l) { + size_t need = listLength_[l] + addPerList[l]; + if (need <= listCapacity_[l]) { + continue; + } + size_t cap = std::max(1, listCapacity_[l]); + while (cap < need) { + cap *= 2; + } + FAISS_THROW_IF_NOT_MSG( + cap <= kMaxU32SizeT, + "MetalIVFFlatImpl: list capacity exceeds uint32 range"); + newCap[l] = cap; + anyMoved = true; + if (movedLists) { + (*movedLists)[l] = 1; + } + } + + if (!anyMoved) { + return false; + } + appendStats_.relayoutEvents += 1; + + for (size_t l = 0; l < (size_t)nlist_; ++l) { + if (newCap[l] == 0) { + continue; + } + + size_t len = listLength_[l]; + size_t oldOff = listOffset_[l]; + size_t oldCap = listCapacity_[l]; + size_t newOff = allocSegment_(newCap[l]); + appendStats_.movedLists += 1; + appendStats_.movedVectors += len; + + if (len > 0) { + std::memcpy( + hostCodes_.data() + newOff * (size_t)dim_, + hostCodes_.data() + oldOff * (size_t)dim_, + len * (size_t)dim_ * sizeof(float)); + std::memcpy( + hostIds_.data() + newOff, + hostIds_.data() + oldOff, + len * sizeof(idx_t)); + } + + listOffset_[l] = newOff; + listCapacity_[l] = newCap[l]; + freeSegment_(oldOff, oldCap, true); + } + + return true; +} + +size_t MetalIVFFlatImpl::allocSegment_(size_t length) { + if (length == 0) { + return 0; + } + + size_t bestIdx = (size_t)-1; + size_t bestLen = std::numeric_limits::max(); + for (size_t i = 0; i < freeSegments_.size(); ++i) { + const auto& seg = freeSegments_[i]; + if (seg.length < length) { + continue; + } + if (seg.length < bestLen) { + bestLen = seg.length; + bestIdx = i; + } + } + + if (bestIdx != (size_t)-1) { + auto seg = freeSegments_[bestIdx]; + size_t out = seg.offset; + if (seg.length == length) { + freeSegments_.erase(freeSegments_.begin() + (long)bestIdx); + } else { + freeSegments_[bestIdx].offset += length; + freeSegments_[bestIdx].length -= length; + } + appendStats_.reusedSegmentAllocs += 1; + appendStats_.reusedCapacityVecs += length; + return out; + } + + size_t out = totalCapacityVecs_; + FAISS_THROW_IF_NOT_MSG( + out <= kMaxU32SizeT && length <= (kMaxU32SizeT - out), + "MetalIVFFlatImpl: total capacity exceeds uint32 list-offset range"); + totalCapacityVecs_ += length; + hostCodes_.resize(totalCapacityVecs_ * (size_t)dim_, 0.0f); + hostIds_.resize(totalCapacityVecs_, (idx_t)-1); + appendStats_.tailSegmentAllocs += 1; + appendStats_.tailCapacityVecs += length; + return out; +} + +void MetalIVFFlatImpl::freeSegment_( + size_t offset, + size_t length, + bool allowTailShrink) { + if (length == 0) { + return; + } + freeSegments_.push_back({offset, length}); + coalesceFreeSegments_(); + if (allowTailShrink) { + tryShrinkTail_(); + } +} + +void MetalIVFFlatImpl::coalesceFreeSegments_() { + if (freeSegments_.empty()) { + return; + } + std::sort( + freeSegments_.begin(), + freeSegments_.end(), + [](const FreeSegment& a, const FreeSegment& b) { + return a.offset < b.offset; + }); + + std::vector merged; + merged.reserve(freeSegments_.size()); + merged.push_back(freeSegments_[0]); + for (size_t i = 1; i < freeSegments_.size(); ++i) { + auto& back = merged.back(); + const auto& cur = freeSegments_[i]; + if (back.offset + back.length == cur.offset) { + back.length += cur.length; + } else { + merged.push_back(cur); + } + } + freeSegments_.swap(merged); +} + +void MetalIVFFlatImpl::tryShrinkTail_() { + const size_t minTailShrinkVecs = getTailShrinkMinVecs(); + if (minTailShrinkVecs == 0) { + return; + } + if (freeSegments_.empty()) { + return; + } + while (!freeSegments_.empty()) { + const auto& tail = freeSegments_.back(); + if (tail.offset + tail.length != totalCapacityVecs_) { + break; + } + if (tail.length < minTailShrinkVecs) { + break; + } + totalCapacityVecs_ = tail.offset; + appendStats_.tailShrinkEvents += 1; + appendStats_.tailShrunkVecs += tail.length; + freeSegments_.pop_back(); + } + hostCodes_.resize(totalCapacityVecs_ * (size_t)dim_); + hostIds_.resize(totalCapacityVecs_); +} + +void MetalIVFFlatImpl::uploadToGpu_( + const std::vector& oldLength, + const std::vector& addPerList, + const std::vector& movedLists, + bool forceFullUpload) { + if (!resources_ || !resources_->isAvailable()) { + return; + } + + size_t codesBytes = totalCapacityVecs_ * (size_t)dim_ * sizeof(float); + size_t idsBytes = totalCapacityVecs_ * sizeof(idx_t); + size_t metaBytes = (size_t)nlist_ * sizeof(uint32_t); + FAISS_THROW_IF_NOT_MSG( + (size_t)nlist_ <= (size_t)std::numeric_limits::max(), + "MetalIVFFlatImpl: nlist exceeds uint32 metadata range"); + + // Always update metadata buffers — they reflect current list layout. + if (metaBytes > 0) { + if (listOffsetBuf_ == nil || [listOffsetBuf_ length] < metaBytes) { + if (listOffsetBuf_ != nil) { + resources_->deallocBuffer( + listOffsetBuf_, MetalAllocType::IVFLists); + } + listOffsetBuf_ = resources_->allocBuffer( + metaBytes, MetalAllocType::IVFLists); + } + if (listLengthBuf_ == nil || [listLengthBuf_ length] < metaBytes) { + if (listLengthBuf_ != nil) { + resources_->deallocBuffer( + listLengthBuf_, MetalAllocType::IVFLists); + } + listLengthBuf_ = resources_->allocBuffer( + metaBytes, MetalAllocType::IVFLists); + } + FAISS_THROW_IF_NOT_MSG( + listOffsetBuf_ && listLengthBuf_, + "MetalIVFFlatImpl: failed to allocate metadata GPU buffers"); + + auto* offPtr = reinterpret_cast([listOffsetBuf_ contents]); + auto* lenPtr = reinterpret_cast([listLengthBuf_ contents]); + for (size_t i = 0; i < (size_t)nlist_; ++i) { + offPtr[i] = checkedToU32( + listOffset_[i], + "MetalIVFFlatImpl: list offset exceeds uint32 range"); + lenPtr[i] = checkedToU32( + listLength_[i], + "MetalIVFFlatImpl: list length exceeds uint32 range"); + } + } + + if (codesBytes == 0 || idsBytes == 0) { + return; + } + + if (idsBuffer_ == nil || [idsBuffer_ length] < idsBytes) { + if (idsBuffer_ != nil) { + resources_->deallocBuffer(idsBuffer_, MetalAllocType::IVFLists); + } + idsBuffer_ = + resources_->allocBuffer(idsBytes, MetalAllocType::IVFLists); + forceFullUpload = true; + } + + FAISS_THROW_IF_NOT_MSG( + idsBuffer_, "MetalIVFFlatImpl: failed to allocate IVF ids buffer"); + + if (codesBuffer_ == nil || [codesBuffer_ length] < codesBytes) { + if (codesBuffer_ != nil) { + resources_->deallocBuffer(codesBuffer_, MetalAllocType::IVFLists); + } + codesBuffer_ = + resources_->allocBuffer(codesBytes, MetalAllocType::IVFLists); + forceFullUpload = true; + } + FAISS_THROW_IF_NOT_MSG( + codesBuffer_, + "MetalIVFFlatImpl: failed to allocate IVF codes buffer"); + + if (forceFullUpload) { + std::memcpy([idsBuffer_ contents], hostIds_.data(), idsBytes); + std::memcpy([codesBuffer_ contents], hostCodes_.data(), codesBytes); + } else { + for (size_t l = 0; l < (size_t)nlist_; ++l) { + size_t add = addPerList[l]; + bool moved = (l < movedLists.size() && movedLists[l] != 0); + if (add == 0 && !moved) { + continue; + } + size_t start = + moved ? listOffset_[l] : (listOffset_[l] + oldLength[l]); + size_t count = moved ? listLength_[l] : add; + if (count == 0) { + continue; + } + + std::memcpy( + reinterpret_cast([idsBuffer_ contents]) + start, + hostIds_.data() + start, + count * sizeof(idx_t)); + std::memcpy( + reinterpret_cast([codesBuffer_ contents]) + + start * (size_t)dim_, + hostCodes_.data() + start * (size_t)dim_, + count * (size_t)dim_ * sizeof(float)); + } + } + + if (interleavedLayout_) { + interleavedDirty_ = true; + } else { + if (interleavedCodesBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesBuf_, MetalAllocType::IVFLists); + interleavedCodesBuf_ = nil; + } + if (interleavedCodesOffsetBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesOffsetBuf_, MetalAllocType::IVFLists); + interleavedCodesOffsetBuf_ = nil; + } + interleavedDirty_ = false; + } +} + +void MetalIVFFlatImpl::rebuildInterleavedBuffers_() { + if (!interleavedLayout_ || !resources_ || !resources_->isAvailable()) { + return; + } + + const size_t metaBytes = (size_t)nlist_ * sizeof(uint32_t); + constexpr int G = kInterleavedGroupSize; // 32 + + std::vector ilOffsets((size_t)nlist_); + size_t totalIlFloats = 0; + for (size_t l = 0; l < (size_t)nlist_; ++l) { + ilOffsets[l] = checkedToU32( + totalIlFloats, + "MetalIVFFlatImpl: interleaved offset exceeds uint32 range"); + size_t len = listLength_[l]; + size_t numBlocks = (len + G - 1) / G; + totalIlFloats += numBlocks * G * (size_t)dim_; + } + + if (totalIlFloats == 0) { + if (interleavedCodesBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesBuf_, MetalAllocType::IVFLists); + interleavedCodesBuf_ = nil; + } + if (interleavedCodesOffsetBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesOffsetBuf_, MetalAllocType::IVFLists); + interleavedCodesOffsetBuf_ = nil; + } + interleavedDirty_ = false; + return; + } + + const size_t ilBytes = totalIlFloats * sizeof(float); + if (interleavedCodesBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesBuf_, MetalAllocType::IVFLists); + interleavedCodesBuf_ = nil; + } + interleavedCodesBuf_ = + resources_->allocBuffer(ilBytes, MetalAllocType::IVFLists); + + if (interleavedCodesOffsetBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesOffsetBuf_, MetalAllocType::IVFLists); + interleavedCodesOffsetBuf_ = nil; + } + interleavedCodesOffsetBuf_ = + resources_->allocBuffer(metaBytes, MetalAllocType::IVFLists); + + if (!interleavedCodesBuf_ || !interleavedCodesOffsetBuf_) { + if (interleavedCodesBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesBuf_, MetalAllocType::IVFLists); + interleavedCodesBuf_ = nil; + } + if (interleavedCodesOffsetBuf_ != nil) { + resources_->deallocBuffer( + interleavedCodesOffsetBuf_, MetalAllocType::IVFLists); + interleavedCodesOffsetBuf_ = nil; + } + interleavedDirty_ = true; + return; + } + + auto* dst = reinterpret_cast([interleavedCodesBuf_ contents]); + std::memset(dst, 0, ilBytes); + + for (size_t l = 0; l < (size_t)nlist_; ++l) { + size_t len = listLength_[l]; + if (len == 0) + continue; + size_t srcOff = listOffset_[l]; + size_t dstOff = ilOffsets[l]; + size_t numBlocks = (len + G - 1) / G; + + for (size_t b = 0; b < numBlocks; ++b) { + for (int dd = 0; dd < dim_; ++dd) { + for (int g = 0; g < G; ++g) { + size_t vi = b * G + g; + float val = (vi < len) + ? hostCodes_[(srcOff + vi) * (size_t)dim_ + dd] + : 0.0f; + dst[dstOff + b * G * (size_t)dim_ + dd * G + g] = val; + } + } + } + } + + auto* ptr = + reinterpret_cast([interleavedCodesOffsetBuf_ contents]); + for (size_t i = 0; i < (size_t)nlist_; ++i) { + ptr[i] = ilOffsets[i]; + } + interleavedDirty_ = false; +} + +void MetalIVFFlatImpl::ensureInterleavedLayoutUpToDate() { + if (!interleavedLayout_) { + return; + } + if (!interleavedDirty_ && interleavedCodesBuf_ && + interleavedCodesOffsetBuf_) { + return; + } + rebuildInterleavedBuffers_(); +} + +} // namespace gpu_metal +} // namespace faiss diff --git a/faiss/gpu_metal/test/CMakeLists.txt b/faiss/gpu_metal/test/CMakeLists.txt index 31312de7f..de192f2d4 100644 --- a/faiss/gpu_metal/test/CMakeLists.txt +++ b/faiss/gpu_metal/test/CMakeLists.txt @@ -25,3 +25,19 @@ set_target_properties(TestMetalIndexFlat PROPERTIES OBJCXX_STANDARD_REQUIRED ON ) gtest_discover_tests(TestMetalIndexFlat) + +add_executable(TestMetalIndexIVFFlat TestMetalIndexIVFFlat.mm) +target_link_libraries(TestMetalIndexIVFFlat PRIVATE + faiss_metal + faiss + GTest::gtest_main +) +target_include_directories(TestMetalIndexIVFFlat PRIVATE + ${PROJECT_SOURCE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/../.. +) +set_target_properties(TestMetalIndexIVFFlat PROPERTIES + OBJCXX_STANDARD 17 + OBJCXX_STANDARD_REQUIRED ON +) +gtest_discover_tests(TestMetalIndexIVFFlat) diff --git a/faiss/gpu_metal/test/TestMetalIndexIVFFlat.mm b/faiss/gpu_metal/test/TestMetalIndexIVFFlat.mm new file mode 100644 index 000000000..1336cbf10 --- /dev/null +++ b/faiss/gpu_metal/test/TestMetalIndexIVFFlat.mm @@ -0,0 +1,200 @@ +// @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. + * + * Tests for MetalIndexIVFFlat: train, add, search, reset, cloner round-trip. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#import +#import +#import + +namespace { + +constexpr float kTolerance = 1e-4f; + +void compareSearchResults( + int nq, + int k, + const float* refDist, + const faiss::idx_t* refLab, + const float* testDist, + const faiss::idx_t* testLab) { + for (int i = 0; i < nq * k; ++i) { + EXPECT_NEAR( + refDist[i], + testDist[i], + kTolerance * (std::fabs(refDist[i]) + 1.0f)) + << "i=" << i; + EXPECT_EQ(refLab[i], testLab[i]) << "i=" << i; + } +} + +} // namespace + +class TestMetalIndexIVFFlat : public ::testing::Test { + protected: + void SetUp() override { + resources_ = std::make_shared(); + if (!resources_->isAvailable()) { + GTEST_SKIP() << "Metal not available"; + } + } + std::shared_ptr resources_; +}; + +TEST_F(TestMetalIndexIVFFlat, L2_TrainAddSearch) { + const int dim = 64; + const int nb = 1000; + const int nq = 10; + const int nlist = 8; + const int nprobe = 4; + const int k = 5; + + std::vector vecs((size_t)nb * dim); + faiss::float_rand(vecs.data(), vecs.size(), 1234); + std::vector queries((size_t)nq * dim); + faiss::float_rand(queries.data(), queries.size(), 5678); + + faiss::IndexFlatL2 quantizer(dim); + faiss::IndexIVFFlat cpuIndex(&quantizer, dim, nlist); + cpuIndex.nprobe = nprobe; + cpuIndex.train(nb, vecs.data()); + cpuIndex.add(nb, vecs.data()); + + faiss::gpu_metal::MetalIndexIVFFlat metalIndex(resources_, &cpuIndex); + ASSERT_EQ(metalIndex.ntotal, nb); + + std::vector refDist((size_t)nq * k); + std::vector refLab((size_t)nq * k, -1); + std::vector testDist((size_t)nq * k); + std::vector testLab((size_t)nq * k, -1); + + cpuIndex.search(nq, queries.data(), k, refDist.data(), refLab.data()); + metalIndex.search(nq, queries.data(), k, testDist.data(), testLab.data()); + + compareSearchResults( + nq, + k, + refDist.data(), + refLab.data(), + testDist.data(), + testLab.data()); +} + +TEST_F(TestMetalIndexIVFFlat, IP_TrainAddSearch) { + const int dim = 64; + const int nb = 1000; + const int nq = 10; + const int nlist = 8; + const int nprobe = 4; + const int k = 5; + + std::vector vecs((size_t)nb * dim); + faiss::float_rand(vecs.data(), vecs.size(), 1234); + std::vector queries((size_t)nq * dim); + faiss::float_rand(queries.data(), queries.size(), 5678); + + faiss::IndexFlatIP quantizer(dim); + faiss::IndexIVFFlat cpuIndex( + &quantizer, dim, nlist, faiss::METRIC_INNER_PRODUCT); + cpuIndex.nprobe = nprobe; + cpuIndex.train(nb, vecs.data()); + cpuIndex.add(nb, vecs.data()); + + faiss::gpu_metal::MetalIndexIVFFlat metalIndex(resources_, &cpuIndex); + + std::vector refDist((size_t)nq * k); + std::vector refLab((size_t)nq * k, -1); + std::vector testDist((size_t)nq * k); + std::vector testLab((size_t)nq * k, -1); + + cpuIndex.search(nq, queries.data(), k, refDist.data(), refLab.data()); + metalIndex.search(nq, queries.data(), k, testDist.data(), testLab.data()); + + compareSearchResults( + nq, + k, + refDist.data(), + refLab.data(), + testDist.data(), + testLab.data()); +} + +TEST_F(TestMetalIndexIVFFlat, Reset) { + const int dim = 32; + const int nb = 500; + const int nlist = 4; + + std::vector vecs((size_t)nb * dim); + faiss::float_rand(vecs.data(), vecs.size(), 42); + + faiss::IndexFlatL2 quantizer(dim); + faiss::IndexIVFFlat cpuIndex(&quantizer, dim, nlist); + cpuIndex.train(nb, vecs.data()); + cpuIndex.add(nb, vecs.data()); + + faiss::gpu_metal::MetalIndexIVFFlat metalIndex(resources_, &cpuIndex); + EXPECT_EQ(metalIndex.ntotal, nb); + + metalIndex.reset(); + EXPECT_EQ(metalIndex.ntotal, 0); +} + +TEST_F(TestMetalIndexIVFFlat, CpuToGpuRoundTrip) { + const int dim = 64; + const int nb = 800; + const int nq = 5; + const int nlist = 8; + const int nprobe = 4; + const int k = 5; + + std::vector vecs((size_t)nb * dim); + faiss::float_rand(vecs.data(), vecs.size(), 300); + std::vector queries((size_t)nq * dim); + faiss::float_rand(queries.data(), queries.size(), 301); + + faiss::IndexFlatL2 quantizer(dim); + faiss::IndexIVFFlat cpuOrig(&quantizer, dim, nlist); + cpuOrig.nprobe = nprobe; + cpuOrig.train(nb, vecs.data()); + cpuOrig.add(nb, vecs.data()); + + faiss::gpu_metal::StandardMetalResources res; + std::unique_ptr metalIndex( + faiss::gpu_metal::index_cpu_to_metal_gpu(&res, 0, &cpuOrig)); + ASSERT_NE(metalIndex, nullptr); + EXPECT_EQ(metalIndex->ntotal, nb); + + std::unique_ptr cpuBack( + faiss::gpu_metal::index_metal_gpu_to_cpu(metalIndex.get())); + ASSERT_NE(cpuBack, nullptr); + EXPECT_EQ(cpuBack->ntotal, nb); + + std::vector refDist((size_t)nq * k); + std::vector refLab((size_t)nq * k, -1); + std::vector testDist((size_t)nq * k); + std::vector testLab((size_t)nq * k, -1); + + cpuOrig.search(nq, queries.data(), k, refDist.data(), refLab.data()); + cpuBack->search(nq, queries.data(), k, testDist.data(), testLab.data()); + + compareSearchResults( + nq, + k, + refDist.data(), + refLab.data(), + testDist.data(), + testLab.data()); +} diff --git a/faiss/gpu_metal/test/test_metal_python.py b/faiss/gpu_metal/test/test_metal_python.py new file mode 100644 index 000000000..c01923ac1 --- /dev/null +++ b/faiss/gpu_metal/test/test_metal_python.py @@ -0,0 +1,186 @@ +# 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. + +import numpy as np +import unittest +import faiss + + +class TestMetalFlat(unittest.TestCase): + + def test_get_num_gpus(self): + n = faiss.get_num_gpus() + self.assertGreaterEqual(n, 0) + if n == 0: + self.skipTest("No Metal device") + + def test_standard_gpu_resources(self): + if faiss.get_num_gpus() == 0: + self.skipTest("No Metal device") + res = faiss.StandardGpuResources() + self.assertIsNotNone(res) + + def test_profiler_and_sync(self): + faiss.gpu_profiler_start() + faiss.gpu_sync_all_devices() + faiss.gpu_profiler_stop() + + def test_flat_l2(self): + if faiss.get_num_gpus() == 0: + self.skipTest("No Metal device") + d, nb, nq, k = 64, 1000, 10, 5 + np.random.seed(1234) + xb = np.random.randn(nb, d).astype(np.float32) + xq = np.random.randn(nq, d).astype(np.float32) + + cpu_index = faiss.IndexFlatL2(d) + cpu_index.add(xb) + + res = faiss.StandardGpuResources() + gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) + + D_gpu, I_gpu = gpu_index.search(xq, k) + D_cpu, I_cpu = cpu_index.search(xq, k) + np.testing.assert_allclose(D_gpu, D_cpu, rtol=1e-5, atol=1e-5) + np.testing.assert_array_equal(I_gpu, I_cpu) + + def test_flat_ip(self): + if faiss.get_num_gpus() == 0: + self.skipTest("No Metal device") + d, nb, nq, k = 32, 500, 5, 3 + np.random.seed(5678) + xb = np.random.randn(nb, d).astype(np.float32) + xq = np.random.randn(nq, d).astype(np.float32) + + cpu_index = faiss.IndexFlatIP(d) + cpu_index.add(xb) + + res = faiss.StandardGpuResources() + gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) + + D_gpu, I_gpu = gpu_index.search(xq, k) + D_cpu, I_cpu = cpu_index.search(xq, k) + np.testing.assert_allclose(D_gpu, D_cpu, rtol=1e-5, atol=1e-5) + np.testing.assert_array_equal(I_gpu, I_cpu) + + def test_flat_gpu_to_cpu(self): + if faiss.get_num_gpus() == 0: + self.skipTest("No Metal device") + d, nb, nq, k = 32, 200, 5, 3 + np.random.seed(9999) + xb = np.random.randn(nb, d).astype(np.float32) + xq = np.random.randn(nq, d).astype(np.float32) + + cpu_index = faiss.IndexFlatL2(d) + cpu_index.add(xb) + + res = faiss.StandardGpuResources() + gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) + cpu_back = faiss.index_gpu_to_cpu(gpu_index) + + self.assertEqual(cpu_back.ntotal, nb) + D_orig, I_orig = cpu_index.search(xq, k) + D_back, I_back = cpu_back.search(xq, k) + np.testing.assert_array_almost_equal(D_back, D_orig) + np.testing.assert_array_equal(I_back, I_orig) + + +class TestMetalIVFFlat(unittest.TestCase): + + def test_ivfflat_l2(self): + if faiss.get_num_gpus() == 0: + self.skipTest("No Metal device") + d, nb, nq, k = 64, 1000, 10, 5 + nlist, nprobe = 8, 4 + np.random.seed(1234) + xb = np.random.randn(nb, d).astype(np.float32) + xq = np.random.randn(nq, d).astype(np.float32) + + quantizer = faiss.IndexFlatL2(d) + cpu_index = faiss.IndexIVFFlat(quantizer, d, nlist) + cpu_index.nprobe = nprobe + cpu_index.train(xb) + cpu_index.add(xb) + + res = faiss.StandardGpuResources() + gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) + + D_gpu, I_gpu = gpu_index.search(xq, k) + D_cpu, I_cpu = cpu_index.search(xq, k) + np.testing.assert_allclose(D_gpu, D_cpu, rtol=1e-4, atol=1e-4) + np.testing.assert_array_equal(I_gpu, I_cpu) + + def test_ivfflat_ip(self): + if faiss.get_num_gpus() == 0: + self.skipTest("No Metal device") + d, nb, nq, k = 64, 1000, 10, 5 + nlist, nprobe = 8, 4 + np.random.seed(5678) + xb = np.random.randn(nb, d).astype(np.float32) + xq = np.random.randn(nq, d).astype(np.float32) + + quantizer = faiss.IndexFlatIP(d) + cpu_index = faiss.IndexIVFFlat( + quantizer, d, nlist, faiss.METRIC_INNER_PRODUCT) + cpu_index.nprobe = nprobe + cpu_index.train(xb) + cpu_index.add(xb) + + res = faiss.StandardGpuResources() + gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) + + D_gpu, I_gpu = gpu_index.search(xq, k) + D_cpu, I_cpu = cpu_index.search(xq, k) + np.testing.assert_allclose(D_gpu, D_cpu, rtol=1e-4, atol=1e-4) + np.testing.assert_array_equal(I_gpu, I_cpu) + + def test_ivfflat_reset(self): + if faiss.get_num_gpus() == 0: + self.skipTest("No Metal device") + d, nb = 32, 500 + nlist = 4 + np.random.seed(42) + xb = np.random.randn(nb, d).astype(np.float32) + + quantizer = faiss.IndexFlatL2(d) + cpu_index = faiss.IndexIVFFlat(quantizer, d, nlist) + cpu_index.train(xb) + cpu_index.add(xb) + + res = faiss.StandardGpuResources() + gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) + self.assertEqual(gpu_index.ntotal, nb) + + gpu_index.reset() + self.assertEqual(gpu_index.ntotal, 0) + + def test_ivfflat_gpu_to_cpu_roundtrip(self): + if faiss.get_num_gpus() == 0: + self.skipTest("No Metal device") + d, nb, nq, k = 64, 800, 5, 5 + nlist, nprobe = 8, 4 + np.random.seed(300) + xb = np.random.randn(nb, d).astype(np.float32) + xq = np.random.randn(nq, d).astype(np.float32) + + quantizer = faiss.IndexFlatL2(d) + cpu_index = faiss.IndexIVFFlat(quantizer, d, nlist) + cpu_index.nprobe = nprobe + cpu_index.train(xb) + cpu_index.add(xb) + + res = faiss.StandardGpuResources() + gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) + cpu_back = faiss.index_gpu_to_cpu(gpu_index) + + self.assertEqual(cpu_back.ntotal, nb) + D_orig, I_orig = cpu_index.search(xq, k) + D_back, I_back = cpu_back.search(xq, k) + np.testing.assert_allclose(D_back, D_orig, rtol=1e-5, atol=1e-5) + np.testing.assert_array_equal(I_back, I_orig) + + +if __name__ == "__main__": + unittest.main() diff --git a/faiss/invlists/InvertedLists.cpp b/faiss/invlists/InvertedLists.cpp index 6dff9d5b1..cd0912701 100644 --- a/faiss/invlists/InvertedLists.cpp +++ b/faiss/invlists/InvertedLists.cpp @@ -498,12 +498,12 @@ void ReadOnlyInvertedLists::resize(size_t, size_t) { * HStackInvertedLists implementation ******************************************/ -HStackInvertedLists::HStackInvertedLists(int nil, const InvertedLists** ils_in) +HStackInvertedLists::HStackInvertedLists(int n_il, const InvertedLists** ils_in) : ReadOnlyInvertedLists( - nil > 0 ? ils_in[0]->nlist : 0, - nil > 0 ? ils_in[0]->code_size : 0) { - FAISS_THROW_IF_NOT(nil > 0); - for (int i = 0; i < nil; i++) { + n_il > 0 ? ils_in[0]->nlist : 0, + n_il > 0 ? ils_in[0]->code_size : 0) { + FAISS_THROW_IF_NOT(n_il > 0); + for (int i = 0; i < n_il; i++) { ils.push_back(ils_in[i]); FAISS_THROW_IF_NOT( ils_in[i]->code_size == code_size && ils_in[i]->nlist == nlist); @@ -683,9 +683,9 @@ int translate_list_no(const VStackInvertedLists* vil, idx_t list_no) { return i0; } -idx_t sum_il_sizes(int nil, const InvertedLists** ils_in) { +idx_t sum_il_sizes(int n_il, const InvertedLists** ils_in) { idx_t tot = 0; - for (int i = 0; i < nil; i++) { + for (int i = 0; i < n_il; i++) { tot += ils_in[i]->nlist; } return tot; @@ -693,13 +693,13 @@ idx_t sum_il_sizes(int nil, const InvertedLists** ils_in) { } // namespace -VStackInvertedLists::VStackInvertedLists(int nil, const InvertedLists** ils_in) +VStackInvertedLists::VStackInvertedLists(int n_il, const InvertedLists** ils_in) : ReadOnlyInvertedLists( - sum_il_sizes(nil, ils_in), - nil > 0 ? ils_in[0]->code_size : 0) { - FAISS_THROW_IF_NOT(nil > 0); - cumsz.resize(nil + 1); - for (int i = 0; i < nil; i++) { + sum_il_sizes(n_il, ils_in), + n_il > 0 ? ils_in[0]->code_size : 0) { + FAISS_THROW_IF_NOT(n_il > 0); + cumsz.resize(n_il + 1); + for (int i = 0; i < n_il; i++) { ils.push_back(ils_in[i]); FAISS_THROW_IF_NOT(ils_in[i]->code_size == code_size); cumsz[i + 1] = cumsz[i] + ils_in[i]->nlist; diff --git a/faiss/invlists/InvertedLists.h b/faiss/invlists/InvertedLists.h index f0c7b450d..bcbe9d32a 100644 --- a/faiss/invlists/InvertedLists.h +++ b/faiss/invlists/InvertedLists.h @@ -376,7 +376,7 @@ struct HStackInvertedLists : ReadOnlyInvertedLists { std::vector ils; /// build InvertedLists by concatenating nil of them - HStackInvertedLists(int nil, const InvertedLists** ils); + HStackInvertedLists(int n_il, const InvertedLists** ils); size_t list_size(size_t list_no) const override; const uint8_t* get_codes(size_t list_no) const override; @@ -422,7 +422,7 @@ struct VStackInvertedLists : ReadOnlyInvertedLists { std::vector cumsz; /// build InvertedLists by concatenating nil of them - VStackInvertedLists(int nil, const InvertedLists** ils); + VStackInvertedLists(int n_il, const InvertedLists** ils); size_t list_size(size_t list_no) const override; const uint8_t* get_codes(size_t list_no) const override; diff --git a/faiss/python/CMakeLists.txt b/faiss/python/CMakeLists.txt index 45ca8d4ef..fdad19780 100644 --- a/faiss/python/CMakeLists.txt +++ b/faiss/python/CMakeLists.txt @@ -58,6 +58,11 @@ macro(configure_swigfaiss source) ) endif() endif() + if(FAISS_ENABLE_METAL) + set_property(SOURCE ${source} APPEND PROPERTY + COMPILE_DEFINITIONS FAISS_ENABLE_METAL + ) + endif() endmacro() # CMake's SWIG wrappers only allow tweaking certain settings at source level, so @@ -238,6 +243,16 @@ if(FAISS_ENABLE_GPU) endif() endif() +if(FAISS_ENABLE_METAL) + target_link_libraries(swigfaiss PRIVATE + faiss_metal + "-framework Metal" + "-framework MetalKit" + "-framework MetalPerformanceShaders" + "-framework Foundation" + ) +endif() + find_package(OpenMP REQUIRED) # When scikit-build-core requests abi3 (Stable ABI), it passes diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index 4ebd8b7af..e3d4fc021 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -461,7 +461,45 @@ int64_t cast_cudastream_t_to_integer(cudaStream_t x) { %} #endif // FAISS_ENABLE_ROCM -#else // GPU_WRAPPER +#elif defined(FAISS_ENABLE_METAL) + +%{ +#include + +int get_num_gpus() +{ + return faiss::gpu_metal::get_num_gpus(); +} + +void gpu_profiler_start() +{ + return faiss::gpu_metal::gpu_profiler_start(); +} + +void gpu_profiler_stop() +{ + return faiss::gpu_metal::gpu_profiler_stop(); +} + +void gpu_sync_all_devices() +{ + return faiss::gpu_metal::gpu_sync_all_devices(); +} +%} + +%ignore faiss::gpu_metal::get_num_gpus; +%ignore faiss::gpu_metal::gpu_profiler_start; +%ignore faiss::gpu_metal::gpu_profiler_stop; +%ignore faiss::gpu_metal::gpu_sync_all_devices; + +%rename(StandardGpuResources) faiss::gpu_metal::StandardMetalResourcesHolder; + +%newobject faiss::gpu_metal::index_cpu_to_gpu; +%newobject faiss::gpu_metal::index_gpu_to_cpu; + +%include + +#else // no GPU backend %{ int get_num_gpus() @@ -483,7 +521,7 @@ void gpu_sync_all_devices() %} -#endif // GPU_WRAPPER +#endif // GPU_WRAPPER / FAISS_ENABLE_METAL // order matters because includes are not recursive