File GpuIndexIVF.h

namespace faiss

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

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

此源代码已获得 MIT 许可证的许可，该许可证位于此源树的根目录下的 LICENSE 文件中。

IDSelector旨在定义要处理的向量子集（用于删除或作为搜索的子集）

PQ4 SIMD 打包和累积函数

基本内核使用 bbs = nb * 2 * 16 个向量累积 nq 个查询向量，并为此生成一个输出矩阵。 对于 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
struct GpuIndexIVFConfig : public faiss::gpu::GpuIndexConfig

被 faiss::gpu::GpuIndexIVFFlatConfig、faiss::gpu::GpuIndexIVFPQConfig、faiss::gpu::GpuIndexIVFScalarQuantizerConfig 继承

公共成员

IndicesOptions indicesOptions = INDICES_64_BIT

GPU 的 Index 存储选项。

GpuIndexFlatConfig flatConfig

粗量化器对象的配置。

bool allowCpuCoarseQuantizer = false

该标志控制索引的粗量化器组件的 CPU 回退逻辑。 当设置为 false（默认值）时，克隆器将为 GPU 上未实现的索引抛出异常。 当设置为 true 时，它将回退到 CPU 实现。

class GpuIndexIVF : public faiss::gpu::GpuIndex, public faiss::IndexIVFInterface
#include <GpuIndexIVF.h>

所有 GPU IVF 索引类型的基类。 这种（目前）故意不继承自 IndexIVF，因为 IndexIVF 中的许多公共数据成员和功能在 GPU 上不支持以相同的方式实现。

由 faiss::gpu::GpuIndexIVFFlat, faiss::gpu::GpuIndexIVFPQ, faiss::gpu::GpuIndexIVFScalarQuantizer 继承

公共函数

GpuIndexIVF(GpuResourcesProvider *provider, int dims, faiss::MetricType metric, float metricArg, idx_t nlist, GpuIndexIVFConfig config = GpuIndexIVFConfig())

基于所需的指标自动构建扁平粗量化器的版本

GpuIndexIVF(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, faiss::MetricType metric, float metricArg, idx_t nlist, GpuIndexIVFConfig config = GpuIndexIVFConfig())

采用粗量化器实例的版本。 默认情况下，GpuIndexIVF 不拥有 coarseQuantizer 实例（如 IndexIVF 等函数）。

~GpuIndexIVF() override
void copyFrom(const faiss::IndexIVF *index)

从 CPU 等效项复制我们需要的内容。

void copyTo(faiss::IndexIVF *index) const

将我们拥有的内容复制到 CPU 等效项。

virtual void updateQuantizer() = 0

如果用户手动更改 IVF 粗量化器的状态（例如，替换新实例或在训练范围之外更改粗量化器中的向量），则应调用该方法。

virtual idx_t getNumLists() const

返回我们正在管理的倒排列表的数量。

virtual idx_t getListLength(idx_t listId) const

返回特定倒排列表中存在的向量数量。

virtual std::vector<uint8_t> getListVectorData(idx_t listId, bool gpuFormat = false) const

返回特定倒排列表中包含的编码向量数据，用于调试目的。如果 gpuFormat 为 true，则数据将以 GPU 端表示形式进行编码后返回。 否则，它将被转换为 CPU 格式。兼容格式，而原生 GPU 格式可能有所不同。

virtual std::vector<idx_t> getListIndices(idx_t listId) const

返回特定倒排列表中包含的向量索引，用于调试目的。

virtual void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const SearchParametersIVF *params = nullptr, IndexIVFStats *stats = nullptr) const override

搜索一组向量，这些向量已由 IVF 量化器预量化。 使用查询结果填充相应的堆。 默认实现使用 InvertedListScanners 进行搜索。

参数:

  • n – 要查询的向量数

  • x – 查询向量，大小为 nx * d

  • assign – 粗量化索引，大小为 nx * nprobe

  • centroid_dis – 到粗糙质心的距离，大小为 nx * nprobe

  • distance – 输出距离，大小为 n * k

  • labels – 输出标签，大小为 n * k

  • store_pairs – 在结果的上/下 32 位中存储 inv 列表索引 + inv 列表偏移量，而不是 ids（用于重新排序）。

  • params – 用于覆盖对象的搜索参数

  • stats – 要更新的搜索统计信息（可以为空）

virtual void range_search_preassigned(idx_t nx, const float *x, float radius, const idx_t *keys, const float *coarse_dis, RangeSearchResult *result, bool store_pairs = false, const IVFSearchParameters *params = nullptr, IndexIVFStats *stats = nullptr) const override

范围搜索一组向量，这些向量已由 IVF 量化器预量化。 填写 RangeSearchResults 结果。 默认实现使用 InvertedListScanners 进行搜索。

参数:

  • n – 要查询的向量数

  • x – 查询向量，大小为 nx * d

  • assign – 粗量化索引，大小为 nx * nprobe

  • centroid_dis – 到粗糙质心的距离，大小为 nx * nprobe

  • result – 输出结果

  • store_pairs – 在结果的上/下 32 位中存储 inv 列表索引 + inv 列表偏移量，而不是 ids（用于重新排序）。

  • params – 用于覆盖对象的搜索参数

  • stats – 要更新的搜索统计信息（可以为空）

保护函数

int getCurrentNProbe_(const SearchParameters *params) const

从当前设置的 nprobe 或 SearchParameters（如果可用）中，返回我们应该用于当前搜索的 nprobe

void verifyIVFSettings_() const
virtual bool addImplRequiresIDs_() const override

addImpl_ 是否需要 ID？如果需要，且没有提供 ID，我们将根据添加 ID 的顺序依次生成 ID。

virtual void trainQuantizer_(idx_t n, const float *x)
virtual void addImpl_(idx_t n, const float *x, const idx_t *ids) override

从 GpuIndex 调用，用于 add/add_with_ids。

virtual void searchImpl_(idx_t n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override

从 GpuIndex 调用，用于搜索。

受保护的属性

const GpuIndexIVFConfig ivfConfig_

我们的配置选项。

std::shared_ptr<IVFBase> baseIndex_

对于经过训练/初始化的索引，这是对基类的引用。

私有函数

void init_()

共享初始化函数。