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, + )