From 76c0e01e0f4db5f943b44cf85530e68efe29f355 Mon Sep 17 00:00:00 2001 From: Evan Haque Date: Thu, 1 Oct 2026 21:07:27 -0700 Subject: [PATCH] Add GPU-accelerated IVFPQ search to the Metal backend (#5449) Summary: This is a follow-up to https://github.com/facebookresearch/faiss/issues/5288, which added `MetalIndexIVFPQ` with GPU-resident storage while delegating search to the CPU. Adds GPU-accelerated IVFPQ search to the Metal backend, including optimized scanning through `M <= 64`. - Add Metal kernels for 8-bit PQ distance computation, inverted-list scanning, and top-k merging - Support optimized GPU scanning through `M <= 64` - Compute distance contributions only for PQ codes encountered in each list when `d <= 256` - Add a 32-thread cooperative kernel for lists containing at most 128 vectors - Retain the cached precomputed-LUT kernel for suitable `M <= 16` workloads - Stream precomputed terms for larger dimensions without allocating an `M * 256` threadgroup LUT - Support both L2 and inner product metrics - Add exact segmented top-k selection for inverted lists of any length - Add grouped multi-round merging without a fixed `nprobe * k` candidate limit - Compact invalid per-list padding before merging when lists contain fewer than `k` results - Reuse dynamically sized Metal scratch buffers across searches - Retain the legacy per-probe LUT path with optional FP16 lookup tables - Fall back to CPU search for unsupported configurations, search parameters, or Metal failures ## Changes - **Modified:** `MetalDistance.metal` - on demand PQ scans, large-M precomputed scan, short-list specialization, segmented top-k, and compacting grouped merge - **Modified:** `MetalDistance.h/.mm` - IVFPQ search orchestration and support for `M <= 64` - **Modified:** `MetalKernels.h/.mm` - dispatch for cached, on-demand, short-list, and large-M scan kernels - **Modified:** `MetalIndexIVFPQ.h/.mm` - GPU scan selection, reusable buffers, centroid uploads, precomputed terms, and CPU fallback - **Modified:** `MetalIndex.h` - add the `useFloat16` LUT configuration - **Modified:** `TestMetalIndexIVFPQ.mm` - direct-GPU scan validation and search-parameter fallback coverage ## Differences from CUDA IVFPQ **Training:** `MetalIndexIVFPQ::train` delegates training to its CPU index, as introduced in https://github.com/facebookresearch/faiss/issues/5288. CUDA can train the coarse and product quantizers on GPU. Training is a one-time cost. **Add path:** Coarse assignment, residual computation, and PQ encoding remain on the CPU. The encoded PQ codes are stored in GPU-resident Metal buffers. CUDA performs these operations on GPU. **Coarse quantization:** The coarse quantizer search runs on the CPU. The selected lists and coarse distances are passed to the Metal scan. **Distance computation:** For `d <= 256`, the preferred Metal path evaluates only the PQ codes encountered while scanning a list. This avoids constructing or loading all 256 lookup entries for every subquantizer, which is especially beneficial for short IVF lists. For other supported dimensions, Metal uses the CPU IVFPQ precomputed-table decomposition. The query-independent centroid/PQ term is computed once per trained index and the query term once per batch. **Threadgroup memory:** The existing cached lookup-table kernel remains available for suitable `M <= 16` workloads. Larger `M` values do not expand its `M * 256` threadgroup LUT, which would exceed Metal threadgroup-memory limits. Instead, they use the on-demand or streamed-precomputed scan paths. **List scanning:** Short lists use a 32-thread cooperative kernel with approximately 2 KiB of threadgroup scratch. Longer lists use segmented selection to maintain an exact running top-k. **Top-k merge:** Per-list results are merged in groups over multiple rounds. Invalid padding is compacted when lists contain fewer than `k` results. **Fallback:** The optimized scan supports `M <= 64`, `d / M <= 256`, and `k <= 512`. Unsupported configurations fall back to the legacy Metal path or CPU search. Search options that the GPU path cannot preserve, including selectors, scan budgets, and polysemous filtering, are forwarded to the CPU with their original parameters. ## Build and test ```bash cmake -B build \ -DFAISS_ENABLE_GPU=OFF \ -DFAISS_ENABLE_METAL=ON \ -DBUILD_TESTING=ON \ -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_PREFIX_PATH="$(brew --prefix libomp)" \ . cmake --build build \ --target faiss faiss_metal TestMetalIndexIVFPQ \ -j$(sysctl -n hw.logicalcpu) cd build && ctest -R TestMetalIndexIVFPQ --output-on-failure ``` All six IVFPQ tests pass. Direct-GPU regression coverage exercises 800 combinations across `M=8,16,17,32,48,64`, L2 and inner product, `k` through 512, short and long lists, both merge modes, and 64-bit IDs. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5449 Reviewed By: junjieqi Differential Revision: D113864779 Pulled By: mnorris11 fbshipit-source-id: 741c12b2b014ef28a82569c6b0344352351a7325 --- faiss/gpu_metal/MetalDistance.h | 92 ++ faiss/gpu_metal/MetalDistance.metal | 1378 ++++++++++++++++++- faiss/gpu_metal/MetalDistance.mm | 323 +++++ faiss/gpu_metal/MetalIndex.h | 4 + faiss/gpu_metal/MetalIndexIVFPQ.h | 38 + faiss/gpu_metal/MetalIndexIVFPQ.mm | 401 +++++- faiss/gpu_metal/MetalKernels.h | 93 ++ faiss/gpu_metal/MetalKernels.mm | 212 +++ faiss/gpu_metal/test/TestMetalIndexIVFPQ.mm | 260 ++++ tests/test_metal_ivfpq.py | 103 ++ 10 files changed, 2878 insertions(+), 26 deletions(-) create mode 100644 tests/test_metal_ivfpq.py diff --git a/faiss/gpu_metal/MetalDistance.h b/faiss/gpu_metal/MetalDistance.h index 2785f5d47..ffa9a1ff4 100644 --- a/faiss/gpu_metal/MetalDistance.h +++ b/faiss/gpu_metal/MetalDistance.h @@ -83,5 +83,97 @@ bool runMetalIVFFlatFullSearch( bool centroidsAreFP16 = false, bool waitForCompletion = true); +/// Full IVF-PQ search on the GPU: builds per-(query, probe) PQ lookup tables, +/// scans the assigned inverted lists with 8-bit ADC, and merges the per-list +/// top-k into the final top-k. Coarse quantization is performed on the CPU and +/// the assignments are passed in via @p coarseAssign. @p coarseDist carries +/// the per-(query, probe) coarse scores, added for inner product (the IP +/// tables hold only); unused for L2. +bool runMetalIVFPQFullSearch( + id device, + id queue, + id queries, + id coarseAssign, + id coarseDist, + id coarseCentroids, + id pqCentroids, + id lookupTable, + id codes, + id ids, + id listOffset, + id listLength, + int nq, + int d, + int M, + int k, + int nprobe, + int nlist, + int avgListLen, + bool lookupFp16, + bool isL2, + id outDistances, + id outIndices, + id perListDistBuf, + id perListIdxBuf, + bool waitForCompletion = true); + +/// Build the query-independent precomputed term of the IVFPQ L2 distance +/// decomposition on the GPU: term2[l][m][c] = ||pq[m][c]||^2 + +/// 2 , laid out as nlist * M * 256 floats. Run once +/// per trained index; consumed by runMetalIVFPQPrecompSearch. +bool runMetalIVFPQPrecomputeTerm2( + id device, + id queue, + id coarseCentroids, + id pqCentroids, + id outTerm2, + int nlist, + int d, + int M); + +/// IVF-PQ search using the precomputed-table decomposition: a per-query +/// M*256 term is built once per batch (instead of a full residual LUT per +/// (query, probe)), combined with @p term2 in threadgroup memory by the scan +/// kernel, and the coarse distance is added as the constant ||x - c||^2 term. +/// The scan keeps an exact running top-k over arbitrarily long lists and the +/// merge runs in rounds, so there is no list-length or nprobe*k cap; the +/// remaining requirements are k <= 512, M <= 64 and d/M <= 256. +/// onTheFly evaluates encountered codes directly (d <= 256), using the +/// supplied coarse centroids; shortLists requires every list to have <= 128 +/// entries and selects a smaller threadgroup and scratch allocation. +/// @p term2 may be nil for inner product (qterm is the whole table); +/// @p useDis0 adds the coarse distance (required for L2, by_residual for IP). +bool runMetalIVFPQPrecompSearch( + id device, + id queue, + id queries, + id coarseAssign, + id coarseDist, + id term2, + id qtermScratch, + id pqCentroids, + id codes, + id ids, + id listOffset, + id listLength, + int nq, + int d, + int M, + int k, + int nprobe, + bool isL2, + bool useDis0, + id outDistances, + id outIndices, + id perListDistBuf, + id perListIdxBuf, + id mergeScratchDistBuf, + id mergeScratchIdxBuf, + bool waitForCompletion = true, + id coarseCentroids = nil, + bool onTheFly = false, + bool shortLists = false, + bool compactMerge = true); + } // namespace gpu_metal } // namespace faiss diff --git a/faiss/gpu_metal/MetalDistance.metal b/faiss/gpu_metal/MetalDistance.metal index d2f96de47..2281754c2 100644 --- a/faiss/gpu_metal/MetalDistance.metal +++ b/faiss/gpu_metal/MetalDistance.metal @@ -434,16 +434,13 @@ kernel void l2_with_norms( // 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) { +template +inline bool ivf_better( + float aDist, + IndexT aIdx, + float bDist, + IndexT bIdx, + uint want_min) { constexpr float tie_eps = 1e-6f; float delta = aDist - bDist; bool strictly = want_min ? (delta < -tie_eps) : (delta > tie_eps); @@ -557,7 +554,7 @@ kernel void ivf_scan_list( localDist[pos] = dist; localIdx [pos] = vi; while (pos > 0) { - bool sw = ivf_better_int( + bool sw = ivf_better( localDist[pos], localIdx[pos], localDist[pos - 1], localIdx[pos - 1], want_min); @@ -586,11 +583,11 @@ kernel void ivf_scan_list( uint partner = idx ^ j; if (partner < CAND && partner > idx) { bool ascending = ((idx & k2) == 0); - bool pB = ivf_better_int( + bool pB = ivf_better( tgDist[partner], tgIdx[partner], tgDist[idx], tgIdx[idx], want_min); - bool iB = ivf_better_int( + bool iB = ivf_better( tgDist[idx], tgIdx[idx], tgDist[partner], tgIdx[partner], want_min); @@ -709,7 +706,7 @@ kernel void ivf_scan_list_small( int vi = (int)vecIdx; - bool better = ivf_better_int(dist, vi, bestDist, bestIdx, want_min); + bool better = ivf_better(dist, vi, bestDist, bestIdx, want_min); if (better) { bestDist = dist; bestIdx = vi; @@ -728,11 +725,11 @@ kernel void ivf_scan_list_small( uint partner = tid ^ j; if (partner < CAND && partner > tid) { bool ascending = ((tid & k2) == 0); - bool pB = ivf_better_int( + bool pB = ivf_better( tgDist[partner], tgIdx[partner], tgDist[tid], tgIdx[tid], want_min); - bool iB = ivf_better_int( + bool iB = ivf_better( tgDist[tid], tgIdx[tid], tgDist[partner], tgIdx[partner], want_min); @@ -875,7 +872,7 @@ kernel void ivf_scan_list_interleaved( localDist[pos] = dist; localIdx [pos] = vi; while (pos > 0) { - bool sw = ivf_better_int( + bool sw = ivf_better( localDist[pos], localIdx[pos], localDist[pos - 1], localIdx[pos - 1], want_min); @@ -904,11 +901,11 @@ kernel void ivf_scan_list_interleaved( uint partner = idx ^ j; if (partner < CAND && partner > idx) { bool ascending = ((idx & k2) == 0); - bool pB = ivf_better_int( + bool pB = ivf_better( tgDist[partner], tgIdx[partner], tgDist[idx], tgIdx[idx], want_min); - bool iB = ivf_better_int( + bool iB = ivf_better( tgDist[idx], tgIdx[idx], tgDist[partner], tgIdx[partner], want_min); @@ -982,7 +979,7 @@ kernel void ivf_merge_lists( localDist[pos] = d; localIdx [pos] = v; while (pos > 0) { - bool sw = ivf_better_long( + bool sw = ivf_better( localDist[pos], localIdx[pos], localDist[pos - 1], localIdx[pos - 1], want_min); @@ -1010,11 +1007,11 @@ kernel void ivf_merge_lists( uint partner = idx ^ j; if (partner < CAND && partner > idx) { bool ascending = ((idx & k2) == 0); - bool pB = ivf_better_long( + bool pB = ivf_better( tgDist[partner], tgIdx[partner], tgDist[idx], tgIdx[idx], want_min); - bool iB = ivf_better_long( + bool iB = ivf_better( tgDist[idx], tgIdx[idx], tgDist[partner], tgIdx[partner], want_min); @@ -1077,3 +1074,1340 @@ inline float sq6_decode_component(device const uchar* code, uint i) { return (float(bits) + 0.5f) / 63.0f; } + +// IVF-PQ scan kernel (8-bit codes) +// ============================================================ +// +// One threadgroup per (query, probe). Layout matches ivf_scan_list: +// output is per-list top-k in perListDist/perListIdx. +// +// lookupTable layout: (nq * nprobe) * M * 256 floats +// For entry (qi, probe), offset = (qi * nprobe + probe) * M * 256 +// table[m * 256 + code] gives the partial distance contribution. +// +// params: [nq, M, k, nprobe, want_min] + +// ============================================================ +// IVFPQ LUT build kernels +// ============================================================ +// +// Build lookup table layout: (nq * nprobe * M * 256), contiguous. +// For each (qi, pi, m, c): +// L2: tab = || (q_m - coarse_m) - pq[m][c] ||^2 +// IP: tab = +// +// params: [nq, d, M, nprobe] + +kernel void ivfpq_build_lut_l2( + device const float* queries [[buffer(0)]], + device const int* coarseAssign [[buffer(1)]], + device const float* coarseCentroids [[buffer(2)]], + device const float* pqCentroids [[buffer(3)]], + device float* outLookup [[buffer(4)]], + constant uint* params [[buffer(5)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + uint nq = params[0]; + uint d = params[1]; + uint M = params[2]; + uint nprobe = params[3]; + if (M == 0 || nprobe == 0) return; + uint dsub = d / M; + if (dsub == 0) return; + + uint totalTG = nq * nprobe * M; + if (tgid >= totalTG || tid >= 256) return; + uint m = tgid % M; + uint qp = tgid / M; + uint pi = qp % nprobe; + uint qi = qp / nprobe; + if (qi >= nq) return; + uint c = tid; + + int listNo = coarseAssign[qi * nprobe + pi]; + if (listNo < 0) { + outLookup[(tgid * 256) + c] = 1e38f; + return; + } + + threadgroup float tgResidual[256]; + uint qbase = qi * d + m * dsub; + uint cbase = uint(listNo) * d + m * dsub; + uint pqbase = ((m * 256 + c) * dsub); + for (uint j = tid; j < dsub; j += 256) { + tgResidual[j] = queries[qbase + j] - coarseCentroids[cbase + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float acc = 0.0f; + for (uint j = 0; j < dsub; ++j) { + float diff = tgResidual[j] - pqCentroids[pqbase + j]; + acc += diff * diff; + } + outLookup[(tgid * 256) + c] = acc; +} + +kernel void ivfpq_build_lut_ip( + device const float* queries [[buffer(0)]], + device const int* coarseAssign [[buffer(1)]], + device const float* coarseCentroids [[buffer(2)]], + device const float* pqCentroids [[buffer(3)]], + device float* outLookup [[buffer(4)]], + constant uint* params [[buffer(5)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + (void)coarseAssign; + (void)coarseCentroids; + + uint nq = params[0]; + uint d = params[1]; + uint M = params[2]; + uint nprobe = params[3]; + if (M == 0 || nprobe == 0) return; + uint dsub = d / M; + if (dsub == 0) return; + + uint totalTG = nq * nprobe * M; + if (tgid >= totalTG || tid >= 256) return; + uint m = tgid % M; + uint qp = tgid / M; + uint qi = qp / nprobe; + if (qi >= nq) return; + uint c = tid; + + threadgroup float tgQuery[256]; + uint qbase = qi * d + m * dsub; + uint pqbase = ((m * 256 + c) * dsub); + for (uint j = tid; j < dsub; j += 256) { + tgQuery[j] = queries[qbase + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float acc = 0.0f; + for (uint j = 0; j < dsub; ++j) { + acc += tgQuery[j] * pqCentroids[pqbase + j]; + } + outLookup[(tgid * 256) + c] = acc; +} + +kernel void ivfpq_build_lut_l2_f16( + device const float* queries [[buffer(0)]], + device const int* coarseAssign [[buffer(1)]], + device const float* coarseCentroids [[buffer(2)]], + device const float* pqCentroids [[buffer(3)]], + device half* outLookup [[buffer(4)]], + constant uint* params [[buffer(5)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + uint nq = params[0]; + uint d = params[1]; + uint M = params[2]; + uint nprobe = params[3]; + if (M == 0 || nprobe == 0) return; + uint dsub = d / M; + if (dsub == 0) return; + + uint totalTG = nq * nprobe * M; + if (tgid >= totalTG || tid >= 256) return; + uint m = tgid % M; + uint qp = tgid / M; + uint pi = qp % nprobe; + uint qi = qp / nprobe; + if (qi >= nq) return; + uint c = tid; + + int listNo = coarseAssign[qi * nprobe + pi]; + if (listNo < 0) { + outLookup[(tgid * 256) + c] = half(65504.0f); + return; + } + + threadgroup float tgResidual[256]; + uint qbase = qi * d + m * dsub; + uint cbase = uint(listNo) * d + m * dsub; + uint pqbase = ((m * 256 + c) * dsub); + for (uint j = tid; j < dsub; j += 256) { + tgResidual[j] = queries[qbase + j] - coarseCentroids[cbase + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float acc = 0.0f; + for (uint j = 0; j < dsub; ++j) { + float diff = tgResidual[j] - pqCentroids[pqbase + j]; + acc += diff * diff; + } + outLookup[(tgid * 256) + c] = half(acc); +} + +kernel void ivfpq_build_lut_ip_f16( + device const float* queries [[buffer(0)]], + device const int* coarseAssign [[buffer(1)]], + device const float* coarseCentroids [[buffer(2)]], + device const float* pqCentroids [[buffer(3)]], + device half* outLookup [[buffer(4)]], + constant uint* params [[buffer(5)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + (void)coarseAssign; + (void)coarseCentroids; + + uint nq = params[0]; + uint d = params[1]; + uint M = params[2]; + uint nprobe = params[3]; + if (M == 0 || nprobe == 0) return; + uint dsub = d / M; + if (dsub == 0) return; + + uint totalTG = nq * nprobe * M; + if (tgid >= totalTG || tid >= 256) return; + uint m = tgid % M; + uint qp = tgid / M; + uint qi = qp / nprobe; + if (qi >= nq) return; + uint c = tid; + + threadgroup float tgQuery[256]; + uint qbase = qi * d + m * dsub; + uint pqbase = ((m * 256 + c) * dsub); + for (uint j = tid; j < dsub; j += 256) { + tgQuery[j] = queries[qbase + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float acc = 0.0f; + for (uint j = 0; j < dsub; ++j) { + acc += tgQuery[j] * pqCentroids[pqbase + j]; + } + outLookup[(tgid * 256) + c] = half(acc); +} + +kernel void ivf_scan_list_pq8( + device const float* lookupTable [[buffer(0)]], + device const uchar* 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 float* coarseDist [[buffer(9)]], + 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 M = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + uint useDis0 = params[5]; + + 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 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] = -1L; + } + return; + } + + uint tableOff = (qi * nprobe + pi) * M * 256; + + // Inner product needs the coarse score: the table holds only. + // L2 tables hold the complete residual distance, so dis0 stays 0. + float dis0 = useDis0 ? coarseDist[qi * nprobe + pi] : 0.0f; + + 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 li = tid; li < lLen; li += TG_SIZE) { + uint vecIdx = lOff + li; + device const uchar* cvec = codes + vecIdx * M; + float dist = dis0; + for (uint m = 0; m < M; m++) { + dist += lookupTable[tableOff + m * 256 + uint(cvec[m])]; + } + + 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 = want_min + ? (localDist[pos] < localDist[pos-1] || (localDist[pos] == localDist[pos-1] && localIdx[pos] < localIdx[pos-1])) + : (localDist[pos] > localDist[pos-1] || (localDist[pos] == localDist[pos-1] && localIdx[pos] < localIdx[pos-1])); + 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; + 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 = 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 iB = 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])); + 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; + } +} + +kernel void ivf_scan_list_pq8_f16( + device const half* lookupTable [[buffer(0)]], + device const uchar* 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 float* coarseDist [[buffer(9)]], + 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 M = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + uint useDis0 = params[5]; + + 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 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] = -1L; + } + return; + } + + uint tableOff = (qi * nprobe + pi) * M * 256; + + // Inner product needs the coarse score: the table holds only. + // L2 tables hold the complete residual distance, so dis0 stays 0. + float dis0 = useDis0 ? coarseDist[qi * nprobe + pi] : 0.0f; + + 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 li = tid; li < lLen; li += TG_SIZE) { + uint vecIdx = lOff + li; + device const uchar* cvec = codes + vecIdx * M; + float dist = dis0; + for (uint m = 0; m < M; m++) { + dist += float(lookupTable[tableOff + m * 256 + uint(cvec[m])]); + } + + 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 = want_min + ? (localDist[pos] < localDist[pos - 1] || + (localDist[pos] == localDist[pos - 1] && + localIdx[pos] < localIdx[pos - 1])) + : (localDist[pos] > localDist[pos - 1] || + (localDist[pos] == localDist[pos - 1] && + localIdx[pos] < localIdx[pos - 1])); + 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; + 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 = 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 iB = 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])); + 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; + } +} + +// Small-list IVFPQ variant: lower threadgroup/barrier cost when average list +// length is short. +kernel void ivf_scan_list_pq8_small( + device const float* lookupTable [[buffer(0)]], + device const uchar* 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 float* coarseDist [[buffer(9)]], + 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 M = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + uint useDis0 = params[5]; + + 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 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] = -1L; + } + return; + } + + uint tableOff = (qi * nprobe + pi) * M * 256; + + // Inner product needs the coarse score: the table holds only. + // L2 tables hold the complete residual distance, so dis0 stays 0. + float dis0 = useDis0 ? coarseDist[qi * nprobe + pi] : 0.0f; + + float bestDist = sentinel; + int bestIdx = -1; + for (uint li = tid; li < lLen; li += TG_SIZE) { + uint vecIdx = lOff + li; + device const uchar* cvec = codes + vecIdx * M; + float dist = dis0; + for (uint m = 0; m < M; m++) { + dist += lookupTable[tableOff + m * 256 + uint(cvec[m])]; + } + + int vi = (int)vecIdx; + bool better = want_min + ? (dist < bestDist || (dist == bestDist && vi < bestIdx)) + : (dist > bestDist || (dist == bestDist && vi < bestIdx)); + if (better) { + bestDist = dist; + bestIdx = vi; + } + } + + constexpr uint CAND = TG_SIZE * LOCAL_K; + 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 = want_min + ? (tgDist[partner] < tgDist[tid] || + (tgDist[partner] == tgDist[tid] && + tgIdx[partner] < tgIdx[tid])) + : (tgDist[partner] > tgDist[tid] || + (tgDist[partner] == tgDist[tid] && + tgIdx[partner] < tgIdx[tid])); + bool iB = want_min + ? (tgDist[tid] < tgDist[partner] || + (tgDist[tid] == tgDist[partner] && + tgIdx[tid] < tgIdx[partner])) + : (tgDist[tid] > tgDist[partner] || + (tgDist[tid] == tgDist[partner] && + tgIdx[tid] < tgIdx[partner])); + 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; + } +} + +// ============================================================ +// IVF-PQ precomputed-table search path +// ============================================================ +// +// Mirrors faiss::IndexIVFPQ use_precomputed_table for L2/by_residual: +// +// ||x - (c + pq)||^2 = ||x - c||^2 (coarse distance, dis0) +// + (||pq||^2 + 2) (term2: per index, once) +// - 2 (qterm: per query, once) +// +// term2 is built once per index (nlist * M * 256), qterm once per query +// batch (nq * M * 256). The scan kernel adds the two tables into threadgroup +// memory per (query, probe) and streams the list codes against it, so no +// per-(query, probe) lookup table is ever written to device memory. +// +// For inner product the score is + : qterm holds and +// dis0 is the coarse score (term2 unused). +// +// The scan keeps an exact running top-k over segments of PQPRE_SEG codes: +// per segment, candidates that beat the current k-th best are gathered via a +// threadgroup atomic and re-merged with a bitonic sort sized to the actual +// candidate count. Exact for any list length; requires k <= PQPRE_MAX_K. + +constant constexpr uint PQPRE_TG = 256; +constant constexpr uint PQPRE_SEG = 512; +constant constexpr uint PQPRE_MAX_K = 512; +constant constexpr uint PQPRE_CAND = PQPRE_MAX_K + PQPRE_SEG; // 1024 +constant constexpr uint PQPRE_LUT_MAX = 16 * 256; // M <= 16 + +// Bitonic sort over tgDist/tgIdx[0..n), n a power of two. Sorted best-first. +template +inline void ivf_tg_bitonic( + threadgroup float* tgDist, + threadgroup IndexT* tgIdx, + uint n, + uint tid, + uint want_min +) { + for (uint k2 = 2; k2 <= n; k2 *= 2) { + for (uint j = k2 >> 1; j > 0; j >>= 1) { + for (uint idx = tid; idx < n; idx += ThreadgroupSize) { + uint partner = idx ^ j; + if (partner < n && partner > idx) { + bool ascending = ((idx & k2) == 0); + bool pB = ivf_better( + tgDist[partner], tgIdx[partner], + tgDist[idx], tgIdx[idx], want_min); + bool iB = ivf_better( + 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; + IndexT ti = tgIdx[idx]; + tgIdx[idx] = tgIdx[partner]; + tgIdx[partner] = ti; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + } +} + +inline uint ivf_next_pow2(uint v) { + uint p = 1; + while (p < v) p <<= 1; + return p; +} + +// term2[(l * M + m) * 256 + c] = ||pq[m][c]||^2 + 2 +// One threadgroup per (l, m), one thread per code c. Run once per index. +// params: [nlist, d, M] +kernel void ivfpq_precompute_term2( + device const float* coarseCentroids [[buffer(0)]], + device const float* pqCentroids [[buffer(1)]], + device float* outTerm2 [[buffer(2)]], + constant uint* params [[buffer(3)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + uint nlist = params[0]; + uint d = params[1]; + uint M = params[2]; + if (M == 0) return; + uint dsub = d / M; + if (dsub == 0 || dsub > 256) return; + + uint m = tgid % M; + uint l = tgid / M; + if (l >= nlist || tid >= 256) return; + uint c = tid; + + threadgroup float tgC[256]; + for (uint j = tid; j < dsub; j += 256) { + tgC[j] = coarseCentroids[l * d + m * dsub + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint pqbase = (m * 256 + c) * dsub; + float acc = 0.0f; + for (uint j = 0; j < dsub; ++j) { + float p = pqCentroids[pqbase + j]; + acc += p * p + 2.0f * tgC[j] * p; + } + outTerm2[(l * M + m) * 256 + c] = acc; +} + +// qterm[(qi * M + m) * 256 + c] = -2 (L2) +// = (IP) +// One threadgroup per (qi, m), one thread per code c. Run once per batch. +// params: [nq, d, M, isL2] +kernel void ivfpq_build_query_term( + device const float* queries [[buffer(0)]], + device const float* pqCentroids [[buffer(1)]], + device float* outQTerm [[buffer(2)]], + constant uint* params [[buffer(3)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + uint nq = params[0]; + uint d = params[1]; + uint M = params[2]; + uint isL2 = params[3]; + if (M == 0) return; + uint dsub = d / M; + if (dsub == 0 || dsub > 256) return; + + uint m = tgid % M; + uint qi = tgid / M; + if (qi >= nq || tid >= 256) return; + uint c = tid; + + threadgroup float tgQ[256]; + for (uint j = tid; j < dsub; j += 256) { + tgQ[j] = queries[qi * d + m * dsub + j]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint pqbase = (m * 256 + c) * dsub; + float acc = 0.0f; + for (uint j = 0; j < dsub; ++j) { + acc += tgQ[j] * pqCentroids[pqbase + j]; + } + outQTerm[(qi * M + m) * 256 + c] = isL2 ? (-2.0f * acc) : acc; +} + +// Scan with precomputed tables. One threadgroup per (query, probe). +// The M*256 lookup table is assembled in threadgroup memory as +// qterm[qi] (+ term2[list_no] when useTerm2), then the list is streamed in +// segments of PQPRE_SEG codes. A running, sorted top-k is kept in tgDist[0..k); +// per segment only candidates beating the current k-th best are gathered +// (threadgroup atomic) and re-merged with a bitonic sort sized to the actual +// candidate count. Exact for any list length when k <= PQPRE_MAX_K. +// params: [nq, M, k, nprobe, want_min, useTerm2, useDis0] +kernel void ivf_scan_list_pq8_precomp( + device const float* term2 [[buffer(0)]], + device const float* qterm [[buffer(1)]], + device const float* coarseDist [[buffer(2)]], + device const uchar* codes [[buffer(3)]], + device const long* ids [[buffer(4)]], + device const uint* listOffset [[buffer(5)]], + device const uint* listLength [[buffer(6)]], + device const int* coarseAssign [[buffer(7)]], + device float* perListDist [[buffer(8)]], + device long* perListIdx [[buffer(9)]], + constant uint* params [[buffer(10)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + uint nq = params[0]; + uint M = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + uint useTerm2 = params[5]; + uint useDis0 = params[6]; + + 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; + if (k > PQPRE_MAX_K || M > PQPRE_LUT_MAX / 256) { + for (uint i = tid; i < k; i += PQPRE_TG) { + perListDist[outBase + i] = sentinel; + perListIdx[outBase + i] = -1L; + } + return; + } + + int list_no = coarseAssign[qi * nprobe + pi]; + uint lLen = (list_no < 0) ? 0u : listLength[(uint)list_no]; + if (lLen == 0) { + for (uint i = tid; i < k; i += PQPRE_TG) { + perListDist[outBase + i] = sentinel; + perListIdx [outBase + i] = -1L; + } + return; + } + + threadgroup float tgLUT[PQPRE_LUT_MAX]; + threadgroup float tgDist[PQPRE_CAND]; + threadgroup int tgIdx [PQPRE_CAND]; + threadgroup float tgThresh; + threadgroup atomic_uint tgCount; + + // Assemble the per-(query, probe) lookup table in threadgroup memory. + uint lutN = M * 256; + uint qBase = qi * lutN; + uint t2Base = (uint)list_no * lutN; + for (uint i = tid; i < lutN; i += PQPRE_TG) { + tgLUT[i] = qterm[qBase + i] + (useTerm2 ? term2[t2Base + i] : 0.0f); + } + if (tid == 0) { + tgThresh = sentinel; + atomic_store_explicit(&tgCount, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float dis0 = useDis0 ? coarseDist[qi * nprobe + pi] : 0.0f; + uint lOff = listOffset[(uint)list_no]; + uint kk = 0; // current size of the sorted running top-k (uniform) + + for (uint segStart = 0; segStart < lLen; segStart += PQPRE_SEG) { + uint segN = min(PQPRE_SEG, lLen - segStart); + float thresh = tgThresh; + bool haveK = (kk >= k); + + for (uint li = tid; li < segN; li += PQPRE_TG) { + uint vecIdx = lOff + segStart + li; + device const uchar* cv = codes + (ulong)vecIdx * M; + float dist = dis0; + for (uint m = 0; m < M; ++m) { + dist += tgLUT[m * 256 + (uint)cv[m]]; + } + bool accept = !haveK || + (want_min ? (dist < thresh) : (dist > thresh)); + if (accept) { + uint slot = atomic_fetch_add_explicit( + &tgCount, 1u, memory_order_relaxed); + tgDist[kk + slot] = dist; + tgIdx [kk + slot] = (int)vecIdx; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint count = atomic_load_explicit(&tgCount, memory_order_relaxed); + if (count > 0) { + uint total = kk + count; + uint paddedN = ivf_next_pow2(total); + for (uint i = total + tid; i < paddedN; i += PQPRE_TG) { + tgDist[i] = sentinel; + tgIdx [i] = -1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + ivf_tg_bitonic(tgDist, tgIdx, paddedN, tid, want_min); + kk = min(k, total); + } + if (tid == 0) { + atomic_store_explicit(&tgCount, 0u, memory_order_relaxed); + if (kk >= k) + tgThresh = tgDist[k - 1]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (uint i = tid; i < k; i += PQPRE_TG) { + if (i < kk) { + int vi = tgIdx[i]; + perListDist[outBase + i] = tgDist[i]; + perListIdx [outBase + i] = (vi < 0) ? -1L : ids[vi]; + } else { + perListDist[outBase + i] = sentinel; + perListIdx [outBase + i] = -1L; + } + } +} + +// Large-dimensional fallback: stream the precomputed terms directly rather +// than allocating an M*256 threadgroup LUT. Selection remains segmented/exact. +kernel void ivf_scan_list_pq8_precomp_large( + device const float* term2 [[buffer(0)]], + device const float* qterm [[buffer(1)]], + device const float* coarseDist [[buffer(2)]], + device const uchar* codes [[buffer(3)]], + device const long* ids [[buffer(4)]], + device const uint* listOffset [[buffer(5)]], + device const uint* listLength [[buffer(6)]], + device const int* coarseAssign [[buffer(7)]], + device float* perListDist [[buffer(8)]], + device long* perListIdx [[buffer(9)]], + constant uint* params [[buffer(10)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + uint nq = params[0]; + uint M = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + uint useTerm2 = params[5]; + uint useDis0 = params[6]; + + 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; + if (k > PQPRE_MAX_K || M > 64) { + for (uint i = tid; i < k; i += PQPRE_TG) { + perListDist[outBase + i] = sentinel; + perListIdx[outBase + i] = -1L; + } + return; + } + + int list_no = coarseAssign[qi * nprobe + pi]; + uint lLen = (list_no < 0) ? 0u : listLength[(uint)list_no]; + if (lLen == 0) { + for (uint i = tid; i < k; i += PQPRE_TG) { + perListDist[outBase + i] = sentinel; + perListIdx[outBase + i] = -1L; + } + return; + } + + threadgroup float tgDist[PQPRE_CAND]; + threadgroup int tgIdx[PQPRE_CAND]; + threadgroup float tgThresh; + threadgroup atomic_uint tgCount; + + // Address the query and centroid terms without a threadgroup LUT. + uint lutN = M * 256; + uint qBase = qi * lutN; + uint t2Base = (uint)list_no * lutN; + if (tid == 0) { + tgThresh = sentinel; + atomic_store_explicit(&tgCount, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float dis0 = useDis0 ? coarseDist[qi * nprobe + pi] : 0.0f; + uint lOff = listOffset[(uint)list_no]; + uint kk = 0; // current size of the sorted running top-k (uniform) + + for (uint segStart = 0; segStart < lLen; segStart += PQPRE_SEG) { + uint segN = min(PQPRE_SEG, lLen - segStart); + float thresh = tgThresh; + bool haveK = (kk >= k); + + for (uint li = tid; li < segN; li += PQPRE_TG) { + uint vecIdx = lOff + segStart + li; + device const uchar* cv = codes + (ulong)vecIdx * M; + float dist = dis0; + for (uint m = 0; m < M; ++m) { + uint idx = m * 256 + (uint)cv[m]; + dist += qterm[qBase + idx] + + (useTerm2 ? term2[t2Base + idx] : 0.0f); + } + bool accept = + !haveK || (want_min ? (dist < thresh) : (dist > thresh)); + if (accept) { + uint slot = atomic_fetch_add_explicit( + &tgCount, 1u, memory_order_relaxed); + tgDist[kk + slot] = dist; + tgIdx[kk + slot] = (int)vecIdx; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint count = atomic_load_explicit(&tgCount, memory_order_relaxed); + if (count > 0) { + uint total = kk + count; + uint paddedN = ivf_next_pow2(total); + for (uint i = total + tid; i < paddedN; i += PQPRE_TG) { + tgDist[i] = sentinel; + tgIdx[i] = -1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + ivf_tg_bitonic(tgDist, tgIdx, paddedN, tid, want_min); + kk = min(k, total); + } + if (tid == 0) { + atomic_store_explicit(&tgCount, 0u, memory_order_relaxed); + if (kk >= k) + tgThresh = tgDist[k - 1]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (uint i = tid; i < k; i += PQPRE_TG) { + if (i < kk) { + int vi = tgIdx[i]; + perListDist[outBase + i] = tgDist[i]; + perListIdx[outBase + i] = (vi < 0) ? -1L : ids[vi]; + } else { + perListDist[outBase + i] = sentinel; + perListIdx[outBase + i] = -1L; + } + } +} + +// For d <= 256, the codebook is small enough to reuse through the device cache. +// Evaluate only encountered codes: dis0 + sum(p*p - 2*(q-c)*p) for L2, +// or dis0 + sum(q*p) for IP. No M*256 table is needed, regardless of M. +// Buffers 0/1 contain coarse centroids/queries, not precomputed terms. +// This variant handles arbitrarily long lists in 512-code segments. +kernel void ivf_scan_list_pq8_onthefly( + device const float* coarseCentroids [[buffer(0)]], + device const float* queries [[buffer(1)]], + device const float* coarseDist [[buffer(2)]], + device const uchar* codes [[buffer(3)]], + device const long* ids [[buffer(4)]], + device const uint* listOffset [[buffer(5)]], + device const uint* listLength [[buffer(6)]], + device const int* coarseAssign [[buffer(7)]], + device float* perListDist [[buffer(8)]], + device long* perListIdx [[buffer(9)]], + constant uint* params [[buffer(10)]], + device const float* pqCentroids [[buffer(11)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + uint nq = params[0]; + uint M = params[1]; + uint k = params[2]; + uint nprobe = params[3]; + uint want_min = params[4]; + uint isL2 = params[5]; + uint useDis0 = params[6]; + + 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; + if (k > PQPRE_MAX_K || M == 0 || M > 64 || params[7] == 0 || + params[7] > 256 || params[7] % M != 0) { + for (uint i = tid; i < k; i += PQPRE_TG) { + perListDist[outBase + i] = sentinel; + perListIdx[outBase + i] = -1L; + } + return; + } + + int list_no = coarseAssign[qi * nprobe + pi]; + uint lLen = (list_no < 0) ? 0u : listLength[(uint)list_no]; + if (lLen == 0) { + for (uint i = tid; i < k; i += PQPRE_TG) { + perListDist[outBase + i] = sentinel; + perListIdx[outBase + i] = -1L; + } + return; + } + + threadgroup float tgDist[PQPRE_CAND]; + threadgroup int tgIdx[PQPRE_CAND]; + threadgroup float tgThresh; + threadgroup atomic_uint tgCount; + + uint d = params[7]; + uint dsub = d / M; + threadgroup float residual[256]; + for (uint j = tid; j < d; j += PQPRE_TG) + residual[j] = queries[qi * d + j] - + (isL2 ? coarseCentroids[(uint)list_no * d + j] : 0.0f); + if (tid == 0) { + tgThresh = sentinel; + atomic_store_explicit(&tgCount, 0u, memory_order_relaxed); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + float dis0 = useDis0 ? coarseDist[qi * nprobe + pi] : 0.0f; + uint lOff = listOffset[(uint)list_no]; + uint kk = 0; // current size of the sorted running top-k (uniform) + + for (uint segStart = 0; segStart < lLen; segStart += PQPRE_SEG) { + uint segN = min(PQPRE_SEG, lLen - segStart); + float thresh = tgThresh; + bool haveK = (kk >= k); + + for (uint li = tid; li < segN; li += PQPRE_TG) { + uint vecIdx = lOff + segStart + li; + device const uchar* cv = codes + (ulong)vecIdx * M; + float dist = dis0; + for (uint m = 0; m < M; ++m) { + uint base = (m * 256 + (uint)cv[m]) * dsub; + for (uint j = 0; j < dsub; ++j) { + float p = pqCentroids[base + j]; + dist += isL2 ? p * p - 2.0f * residual[m * dsub + j] * p + : residual[m * dsub + j] * p; + } + } + bool accept = + !haveK || (want_min ? (dist < thresh) : (dist > thresh)); + if (accept) { + uint slot = atomic_fetch_add_explicit( + &tgCount, 1u, memory_order_relaxed); + tgDist[kk + slot] = dist; + tgIdx[kk + slot] = (int)vecIdx; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint count = atomic_load_explicit(&tgCount, memory_order_relaxed); + if (count > 0) { + uint total = kk + count; + uint paddedN = ivf_next_pow2(total); + for (uint i = total + tid; i < paddedN; i += PQPRE_TG) { + tgDist[i] = sentinel; + tgIdx[i] = -1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + ivf_tg_bitonic(tgDist, tgIdx, paddedN, tid, want_min); + kk = min(k, total); + } + if (tid == 0) { + atomic_store_explicit(&tgCount, 0u, memory_order_relaxed); + if (kk >= k) + tgThresh = tgDist[k - 1]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + for (uint i = tid; i < k; i += PQPRE_TG) { + if (i < kk) { + int vi = tgIdx[i]; + perListDist[outBase + i] = tgDist[i]; + perListIdx[outBase + i] = (vi < 0) ? -1L : ids[vi]; + } else { + perListDist[outBase + i] = sentinel; + perListIdx[outBase + i] = -1L; + } + } +} + +// All lists must have <= 128 codes (checked by the host). Four vectors share +// a 32-thread group, eight lanes per vector, to parallelize over subquantizers. +// The residual and sorting scratch occupy only 2 KiB of threadgroup memory. +// Buffers 0/1 contain coarse centroids/queries; params[7] is d (<= 256). +kernel void ivf_scan_list_pq8_onthefly_small( + device const float* coarseCentroids [[buffer(0)]], + device const float* queries [[buffer(1)]], + device const float* coarseDist [[buffer(2)]], + device const uchar* codes [[buffer(3)]], + device const long* ids [[buffer(4)]], + device const uint* listOffset [[buffer(5)]], + device const uint* listLength [[buffer(6)]], + device const int* coarseAssign [[buffer(7)]], + device float* perListDist [[buffer(8)]], + device long* perListIdx [[buffer(9)]], + constant uint* params [[buffer(10)]], + device const float* pqCentroids [[buffer(11)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) { + uint M = params[1], k = params[2], nprobe = params[3], want_min = params[4]; + uint isL2 = params[5], useDis0 = params[6], qi = tgid / nprobe; + if (qi >= params[0] || k == 0) + return; + int list = coarseAssign[tgid]; + uint len = list < 0 ? 0u : listLength[list]; + float sentinel = want_min ? 1e38f : -1e38f; + threadgroup float ds[128]; + threadgroup int ix[128]; + if (len > 128 || M > 64) + len = 0; + uint off = list < 0 ? 0u : listOffset[list]; + float dis0 = useDis0 ? coarseDist[tgid] : 0.0f; + uint d = params[7]; + if (M == 0 || M > 64 || k > PQPRE_MAX_K || d == 0 || d > 256 || + d % M != 0) { + for (uint i = tid; i < k; i += 32) { + perListDist[tgid * k + i] = sentinel; + perListIdx[tgid * k + i] = -1L; + } + return; + } + uint dsub = d / M; + threadgroup float residual[256]; + for (uint j = tid; j < d; j += 32) + residual[j] = queries[qi * d + j] - + (isL2 && list >= 0 ? coarseCentroids[uint(list) * d + j] + : 0.0f); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint li = tid / 8; li < ((len + 3) / 4) * 4; li += 4) { + uint vec = off + li; + device const uchar* cv = codes + (ulong)vec * M; + float dist = 0; + for (uint m = tid % 8; li < len && m < M; m += 8) { + uint base = (m * 256 + uint(cv[m])) * dsub; + for (uint j = 0; j < dsub; ++j) { + float p = pqCentroids[base + j]; + dist += isL2 ? p * p - 2.0f * residual[m * dsub + j] * p + : residual[m * dsub + j] * p; + } + } + dist += simd_shuffle_down(dist, 4); + dist += simd_shuffle_down(dist, 2); + dist += simd_shuffle_down(dist, 1); + if (li < len && tid % 8 == 0) { + ds[li] = dist + dis0; + ix[li] = int(vec); + } + } + uint padded = ivf_next_pow2(len); + for (uint i = len + tid; i < padded; i += 32) { + ds[i] = sentinel; + ix[i] = -1; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + ivf_tg_bitonic(ds, ix, padded, tid, want_min); + for (uint i = tid; i < k; i += 32) { + perListDist[tgid * k + i] = i < len ? ds[i] : sentinel; + perListIdx[tgid * k + i] = i < len ? ids[ix[i]] : -1L; + } +} + +// Merge numLists per-query candidate lists (k entries each, best-first or +// sentinel-padded) in groups of groupSize lists per threadgroup. Output is +// nGroups = ceil(numLists / groupSize) lists of k per query. Run repeatedly +// (ping-pong) until one list remains; exact for any nprobe * k as long as +// groupSize * k <= MERGE_GROUP_CAND. +// params: [nq, numLists, groupSize, k, want_min, compact] +constant constexpr uint MERGE_GROUP_CAND = 2048; + +kernel void ivf_merge_lists_grouped( + device const float* inDist [[buffer(0)]], + device const long* inIdx [[buffer(1)]], + device float* outDist [[buffer(2)]], + device long* outIdx [[buffer(3)]], + constant uint* params [[buffer(4)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]] +) { + uint nq = params[0]; + uint numLists = params[1]; + uint groupSize = params[2]; + uint k = params[3]; + uint want_min = params[4]; + if (numLists == 0 || groupSize == 0 || k == 0) + return; + if (groupSize * k > MERGE_GROUP_CAND) + return; + + uint nGroups = (numLists + groupSize - 1) / groupSize; + uint qi = tgid / nGroups; + uint g = tgid % nGroups; + if (qi >= nq) return; + + float sentinel = want_min ? 1e38f : -1e38f; + uint listStart = g * groupSize; + uint nLists = min(groupSize, numLists - listStart); + uint count = nLists * k; + + threadgroup float tgDist[MERGE_GROUP_CAND]; + threadgroup long tgIdx [MERGE_GROUP_CAND]; + threadgroup atomic_uint validCount; + + uint inBase = (qi * numLists + listStart) * k; + if (params[5]) { + // Short IVF lists often contribute far fewer than k results. Compact + // away their padding before sorting, especially for large k. + if (tid == 0) + atomic_store_explicit(&validCount, 0u, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = tid; i < count; i += PQPRE_TG) { + long id = inIdx[inBase + i]; + if (id >= 0) { + uint slot = atomic_fetch_add_explicit( + &validCount, 1u, memory_order_relaxed); + tgDist[slot] = inDist[inBase + i]; + tgIdx[slot] = id; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + count = atomic_load_explicit(&validCount, memory_order_relaxed); + } else { + for (uint i = tid; i < count; i += PQPRE_TG) { + tgDist[i] = inDist[inBase + i]; + tgIdx[i] = inIdx[inBase + i]; + } + } + uint paddedN = ivf_next_pow2(count); + for (uint i = count + tid; i < paddedN; i += PQPRE_TG) { + tgDist[i] = sentinel; + tgIdx [i] = -1L; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + ivf_tg_bitonic(tgDist, tgIdx, paddedN, tid, want_min); + + uint outBase = (qi * nGroups + g) * k; + for (uint i = tid; i < k; i += PQPRE_TG) { + outDist[outBase + i] = i < count ? tgDist[i] : sentinel; + outIdx [outBase + i] = i < count ? tgIdx[i] : -1L; + } +} diff --git a/faiss/gpu_metal/MetalDistance.mm b/faiss/gpu_metal/MetalDistance.mm index a17e9d9ae..4540523aa 100644 --- a/faiss/gpu_metal/MetalDistance.mm +++ b/faiss/gpu_metal/MetalDistance.mm @@ -355,5 +355,328 @@ bool runMetalIVFFlatFullSearch( return cmdBuf.status == MTLCommandBufferStatusCompleted; } +// ============================================================ +// runMetalIVFPQFullSearch +// ============================================================ + +bool runMetalIVFPQFullSearch( + id device, + id queue, + id queries, + id coarseAssign, + id coarseDist, + id coarseCentroids, + id pqCentroids, + id lookupTable, + id codes, + id ids, + id listOffset, + id listLength, + int nq, + int d, + int M, + int k, + int nprobe, + int nlist, + int avgListLen, + bool lookupFp16, + bool isL2, + id outDistances, + id outIndices, + id perListDistBuf, + id perListIdxBuf, + bool waitForCompletion) { + if (!device || !queue || !queries || !coarseAssign || !coarseDist || + !pqCentroids || !lookupTable || !codes || !ids || !listOffset || + !listLength || !outDistances || !outIndices || !perListDistBuf || + !perListIdxBuf) { + return false; + } + if (nq <= 0 || d <= 0 || M <= 0 || k <= 0 || nprobe <= 0 || nlist <= 0) + return false; + // LUT build stages subvectors in 256-wide threadgroup scratch: + // d / M > 256 (degenerate super-quantized shape) is unsupported. + if (isL2 && !coarseCentroids) + return false; + if (!fitsU32(nq) || !fitsU32(M) || !fitsU32(k) || !fitsU32(nprobe)) + return false; + + // The scan and merge kernels retain only kIvfReduceExactCandidates + // candidates (TG_SIZE * LOCAL_K). The result is exact only when both the + // merge input (nprobe * k) and every scanned list fit within that bound; + // otherwise fall back to the caller's CPU path. + if (!ivfMergeExactnessHolds(nprobe, k)) + return false; + if (!ivfScanExactnessHoldsForAllLists( + listLength, nlist, kIvfReduceExactCandidates)) { + return false; + } + + MetalKernels& K = getMetalKernels(device); + if (!K.isValid()) + return false; + + // params[5] selects the coarse-score term: inner-product tables hold + // only, so the per-(query, probe) coarse score is added + // during the scan. L2 tables hold the complete residual distance. + uint32_t sp[6] = { + (uint32_t)nq, + (uint32_t)M, + (uint32_t)k, + (uint32_t)nprobe, + isL2 ? 1u : 0u, + isL2 ? 0u : 1u}; + id paramsBuf = + [device newBufferWithBytes:sp + length:sizeof(sp) + options:MTLResourceStorageModeShared]; + if (!paramsBuf) + return false; + + id cmdBuf = [queue commandBuffer]; + id enc = [cmdBuf computeCommandEncoder]; + + // Step 1: build per-(query, probe) PQ lookup tables. + K.encodeIVFPQBuildLookupTables( + enc, + isL2, + lookupFp16, + queries, + coarseAssign, + coarseCentroids, + pqCentroids, + lookupTable, + nq, + d, + M, + nprobe); + + // Step 2: scan the assigned inverted lists with 8-bit ADC. The small-list + // kernel is exact only when every list is provably bounded. + (void)avgListLen; + const bool useSmall = !lookupFp16 && k <= (int)kIvfSmallExactCandidates && + ivfScanExactnessHoldsForAllLists( + listLength, nlist, kIvfSmallExactCandidates); + K.encodeIVFPQScanList( + enc, + useSmall, + lookupFp16, + lookupTable, + codes, + ids, + listOffset, + listLength, + coarseAssign, + perListDistBuf, + perListIdxBuf, + paramsBuf, + coarseDist, + nq, + nprobe); + + // Step 3: merge per-list top-k into the final top-k. + K.encodeIVFMergeLists( + enc, + perListDistBuf, + perListIdxBuf, + outDistances, + outIndices, + paramsBuf, + nq); + + [enc endEncoding]; + [cmdBuf commit]; + if (!waitForCompletion) + return true; + [cmdBuf waitUntilCompleted]; + return cmdBuf.status == MTLCommandBufferStatusCompleted; +} + +// ============================================================ +// IVF-PQ precomputed-table path +// ============================================================ + +// Must match the limits in MetalDistance.metal (PQPRE_* / MERGE_GROUP_CAND). +constexpr int kIvfPQPrecompMaxK = 512; +constexpr int kIvfPQPrecompMaxM = 64; +constexpr int kIvfPQPrecompMaxDsub = 256; +constexpr uint32_t kIvfMergeGroupCand = 2048; + +bool runMetalIVFPQPrecomputeTerm2( + id device, + id queue, + id coarseCentroids, + id pqCentroids, + id outTerm2, + int nlist, + int d, + int M) { + if (!device || !queue || !coarseCentroids || !pqCentroids || !outTerm2) + return false; + if (nlist <= 0 || d <= 0 || M <= 0 || (d % M) != 0) + return false; + if (M > kIvfPQPrecompMaxM || (d / M) > kIvfPQPrecompMaxDsub) + return false; + if (!fitsU32(nlist) || !fitsU32(d) || !fitsU32(M)) + return false; + + MetalKernels& K = getMetalKernels(device); + if (!K.isValid()) + return false; + + id cmdBuf = [queue commandBuffer]; + id enc = [cmdBuf computeCommandEncoder]; + K.encodeIVFPQPrecomputeTerm2( + enc, coarseCentroids, pqCentroids, outTerm2, nlist, d, M); + [enc endEncoding]; + [cmdBuf commit]; + [cmdBuf waitUntilCompleted]; + return cmdBuf.status == MTLCommandBufferStatusCompleted; +} + +bool runMetalIVFPQPrecompSearch( + id device, + id queue, + id queries, + id coarseAssign, + id coarseDist, + id term2, + id qtermScratch, + id pqCentroids, + id codes, + id ids, + id listOffset, + id listLength, + int nq, + int d, + int M, + int k, + int nprobe, + bool isL2, + bool useDis0, + id outDistances, + id outIndices, + id perListDistBuf, + id perListIdxBuf, + id mergeScratchDistBuf, + id mergeScratchIdxBuf, + bool waitForCompletion, + id coarseCentroids, + bool onTheFly, + bool shortLists, + bool compactMerge) { + if (!device || !queue || !queries || !coarseAssign || !coarseDist || + (!onTheFly && !qtermScratch) || !pqCentroids || !codes || !ids || + !listOffset || !listLength || !outDistances || !outIndices) { + return false; + } + if (nq <= 0 || d <= 0 || M <= 0 || k <= 0 || nprobe <= 0) + return false; + if ((d % M) != 0) + return false; + if (M > kIvfPQPrecompMaxM || (d / M) > kIvfPQPrecompMaxDsub) + return false; + if (k > kIvfPQPrecompMaxK) + return false; + if (onTheFly && (d > 256 || !coarseCentroids)) + return false; + if (isL2 && !onTheFly && !term2) + return false; + if (nprobe > 1 && + (!perListDistBuf || !perListIdxBuf || !mergeScratchDistBuf || + !mergeScratchIdxBuf)) { + return false; + } + if (!fitsU32(nq) || !fitsU32(d) || !fitsU32(M) || !fitsU32(k) || + !fitsU32(nprobe)) { + return false; + } + + MetalKernels& K = getMetalKernels(device); + if (!K.isValid()) + return false; + + const bool wantMin = isL2; + const bool useTerm2 = isL2; + + id cmdBuf = [queue commandBuffer]; + id enc = [cmdBuf computeCommandEncoder]; + + // Step 1: per-query M*256 term (once per batch, not per probe). + if (!onTheFly) + K.encodeIVFPQBuildQueryTerm( + enc, queries, pqCentroids, qtermScratch, nq, d, M, isL2); + + // Step 2: fused LUT-combine + list scan with exact running top-k. + // With a single probe the per-list result is the final result. + id scanDist = (nprobe == 1) ? outDistances : perListDistBuf; + id scanIdx = (nprobe == 1) ? outIndices : perListIdxBuf; + K.encodeIVFPQScanListPrecomp( + enc, + useTerm2 ? term2 : nil, + qtermScratch, + coarseDist, + codes, + ids, + listOffset, + listLength, + coarseAssign, + scanDist, + scanIdx, + nq, + M, + k, + nprobe, + wantMin, + useTerm2, + useDis0, + queries, + coarseCentroids, + pqCentroids, + d, + onTheFly, + shortLists); + + // Step 3: merge the nprobe per-list top-k in rounds of groupSize lists + // per threadgroup, ping-ponging between the per-list and scratch buffers. + if (nprobe > 1) { + const int groupSize = + std::max(2, (int)(kIvfMergeGroupCand / (uint32_t)k)); + int numLists = nprobe; + id curD = perListDistBuf; + id curI = perListIdxBuf; + id altD = mergeScratchDistBuf; + id altI = mergeScratchIdxBuf; + while (numLists > 1) { + const int nGroups = (numLists + groupSize - 1) / groupSize; + const bool last = (nGroups == 1); + K.encodeIVFMergeListsGrouped( + enc, + curD, + curI, + last ? outDistances : altD, + last ? outIndices : altI, + nq, + numLists, + groupSize, + k, + wantMin, + compactMerge); + if (!last) { + std::swap(curD, altD); + std::swap(curI, altI); + } + numLists = nGroups; + } + } + + [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 f40a4ca49..0b5d6ab6f 100644 --- a/faiss/gpu_metal/MetalIndex.h +++ b/faiss/gpu_metal/MetalIndex.h @@ -24,6 +24,10 @@ struct MetalIndexConfig { bool useFloat16CoarseQuantizer = false; + /// Build the IVF-PQ distance lookup tables in half precision. Halves LUT + /// bandwidth during the GPU scan at a small accuracy cost. + bool useFloat16 = false; + faiss::gpu::IndicesOptions indicesOptions = faiss::gpu::INDICES_64_BIT; bool interleavedLayout = true; diff --git a/faiss/gpu_metal/MetalIndexIVFPQ.h b/faiss/gpu_metal/MetalIndexIVFPQ.h index 7b59be6ec..193fe520b 100644 --- a/faiss/gpu_metal/MetalIndexIVFPQ.h +++ b/faiss/gpu_metal/MetalIndexIVFPQ.h @@ -75,13 +75,51 @@ class MetalIndexIVFPQ : public MetalIndex { std::unique_ptr cpuIndex_; std::unique_ptr gpuIvf_; + // Reusable GPU scratch buffers for the search path, grown on demand. + mutable id searchQueriesBuf_ = nil; + mutable id searchCoarseBuf_ = nil; + mutable id searchOutDistBuf_ = nil; + mutable id searchOutIdxBuf_ = nil; + mutable id searchPerListDistBuf_ = nil; + mutable id searchPerListIdxBuf_ = nil; + mutable id lookupTableBuf_ = nil; + mutable id lookupTableHalfBuf_ = nil; + mutable id searchCoarseDistBuf_ = nil; + mutable id searchQTermBuf_ = nil; + mutable id searchMergeDistBuf_ = nil; + mutable id searchMergeIdxBuf_ = nil; + mutable size_t searchQueriesCap_ = 0; + mutable size_t searchCoarseCap_ = 0; + mutable size_t searchOutDistCap_ = 0; + mutable size_t searchOutIdxCap_ = 0; + mutable size_t searchPerListDistCap_ = 0; + mutable size_t searchPerListIdxCap_ = 0; + mutable size_t lookupTableCap_ = 0; + mutable size_t lookupTableHalfCap_ = 0; + mutable size_t searchCoarseDistCap_ = 0; + mutable size_t searchQTermCap_ = 0; + mutable size_t searchMergeDistCap_ = 0; + mutable size_t searchMergeIdxCap_ = 0; + + // Coarse centroids mirrored on the GPU for the L2 lookup-table build. + mutable id centroidBuf_ = nil; + + // Query-independent term of the L2 precomputed-table decomposition + // (nlist * M * 256 floats), built once per trained index on the GPU. + mutable id term2Buf_ = nil; + void verifyPQSettings_() const; + void precomputeTerm2_() const; void encodeResidualAndAppend_( idx_t n, const float* x, const idx_t* list_nos, const idx_t* xids); + + void ensureSearchBuf_(id& buf, size_t& cap, size_t needed) const; + void uploadCentroids_() const; + void uploadPQCentroids_() const; }; } // namespace gpu_metal diff --git a/faiss/gpu_metal/MetalIndexIVFPQ.mm b/faiss/gpu_metal/MetalIndexIVFPQ.mm index a4b2606ec..3ff5b76db 100644 --- a/faiss/gpu_metal/MetalIndexIVFPQ.mm +++ b/faiss/gpu_metal/MetalIndexIVFPQ.mm @@ -9,6 +9,7 @@ #import "MetalIndexIVFPQ.h" #include +#include #include #include @@ -21,6 +22,22 @@ namespace faiss { namespace gpu_metal { +namespace { +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; +} +} // namespace + // ============================================================ // Constructors // ============================================================ @@ -113,6 +130,104 @@ MetalIndexIVFPQ::MetalIndexIVFPQ( MetalIndexIVFPQ::~MetalIndexIVFPQ() = default; +// ============================================================ +// GPU buffer helpers +// ============================================================ + +void MetalIndexIVFPQ::ensureSearchBuf_( + id& buf, + size_t& cap, + size_t needed) const { + if (buf != nil && cap >= needed) + return; + size_t newCap = std::max(needed, cap * 2); + id device = resources_->getDevice(); + // Commit only on success so a failed over-large allocation keeps the + // existing valid buffer instead of forcing later reallocations. + id grown = + [device newBufferWithLength:newCap + options:MTLResourceStorageModeShared]; + if (grown) { + buf = grown; + cap = newCap; + } +} + +void MetalIndexIVFPQ::uploadCentroids_() const { + if (!cpuIndex_ || !cpuIndex_->quantizer || !resources_) + return; + auto* flatQ = dynamic_cast(cpuIndex_->quantizer); + if (!flatQ || flatQ->ntotal == 0) { + centroidBuf_ = nil; + return; + } + size_t bytes = (size_t)flatQ->ntotal * (size_t)d * sizeof(float); + id device = resources_->getDevice(); + if (!device) + return; + centroidBuf_ = [device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (centroidBuf_) { + std::memcpy([centroidBuf_ contents], flatQ -> get_xb(), bytes); + } +} + +void MetalIndexIVFPQ::uploadPQCentroids_() const { + if (!cpuIndex_ || !gpuIvf_) + return; + const auto& pq = cpuIndex_->pq; + if (pq.centroids.empty()) + return; + gpuIvf_->setPQCentroids(pq.centroids.data()); +} + +void MetalIndexIVFPQ::precomputeTerm2_() const { + // GPU analogue of faiss::IndexIVFPQ::precompute_table(): the + // query-independent term ||pq||^2 + 2 of the L2 by_residual + // decomposition, built once per trained index. + term2Buf_ = nil; + if (!cpuIndex_ || !resources_ || !gpuIvf_) + return; + if (cpuIndex_->metric_type != METRIC_L2 || !cpuIndex_->by_residual) + return; + if (!centroidBuf_ || !gpuIvf_->pqCentroidsBuffer()) + return; + + const int M = (int)cpuIndex_->pq.M; + const idx_t nl = cpuIndex_->nlist; + if (M <= 0 || M > 64 || (d % M) != 0 || (d / M) > 256 || nl <= 0) + return; + // These indexes evaluate only the encountered codes from the codebook. + if (M > 16 && d <= 256) + return; + + // Mirror the CPU's precomputed_table_max_bytes default (2 GiB). + const size_t bytes = (size_t)nl * (size_t)M * 256 * sizeof(float); + if (bytes > ((size_t)1 << 31)) + return; + + id device = resources_->getDevice(); + id queue = resources_->getCommandQueue(); + if (!device || !queue) + return; + + term2Buf_ = [device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + if (!term2Buf_) + return; + if (!runMetalIVFPQPrecomputeTerm2( + device, + queue, + centroidBuf_, + gpuIvf_->pqCentroidsBuffer(), + term2Buf_, + (int)nl, + d, + M)) { + term2Buf_ = nil; + } +} + // ============================================================ // Train / add / reset // ============================================================ @@ -125,6 +240,11 @@ void MetalIndexIVFPQ::train(idx_t n, const float* x) { cpuIndex_->precompute_table(); } is_trained = cpuIndex_->is_trained; + + // Mirror coarse and PQ centroids onto the GPU for the lookup-table build. + uploadCentroids_(); + uploadPQCentroids_(); + precomputeTerm2_(); } void MetalIndexIVFPQ::encodeResidualAndAppend_( @@ -217,12 +337,270 @@ void MetalIndexIVFPQ::search( FAISS_THROW_IF_NOT(cpuIndex_); FAISS_THROW_IF_NOT(k > 0); - if (auto* ivfParams = dynamic_cast(params)) { - if (ivfParams->nprobe > 0) - cpuIndex_->nprobe = ivfParams->nprobe; + 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; } - cpuIndex_->search(n, x, k, distances, labels); + id device = resources_ ? resources_->getDevice() : nil; + id queue = + resources_ ? resources_->getCommandQueue() : nil; + + // Fall back to CPU when the GPU-resident storage or device is unavailable, + // when k exceeds what the Metal top-k kernels support, or when the PQ + // centroids have not been mirrored onto the GPU. + const int maxK = getMetalDistanceMaxK(); + if (!device || !queue || !gpuIvf_ || !gpuIvf_->codesBuffer() || + !gpuIvf_->idsBuffer() || !gpuIvf_->listOffsetGpuBuffer() || + !gpuIvf_->listLengthGpuBuffer() || !gpuIvf_->pqCentroidsBuffer() || + k > maxK) { + cpuIndex_->search(n, x, k, distances, labels, params); + return; + } + + size_t nprobe = cpuIndex_->nprobe; + if (params) { + const auto* ivfParams = + dynamic_cast(params); + const auto* pqParams = + dynamic_cast(params); + // The GPU scans every code in each probe and cannot apply selectors, + // scan budgets, or polysemous filters. Preserve CPU semantics when + // extending GPU dispatch to configurations that previously fell back. + if (!ivfParams || ivfParams->nprobe == 0 || ivfParams->sel || + ivfParams->max_codes || ivfParams->max_lists_num || + ivfParams->ensure_topk_full || + ivfParams->max_empty_result_buckets || + ivfParams->quantizer_params || ivfParams->inverted_list_context || + (pqParams && + (pqParams->polysemous_ht || pqParams->scan_table_threshold))) { + cpuIndex_->search(n, x, k, distances, labels, params); + return; + } + nprobe = ivfParams->nprobe; + } + nprobe = std::min(nprobe, (size_t)cpuIndex_->nlist); + if (nprobe == 0) { + for (idx_t i = 0; i < n * k; ++i) { + labels[i] = -1; + distances[i] = isL2 ? inf : negInf; + } + return; + } + + // Coarse quantization on the CPU. + std::vector coarseDistVec((size_t)n * nprobe); + std::vector coarseAssignVec((size_t)n * nprobe); + cpuIndex_->quantizer->search( + n, x, (idx_t)nprobe, coarseDistVec.data(), coarseAssignVec.data()); + + const int M = (int)cpuIndex_->pq.M; + const int ksub = (int)cpuIndex_->pq.ksub; + const bool useFp16Lut = config_.useFloat16; + + const size_t outDistBytes = (size_t)n * (size_t)k * sizeof(float); + const size_t outIdxBytes = (size_t)n * (size_t)k * sizeof(int64_t); + const size_t queryBytes = (size_t)n * (size_t)d * sizeof(float); + const size_t perListBytes = (size_t)n * nprobe * (size_t)k * sizeof(float); + const size_t perListIdxB = (size_t)n * nprobe * (size_t)k * sizeof(int64_t); + const size_t coarseBytes = (size_t)n * nprobe * sizeof(int32_t); + + ensureSearchBuf_(searchQueriesBuf_, searchQueriesCap_, queryBytes); + ensureSearchBuf_(searchOutDistBuf_, searchOutDistCap_, outDistBytes); + ensureSearchBuf_(searchOutIdxBuf_, searchOutIdxCap_, outIdxBytes); + ensureSearchBuf_( + searchPerListDistBuf_, searchPerListDistCap_, perListBytes); + ensureSearchBuf_(searchPerListIdxBuf_, searchPerListIdxCap_, perListIdxB); + ensureSearchBuf_(searchCoarseBuf_, searchCoarseCap_, coarseBytes); + + // A failed growth keeps the old (too-small) buffer, so check capacities + // as well as nil before writing into the buffers. + if (!searchQueriesBuf_ || searchQueriesCap_ < queryBytes || + !searchOutDistBuf_ || searchOutDistCap_ < outDistBytes || + !searchOutIdxBuf_ || searchOutIdxCap_ < outIdxBytes || + !searchPerListDistBuf_ || searchPerListDistCap_ < perListBytes || + !searchPerListIdxBuf_ || searchPerListIdxCap_ < perListIdxB || + !searchCoarseBuf_ || searchCoarseCap_ < coarseBytes) { + cpuIndex_->search(n, x, k, distances, labels, params); + return; + } + std::memcpy([searchQueriesBuf_ contents], x, queryBytes); + + // Upload coarse assignments as int32. + auto* coarseDst = reinterpret_cast([searchCoarseBuf_ contents]); + for (size_t i = 0; i < (size_t)n * nprobe; ++i) { + FAISS_THROW_IF_NOT_MSG( + coarseAssignVec[i] >= + (idx_t)std::numeric_limits::min() && + coarseAssignVec[i] <= + (idx_t)std::numeric_limits::max(), + "MetalIndexIVFPQ: coarse assignment exceeds int32 range"); + coarseDst[i] = (int32_t)coarseAssignVec[i]; + } + + const int nI = checkedIdxToInt(n, "MetalIndexIVFPQ: n exceeds int range"); + const int kI = checkedIdxToInt(k, "MetalIndexIVFPQ: k exceeds int range"); + const int nprobeI = checkedSizeToInt( + nprobe, "MetalIndexIVFPQ: nprobe exceeds int range"); + const int avgListLen = cpuIndex_->nlist > 0 + ? (int)(cpuIndex_->ntotal / cpuIndex_->nlist) + : 0; + + bool ok = false; + + // Preferred path: evaluate encountered codes directly for small dimensions, + // or use the precomputed decomposition. Both scans retain an exact top-k; + // merging runs in rounds for any nprobe * k. Keep the cached LUT kernel for + // M <= 16 when the small-list specialization cannot be used. + const int dsub = d / M; + const bool byResidual = cpuIndex_->by_residual; + const auto& listLengths = gpuIvf_->listLength(); + const bool shortLists = + std::all_of(listLengths.begin(), listLengths.end(), [](size_t len) { + return len <= 128; + }); + const bool onTheFly = + d <= 256 && M >= 8 && (M > 16 || shortLists) && centroidBuf_ != nil; + const bool compactMerge = std::any_of( + listLengths.begin(), listLengths.end(), [k](size_t len) { + return len < (size_t)k; + }); + const bool canPrecomp = M <= 64 && dsub <= 256 && kI <= 512 && + (!isL2 || (byResidual && (onTheFly || term2Buf_ != nil))) && + gpuIvf_->totalVecs() <= (size_t)std::numeric_limits::max(); + if (canPrecomp) { + const size_t coarseDistBytes = (size_t)n * nprobe * sizeof(float); + const size_t qtermBytes = + (size_t)n * (size_t)M * (size_t)ksub * sizeof(float); + ensureSearchBuf_( + searchCoarseDistBuf_, searchCoarseDistCap_, coarseDistBytes); + if (!onTheFly) + ensureSearchBuf_(searchQTermBuf_, searchQTermCap_, qtermBytes); + if (nprobeI > 1) { + ensureSearchBuf_( + searchMergeDistBuf_, searchMergeDistCap_, perListBytes); + ensureSearchBuf_( + searchMergeIdxBuf_, searchMergeIdxCap_, perListIdxB); + } + if (searchCoarseDistBuf_ && searchCoarseDistCap_ >= coarseDistBytes && + (onTheFly || (searchQTermBuf_ && searchQTermCap_ >= qtermBytes)) && + (nprobeI == 1 || + (searchMergeDistBuf_ && searchMergeDistCap_ >= perListBytes && + searchMergeIdxBuf_ && searchMergeIdxCap_ >= perListIdxB))) { + std::memcpy( + [searchCoarseDistBuf_ contents], + coarseDistVec.data(), + coarseDistBytes); + ok = runMetalIVFPQPrecompSearch( + device, + queue, + searchQueriesBuf_, + searchCoarseBuf_, + searchCoarseDistBuf_, + isL2 ? term2Buf_ : nil, + searchQTermBuf_, + gpuIvf_->pqCentroidsBuffer(), + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + gpuIvf_->listOffsetGpuBuffer(), + gpuIvf_->listLengthGpuBuffer(), + nI, + d, + M, + kI, + nprobeI, + isL2, + byResidual, + searchOutDistBuf_, + searchOutIdxBuf_, + searchPerListDistBuf_, + searchPerListIdxBuf_, + searchMergeDistBuf_, + searchMergeIdxBuf_, + true, + centroidBuf_, + onTheFly, + shortLists, + compactMerge); + } + } + + // Legacy path: per-(query, probe) lookup tables. Only allocate the large + // LUT buffer when this path is actually taken. The coarse scores are + // needed as well: inner-product tables hold only. + if (!ok) { + const size_t tableSize = (size_t)n * nprobe * M * ksub; + const size_t tableBytes = tableSize * sizeof(float); + const size_t tableHalfBytes = tableSize * sizeof(uint16_t); + const size_t coarseDistBytes = (size_t)n * nprobe * sizeof(float); + ensureSearchBuf_(lookupTableBuf_, lookupTableCap_, tableBytes); + ensureSearchBuf_( + searchCoarseDistBuf_, searchCoarseDistCap_, coarseDistBytes); + if (useFp16Lut) { + ensureSearchBuf_( + lookupTableHalfBuf_, lookupTableHalfCap_, tableHalfBytes); + } + if (lookupTableBuf_ && lookupTableCap_ >= tableBytes && + searchCoarseDistBuf_ && searchCoarseDistCap_ >= coarseDistBytes && + (!useFp16Lut || + (lookupTableHalfBuf_ && lookupTableHalfCap_ >= tableHalfBytes))) { + std::memcpy( + [searchCoarseDistBuf_ contents], + coarseDistVec.data(), + coarseDistBytes); + ok = runMetalIVFPQFullSearch( + device, + queue, + searchQueriesBuf_, + searchCoarseBuf_, + searchCoarseDistBuf_, + centroidBuf_, + gpuIvf_->pqCentroidsBuffer(), + useFp16Lut ? lookupTableHalfBuf_ : lookupTableBuf_, + gpuIvf_->codesBuffer(), + gpuIvf_->idsBuffer(), + gpuIvf_->listOffsetGpuBuffer(), + gpuIvf_->listLengthGpuBuffer(), + nI, + d, + M, + kI, + nprobeI, + (int)cpuIndex_->nlist, + avgListLen, + useFp16Lut, + isL2, + searchOutDistBuf_, + searchOutIdxBuf_, + searchPerListDistBuf_, + searchPerListIdxBuf_); + } + } + + if (!ok) { + cpuIndex_->search(n, x, k, distances, labels, params); + return; + } + + const float* outDistPtr = + reinterpret_cast([searchOutDistBuf_ contents]); + const int64_t* outIdxPtr = + reinterpret_cast([searchOutIdxBuf_ contents]); + for (idx_t qi = 0; qi < n; ++qi) { + for (idx_t j = 0; j < k; ++j) { + const size_t pos = (size_t)qi * (size_t)k + (size_t)j; + const int64_t globalId = outIdxPtr[pos]; + labels[pos] = (globalId < 0) ? -1 : (idx_t)globalId; + distances[pos] = outDistPtr[pos]; + } + } } // ============================================================ @@ -241,6 +619,16 @@ void MetalIndexIVFPQ::copyFrom(const faiss::IndexIVFPQ* src) { FAISS_THROW_IF_NOT_MSG( src->polysemous_ht == 0, "copyFrom: polysemous codes are not supported"); + FAISS_THROW_IF_NOT_FMT( + src->d == d, + "copyFrom: dimension mismatch (%zd vs %zd)", + (size_t)src->d, + (size_t)d); + FAISS_THROW_IF_NOT_FMT( + src->pq.M == cpuIndex_->pq.M, + "copyFrom: PQ sub-quantizer mismatch (%zd vs %zd)", + (size_t)src->pq.M, + (size_t)cpuIndex_->pq.M); FAISS_THROW_IF_NOT_FMT( src->nlist == cpuIndex_->nlist, "copyFrom: nlist mismatch (%zd vs %zd)", @@ -280,6 +668,11 @@ void MetalIndexIVFPQ::copyFrom(const faiss::IndexIVFPQ* src) { } is_trained = true; + // Mirror coarse and PQ centroids onto the GPU for the lookup-table build. + uploadCentroids_(); + uploadPQCentroids_(); + precomputeTerm2_(); + // Gather IVF list data. size_t totalN = 0; for (size_t l = 0; l < (size_t)src->nlist; ++l) diff --git a/faiss/gpu_metal/MetalKernels.h b/faiss/gpu_metal/MetalKernels.h index 92fb79df5..82f118221 100644 --- a/faiss/gpu_metal/MetalKernels.h +++ b/faiss/gpu_metal/MetalKernels.h @@ -91,6 +91,99 @@ class MetalKernels { id paramsBuf, int nq); + // ---- IVF-PQ ---- + + void encodeIVFPQBuildLookupTables( + id enc, + bool isL2, + bool outFp16, + id queries, + id coarseAssign, + id coarseCentroids, + id pqCentroids, + id outLookup, + int nq, + int d, + int M, + int nprobe); + + void encodeIVFPQScanList( + id enc, + bool useSmall, + bool useFp16Lookup, + id lookupTable, + id codes, + id ids, + id listOffset, + id listLength, + id coarseAssign, + id perListDist, + id perListIdx, + id paramsBuf, + id coarseDist, + int nq, + int nprobe); + + // ---- IVF-PQ precomputed-table path ---- + + void encodeIVFPQPrecomputeTerm2( + id enc, + id coarseCentroids, + id pqCentroids, + id outTerm2, + int nlist, + int d, + int M); + + void encodeIVFPQBuildQueryTerm( + id enc, + id queries, + id pqCentroids, + id outQTerm, + int nq, + int d, + int M, + bool isL2); + + void encodeIVFPQScanListPrecomp( + id enc, + id term2, // nil when useTerm2 is false (IP path) + id qterm, + id coarseDist, + id codes, + id ids, + id listOffset, + id listLength, + id coarseAssign, + id perListDist, + id perListIdx, + int nq, + int M, + int k, + int nprobe, + bool wantMin, + bool useTerm2, + bool useDis0, + id queries, + id coarseCentroids, + id pqCentroids, + int d, + bool onTheFly, + bool shortLists); + + void encodeIVFMergeListsGrouped( + id enc, + id inDist, + id inIdx, + id outDist, + id outIdx, + int nq, + int numLists, + int groupSize, + int k, + bool wantMin, + bool compact); + static int selectTopKVariantIndex(int k); private: diff --git a/faiss/gpu_metal/MetalKernels.mm b/faiss/gpu_metal/MetalKernels.mm index 84ec24496..391582b6e 100644 --- a/faiss/gpu_metal/MetalKernels.mm +++ b/faiss/gpu_metal/MetalKernels.mm @@ -274,6 +274,218 @@ void MetalKernels::encodeIVFMergeLists( threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; } +void MetalKernels::encodeIVFPQBuildLookupTables( + id enc, + bool isL2, + bool outFp16, + id queries, + id coarseAssign, + id coarseCentroids, + id pqCentroids, + id outLookup, + int nq, + int d, + int M, + int nprobe) { + const char* name = nullptr; + if (isL2) { + name = outFp16 ? "ivfpq_build_lut_l2_f16" : "ivfpq_build_lut_l2"; + } else { + name = outFp16 ? "ivfpq_build_lut_ip_f16" : "ivfpq_build_lut_ip"; + } + [enc setComputePipelineState:pipeline(name)]; + [enc setBuffer:queries offset:0 atIndex:0]; + [enc setBuffer:coarseAssign offset:0 atIndex:1]; + [enc setBuffer:coarseCentroids offset:0 atIndex:2]; + [enc setBuffer:pqCentroids offset:0 atIndex:3]; + [enc setBuffer:outLookup offset:0 atIndex:4]; + uint32_t args[4] = { + (uint32_t)nq, (uint32_t)d, (uint32_t)M, (uint32_t)nprobe}; + [enc setBytes:args length:sizeof(args) atIndex:5]; + const NSUInteger totalTG = + (NSUInteger)nq * (NSUInteger)nprobe * (NSUInteger)M; + [enc dispatchThreadgroups:MTLSizeMake(totalTG, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; +} + +void MetalKernels::encodeIVFPQScanList( + id enc, + bool useSmall, + bool useFp16Lookup, + id lookupTable, + id codes, + id ids, + id listOffset, + id listLength, + id coarseAssign, + id perListDist, + id perListIdx, + id paramsBuf, + id coarseDist, + int nq, + int nprobe) { + const char* name = "ivf_scan_list_pq8"; + if (useFp16Lookup) { + name = "ivf_scan_list_pq8_f16"; + } else if (useSmall) { + name = "ivf_scan_list_pq8_small"; + } + [enc setComputePipelineState:pipeline(name)]; + [enc setBuffer:lookupTable 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]; + [enc setBuffer:coarseDist offset:0 atIndex:9]; + [enc dispatchThreadgroups:MTLSizeMake( + (NSUInteger)nq * (NSUInteger)nprobe, 1, 1) + threadsPerThreadgroup:MTLSizeMake( + (useSmall && !useFp16Lookup) ? 32 + : 256, + 1, + 1)]; +} + +void MetalKernels::encodeIVFPQPrecomputeTerm2( + id enc, + id coarseCentroids, + id pqCentroids, + id outTerm2, + int nlist, + int d, + int M) { + [enc setComputePipelineState:pipeline("ivfpq_precompute_term2")]; + [enc setBuffer:coarseCentroids offset:0 atIndex:0]; + [enc setBuffer:pqCentroids offset:0 atIndex:1]; + [enc setBuffer:outTerm2 offset:0 atIndex:2]; + uint32_t args[3] = {(uint32_t)nlist, (uint32_t)d, (uint32_t)M}; + [enc setBytes:args length:sizeof(args) atIndex:3]; + const NSUInteger totalTG = (NSUInteger)nlist * (NSUInteger)M; + [enc dispatchThreadgroups:MTLSizeMake(totalTG, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; +} + +void MetalKernels::encodeIVFPQBuildQueryTerm( + id enc, + id queries, + id pqCentroids, + id outQTerm, + int nq, + int d, + int M, + bool isL2) { + [enc setComputePipelineState:pipeline("ivfpq_build_query_term")]; + [enc setBuffer:queries offset:0 atIndex:0]; + [enc setBuffer:pqCentroids offset:0 atIndex:1]; + [enc setBuffer:outQTerm offset:0 atIndex:2]; + uint32_t args[4] = {(uint32_t)nq, (uint32_t)d, (uint32_t)M, isL2 ? 1u : 0u}; + [enc setBytes:args length:sizeof(args) atIndex:3]; + const NSUInteger totalTG = (NSUInteger)nq * (NSUInteger)M; + [enc dispatchThreadgroups:MTLSizeMake(totalTG, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; +} + +void MetalKernels::encodeIVFPQScanListPrecomp( + id enc, + id term2, + id qterm, + id coarseDist, + id codes, + id ids, + id listOffset, + id listLength, + id coarseAssign, + id perListDist, + id perListIdx, + int nq, + int M, + int k, + int nprobe, + bool wantMin, + bool useTerm2, + bool useDis0, + id queries, + id coarseCentroids, + id pqCentroids, + int d, + bool onTheFly, + bool shortLists) { + const char* name = onTheFly + ? (shortLists ? "ivf_scan_list_pq8_onthefly_small" + : "ivf_scan_list_pq8_onthefly") + : (M > 16 ? "ivf_scan_list_pq8_precomp_large" + : "ivf_scan_list_pq8_precomp"); + [enc setComputePipelineState:pipeline(name)]; + // The kernel never reads term2 when useTerm2 is 0, but Metal requires a + // valid binding; rebind qterm in that case. + [enc setBuffer:(onTheFly ? coarseCentroids : (term2 ? term2 : qterm)) + offset:0 + atIndex:0]; + [enc setBuffer:(onTheFly ? queries : qterm) offset:0 atIndex:1]; + [enc setBuffer:coarseDist offset:0 atIndex:2]; + [enc setBuffer:codes offset:0 atIndex:3]; + [enc setBuffer:ids offset:0 atIndex:4]; + [enc setBuffer:listOffset offset:0 atIndex:5]; + [enc setBuffer:listLength offset:0 atIndex:6]; + [enc setBuffer:coarseAssign offset:0 atIndex:7]; + [enc setBuffer:perListDist offset:0 atIndex:8]; + [enc setBuffer:perListIdx offset:0 atIndex:9]; + uint32_t args[8] = { + (uint32_t)nq, + (uint32_t)M, + (uint32_t)k, + (uint32_t)nprobe, + wantMin ? 1u : 0u, + useTerm2 ? 1u : 0u, + useDis0 ? 1u : 0u, + (uint32_t)d}; + [enc setBytes:args length:sizeof(args) atIndex:10]; + if (onTheFly) + [enc setBuffer:pqCentroids offset:0 atIndex:11]; + [enc dispatchThreadgroups:MTLSizeMake( + (NSUInteger)nq * (NSUInteger)nprobe, 1, 1) + threadsPerThreadgroup:MTLSizeMake( + onTheFly && shortLists ? 32 : 256, + 1, + 1)]; +} + +void MetalKernels::encodeIVFMergeListsGrouped( + id enc, + id inDist, + id inIdx, + id outDist, + id outIdx, + int nq, + int numLists, + int groupSize, + int k, + bool wantMin, + bool compact) { + [enc setComputePipelineState:pipeline("ivf_merge_lists_grouped")]; + [enc setBuffer:inDist offset:0 atIndex:0]; + [enc setBuffer:inIdx offset:0 atIndex:1]; + [enc setBuffer:outDist offset:0 atIndex:2]; + [enc setBuffer:outIdx offset:0 atIndex:3]; + uint32_t args[6] = { + (uint32_t)nq, + (uint32_t)numLists, + (uint32_t)groupSize, + (uint32_t)k, + wantMin ? 1u : 0u, + compact ? 1u : 0u}; + [enc setBytes:args length:sizeof(args) atIndex:4]; + const NSUInteger nGroups = + ((NSUInteger)numLists + (NSUInteger)groupSize - 1) / + (NSUInteger)groupSize; + [enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nq * nGroups, 1, 1) + threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; +} + MetalKernels& getMetalKernels(id device) { static std::mutex mu; static std::unordered_map> map; diff --git a/faiss/gpu_metal/test/TestMetalIndexIVFPQ.mm b/faiss/gpu_metal/test/TestMetalIndexIVFPQ.mm index 52f725a25..48d2b8d09 100644 --- a/faiss/gpu_metal/test/TestMetalIndexIVFPQ.mm +++ b/faiss/gpu_metal/test/TestMetalIndexIVFPQ.mm @@ -11,12 +11,16 @@ #include #include #include +#include #include #include #include +#include #include #include +#include #import +#include #import #import #import @@ -170,3 +174,259 @@ TEST_F(TestMetalIndexIVFPQ, CpuToGpuRoundTrip) { std::unique_ptr cpuBackPtr(cpuBack); ASSERT_EQ(cpuBack->ntotal, nb); } + +TEST_F(TestMetalIndexIVFPQ, LargeMSearchParametersKeepCpuSemantics) { + constexpr int d = 128, nb = 2000, nq = 4, k = 20; + std::vector xb(nb * d), xq(nq * d); + faiss::float_rand(xb.data(), xb.size(), 123); + faiss::float_rand(xq.data(), xq.size(), 456); + faiss::IndexFlatL2 quantizer(d); + faiss::IndexIVFPQ cpu(&quantizer, d, 8, 32, 8); + cpu.nprobe = 2; + cpu.train(nb, xb.data()); + cpu.add(nb, xb.data()); + faiss::gpu_metal::MetalIndexIVFPQ metal(resources_, &cpu); + faiss::IDSelectorRange selector(0, 10); + for (bool useSelector : {true, false}) { + faiss::IVFPQSearchParameters params; + params.nprobe = 8; + params.sel = useSelector ? &selector : nullptr; + params.max_codes = useSelector ? 0 : 12; + std::vector cpuD(nq * k), gpuD(nq * k); + std::vector cpuI(nq * k), gpuI(nq * k); + cpu.search(nq, xq.data(), k, cpuD.data(), cpuI.data(), ¶ms); + metal.search(nq, xq.data(), k, gpuD.data(), gpuI.data(), ¶ms); + EXPECT_EQ(cpuI, gpuI); + EXPECT_EQ(cpuD, gpuD); + EXPECT_EQ(metal.nprobe(), 2); + } +} + +// Exercise the Metal entry point directly: a CPU fallback cannot satisfy this +// test. The oracle reconstructs PQ vectors and computes their distances, +// without using the GPU's lookup-table decomposition or selection +// implementation. +TEST_F(TestMetalIndexIVFPQ, ScanVariantsAgainstReconstruction) { + struct Config { + int M; + int d; + bool onTheFly; + }; + const Config configs[] = { + {8, 128, false}, + {16, 128, false}, + {8, 128, true}, + {16, 128, true}, + {17, 136, true}, + {32, 128, true}, + {48, 192, true}, + {64, 128, true}, + {32, 512, false}, + {64, 512, false}}; + for (const auto& config : configs) { + for (int shape : {0, 1, 2, 3, 4}) { + const bool shortLists = shape != 1 && shape != 3; + for (bool isL2 : {true, false}) { + @autoreleasepool { + const int M = config.M, d = config.d, dsub = d / M; + const int nq = 3; + std::vector lengths = shortLists + ? std::vector{0, 1, 6, 24, 127, 128} + : std::vector{ + 0, + 1, + 6, + 24, + 127, + 128, + 129, + 511, + 512, + 513, + 1025, + 2051, + 3}; + if (shape == 2) + lengths = {0, 0, 0}; + else if (shape == 3) + lengths = {2051}; + else if (shape == 4) + lengths = {128}; + const int nl = lengths.size(), np = nl; + std::vector offsets(nl); + size_t nb = 0; + for (int l = 0; l < nl; ++l) { + offsets[l] = nb; + nb += lengths[l]; + } + std::vector queries(nq * d), coarse(nl * d), + pq(d * 256); + faiss::float_rand(queries.data(), queries.size(), 43); + faiss::float_rand(coarse.data(), coarse.size(), 44); + faiss::float_rand(pq.data(), pq.size(), 45); + std::vector codes(nb * M); + for (size_t i = 0; i < codes.size(); ++i) + codes[i] = (i * 131 + i / 7 + i / 257) % 256; + std::vector ids(nb); + for (size_t i = 0; i < nb; ++i) + ids[i] = (int64_t(1) << 40) + i * 17; + std::vector assign(nq * np); + std::vector coarseDist(nq * np); + std::vector>> ref(nq); + std::vector> score(nq); + for (int q = 0; q < nq; ++q) { + for (int p = 0; p < np; ++p) { + int l = (p + q) % nl; + assign[q * np + p] = np > 1 && p == np - 1 ? -1 : l; + if (assign[q * np + p] < 0) + continue; + for (int j = 0; j < d; ++j) { + float x = queries[q * d + j], + c = coarse[l * d + j]; + coarseDist[q * np + p] += + isL2 ? (x - c) * (x - c) : x * c; + } + for (uint32_t v = offsets[l]; + v < offsets[l] + lengths[l]; + ++v) { + double distance = 0; + for (int j = 0; j < d; ++j) { + float reconstructed = coarse[l * d + j] + + pq[((j / dsub) * 256 + + codes[size_t(v) * M + + j / dsub]) * + dsub + + j % dsub]; + double x = queries[q * d + j]; + distance += isL2 ? (x - reconstructed) * + (x - reconstructed) + : x * reconstructed; + } + score[q][ids[v]] = distance; + ref[q].emplace_back(distance, ids[v]); + } + } + std::sort( + ref[q].begin(), + ref[q].end(), + [isL2](const auto& a, const auto& b) { + return a.first == b.first + ? a.second < b.second + : (isL2 ? a.first < b.first + : a.first > b.first); + }); + } + auto device = resources_->getDevice(); + auto buffer = [&](const void* src, size_t bytes) { + if (bytes == 0) + return [device + newBufferWithLength:1 + options:MTLResourceStorageModeShared]; + return src + ? [device newBufferWithBytes:src + length:bytes + options:MTLResourceStorageModeShared] + : [device newBufferWithLength:bytes + options:MTLResourceStorageModeShared]; + }; + auto qb = buffer(queries.data(), queries.size() * 4); + auto cb = buffer(coarse.data(), coarse.size() * 4); + auto pb = buffer(pq.data(), pq.size() * 4); + auto codesb = buffer(codes.data(), codes.size()); + auto ib = buffer(ids.data(), ids.size() * 8); + auto ob = buffer(offsets.data(), offsets.size() * 4); + auto lb = buffer(lengths.data(), lengths.size() * 4); + auto ab = buffer(assign.data(), assign.size() * 4); + auto db = buffer(coarseDist.data(), coarseDist.size() * 4); + auto qt = buffer(nullptr, size_t(nq) * M * 256 * 4); + auto t2 = buffer(nullptr, size_t(nl) * M * 256 * 4); + if (isL2 && !config.onTheFly) { + ASSERT_TRUE( + faiss::gpu_metal::runMetalIVFPQPrecomputeTerm2( + device, + resources_->getCommandQueue(), + cb, + pb, + t2, + nl, + d, + M)); + } + for (int k : {1, 20, 128, 512}) { + for (bool compact : {false, true}) { + SCOPED_TRACE( + ::testing::Message() + << "M=" << M << " d=" << d + << " short=" << shortLists << " L2=" << isL2 + << " onTheFly=" << config.onTheFly + << " k=" << k); + auto outD = buffer(nullptr, nq * k * 4), + outI = buffer(nullptr, nq * k * 8); + auto perD = buffer(nullptr, nq * np * k * 4), + perI = buffer(nullptr, nq * np * k * 8); + auto tmpD = buffer(nullptr, nq * np * k * 4), + tmpI = buffer(nullptr, nq * np * k * 8); + ASSERT_TRUE( + faiss::gpu_metal:: + runMetalIVFPQPrecompSearch( + device, + resources_ + ->getCommandQueue(), + qb, + ab, + db, + t2, + qt, + pb, + codesb, + ib, + ob, + lb, + nq, + d, + M, + k, + np, + isL2, + true, + outD, + outI, + perD, + perI, + tmpD, + tmpI, + true, + cb, + config.onTheFly, + shortLists, + compact)); + const auto* gd = + static_cast(outD.contents); + const auto* gi = + static_cast(outI.contents); + for (int q = 0; q < nq; ++q) { + std::set seen; + for (int j = 0; j < k; ++j) { + int pos = q * k + j; + if (j >= ref[q].size()) { + EXPECT_EQ(gi[pos], -1); + EXPECT_GE( + isL2 ? gd[pos] : -gd[pos], + 1e38f); + continue; + } + ASSERT_TRUE(score[q].count(gi[pos])); + EXPECT_TRUE(seen.insert(gi[pos]).second); + EXPECT_NEAR( + gd[pos], score[q][gi[pos]], 0.002); + EXPECT_NEAR( + gd[pos], ref[q][j].first, 0.002); + } + } + } + } + } + } + } + } +} diff --git a/tests/test_metal_ivfpq.py b/tests/test_metal_ivfpq.py new file mode 100644 index 000000000..a3dcc1385 --- /dev/null +++ b/tests/test_metal_ivfpq.py @@ -0,0 +1,103 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +"""Regression tests for the Metal IVF-PQ search backend. + +The per-(query, probe) lookup-table path must add the coarse-quantizer score +for inner product: its tables hold only, while the CPU reference +adds dis0 = (see QueryTables::precompute_list_tables_IP). +These tests force the lookup-table path with configurations the +precomputed-table path declines, and compare GPU results against the CPU +index. +""" + +import functools +import unittest + +import numpy as np + +import faiss + + +_HAS_METAL = "MAC_METAL" in faiss.get_compile_options() +needs_metal_build = unittest.skipUnless( + _HAS_METAL, "needs a faiss Metal build" +) + + +def needs_metal_device(test_func): + """Skip the test when no Metal device is available.""" + + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + if faiss.get_num_gpus() == 0: + raise unittest.SkipTest("no Metal device available") + return test_func(*args, **kwargs) + + return wrapper + + +@needs_metal_build +class TestMetalIVFPQ(unittest.TestCase): + + def check_matches_cpu( + self, metric, d, M, nlist, nb, nq, nprobe, k + ): + rs = np.random.RandomState(1234) + xb = rs.rand(nb, d).astype(np.float32) + xq = rs.rand(nq, d).astype(np.float32) + + quantizer = ( + faiss.IndexFlatIP(d) + if metric == faiss.METRIC_INNER_PRODUCT + else faiss.IndexFlatL2(d) + ) + cpu_index = faiss.IndexIVFPQ(quantizer, d, nlist, M, 8, metric) + cpu_index.nprobe = nprobe + cpu_index.train(xb) + cpu_index.add(xb) + ref_D, ref_I = cpu_index.search(xq, k) + + res = faiss.StandardGpuResources() + gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index) + got_D, got_I = gpu_index.search(xq, k) + + for q in range(nq): + valid = ref_I[q] != -1 + self.assertTrue(np.any(valid)) + self.assertEqual(set(got_I[q][valid]), set(ref_I[q][valid])) + np.testing.assert_allclose( + got_D[q][valid], ref_D[q][valid], atol=1e-2 + ) + + @needs_metal_device + def test_lookup_table_path_ip_large_k(self): + # k > 512 declines the precomputed-table path. With nprobe == 1 the + # missing coarse term keeps the ranking but shifts every score. + self.check_matches_cpu( + faiss.METRIC_INNER_PRODUCT, + d=64, + M=8, + nlist=8, + nb=500, + nq=5, + nprobe=1, + k=600, + ) + + @needs_metal_device + def test_lookup_table_path_l2_large_k(self): + # L2 tables hold the complete residual distance, so the coarse term + # must stay off: guard the flag plumbing in both directions. + self.check_matches_cpu( + faiss.METRIC_L2, + d=64, + M=8, + nlist=8, + nb=500, + nq=5, + nprobe=1, + k=600, + )