文件 GpuDistance.h

namespace faiss

实现了多种变体的 k-means 聚类。

版权所有 (c) Facebook, Inc. 及其附属公司。

此源代码根据 MIT 许可获得许可,该许可位于此源树的根目录中的 LICENSE 文件中。

IDSelector 旨在定义要处理的向量子集(用于移除或作为搜索的子集)

PQ4 SIMD 压缩和累加函数

基本内核将 nq 查询向量与 bbs = nb * 2 * 16 个向量累加,并生成该向量的输出矩阵。 对于 nq * nb <= 4 很有用,否则寄存器溢出变得太大。

这些函数的实现分布在 3 个 cpp 文件中,以减少并行编译时间。 模板被显式实例化。

此文件包含用于计算距离的内核的回调。

在整个库中,向量以 float * 指针的形式提供。 当批量处理(添加/搜索)多个向量时,大多数算法都可以得到优化。 在这种情况下,它们作为矩阵传入。 当大小为 d 的 n 个向量作为 float * x 提供时,向量 i 的分量 j 是

x[ i * d + j ]

其中 0 <= i < n 且 0 <= j < d。 换句话说,矩阵始终是紧凑的。 在指定矩阵的大小时,我们称其为 n*d 矩阵,这意味着行优先存储。

I/O 函数可以读取/写入到文件名、文件句柄或抽象介质的对象。

读取函数返回应使用 delete 释放的对象。 这些对象中的所有引用都归该对象所有。

倒排列表的定义 + 一些实现该接口的常见类。

由于 IVF(倒排文件)索引对于大规模用例非常有用,因此我们将与它们相关的一些函数组合在这个小型库中。 大多数函数都可以在 IndexIVF 和嵌入在 IndexPreTransform 中的 IndexIVF 上工作。

此文件包含 L2 和内积之外的额外指标的实现

实现一些神经网络层,主要用于支持 QINCo

定义一些将变换应用于一组向量的对象 通常这些是预处理步骤。

namespace gpu

枚举

enum class DistanceDataType

enumerator F32
枚举值 F16
枚举值 BF16
枚举 IndicesDataType

枚举值 I64
枚举值 I32

函数

bool should_use_cuvs(GpuDistanceParams args)

一个根据各种条件(例如,不支持的架构)确定是否应使用 cuVS 的函数

void bfKnn(GpuResourcesProvider *resources, const GpuDistanceParams &args)

gpu/impl/Distance.cuh 的一个包装器,用于在外部提供的内存区域(例如,来自 pytorch 张量)上公开直接的暴力 k 近邻搜索。 数据(向量、查询、outDistances、outIndices)可以驻留在 GPU 或 CPU 上,但所有计算都在 GPU 上执行。 如果结果缓冲区位于 CPU 上,则完成后会将结果复制回来。

所有 GPU 计算都在当前的 CUDA 设备上执行,并且相对于 resources->getDefaultStreamCurrentDevice() 排序。

对于 queries 中的每个向量,搜索所有 vectors 以找到其 k 个最近邻,并以给定的度量为准

void bfKnn_tiling(GpuResourcesProvider *resources, const GpuDistanceParams &args, size_t vectorsMemoryLimit, size_t queriesMemoryLimit)
void bruteForceKnn(GpuResourcesProvider *resources, faiss::MetricType metric, const float *vectors, bool vectorsRowMajor, idx_t numVectors, const float *queries, bool queriesRowMajor, idx_t numQueries, int dims, int k, float *outDistances, idx_t *outIndices)

已弃用的旧版实现。

结构体 GpuDistanceParams
#include <GpuDistance.h>

暴力 GPU k 近邻搜索的参数。

公共成员

faiss::MetricType metric = METRIC_L2

搜索参数:距离度量。

float metricArg = 0

搜索参数:距离度量参数(如果适用)。 对于 metric == METRIC_Lp,这是 p 值

int k = 0

搜索参数:返回 k 个最近邻 如果提供的值为 -1,则我们报告所有成对距离,而不进行 top-k 过滤

int dims = 0

向量维度。

const void *vectors = nullptr

如果 vectorsRowMajor 为 true,则这是 numVectors x dims,其中 dims 是最里面的;否则,dims x numVectors,其中 numVectors 是最里面的

DistanceDataType vectorType = DistanceDataType::F32
bool vectorsRowMajor = true
idx_t numVectors = 0
const float *vectorNorms = nullptr

vectors 中的每个向量预先计算的 L2 范数,可以选择提前提供以加速 METRIC_L2 的计算

const void *queries = nullptr

如果 queriesRowMajor 为 true,则这是 numQueries x dims,dims 在最内层;否则,dims x numQueries,numQueries 在最内层

DistanceDataType queryType = DistanceDataType::F32
bool queriesRowMajor = true
idx_t numQueries = 0
float *outDistances = nullptr

大小为 numQueries x k 的内存区域,如果 k > 0,则 k 在最内层(行优先);如果 k == -1,则大小为 numQueries x numVectors 的内存区域

bool ignoreOutDistances = false

我们是否只关心报告的索引,而不是输出距离? 如果 k == -1(所有成对距离)则不使用

IndicesDataType outIndicesType = IndicesDataType::I64

大小为 numQueries x k 的内存区域,k 在最内层(行优先)。 如果 k == -1(所有成对距离)则不使用

void *outIndices = nullptr
int device = -1

搜索应在哪个 GPU 设备上运行? -1 表示使用当前 CUDA 线程局部设备(通过 cudaGetDevice/cudaSetDevice) 否则,整数 0 <= device < numDevices 表示执行设备

bool use_cuvs = false

索引是否应分派到 cuVS?