facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add GPU-accelerated IVFPQ search to the Metal backend (#5449)

Browse Source 
Summary:
This is a follow-up to https://github.com/facebookresearch/faiss/issues/5288, which added `MetalIndexIVFPQ` with GPU-resident storage while delegating search to the CPU.

Adds GPU-accelerated IVFPQ search to the Metal backend, including optimized scanning through `M <= 64`.

- Add Metal kernels for 8-bit PQ distance computation, inverted-list scanning, and top-k merging
- Support optimized GPU scanning through `M <= 64`
- Compute distance contributions only for PQ codes encountered in each list when `d <= 256`
- Add a 32-thread cooperative kernel for lists containing at most 128 vectors
- Retain the cached precomputed-LUT kernel for suitable `M <= 16` workloads
- Stream precomputed terms for larger dimensions without allocating an `M * 256` threadgroup LUT
- Support both L2 and inner product metrics
- Add exact segmented top-k selection for inverted lists of any length
- Add grouped multi-round merging without a fixed `nprobe * k` candidate limit
- Compact invalid per-list padding before merging when lists contain fewer than `k` results
- Reuse dynamically sized Metal scratch buffers across searches
- Retain the legacy per-probe LUT path with optional FP16 lookup tables
- Fall back to CPU search for unsupported configurations, search parameters, or Metal failures

## Changes

- **Modified:** `MetalDistance.metal` - on demand PQ scans, large-M precomputed scan, short-list specialization, segmented top-k, and compacting grouped merge
- **Modified:** `MetalDistance.h/.mm` - IVFPQ search orchestration and support for `M <= 64`
- **Modified:** `MetalKernels.h/.mm` - dispatch for cached, on-demand, short-list, and large-M scan kernels
- **Modified:** `MetalIndexIVFPQ.h/.mm` - GPU scan selection, reusable buffers, centroid uploads, precomputed terms, and CPU fallback
- **Modified:** `MetalIndex.h` - add the `useFloat16` LUT configuration
- **Modified:** `TestMetalIndexIVFPQ.mm` - direct-GPU scan validation and search-parameter fallback coverage

## Differences from CUDA IVFPQ

**Training:** `MetalIndexIVFPQ::train` delegates training to its CPU index, as introduced in https://github.com/facebookresearch/faiss/issues/5288. CUDA can train the coarse and product quantizers on GPU. Training is a one-time cost.

**Add path:** Coarse assignment, residual computation, and PQ encoding remain on the CPU. The encoded PQ codes are stored in GPU-resident Metal buffers. CUDA performs these operations on GPU.

**Coarse quantization:** The coarse quantizer search runs on the CPU. The selected lists and coarse distances are passed to the Metal scan.

**Distance computation:** For `d <= 256`, the preferred Metal path evaluates only the PQ codes encountered while scanning a list. This avoids constructing or loading all 256 lookup entries for every subquantizer, which is especially beneficial for short IVF lists.

For other supported dimensions, Metal uses the CPU IVFPQ precomputed-table decomposition. The query-independent centroid/PQ term is computed once per trained index and the query term once per batch.

**Threadgroup memory:** The existing cached lookup-table kernel remains available for suitable `M <= 16` workloads. Larger `M` values do not expand its `M * 256` threadgroup LUT, which would exceed Metal threadgroup-memory limits. Instead, they use the on-demand or streamed-precomputed scan paths.

**List scanning:** Short lists use a 32-thread cooperative kernel with approximately 2 KiB of threadgroup scratch. Longer lists use segmented selection to maintain an exact running top-k.

**Top-k merge:** Per-list results are merged in groups over multiple rounds. Invalid padding is compacted when lists contain fewer than `k` results.

**Fallback:** The optimized scan supports `M <= 64`, `d / M <= 256`, and `k <= 512`. Unsupported configurations fall back to the legacy Metal path or CPU search. Search options that the GPU path cannot preserve, including selectors, scan budgets, and polysemous filtering, are forwarded to the CPU with their original parameters.

## Build and test

```bash
cmake -B build \
  -DFAISS_ENABLE_GPU=OFF \
  -DFAISS_ENABLE_METAL=ON \
  -DBUILD_TESTING=ON \
  -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_PREFIX_PATH="$(brew --prefix libomp)" \
  .

cmake --build build \
  --target faiss faiss_metal TestMetalIndexIVFPQ \
  -j$(sysctl -n hw.logicalcpu)

cd build && ctest -R TestMetalIndexIVFPQ --output-on-failure
```
All six IVFPQ tests pass. Direct-GPU regression coverage exercises 800 combinations across `M=8,16,17,32,48,64`, L2 and inner product, `k` through 512, short and long lists, both merge modes, and 64-bit IDs.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5449

Reviewed By: junjieqi

Differential Revision: D113864779

Pulled By: mnorris11

fbshipit-source-id: 741c12b2b014ef28a82569c6b0344352351a7325
This commit is contained in:
Evan Haque
2026-10-01 21:07:27 -07:00
committed by meta-codesync[bot]
parent 0db2f30886
commit 76c0e01e0f
10 changed files with 2878 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/gpu_metal/MetalDistance.h
+92
View File
@@ -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 <q_m, pq_m> only); unused for L2.
bool runMetalIVFPQFullSearch(
id<MTLDevice> device,
id<MTLCommandQueue> queue,
id<MTLBuffer> queries,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> coarseDist,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> lookupTable,
id<MTLBuffer> codes,
id<MTLBuffer> ids,
id<MTLBuffer> listOffset,
id<MTLBuffer> listLength,
int nq,
int d,
int M,
int k,
int nprobe,
int nlist,
int avgListLen,
bool lookupFp16,
bool isL2,
id<MTLBuffer> outDistances,
id<MTLBuffer> outIndices,
id<MTLBuffer> perListDistBuf,
id<MTLBuffer> 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 <coarse_l[m], pq[m][c]>, laid out as nlist * M * 256 floats. Run once
/// per trained index; consumed by runMetalIVFPQPrecompSearch.
bool runMetalIVFPQPrecomputeTerm2(
id<MTLDevice> device,
id<MTLCommandQueue> queue,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> 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<MTLDevice> device,
id<MTLCommandQueue> queue,
id<MTLBuffer> queries,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> coarseDist,
id<MTLBuffer> term2,
id<MTLBuffer> qtermScratch,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> codes,
id<MTLBuffer> ids,
id<MTLBuffer> listOffset,
id<MTLBuffer> listLength,
int nq,
int d,
int M,
int k,
int nprobe,
bool isL2,
bool useDis0,
id<MTLBuffer> outDistances,
id<MTLBuffer> outIndices,
id<MTLBuffer> perListDistBuf,
id<MTLBuffer> perListIdxBuf,
id<MTLBuffer> mergeScratchDistBuf,
id<MTLBuffer> mergeScratchIdxBuf,
bool waitForCompletion = true,
id<MTLBuffer> coarseCentroids = nil,
bool onTheFly = false,
bool shortLists = false,
bool compactMerge = true);
} // namespace gpu_metal
} // namespace faiss
faiss/gpu_metal/MetalDistance.metal
+1356 -22
View File
File diff suppressed because it is too large Load Diff
faiss/gpu_metal/MetalDistance.mm
+323
View File
@@ -355,5 +355,328 @@ bool runMetalIVFFlatFullSearch(
return cmdBuf.status == MTLCommandBufferStatusCompleted;
}
// ============================================================
// runMetalIVFPQFullSearch
// ============================================================
bool runMetalIVFPQFullSearch(
id<MTLDevice> device,
id<MTLCommandQueue> queue,
id<MTLBuffer> queries,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> coarseDist,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> lookupTable,
id<MTLBuffer> codes,
id<MTLBuffer> ids,
id<MTLBuffer> listOffset,
id<MTLBuffer> listLength,
int nq,
int d,
int M,
int k,
int nprobe,
int nlist,
int avgListLen,
bool lookupFp16,
bool isL2,
id<MTLBuffer> outDistances,
id<MTLBuffer> outIndices,
id<MTLBuffer> perListDistBuf,
id<MTLBuffer> 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
// <q_m, pq_m> 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<MTLBuffer> paramsBuf =
[device newBufferWithBytes:sp
length:sizeof(sp)
options:MTLResourceStorageModeShared];
if (!paramsBuf)
return false;
id<MTLCommandBuffer> cmdBuf = [queue commandBuffer];
id<MTLComputeCommandEncoder> 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<MTLDevice> device,
id<MTLCommandQueue> queue,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> 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<MTLCommandBuffer> cmdBuf = [queue commandBuffer];
id<MTLComputeCommandEncoder> 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<MTLDevice> device,
id<MTLCommandQueue> queue,
id<MTLBuffer> queries,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> coarseDist,
id<MTLBuffer> term2,
id<MTLBuffer> qtermScratch,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> codes,
id<MTLBuffer> ids,
id<MTLBuffer> listOffset,
id<MTLBuffer> listLength,
int nq,
int d,
int M,
int k,
int nprobe,
bool isL2,
bool useDis0,
id<MTLBuffer> outDistances,
id<MTLBuffer> outIndices,
id<MTLBuffer> perListDistBuf,
id<MTLBuffer> perListIdxBuf,
id<MTLBuffer> mergeScratchDistBuf,
id<MTLBuffer> mergeScratchIdxBuf,
bool waitForCompletion,
id<MTLBuffer> 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<MTLCommandBuffer> cmdBuf = [queue commandBuffer];
id<MTLComputeCommandEncoder> 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<MTLBuffer> scanDist = (nprobe == 1) ? outDistances : perListDistBuf;
id<MTLBuffer> 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<MTLBuffer> curD = perListDistBuf;
id<MTLBuffer> curI = perListIdxBuf;
id<MTLBuffer> altD = mergeScratchDistBuf;
id<MTLBuffer> 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
faiss/gpu_metal/MetalIndex.h
+4
View File
@@ -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;
faiss/gpu_metal/MetalIndexIVFPQ.h
+38
View File
@@ -75,13 +75,51 @@ class MetalIndexIVFPQ : public MetalIndex {
std::unique_ptr<faiss::IndexIVFPQ> cpuIndex_;
std::unique_ptr<MetalIVFPQImpl> gpuIvf_;
// Reusable GPU scratch buffers for the search path, grown on demand.
mutable id<MTLBuffer> searchQueriesBuf_ = nil;
mutable id<MTLBuffer> searchCoarseBuf_ = nil;
mutable id<MTLBuffer> searchOutDistBuf_ = nil;
mutable id<MTLBuffer> searchOutIdxBuf_ = nil;
mutable id<MTLBuffer> searchPerListDistBuf_ = nil;
mutable id<MTLBuffer> searchPerListIdxBuf_ = nil;
mutable id<MTLBuffer> lookupTableBuf_ = nil;
mutable id<MTLBuffer> lookupTableHalfBuf_ = nil;
mutable id<MTLBuffer> searchCoarseDistBuf_ = nil;
mutable id<MTLBuffer> searchQTermBuf_ = nil;
mutable id<MTLBuffer> searchMergeDistBuf_ = nil;
mutable id<MTLBuffer> 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<MTLBuffer> 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<MTLBuffer> 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<MTLBuffer>& buf, size_t& cap, size_t needed) const;
void uploadCentroids_() const;
void uploadPQCentroids_() const;
};
} // namespace gpu_metal
faiss/gpu_metal/MetalIndexIVFPQ.mm
+397 -4
View File
@@ -9,6 +9,7 @@
#import "MetalIndexIVFPQ.h"
#include <faiss/IndexFlat.h>
#include <faiss/gpu_metal/MetalDistance.h>
#include <faiss/gpu_metal/impl/MetalIVFPQ.h>
#include <faiss/impl/FaissAssert.h>
@@ -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<int>::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<int>::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<MTLBuffer>& buf,
size_t& cap,
size_t needed) const {
if (buf != nil && cap >= needed)
return;
size_t newCap = std::max(needed, cap * 2);
id<MTLDevice> device = resources_->getDevice();
// Commit only on success so a failed over-large allocation keeps the
// existing valid buffer instead of forcing later reallocations.
id<MTLBuffer> 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<faiss::IndexFlat*>(cpuIndex_->quantizer);
if (!flatQ || flatQ->ntotal == 0) {
centroidBuf_ = nil;
return;
}
size_t bytes = (size_t)flatQ->ntotal * (size_t)d * sizeof(float);
id<MTLDevice> 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<c, pq> 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<MTLDevice> device = resources_->getDevice();
id<MTLCommandQueue> 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<const IVFSearchParameters*>(params)) {
if (ivfParams->nprobe > 0)
cpuIndex_->nprobe = ivfParams->nprobe;
const float inf = std::numeric_limits<float>::infinity();
const float negInf = -std::numeric_limits<float>::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<MTLDevice> device = resources_ ? resources_->getDevice() : nil;
id<MTLCommandQueue> 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<const IVFSearchParameters*>(params);
const auto* pqParams =
dynamic_cast<const IVFPQSearchParameters*>(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<float> coarseDistVec((size_t)n * nprobe);
std::vector<idx_t> 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<int32_t*>([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<int32_t>::min() &&
coarseAssignVec[i] <=
(idx_t)std::numeric_limits<int32_t>::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<int32_t>::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 <q_m, pq_m> 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<const float*>([searchOutDistBuf_ contents]);
const int64_t* outIdxPtr =
reinterpret_cast<const int64_t*>([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)
faiss/gpu_metal/MetalKernels.h
+93
View File
@@ -91,6 +91,99 @@ class MetalKernels {
id<MTLBuffer> paramsBuf,
int nq);
// ---- IVF-PQ ----
void encodeIVFPQBuildLookupTables(
id<MTLComputeCommandEncoder> enc,
bool isL2,
bool outFp16,
id<MTLBuffer> queries,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> outLookup,
int nq,
int d,
int M,
int nprobe);
void encodeIVFPQScanList(
id<MTLComputeCommandEncoder> enc,
bool useSmall,
bool useFp16Lookup,
id<MTLBuffer> lookupTable,
id<MTLBuffer> codes,
id<MTLBuffer> ids,
id<MTLBuffer> listOffset,
id<MTLBuffer> listLength,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> perListDist,
id<MTLBuffer> perListIdx,
id<MTLBuffer> paramsBuf,
id<MTLBuffer> coarseDist,
int nq,
int nprobe);
// ---- IVF-PQ precomputed-table path ----
void encodeIVFPQPrecomputeTerm2(
id<MTLComputeCommandEncoder> enc,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> outTerm2,
int nlist,
int d,
int M);
void encodeIVFPQBuildQueryTerm(
id<MTLComputeCommandEncoder> enc,
id<MTLBuffer> queries,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> outQTerm,
int nq,
int d,
int M,
bool isL2);
void encodeIVFPQScanListPrecomp(
id<MTLComputeCommandEncoder> enc,
id<MTLBuffer> term2, // nil when useTerm2 is false (IP path)
id<MTLBuffer> qterm,
id<MTLBuffer> coarseDist,
id<MTLBuffer> codes,
id<MTLBuffer> ids,
id<MTLBuffer> listOffset,
id<MTLBuffer> listLength,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> perListDist,
id<MTLBuffer> perListIdx,
int nq,
int M,
int k,
int nprobe,
bool wantMin,
bool useTerm2,
bool useDis0,
id<MTLBuffer> queries,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
int d,
bool onTheFly,
bool shortLists);
void encodeIVFMergeListsGrouped(
id<MTLComputeCommandEncoder> enc,
id<MTLBuffer> inDist,
id<MTLBuffer> inIdx,
id<MTLBuffer> outDist,
id<MTLBuffer> outIdx,
int nq,
int numLists,
int groupSize,
int k,
bool wantMin,
bool compact);
static int selectTopKVariantIndex(int k);
private:
faiss/gpu_metal/MetalKernels.mm
+212
View File
@@ -274,6 +274,218 @@ void MetalKernels::encodeIVFMergeLists(
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
}
void MetalKernels::encodeIVFPQBuildLookupTables(
id<MTLComputeCommandEncoder> enc,
bool isL2,
bool outFp16,
id<MTLBuffer> queries,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> 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<MTLComputeCommandEncoder> enc,
bool useSmall,
bool useFp16Lookup,
id<MTLBuffer> lookupTable,
id<MTLBuffer> codes,
id<MTLBuffer> ids,
id<MTLBuffer> listOffset,
id<MTLBuffer> listLength,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> perListDist,
id<MTLBuffer> perListIdx,
id<MTLBuffer> paramsBuf,
id<MTLBuffer> 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<MTLComputeCommandEncoder> enc,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> 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<MTLComputeCommandEncoder> enc,
id<MTLBuffer> queries,
id<MTLBuffer> pqCentroids,
id<MTLBuffer> 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<MTLComputeCommandEncoder> enc,
id<MTLBuffer> term2,
id<MTLBuffer> qterm,
id<MTLBuffer> coarseDist,
id<MTLBuffer> codes,
id<MTLBuffer> ids,
id<MTLBuffer> listOffset,
id<MTLBuffer> listLength,
id<MTLBuffer> coarseAssign,
id<MTLBuffer> perListDist,
id<MTLBuffer> perListIdx,
int nq,
int M,
int k,
int nprobe,
bool wantMin,
bool useTerm2,
bool useDis0,
id<MTLBuffer> queries,
id<MTLBuffer> coarseCentroids,
id<MTLBuffer> 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<MTLComputeCommandEncoder> enc,
id<MTLBuffer> inDist,
id<MTLBuffer> inIdx,
id<MTLBuffer> outDist,
id<MTLBuffer> 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<MTLDevice> device) {
static std::mutex mu;
static std::unordered_map<uintptr_t, std::unique_ptr<MetalKernels>> map;
faiss/gpu_metal/test/TestMetalIndexIVFPQ.mm
+260
View File
@@ -11,12 +11,16 @@
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/gpu_metal/MetalCloner.h>
#include <faiss/gpu_metal/MetalDistance.h>
#include <faiss/gpu_metal/MetalIndexIVFPQ.h>
#include <faiss/gpu_metal/MetalResources.h>
#include <faiss/gpu_metal/StandardMetalResources.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/utils/random.h>
#include <gtest/gtest.h>
#include <algorithm>
#import <cmath>
#include <map>
#import <memory>
#import <set>
#import <vector>
@@ -170,3 +174,259 @@ TEST_F(TestMetalIndexIVFPQ, CpuToGpuRoundTrip) {
std::unique_ptr<faiss::Index> cpuBackPtr(cpuBack);
ASSERT_EQ(cpuBack->ntotal, nb);
}
TEST_F(TestMetalIndexIVFPQ, LargeMSearchParametersKeepCpuSemantics) {
constexpr int d = 128, nb = 2000, nq = 4, k = 20;
std::vector<float> 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<float> cpuD(nq * k), gpuD(nq * k);
std::vector<faiss::idx_t> cpuI(nq * k), gpuI(nq * k);
cpu.search(nq, xq.data(), k, cpuD.data(), cpuI.data(), &params);
metal.search(nq, xq.data(), k, gpuD.data(), gpuI.data(), &params);
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<uint32_t> lengths = shortLists
? std::vector<uint32_t>{0, 1, 6, 24, 127, 128}
: std::vector<uint32_t>{
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<uint32_t> offsets(nl);
size_t nb = 0;
for (int l = 0; l < nl; ++l) {
offsets[l] = nb;
nb += lengths[l];
}
std::vector<float> 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<uint8_t> codes(nb * M);
for (size_t i = 0; i < codes.size(); ++i)
codes[i] = (i * 131 + i / 7 + i / 257) % 256;
std::vector<int64_t> ids(nb);
for (size_t i = 0; i < nb; ++i)
ids[i] = (int64_t(1) << 40) + i * 17;
std::vector<int32_t> assign(nq * np);
std::vector<float> coarseDist(nq * np);
std::vector<std::vector<std::pair<float, int64_t>>> ref(nq);
std::vector<std::map<int64_t, float>> 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<const float*>(outD.contents);
const auto* gi =
static_cast<const int64_t*>(outI.contents);
for (int q = 0; q < nq; ++q) {
std::set<int64_t> 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);
}
}
}
}
}
}
}
}
}
tests/test_metal_ivfpq.py
+103
View File
@@ -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 <q_m, pq_m> only, while the CPU reference
adds dis0 = <q, centroid> (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,
)