文件 GpuIndexIVFPQ.h

namespace faiss

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

此源代码已获得 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 GpuIndexIVFPQConfig : public faiss::gpu::GpuIndexIVFConfig 

公共成员

bool useFloat16LookupTables = false: 是否在列表扫描内核中使用float16残差距离表。当subQuantizers * 2^bitsPerCode > 16384时，这是必需的。

bool usePrecomputedTables = false: 是否启用搜索的预计算表选项，这会大大增加内存需求。

bool interleavedLayout = false: 使用IVF列表的替代内存布局。警告：这是一个正在开发中的功能，仅在使用cuVS启用索引时才受支持。如果未启用cuVS，请勿使用。

bool useMMCodeDistance = false

对于IVFPQ的无预计算表版本，使用GEMM支持的PQ码距离计算。这是为了调试目的，它不应该实质性地影响结果。

请注意，如果一个人使用每个子量化器的维度数不是本机专门化的（例如7之类的奇数），则会自动启用MM码距离。

class GpuIndexIVFPQ : public faiss::gpu::GpuIndexIVF

#include <GpuIndexIVFPQ.h>

用于GPU的IVFPQ索引。

公共函数

GpuIndexIVFPQ(GpuResourcesProvider *provider, const faiss::IndexIVFPQ *index, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig()): 从预先存在的 faiss::IndexIVFPQ 实例构造，如果输入索引已训练，则将数据复制到给定的GPU。

GpuIndexIVFPQ(GpuResourcesProvider *provider, int dims, idx_t nlist, idx_t subQuantizers, idx_t bitsPerCode, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig()): 构造一个带有空flat量化器的新实例；用户提供所需的IVF列表数量。

GpuIndexIVFPQ(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, idx_t nlist, idx_t subQuantizers, idx_t bitsPerCode, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig()): 使用提供的 CPU 或 GPU 粗量化器构造新实例；用户提供所需的 IVF 列表数量。

~GpuIndexIVFPQ() override

void copyFrom(const faiss::IndexIVFPQ *index): 在 GPU 上为 num 向量的倒排列表保留空间，假设这些向量平均分配。从给定的 CPU 索引初始化自身；将覆盖自身中的所有数据

void copyTo(faiss::IndexIVFPQ *index) const: 将自身复制到给定的 CPU 索引；将覆盖索引实例中的所有数据

void reserveMemory(size_t numVecs): 在倒排列表中为此数量的向量保留 GPU 内存。

void setPrecomputedCodes(bool enable): 启用或禁用预计算代码。

bool getPrecomputedCodes() const: 是否启用了预计算代码？

int getNumSubQuantizers() const: 返回我们正在使用的子量化器的数量。

int getBitsPerCode() const: 返回每个 PQ 代码的位数。

int getCentroidsPerSubQuantizer() const: 返回每个 PQ 代码的质心数量（2^每个代码的位数）

size_t reclaimMemory(): 在添加向量之后，可以调用此函数来回收设备内存，使其精确到所需的量。返回以字节为单位回收的空间

virtual void reset() override: 清除所有倒排列表，但保留粗量化和乘积质心信息

virtual void updateQuantizer() override: 如果用户手动更改 IVF 粗量化器的状态（例如，替换新实例或在训练范围之外更改粗量化器中的向量），则应调用此函数

virtual void train(idx_t n, const float *x) override: 根据给定的向量数据训练粗量化器和乘积量化器。

公共成员

ProductQuantizer pq: 与 CPU 版本一样，我们公开一个可公开访问的 ProductQuantizer 以进行操作

保护函数

void setIndex_(GpuResources *resources, int dim, idx_t nlist, faiss::MetricType metric, float metricArg, int numSubQuantizers, int bitsPerSubQuantizer, bool useFloat16LookupTables, bool useMMCodeDistance, bool interleavedLayout, float *pqCentroidData, IndicesOptions indicesOptions, MemorySpace space): 初始化适当的索引。

void verifyPQSettings_() const: 如果配置设置不正确，则抛出错误。

void trainResidualQuantizer_(idx_t n, const float *x): 基于给定的向量数据训练PQ量化器。

受保护的属性

const GpuIndexIVFPQConfig ivfpqConfig_: 我们初始化时使用的配置选项。

bool usePrecomputedTables_: 运行时覆盖：我们是否使用预计算表。

int subQuantizers_: 每个编码向量的子量化器的数量。

int bitsPerCode_: 每个子量化器代码的位数。

size_t reserveMemoryVecs_: 期望的倒排列表内存预留。

std::shared_ptr<IVFPQ> index_: 我们拥有的乘积量化器实例；包含倒排列表