命名空间 faiss::gpu

namespace gpu

枚举

enum class DistanceDataType

enumerator F32
enumerator F16
enumerator BF16
enum class IndicesDataType

enumerator I64
enumerator I32
enum class graph_build_algo

enumerator IVF_PQ

使用 IVF-PQ 构建全邻近 knn 图。

enumerator NN_DESCENT

使用 NN-Descent 构建全邻近 knn 图。

enum class codebook_gen

一种用于指定如何创建 PQ 码本的类型。

enumerator PER_SUBSPACE
enumerator PER_CLUSTER
enum class search_algo

enumerator SINGLE_CTA

对于大批量大小。

enumerator MULTI_CTA

对于小批量大小。

enumerator MULTI_KERNEL
enumerator AUTO
enum class hash_mode

enumerator HASH
enumerator SMALL
enumerator AUTO
enum IndicesOptions

用户向量索引数据在 GPU 上的存储方式。

enumerator INDICES_CPU

用户索引仅存储在 CPU 上;GPU 返回(倒排列表,偏移量)给 CPU,然后将其转换为实际用户索引。

enumerator INDICES_IVF

索引根本不存储在 CPU 或 GPU 上。 只有(倒排列表,偏移量)作为索引返回给用户。

enumerator INDICES_32_BIT

索引在 GPU 上存储为 32 位整数,但返回为 64 位整数

enumerator INDICES_64_BIT

索引在 GPU 上存储为 64 位整数。

enum AllocType

enumerator Other

未知分配类型或杂项(当前未分类)

enumerator FlatData

GpuIndexFlat 的主数据存储(向量的原始矩阵和如果需要,向量范数)

enumerator IVFLists

GpuIndexIVF* 的主数据存储(每个单独的 IVF 列表的存储)

enumerator Quantizer

Quantizer (PQ, SQ) 字典信息。

enumerator QuantizerPrecomputedCodes

对于 GpuIndexIVFPQ,更高效的 PQ 查找的“预计算代码”需要使用可能很大的表。 这些与 Quantizer 分开标记,因为它们通常可以是 100s - 1000s MiB 大小

枚举器 TemporaryMemoryBuffer

StandardGpuResources 特定于实现的类型。当使用 StandardGpuResources 时,临时内存分配 (MemorySpace::Temporary) 来自于为每个 GPU 预先分配的堆栈内存区域(例如,初始化时为 1.5 GiB)。StandardGpuResources 的此分配被标记为此 AllocType。

枚举器 TemporaryMemoryOverflow

当使用 StandardGpuResources 时,任何在 TemporaryMemoryBuffer 区域内无法满足的 MemorySpace::Temporary 分配都会回退到调用 cudaMalloc,其大小仅为当前请求的大小。这些“溢出”临时分配被标记为此 AllocType。

枚举 MemorySpace

GPU 可访问的内存区域。

枚举器 Temporary

临时设备内存(保证在顶级索引调用退出后不再使用,并且使用它的流已完成 GPU 工作)。通常由设备内存(cudaMalloc/cudaFree)支持。

枚举器 Device

使用 cudaMalloc/cudaFree 管理(典型的 GPU 设备内存)

枚举器 Unified

使用 cudaMallocManaged/cudaFree 管理(典型的统一 CPU/GPU 内存)

函数

faiss::Index *index_gpu_to_cpu(const faiss::Index *gpu_index)

将 gpu_index 中的任何 GPU 索引转换为 CPU 索引

faiss::Index *index_cpu_to_gpu(GpuResourcesProvider *provider, int device, const faiss::Index *index, const GpuClonerOptions *options = nullptr)

将任何可以转换为 GPU 的 CPU 索引进行转换

faiss::Index *index_cpu_to_gpu_multiple(std::vector<GpuResourcesProvider*> &provider, std::vector<int> &devices, const faiss::Index *index, const GpuMultipleClonerOptions *options = nullptr)
faiss::IndexBinary *index_binary_gpu_to_cpu(const faiss::IndexBinary *gpu_index)
faiss::IndexBinary *index_binary_cpu_to_gpu(GpuResourcesProvider *provider, int device, const faiss::IndexBinary *index, const GpuClonerOptions *options = nullptr)

将任何可以转换为 GPU 的 CPU 索引进行转换

将 CPU 上的二进制索引复制到多个 GPU 上。
确定是否应该使用 cuVS 的函数,基于各种条件(例如不支持的架构)。

A function that determines whether cuVS should be used based on various conditions (such as unsupported architecture)

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

A wrapper for gpu/impl/Distance.cuh to expose direct brute-force k-nearest neighbor searches on an externally-provided region of memory (e.g., from a pytorch tensor). The data (vectors, queries, outDistances, outIndices) can be resident on the GPU or the CPU, but all calculations are performed on the GPU. If the result buffers are on the CPU, results will be copied back when done.

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

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

void bfKnn_tiling(GpuResourcesProvider *resources, const GpuDistanceParams &args, size_t vectorsMemoryLimit, size_t queriesMemoryLimit)
暴力 k 最近邻搜索。

已弃用的旧实现。

bool should_use_cuvs(GpuIndexConfig config_)

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

GpuIndex *tryCastGpuIndex(faiss::Index *index)

如果给定的索引是 GPU 索引,则返回该索引实例。

bool isGpuIndex(faiss::Index *index)

给定的索引实例是否为 GPU 索引?

bool isGpuIndexImplemented(faiss::Index *index)

给定的 CPU 索引实例是否具有对应的 GPU 实现?

std::string allocTypeToString(AllocType t)

将 AllocType 转换为字符串。

std::string memorySpaceToString(MemorySpace s)

将 MemorySpace 转换为字符串。

AllocInfo makeDevAlloc(AllocType at, cudaStream_t st)

为具有 MemorySpace::Device 的当前设备创建 AllocInfo

AllocInfo makeTempAlloc(AllocType at, cudaStream_t st)

为具有 MemorySpace::Temporary 的当前设备创建 AllocInfo

AllocInfo makeSpaceAlloc(AllocType at, MemorySpace sp, cudaStream_t st)

为当前设备创建 AllocInfo

int getMaxKSelection()

索引实现的各种实用程序函数的集合。

返回基于我们编译时使用的 CUDA SDK 所支持的最大 k 选择值。 .cu 文件可以使用 DeviceDefs.cuh,但这是针对非 CUDA 文件

void validateKSelect(int k)
void validateNProbe(size_t nprobe)
std::vector<uint8_t> unpackNonInterleaved(std::vector<uint8_t> data, int numVecs, int dims, int bitsPerCode)
std::vector<uint8_t> unpackInterleaved(std::vector<uint8_t> data, int numVecs, int dims, int bitsPerCode)
std::vector<uint8_t> packNonInterleaved(std::vector<uint8_t> data, int numVecs, int dims, int bitsPerCode)
std::vector<uint8_t> packInterleaved(std::vector<uint8_t> data, int numVecs, int dims, int bitsPerCode)
void ivfOffsetToUserIndex(idx_t *indices, idx_t numLists, idx_t queries, int k, const std::vector<std::vector<idx_t>> &listOffsetToUserIndex)

在 CPU 上将(列表 ID,偏移量)转换为用户索引的实用函数。位于 cpp 中是为了使用 OpenMP。

void newTestSeed()

为测试生成并显示一个新的种子。

void setTestSeed(long seed)

为测试使用显式种子。

float relativeError(float a, float b)

返回 a 和 b 之间差异的相对误差 (|a - b| / (0.5 * (|a| + |b|))

int randVal(int a, int b)

生成范围 [a, b] 内的随机整数。

bool randBool()

生成一个随机布尔值。

template<typename T>
T randSelect(std::initializer_list<T> vals)

从作为 initializer_list 提供的给定值列表中选择一个随机值

std::vector<float> randVecs(size_t num, size_t dim)

生成范围 [0, 1] 内的随机向量集合。

std::vector<unsigned char> randBinaryVecs(size_t num, size_t dim)

生成随机位向量的集合。

std::vector<float> roundToHalf(const std::vector<float> &v)
void compareIndices(const std::vector<float> &queryVecs, faiss::Index &refIndex, faiss::Index &testIndex, int numQuery, int dim, int k, const std::string &configMsg, float maxRelativeError = 6e-5f, float pctMaxDiff1 = 0.1f, float pctMaxDiffN = 0.005f)

通过查询相似度比较两个索引,使用用户指定的一组查询向量

void compareIndices(faiss::Index &refIndex, faiss::Index &testIndex, int numQuery, int dim, int k, const std::string &configMsg, float maxRelativeError = 6e-5f, float pctMaxDiff1 = 0.1f, float pctMaxDiffN = 0.005f)

Compare two indices via query for similarity, generating random query vectors

void compareLists(const float *refDist, const faiss::idx_t *refInd, const float *testDist, const faiss::idx_t *testInd, int dim1, int dim2, const std::string &configMsg, bool printBasicStats, bool printDiffs, bool assertOnErr, float maxRelativeError = 6e-5f, float pctMaxDiff1 = 0.1f, float pctMaxDiffN = 0.005f)

Display specific differences in the two (distance, index) lists.

template<typename A, typename B>
void testIVFEquality(A &cpuIndex, B &gpuIndex)

Compare IVF lists between a CPU and GPU index.

inline cuvsDistanceType metricFaissToCuvs(MetricType metric, bool exactDistance)
void validRowIndices(GpuResources *res, Tensor<float, 2, true> &vecs, bool *validRows)

Identify matrix rows containing non NaN values. validRows[i] is false if row i contains a NaN value and true otherwise.

idx_t inplaceGatherFilteredRows(GpuResources *res, Tensor<float, 2, true> &vecs, Tensor<idx_t, 1, true> &indices)

过滤掉包含 NaN 值的矩阵行。向量和索引会原地更新。

int getCurrentDevice()

返回当前线程局部 GPU 设备。

void setCurrentDevice(int device)

设置当前线程局部 GPU 设备。

int getNumDevices()

返回可用 GPU 设备的数量。

void profilerStart()

启动 CUDA 分析器(通过 SWIG 公开)

void profilerStop()

停止 CUDA 分析器(通过 SWIG 公开)

void synchronizeAllDevices()

将 CPU 与所有设备同步(相当于每个设备的 cudaDeviceSynchronize)

const cudaDeviceProp &getDeviceProperties(int device)

返回给定设备的缓存 cudaDeviceProp。

const cudaDeviceProp &getCurrentDeviceProperties()

返回当前设备的缓存 cudaDeviceProp。

int getMaxThreads(int device)

返回给定 GPU 设备可用的最大线程数

int getMaxThreadsCurrentDevice()

等效于 getMaxThreads(getCurrentDevice())

dim3 getMaxGrid(int device)

返回给定 GPU 设备的最大网格大小。

dim3 getMaxGridCurrentDevice()

等效于 getMaxGrid(getCurrentDevice())

size_t getMaxSharedMemPerBlock(int device)

返回给定 GPU 设备可用的最大 smem。

size_t getMaxSharedMemPerBlockCurrentDevice()

等效于 getMaxSharedMemPerBlock(getCurrentDevice())

int getDeviceForAddress(const void *p)

对于给定的指针,返回它是否位于设备上(deviceId >= 0)或主机上(-1)。

bool getFullUnifiedMemSupport(int device)

给定的设备是否支持完全统一的内存共享主机内存?

bool getFullUnifiedMemSupportCurrentDevice()

等效于 getFullUnifiedMemSupport(getCurrentDevice())

bool getTensorCoreSupport(int device)

给定的设备是否支持张量核心操作?

bool getTensorCoreSupportCurrentDevice()

等效于 getTensorCoreSupport(getCurrentDevice())

int getWarpSize(int device)

返回给定 GPU 设备的 warp 大小。

int getWarpSizeCurrentDevice()

等效于 getWarpSize(getCurrentDevice())

size_t getFreeMemory(int device)

返回给定设备上当前可用的内存量。

size_t getFreeMemoryCurrentDevice()

等效于 getFreeMemory(getCurrentDevice())

template<typename L1, typename L2>
void streamWaitBase(const L1 &listWaiting, const L2 &listWaitOn)

调用一个流集合等待。

template<typename L1>
void streamWait(const L1 &a, const std::initializer_list<cudaStream_t> &b)

这些版本允许使用 initializer_list 作为参数,因为否则 {…} 没有类型

template<typename L2>
void streamWait(const std::initializer_list<cudaStream_t> &a, const L2 &b)
inline void streamWait(const std::initializer_list<cudaStream_t> &a, const std::initializer_list<cudaStream_t> &b)
struct GpuParameterSpace : public faiss::ParameterSpace
#include <GpuAutoTune.h>

GPU 索引的参数空间和设置器

公共函数

virtual void initialize(const faiss::Index *index) override

使用索引的合理参数进行初始化

virtual void set_index_parameter(faiss::Index *index, const std::string &name, double val) const override

设置索引上的一组参数

size_t n_combinations() const

组合数, = 值大小的乘积

bool combination_ge(size_t c1, size_t c2) const

返回元组意义上组合 c1 >= c2 是否成立

std::string combination_name(size_t cno) const

获取组合的字符串表示

void display() const

在 stdout 上打印描述

ParameterRange &add_range(const std::string &name)

添加一个新参数(如果它存在则返回它)

void set_index_parameters(Index *index, size_t cno) const

设置索引上的一组参数

void set_index_parameters(Index *index, const char *param_string) const

设置由字符串描述的参数组合

void update_bounds(size_t cno, const OperatingPoint &op, double *upper_bound_perf, double *lower_bound_t) const

针对给定操作点 op,查找配置 cno 的性能上限和 t 的下限

void explore(Index *index, size_t nq, const float *xq, const AutoTuneCriterion &crit, OperatingPoints *ops) const

探索操作点

参数:
  • index – 要在其上运行的索引

  • xq – 查询向量(大小为 nq * index.d)

  • crit – 选择标准

  • ops – 生成的操作点

公共成员

std::vector<ParameterRange> parameter_ranges

所有可调参数

int verbose

探索期间的详细程度

int n_experiments

优化期间的实验次数(0 = 尝试所有组合)

size_t batchsize

一次提交的最大查询数。

bool thread_over_batches

在批次上使用多线程(对于基准测试独立的单次搜索非常有用)

double min_test_duration

多次运行测试,直到它们至少达到此持续时间(以避免 MT 模式下的抖动)

struct ToCPUCloner : public faiss::Cloner
#include <GpuCloner.h>

为 GPU -> CPU 专门设计的Cloner

公共函数

void merge_index(Index *dst, Index *src, bool successive_ids)
virtual Index *clone_Index(const Index *index) override
virtual VectorTransform *clone_VectorTransform(const VectorTransform*)
virtual IndexIVF *clone_IndexIVF(const IndexIVF*)
struct ToGpuCloner : public faiss::Cloner, public faiss::gpu::GpuClonerOptions
#include <GpuCloner.h>

Cloner 专门用于 CPU -> 1 GPU 的克隆。

公共函数

ToGpuCloner(GpuResourcesProvider *prov, int device, const GpuClonerOptions &options)
virtual Index *clone_Index(const Index *index) override
virtual VectorTransform *clone_VectorTransform(const VectorTransform*)
virtual IndexIVF *clone_IndexIVF(const IndexIVF*)

公共成员

GpuResourcesProvider *provider
int device
IndicesOptions indicesOptions = INDICES_64_BIT

索引应该如何在支持索引的索引类型上存储(除了 GpuIndexFlat* 之外的任何类型)?

bool useFloat16CoarseQuantizer = false

粗量化器是否使用 float16?

bool useFloat16 = false

对于 GpuIndexIVFFlat, 存储是否使用 float16? 对于 GpuIndexIVFPQ, 中间计算是否使用 float16?

bool usePrecomputed = false

使用预计算表吗?

long reserveVecs = 0

在 invfiles 中保留向量?

bool storeTransposed = false

对于 GpuIndexFlat, 以转置布局存储数据?

bool verbose = false

在索引上设置 verbose 选项。

bool use_cuvs = false

使用 cuVS 实现

bool allowCpuCoarseQuantizer = false

此标志控制索引的粗量化器组件的 CPU 回退逻辑。如果设置为 false (默认),克隆器将为 GPU 上未实现的索引引发异常。 如果设置为 true,它将回退到 CPU 实现。

struct ToGpuClonerMultiple : public faiss::Cloner, public faiss::gpu::GpuMultipleClonerOptions
#include <GpuCloner.h>

Cloner 专门用于 CPU -> 多个 GPU。

公共函数

ToGpuClonerMultiple(std::vector<GpuResourcesProvider*> &provider, std::vector<int> &devices, const GpuMultipleClonerOptions &options)
ToGpuClonerMultiple(const std::vector<ToGpuCloner> &sub_cloners, const GpuMultipleClonerOptions &options)
void copy_ivf_shard(const IndexIVF *index_ivf, IndexIVF *idx2, idx_t n, idx_t i)
Index *clone_Index_to_shards(const Index *index)
virtual Index *clone_Index(const Index *index) override

主函数

virtual VectorTransform *clone_VectorTransform(const VectorTransform*)
virtual IndexIVF *clone_IndexIVF(const IndexIVF*)

公共成员

std::vector<ToGpuCloner> sub_cloners
bool shard = false

是否跨 GPU 分片索引,而不是跨 GPU 复制。

int shard_type = 1

IndexIVF::copy_subset_to 子集类型。

bool common_ivf_quantizer = false

如果一个 IndexIVF 要被分派到多个 GPU,且这些 GPU 共享一个通用的 IVF 量化器,则设置为 true。即,只有倒排列表会被分片到子索引上(使用 IndexShardsIVF)。

IndicesOptions indicesOptions = INDICES_64_BIT

索引应该如何在支持索引的索引类型上存储(除了 GpuIndexFlat* 之外的任何类型)?

bool useFloat16CoarseQuantizer = false

粗量化器是否使用 float16?

bool useFloat16 = false

对于 GpuIndexIVFFlat, 存储是否使用 float16? 对于 GpuIndexIVFPQ, 中间计算是否使用 float16?

bool usePrecomputed = false

使用预计算表吗?

long reserveVecs = 0

在 invfiles 中保留向量?

bool storeTransposed = false

对于 GpuIndexFlat, 以转置布局存储数据?

bool verbose = false

在索引上设置 verbose 选项。

bool use_cuvs = false

使用 cuVS 实现

bool allowCpuCoarseQuantizer = false

此标志控制索引的粗量化器组件的 CPU 回退逻辑。如果设置为 false (默认),克隆器将为 GPU 上未实现的索引引发异常。 如果设置为 true,它将回退到 CPU 实现。

struct GpuProgressiveDimIndexFactory : public faiss::ProgressiveDimIndexFactory
#include <GpuCloner.h>

用于 ProgressiveDimClustering 对象的索引工厂

公共函数

explicit GpuProgressiveDimIndexFactory(int ngpu)
virtual Index *operator()(int dim) override

所有权转移给调用者

virtual ~GpuProgressiveDimIndexFactory() override

公共成员

GpuMultipleClonerOptions options
std::vector<GpuResourcesProvider*> vres
std::vector<int> devices
int ncall
struct GpuClonerOptions
#include <GpuClonerOptions.h>

设置如何复制到 GPU 的一些选项

子类包括 faiss::gpu::GpuMultipleClonerOptions, faiss::gpu::ToGpuCloner

公共成员

IndicesOptions indicesOptions = INDICES_64_BIT

索引应该如何在支持索引的索引类型上存储(除了 GpuIndexFlat* 之外的任何类型)?

bool useFloat16CoarseQuantizer = false

粗量化器是否使用 float16?

bool useFloat16 = false

对于 GpuIndexIVFFlat, 存储是否使用 float16? 对于 GpuIndexIVFPQ, 中间计算是否使用 float16?

bool usePrecomputed = false

使用预计算表吗?

long reserveVecs = 0

在 invfiles 中保留向量?

bool storeTransposed = false

对于 GpuIndexFlat, 以转置布局存储数据?

是否启用详细模式,默认为false。

在索引上设置 verbose 选项。

是否使用CUVS(CUDA Unified Virtual Addressing),默认为false。

使用 cuVS 实现

是否允许CPU粗量化器,默认为false。

此标志控制索引的粗量化器组件的 CPU 回退逻辑。如果设置为 false (默认),克隆器将为 GPU 上未实现的索引引发异常。 如果设置为 true,它将回退到 CPU 实现。

`GpuMultipleClonerOptions` 结构体,继承自 `GpuClonerOptions`。

该结构体被 `faiss::gpu::ToGpuClonerMultiple` 类继承。

公共成员

是否进行分片(sharding),默认为false。

是否跨 GPU 分片索引,而不是跨 GPU 复制。

分片类型,默认为1。

IndexIVF::copy_subset_to 子集类型。

是否使用通用的IVF量化器,默认为false。

如果一个 IndexIVF 要被分派到多个 GPU,且这些 GPU 共享一个通用的 IVF 量化器,则设置为 true。即,只有倒排列表会被分片到子索引上(使用 IndexShardsIVF)。

索引选项,默认为 `INDICES_64_BIT` (64位索引)。

索引应该如何在支持索引的索引类型上存储(除了 GpuIndexFlat* 之外的任何类型)?

是否使用Float16粗量化器,默认为false。

粗量化器是否使用 float16?

是否使用Float16,默认为false。

对于 GpuIndexIVFFlat, 存储是否使用 float16? 对于 GpuIndexIVFPQ, 中间计算是否使用 float16?

是否使用预计算的结果,默认为false。

使用预计算表吗?

保留的向量数量,默认为0。

在 invfiles 中保留向量?

是否存储转置矩阵,默认为false。

对于 GpuIndexFlat, 以转置布局存储数据?

是否启用详细模式,默认为false。(继承自GpuClonerOptions)

在索引上设置 verbose 选项。

是否使用CUVS(CUDA Unified Virtual Addressing),默认为false。(继承自GpuClonerOptions)

使用 cuVS 实现

是否允许CPU粗量化器,默认为false。(继承自GpuClonerOptions)

此标志控制索引的粗量化器组件的 CPU 回退逻辑。如果设置为 false (默认),克隆器将为 GPU 上未实现的索引引发异常。 如果设置为 true,它将回退到 CPU 实现。

`GpuDistanceParams` 结构体。
#include <GpuDistance.h>

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

公共成员

距离度量标准,默认为 `METRIC_L2` (欧几里得距离)。

搜索参数:距离度量。

距离度量参数,默认为0。 对于 metric == METRIC_Lp,这是p值

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

返回 k 个最近邻,默认为0. 如果提供的值为 -1,则我们报告所有成对距离而不进行 top-k 过滤

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

向量维度,默认为0。

向量维度。

向量数据指针,默认为nullptr。 如果 vectorsRowMajor 为 true,则是 numVectors x dims,其中 dims 是最内层;否则,dims x numVectors,其中 numVectors 是最内层

如果 vectorsRowMajor 为 true,则是 numVectors x dims,其中 dims 是最内层;否则,dims x numVectors,其中 numVectors 是最内层。

向量数据类型,默认为 `F32` (float)。
向量是否按行优先存储,默认为true。
向量的数量,默认为0。
向量的范数数据指针,默认为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?

class GpuIcmEncoder : public faiss::lsq::IcmEncoder
#include <GpuIcmEncoder.h>

在 GPU 上执行 LSQ 编码。

将输入向量拆分到不同的设备,并调用 IcmEncoderImpl::encode 对其进行编码

公共函数

GpuIcmEncoder(const LocalSearchQuantizer *lsq, const std::vector<GpuResourcesProvider*> &provs, const std::vector<int> &devices)
~GpuIcmEncoder()
GpuIcmEncoder(const GpuIcmEncoder&) = delete
GpuIcmEncoder &operator=(const GpuIcmEncoder&) = delete
virtual void set_binary_term() override
virtual void encode(int32_t *codes, const float *x, std::mt19937 &gen, size_t n, size_t ils_iters) const override

使用给定的码本编码向量

参数:
  • codes – 输出码,大小为 n * M

  • x – 要编码的向量,大小为 n * d

  • gen – 随机数生成器

  • n – 向量数量

  • ils_iters – 迭代局部搜索的迭代次数

公共成员

std::vector<float> binaries
bool verbose
const LocalSearchQuantizer *lsq

私有成员

std::unique_ptr<IcmEncoderShards> shards
struct GpuIcmEncoderFactory : public faiss::lsq::IcmEncoderFactory

公共函数

explicit GpuIcmEncoderFactory(int ngpus = 1)
virtual lsq::IcmEncoder *get(const LocalSearchQuantizer *lsq) override

公共成员

std::vector<GpuResourcesProvider*> provs
std::vector<int> devices
struct GpuIndexConfig

faiss::gpu::GpuIndexBinaryFlatConfig, faiss::gpu::GpuIndexCagraConfig, faiss::gpu::GpuIndexFlatConfig, faiss::gpu::GpuIndexIVFConfig 继承

公共成员

int device = 0

索引所在的 GPU 设备。

MemorySpace memorySpace = MemorySpace::Device

用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。

bool use_cuvs = false

索引是否应调度到 cuVS?

class GpuIndex : public faiss::Index

faiss::gpu::GpuIndexCagra, faiss::gpu::GpuIndexFlat, faiss::gpu::GpuIndexIVF 继承

公共类型

using component_t = float
using distance_t = float

公共函数

GpuIndex(std::shared_ptr<GpuResources> resources, int dims, faiss::MetricType metric, float metricArg, GpuIndexConfig config)
int getDevice() const

返回此索引所在的设备。

std::shared_ptr<GpuResources> getResources()

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

void setMinPagingSize(size_t size)

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

size_t getMinPagingSize() const

返回分页搜索的当前最小数据大小。

virtual void add(idx_t, const float *x) override

x 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

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

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void train(idx_t n, const float *x)

在一组代表性向量上执行训练

参数:
  • n – 训练向量的 nb

  • x – 训练向量,大小为 n * d

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

virtual void reset() = 0

从数据库中删除所有元素。

virtual size_t remove_ids(const IDSelector &sel)

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

virtual void reconstruct(idx_t key, float *recons) const

重建存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • key – 要重建的向量的 id

  • recons – 重建的向量(大小为 d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

重建多个存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • n – 要重建的向量数

  • keys – 要重建的向量的 id(大小为 n)

  • recons – 重建的向量(大小为 n * d)

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

重建向量 i0 到 i0 + ni - 1

对于某些索引,可能未定义此函数

参数:
  • i0 – 序列中第一个向量的索引

  • ni – 序列中向量的数量

  • recons – 重建的向量(大小为 ni * d)

virtual DistanceComputer *get_distance_computer() const

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

virtual size_t sa_code_size() const

生成的代码的大小,以字节为单位

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

virtual void add_sa_codes(idx_t n, const uint8_t *codes, const idx_t *xids)

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

公共成员

int d

向量维度

idx_t ntotal

索引向量的总数

bool verbose

详细级别

bool is_trained

如果 Index 不需要训练,或者如果已经完成训练,则设置

MetricType metric_type

此索引用于搜索的度量类型

float metric_arg

度量类型的参数

受保护函数

void copyFrom(const faiss::Index *index)

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

void copyTo(faiss::Index *index) const

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

virtual bool addImplRequiresIDs_() const = 0

addImpl_ 是否需要 ID? 如果是,并且未提供 ID,我们将根据添加 ID 的顺序依次生成它们

virtual void addImpl_(idx_t n, const float *x, const idx_t *ids) = 0

重写此函数以实际执行添加。所有数据都保证驻留在我们的设备上。

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

重写此函数以实际执行搜索。所有数据都保证驻留在我们的设备上。

受保护的属性

std::shared_ptr<GpuResources> resources_

管理设备上的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

将副本从 CPU 分页到 GPU 的大小阈值。

私有函数

void addPaged_(idx_t n, const float *x, const idx_t *ids)

如果添加集太大,则处理分页添加,传递给 addImpl_ 以实际执行当前页面的添加

void addPage_(idx_t n, const float *x, const idx_t *ids)

为单个 GPU 驻留数据页面调用 addImpl_。

void searchNonPaged_(idx_t n, const float *x, int k, float *outDistancesData, idx_t *outIndicesData, const SearchParameters *params) const

为单个 GPU 驻留数据页面调用 searchImpl_。

void searchFromCpuPaged_(idx_t n, const float *x, int k, float *outDistancesData, idx_t *outIndicesData, const SearchParameters *params) const

为单个 GPU 驻留数据页面调用 searchImpl_,处理数据分页和从 CPU 复制

struct GpuIndexBinaryFlatConfig : public faiss::gpu::GpuIndexConfig

公共成员

设备 ID,默认为 0。

索引所在的 GPU 设备。

内存空间,默认为 Device。

用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。

是否使用 cuvs,默认为 false。

索引是否应调度到 cuVS?

GpuIndexBinaryFlat 类继承自 faiss::IndexBinary
包含头文件 <GpuIndexBinaryFlat.h>。

IndexBinaryFlat 的 GPU 版本,用于通过汉明距离进行位向量的暴力比较。

公共类型

定义别名 component_tuint8_t
定义别名 distance_tint32_t

公共函数

从预先存在的 faiss::IndexBinaryFlat 实例构造,将数据复制到给定的 GPU。

Construct from a pre-existing faiss::IndexBinaryFlat instance, copying data over to the given GPU

构造一个可以添加的空实例。

Construct an empty instance that can be added to.

~GpuIndexBinaryFlat() 析构函数,重写基类析构函数。
getDevice() 函数,返回设备 ID。

返回此索引所在的设备。

getResources() 函数,返回 GpuResources 的 shared_ptr。

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

copyFrom(const faiss::IndexBinaryFlat *index) 函数,从给定的 CPU 索引初始化自身;将覆盖自身中的所有数据。

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

copyTo(faiss::IndexBinaryFlat *index) const 函数,将自身复制到给定的 CPU 索引;将覆盖索引实例中的所有数据。

Copy ourselves to the given CPU index; will overwrite all data in the index instance

virtual void add(faiss::idx_t n, const uint8_t *x) override

将维度为 d 的 n 个向量添加到索引。

向量隐式分配标签 ntotal .. ntotal + n - 1

参数:

x – 输入矩阵,大小为 n * d / 8

virtual void reset() override

从数据库中删除所有元素。

virtual void search(idx_t n, const uint8_t *x, idx_t k, int32_t *distances, faiss::idx_t *labels, const faiss::SearchParameters *params = nullptr) const override

查询索引中维度为 d 的 n 个向量。

最多返回 k 个向量。 如果查询的结果不足,则结果数组将填充 -1。

参数:
  • x – 要搜索的输入向量,大小为 n * d / 8

  • labels – NN 的输出标签,大小为 n*k

  • distances – 输出成对距离,大小为 n*k

virtual void reconstruct(faiss::idx_t key, uint8_t *recons) const override

重建存储的向量。

对于某些索引,可能未定义此函数。

参数:
  • key – 要重建的向量的 id

  • recons – 重建的向量(大小为 d / 8)

virtual void train(idx_t n, const uint8_t *x)

在一组有代表性的向量上执行训练。

参数:
  • n – 训练向量的 nb

  • x – 训练向量,大小为 n * d / 8

virtual void add_with_ids(idx_t n, const uint8_t *x, const idx_t *xids)

与 add 相同,但存储 xids 而不是顺序 id。

默认实现会因断言而失败,因为它并非所有索引都支持。

参数:

xids – 如果非空,则为要为向量存储的 id(大小为 n)

virtual void range_search(idx_t n, const uint8_t *x, int radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

查询索引中维度为 d 的 n 个向量。

返回所有距离 < radius 的向量。 请注意,许多索引未实现 range_search(只有 k-NN 搜索是必需的)。 距离转换为 float 以重用 RangeSearchResult 结构,但它们是整数。 按照惯例,仅返回距离 < radius(严格比较)的距离,即。 radius = 0 不返回任何结果,而 1 仅返回完全相同的向量。

参数:
  • x – 要搜索的输入向量,大小为 n * d / 8

  • radius – 搜索半径

  • result – 结果表

void assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k = 1) const

返回最接近查询 x 的 k 个向量的索引。

此函数与搜索相同,但仅返回邻居的标签。

参数:
  • x – 要搜索的输入向量,大小为 n * d / 8

  • labels – NN 的输出标签,大小为 n*k

virtual size_t remove_ids(const IDSelector &sel)

从索引中移除 ID。并非所有索引都支持。

virtual void reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const

重构向量 i0 到 i0 + ni - 1。

对于某些索引,可能未定义此函数。

参数:

recons – 重构后的向量(大小为 ni * d / 8)

virtual void search_and_reconstruct(idx_t n, const uint8_t *x, idx_t k, int32_t *distances, idx_t *labels, uint8_t *recons, const SearchParameters *params = nullptr) const

类似于 search,但也为搜索结果重构存储的向量(或者在有损编码的情况下进行近似)。

如果查询的结果不足,则生成的数组将用 -1 填充。

参数:

recons – 重构后的向量大小 (n, k, d)

void display() const

显示实际的类名和一些更多信息。

virtual void merge_from(IndexBinary &otherIndex, idx_t add_id = 0)

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

virtual void check_compatible_for_merge(const IndexBinary &otherIndex) const

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

virtual size_t sa_code_size() const

生成的代码的大小,以字节为单位

virtual void add_sa_codes(idx_t n, const uint8_t *codes, const idx_t *xids)

IndexBinary 的 add_with_ids 相同。

公共成员

int d = 0

向量维度

int code_size = 0

每个向量的字节数(= d / 8)

idx_t ntotal = 0

索引向量的总数

bool verbose = false

详细级别

bool is_trained = true

如果 Index 不需要训练,或者如果已经完成训练,则设置

MetricType metric_type = METRIC_L2

此索引用于搜索的度量类型

受保护函数

void searchFromCpuPaged_(idx_t n, const uint8_t *x, int k, int32_t *outDistancesData, idx_t *outIndicesData) const

当输入数据位于 CPU 上时从 search 调用;可能允许使用锁页内存

此函数在不分页的情况下执行搜索。

受保护的属性

GPU 资源。

管理设备的流、cuBLAS句柄和临时内存。

二进制平面配置。

配置选项。

二进制平面索引的数据。

保存包含向量列表的GPU数据。

IVFPQ构建Cagra配置。

公共成员

反向列表(簇)的数量。

反向列表(簇)的数量。

提示:每个簇的向量数 (n_rows/n_lists) 应该大约为 1,000 到 10,000。

搜索kmeans中心(索引构建)的迭代次数。

搜索kmeans中心(索引构建)的迭代次数。

迭代kmeans构建期间要使用的数据比例。

迭代kmeans构建期间要使用的数据比例。

PQ压缩后向量元素的位长。

PQ压缩后向量元素的位长。

可能的值:[4, 5, 6, 7, 8].

提示:'pq_bits' 越小,索引大小越小,搜索性能越好,但召回率越低。

PQ压缩后向量的维度。 当为零时,使用启发式方法选择最佳值。

PQ压缩后向量的维度。 当为零时,使用启发式方法选择最佳值。

注意:pq_dim /// pq_bits 必须是 8 的倍数。

提示:较小的 'pq_dim' 会导致较小的索引大小和更好的搜索性能,但召回率较低。 如果 'pq_bits' 为 8,则 'pq_dim' 可以设置为任何数字,但为了获得良好的性能,最好是 8 的倍数。 如果 'pq_bits' 不是 8,则 'pq_dim' 应该是 8 的倍数。 为了获得良好的性能,理想情况下 'pq_dim' 是 32 的倍数。 理想情况下,'pq_dim' 也应该是数据集维度的除数。

codebook_gen codebook_kind = codebook_gen::PER_SUBSPACE

如何创建 PQ 码本。

bool force_random_rotation = false

即使 dim % pq_dim == 0,也在输入数据和查询上应用随机旋转矩阵。

注意:如果 dim 不是 pq_dim 的倍数,则始终将随机旋转应用于输入数据和查询,以将工作空间从 dim 转换为 rot_dim,后者可能略大于原始空间,并且是 pq_dim 的倍数 (rot_dim % pq_dim == 0)。但是,当 dimpq_dim 的倍数时,此转换不是必需的(dim == rot_dim,因此无需添加“额外”数据列/特征)。

默认情况下,如果 dim == rot_dim,则使用单位矩阵初始化旋转变换。 当 force_random_rotation == true 时,无论 dimpq_dim 的值如何,都会生成随机正交变换矩阵。

bool conservative_memory_allocation = false

默认情况下,该算法为单个集群(list_data)分配比必要更多的空间。这可以分摊内存分配的成本,并减少重复调用 extend(扩展数据库)期间的数据副本数量。

另一种方法是保守分配行为;启用后,该算法始终分配存储给定数量的记录所需的最小内存量。如果您希望尽可能少地使用 GPU 内存用于数据库,请将此标志设置为 true

struct IVFPQSearchCagraConfig

公共成员

uint32_t n_probes = 20

要搜索的聚类数量。

cudaDataType_t lut_dtype = CUDA_R_32F

搜索时要动态创建的查找表的数据类型。

可能的值:[CUDA_R_32F, CUDA_R_16F, CUDA_R_8U]

使用低精度类型可减少搜索时所需的共享内存量,因此即使对于具有大维度的数据集,也可以使用快速共享内存内核。请注意,选择低精度类型时,召回率会略有下降。

cudaDataType_t internal_distance_dtype = CUDA_R_32F

搜索时计算的距离/相似度的存储数据类型。

可能的值:[CUDA_R_16F, CUDA_R_32F]

如果搜索时的性能限制器是设备内存访问,则选择 FP16 会略微提高性能。

double preferred_shmem_carveout = 1.0

首选的 SM 统一内存/L1 缓存的分数,用作共享内存。

可能的值:[0.0 - 1.0],作为 sharedMemPerMultiprocessor 的一部分。

人们希望增加 carveout 以确保主搜索内核具有良好的 GPU 占用率,但不要使其过高,以留出一些内存用作 L1 缓存。请注意,此值仅被解释为提示。此外,GPU 通常只允许一组固定的缓存配置,因此提供的值将向上舍入到最近的配置。请参阅目标 GPU 架构的 NVIDIA 调整指南。

请注意,这是一个低级调整参数,如果调整不正确,可能会对搜索性能产生严重的负面影响。

struct GpuIndexCagraConfig : public faiss::gpu::GpuIndexConfig

公共成员

size_t intermediate_graph_degree = 128

用于修剪的输入图的度数。

size_t graph_degree = 64

输出图的度数。

graph_build_algo build_algo = graph_build_algo::IVF_PQ

用于构建 knn 图的 ANN 算法。

size_t nn_descent_niter = 20

如果使用 NN_DESCENT 构建,则运行的迭代次数。

IVFPQBuildCagraConfig *ivf_pq_params = nullptr
IVFPQSearchCagraConfig *ivf_pq_search_params = nullptr
float refine_rate = 2.0f
bool store_dataset = true
int device = 0

索引所在的 GPU 设备。

内存空间 = 设备

用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。

是否使用 cuvs = 否

索引是否应调度到 cuVS?

结构体 SearchParametersCagra: public faiss::SearchParameters

公共成员

size_t max_queries = 0

同时搜索的最大查询数量(批处理大小)。当为 0 时自动选择。

size_t itopk_size = 64

在搜索期间保留的中间搜索结果的数量。

这是调整准确性和搜索速度之间权衡的主要旋钮。 较高的值会提高搜索准确性。

size_t max_iterations = 0

搜索迭代次数的上限。当为 0 时自动选择。

search_algo algo = search_algo::AUTO

要使用哪种搜索实现。

size_t team_size = 0

用于计算单个距离的线程数。 4、8、16 或 32。

size_t search_width = 1

在每次迭代中选择作为搜索起点的图节点数。 又名搜索宽度?

size_t min_iterations = 0

搜索迭代次数的下限。

size_t thread_block_size = 0

线程块大小。 0、64、128、256、512、1024。当为 0 时自动选择。

hash_mode hashmap_mode = hash_mode::AUTO

哈希映射类型。当为 AUTO 时自动选择。

size_t hashmap_min_bitlen = 0

哈希映射位长度的下限。 大于 8。

float hashmap_max_fill_rate = 0.5

哈希映射填充率的上限。 大于 0.1,小于 0.9。

uint32_t num_random_samplings = 1

初始随机种子节点选择的迭代次数。 1 或更多。

uint64_t seed = 0x128394

用于初始随机种子节点选择的位掩码。

IDSelector * sel = nullptr

如果非空,则在搜索期间仅考虑这些 ID。

结构体 GpuIndexCagra: public faiss::gpu::GpuIndex

公共类型

using component_t = float
using distance_t = float

公共函数

GpuIndexCagra(GpuResourcesProvider *provider, int dims, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexCagraConfig config = GpuIndexCagraConfig())
virtual void train(idx_t n, const float *x) override

基于给定的向量数据训练CAGRA。

void copyFrom(const faiss::IndexHNSWCagra *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexHNSWCagra *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

virtual void reset() override

从数据库中删除所有元素。

std::vector<idx_t> get_knngraph() const
int getDevice() const

返回此索引所在的设备。

std::shared_ptr<GpuResources> getResources()

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

void setMinPagingSize(size_t size)

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

size_t getMinPagingSize() const

返回分页搜索的当前最小数据大小。

virtual void add(idx_t, const float *x) override

x 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

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

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

virtual size_t remove_ids(const IDSelector &sel)

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

virtual void reconstruct(idx_t key, float *recons) const

重建存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • key – 要重建的向量的 id

  • recons – 重建的向量(大小为 d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

重建多个存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • n – 要重建的向量数

  • keys – 要重建的向量的 id(大小为 n)

  • recons – 重建的向量(大小为 n * d)

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

重建向量 i0 到 i0 + ni - 1

对于某些索引,可能未定义此函数

参数:
  • i0 – 序列中第一个向量的索引

  • ni – 序列中向量的数量

  • recons – 重建的向量(大小为 ni * d)

virtual DistanceComputer *get_distance_computer() const

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

virtual size_t sa_code_size() const

生成的代码的大小,以字节为单位

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

virtual void add_sa_codes(idx_t n, const uint8_t *codes, const idx_t *xids)

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

公共成员

int d

向量维度

idx_t ntotal

索引向量的总数

bool verbose

详细级别

bool is_trained

如果 Index 不需要训练,或者如果已经完成训练,则设置

MetricType metric_type

此索引用于搜索的度量类型

float metric_arg

度量类型的参数

受保护函数

virtual bool addImplRequiresIDs_() const override

addImpl_ 是否需要 ID? 如果是,并且未提供 ID,我们将根据添加 ID 的顺序依次生成它们

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

重写此函数以实际执行添加。所有数据都保证驻留在我们的设备上。

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

GpuIndex 调用以进行搜索。

void copyFrom(const faiss::Index *index)

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

void copyTo(faiss::Index *index) const

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

受保护的属性

const GpuIndexCagraConfig cagraConfig_

我们的配置选项。

std::shared_ptr<CuvsCagra> index_

我们拥有的实例;包含倒排列表。

std::shared_ptr<GpuResources> resources_

管理设备上的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

将副本从 CPU 分页到 GPU 的大小阈值。

结构体 GpuIndexFlatConfig : public faiss :: gpu :: GpuIndexConfig

公共函数

bool ALIGNED (8) useFloat16

是否将数据存储为 float16 。

公共成员

bool storeTransposed = false

已弃用:不再使用。以前用于指示向量的内部存储是否已转置

int device = 0

索引所在的 GPU 设备。

MemorySpace memorySpace = MemorySpace :: Device

用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。

bool use_cuvs = false

索引是否应调度到 cuVS?

GpuIndexFlat : public faiss :: gpu :: GpuIndex
#include <GpuIndexFlat.h>

GPU 实现的封装器,看起来像 faiss::IndexFlat ; 从给定的 faiss::IndexFlat 复制质心数据

faiss::gpu::GpuIndexFlatIP , faiss::gpu::GpuIndexFlatL2 继承

公共类型

using component_t = float
using distance_t = float

公共函数

GpuIndexFlat ( GpuResourcesProvider * provider , const faiss :: IndexFlat * index , GpuIndexFlatConfig config = GpuIndexFlatConfig ())

从预先存在的 faiss::IndexFlat 实例构造,将数据复制到给定的 GPU

GpuIndexFlat ( std :: shared_ptr < GpuResources > resources , const faiss :: IndexFlat * index , GpuIndexFlatConfig config = GpuIndexFlatConfig ())
GpuIndexFlat ( GpuResourcesProvider * provider , int dims , faiss :: MetricType metric , GpuIndexFlatConfig config = GpuIndexFlatConfig ())

Construct an empty instance that can be added to.

GpuIndexFlat(std::shared_ptr<GpuResources> resources, int dims, faiss::MetricType metric, GpuIndexFlatConfig config = GpuIndexFlatConfig())
~GpuIndexFlat() override
void copyFrom(const faiss::IndexFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexFlat *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t getNumVecs() const

返回我们包含的向量的数量。

virtual void reset() override

从此索引中清除所有向量。

virtual void train(idx_t n, const float *x) override

此索引未经训练,因此不执行任何操作。

virtual void add(idx_t, const float *x) override

重写以避免过度复制。

virtual void reconstruct(idx_t key, float *out) const override

重建方法;推荐批量重建,因为它会更有效率

virtual void reconstruct_n(idx_t i0, idx_t num, float *out) const override

批量重建方法。

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *out) const override

批量重建方法。

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

计算残差。

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

计算残差 (批量模式)

inline FlatIndex *getGpuData()

用于内部访问。

int getDevice() const

返回此索引所在的设备。

std::shared_ptr<GpuResources> getResources()

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

void setMinPagingSize(size_t size)

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

size_t getMinPagingSize() const

返回分页搜索的当前最小数据大小。

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

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

virtual size_t remove_ids(const IDSelector &sel)

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

virtual DistanceComputer *get_distance_computer() const

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

virtual size_t sa_code_size() const

生成的代码的大小,以字节为单位

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

virtual void add_sa_codes(idx_t n, const uint8_t *codes, const idx_t *xids)

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

公共成员

int d

向量维度

idx_t ntotal

索引向量的总数

bool verbose

详细级别

bool is_trained

如果 Index 不需要训练,或者如果已经完成训练,则设置

MetricType metric_type

此索引用于搜索的度量类型

float metric_arg

度量类型的参数

受保护函数

void resetIndex_(int dims)
virtual bool addImplRequiresIDs_() const override

Flat 索引不需要 ID,因为没有可用于存储 ID 的空间

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

GpuIndex 中调用以进行添加。

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

GpuIndex 调用以进行搜索。

void copyFrom(const faiss::Index *index)

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

void copyTo(faiss::Index *index) const

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

受保护的属性

const GpuIndexFlatConfig flatConfig_

我们的配置选项。

std::unique_ptr<FlatIndex> data_

保存包含向量列表的GPU数据。

std::shared_ptr<GpuResources> resources_

管理设备上的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

将副本从 CPU 分页到 GPU 的大小阈值。

class GpuIndexFlatL2 : public faiss::gpu::GpuIndexFlat
#include <GpuIndexFlat.h>

围绕 GPU 实现的包装器,看起来像 faiss::IndexFlatL2;从给定的 faiss::IndexFlat 复制质心数据

公共类型

using component_t = float
using distance_t = float

公共函数

GpuIndexFlatL2(GpuResourcesProvider *provider, faiss::IndexFlatL2 *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())

从预先存在的 faiss::IndexFlatL2 实例构造,将数据复制到给定的 GPU

GpuIndexFlatL2(resources, index, config = GpuIndexFlatConfig())
GpuIndexFlatL2(provider, dims, config = GpuIndexFlatConfig())

Construct an empty instance that can be added to.

GpuIndexFlatL2(resources, dims, config = GpuIndexFlatConfig())
copyFrom(index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

copyTo(index)

Copy ourselves to the given CPU index; will overwrite all data in the index instance

copyFrom(const index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

copyTo(index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

getNumVecs() const

返回我们包含的向量的数量。

reset() override

从此索引中清除所有向量。

train(n, const x) override

此索引未经训练,因此不执行任何操作。

add(idx_t, const x) override

重写以避免过度复制。

reconstruct(key, out) const override

重建方法;推荐批量重建,因为它会更有效率

reconstruct_n(i0, num, out) const override

批量重建方法。

out: 存储重建向量的输出数组。

批量重建方法。

key: 与输入向量对应的索引键。

计算残差。

keys: 与输入向量对应的索引键数组。

计算残差 (批量模式)

返回值:指向`FlatIndex`对象的指针。

用于内部访问。

返回值:设备ID。

返回此索引所在的设备。

返回值:指向`GpuResources`对象的共享指针。

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

size: 最小分页大小。

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

返回值:最小分页大小。

返回分页搜索的当前最小数据大小。

ids: 与要添加的向量对应的ID数组。

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

k: 要考虑的最近邻的数量 (默认为 1)。

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

params: 搜索参数 (可以为 null)。

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

virtual size_t remove_ids(const IDSelector &sel)

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

virtual DistanceComputer *get_distance_computer() const

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

virtual size_t sa_code_size() const

生成的代码的大小,以字节为单位

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

检查是否可以合并

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

添加短码

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

公共成员

维度

向量维度

向量总数

索引向量的总数

是否详细输出

详细级别

是否已训练

如果 Index 不需要训练,或者如果已经完成训练,则设置

距离类型

此索引用于搜索的度量类型

距离参数

度量类型的参数

受保护函数

从faiss索引复制

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

复制到faiss索引

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

重置索引
添加实现需要ID

Flat 索引不需要 ID,因为没有可用于存储 ID 的空间

添加实现

GpuIndex 中调用以进行添加。

搜索实现

GpuIndex 调用以进行搜索。

受保护的属性

平面配置

我们的配置选项。

平面索引数据

保存包含向量列表的GPU数据。

GPU资源

管理设备上的流、cuBLAS 句柄和暂存内存。

GPU索引配置

我们的配置选项。

最小分页大小

将副本从 CPU 分页到 GPU 的大小阈值。

GpuIndexFlatIP 类:GPU 中 IndexFlatIP 的封装
#include <GpuIndexFlat.h>

GPU 实现的包装器,看起来像 faiss::IndexFlatIP; 从给定的 faiss::IndexFlat 复制质心数据

公共类型

组件类型定义为 float
using distance_t = float

公共函数

GpuIndexFlatIP(GpuResourcesProvider *provider, faiss::IndexFlatIP *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())

从预先存在的 faiss::IndexFlatIP 实例构造,将数据复制到给定的 GPU 上

GpuIndexFlatIP(std::shared_ptr<GpuResources> resources, faiss::IndexFlatIP *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())
GpuIndexFlatIP(GpuResourcesProvider *provider, int dims, GpuIndexFlatConfig config = GpuIndexFlatConfig())

Construct an empty instance that can be added to.

GpuIndexFlatIP(std::shared_ptr<GpuResources> resources, int dims, GpuIndexFlatConfig config = GpuIndexFlatConfig())
void copyFrom(faiss::IndexFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexFlat *index)

Copy ourselves to the given CPU index; will overwrite all data in the index instance

void copyFrom(const faiss::IndexFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexFlat *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t getNumVecs() const

返回我们包含的向量的数量。

virtual void reset() override

从此索引中清除所有向量。

参数:

此索引未经训练,因此不执行任何操作。

参数:

重写以避免过度复制。

参数:

重建方法;推荐批量重建,因为它会更有效率

参数:

批量重建方法。

参数:

批量重建方法。

参数:

计算残差。

参数:

计算残差 (批量模式)

返回值:

用于内部访问。

返回值:

返回此索引所在的设备。

返回值:

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

设置最小分页大小

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

获取最小分页大小

返回分页搜索的当前最小数据大小。

添加带有id的向量。

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

分配向量到最近的聚类中心。

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

搜索k个最近邻。

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

搜索k个最近邻并重构它们。

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

在指定半径内搜索近邻。

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

移除与ID选择器匹配的ID。

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

获取距离计算器。

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

获取SA代码大小。

生成的代码的大小,以字节为单位

对向量进行编码并存储到字节数组中。

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

virtual void add_sa_codes(idx_t n, const uint8_t *codes, const idx_t *xids)

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

公共成员

int d

向量维度

idx_t ntotal

索引向量的总数

bool verbose

详细级别

bool is_trained

如果 Index 不需要训练,或者如果已经完成训练,则设置

MetricType metric_type

此索引用于搜索的度量类型

float metric_arg

度量类型的参数

受保护函数

void copyFrom(const faiss::Index *index)

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

void copyTo(faiss::Index *index) const

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

void resetIndex_(int dims)
virtual bool addImplRequiresIDs_() const override

Flat 索引不需要 ID,因为没有可用于存储 ID 的空间

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

GpuIndex 中调用以进行添加。

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

GpuIndex 调用以进行搜索。

受保护的属性

const GpuIndexFlatConfig flatConfig_

我们的配置选项。

std::unique_ptr<FlatIndex> data_

保存包含向量列表的GPU数据。

std::shared_ptr<GpuResources> resources_

管理设备上的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

将副本从 CPU 分页到 GPU 的大小阈值。

struct GpuIndexIVFConfig : public faiss::gpu::GpuIndexConfig

Subclassed by 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 实现。

int device = 0

索引所在的 GPU 设备。

MemorySpace memorySpace = MemorySpace::Device

用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。

bool use_cuvs = false

索引是否应调度到 cuVS?

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

所有GPU IVF索引类型的基类。 目前,它有意不继承自IndexIVF,因为IndexIVF中的许多公共数据成员和功能在GPU上不支持以相同的方式使用。

faiss::gpu::GpuIndexIVFFlatfaiss::gpu::GpuIndexIVFPQfaiss::gpu::GpuIndexIVFScalarQuantizer继承。

公共类型

using component_t = float
using distance_t = float

公共函数

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 – 将倒排列表索引 + 倒排列表偏移量存储在结果的 upper/lower 32 位中,而不是 ID(用于重排序)。

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

  • stats – 要更新的搜索统计信息(可以为 null)

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 – 将倒排列表索引 + 倒排列表偏移量存储在结果的 upper/lower 32 位中,而不是 ID(用于重排序)。

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

  • stats – 要更新的搜索统计信息(可以为 null)

int getDevice() const

返回此索引所在的设备。

std::shared_ptr<GpuResources> getResources()

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

void setMinPagingSize(size_t size)

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

size_t getMinPagingSize() const

返回分页搜索的当前最小数据大小。

virtual void add(idx_t, const float *x) override

x 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

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

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

计算残差

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

计算n个向量的残差。

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

训练索引。

在一组代表性向量上执行训练

参数:
  • n – 训练向量的 nb

  • x – 训练向量,大小为 n * d

在指定半径内进行范围搜索。

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

重置索引,移除所有数据。

从数据库中删除所有元素。

移除满足ID选择器的向量。

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

重建向量。

重建存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • key – 要重建的向量的 id

  • recons – 重建的向量(大小为 d)

批量重建向量。

重建多个存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • n – 要重建的向量数

  • keys – 要重建的向量的 id(大小为 n)

  • recons – 重建的向量(大小为 n * d)

重建一系列向量。

重建向量 i0 到 i0 + ni - 1

对于某些索引,可能未定义此函数

参数:
  • i0 – 序列中第一个向量的索引

  • ni – 序列中向量的数量

  • recons – 重建的向量(大小为 ni * d)

获取距离计算器。

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

获取标量量化编码的大小。

生成的代码的大小,以字节为单位

sa_encode 的函数签名,用于编码。

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

sa_decode 的函数签名,用于解码。

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

merge_from 的函数签名,用于从另一个索引合并。

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

check_compatible_for_merge 的函数签名,用于检查合并的兼容性。

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

add_sa_codes 的函数签名,用于添加 SA 编码。

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

train_q1 的函数签名,用于训练一级量化器。

训练量化器并调用 train_residual 来训练子量化器。

coarse_code_size 的函数签名,用于获取粗糙编码的大小。

计算存储列表 ID 所需的字节数。

encode_listno 的函数签名,用于编码列表编号。
decode_listno 的函数签名,用于解码列表编号。

公共成员

向量维度。

向量维度

索引中的向量总数。

索引向量的总数

是否启用详细输出。

详细级别

索引是否已训练。

如果 Index 不需要训练,或者如果已经完成训练,则设置

使用的距离度量类型。

此索引用于搜索的度量类型

距离度量的参数 (例如 Minkowski 距离的 p 值)。

度量类型的参数

查询时探针的数量。

查询时探针的数量

查询时要访问的最大代码数量。

查询时要访问的最大代码数量

将向量映射到倒排列表的量化器。

将向量映射到倒排列表的量化器

倒排列表的数量。

倒排列表的数量

量化器训练模式:0: 将量化器用作 kmeans 训练中的索引;1: 将训练集传递给量化器的 train();2: 在平面索引上进行 kmeans 训练 + 将质心添加到量化器。

= 0: use the quantizer as index in a kmeans training = 1: just pass on the training set to the train() of the quantizer = 2: kmeans training on a flat index + add the centroids to the quantizer

对象是否拥有量化器。

对象是否拥有量化器

用于覆盖默认聚类参数。

用于覆盖默认聚类参数

Index *clustering_index = nullptr

覆盖聚类期间使用的索引

受保护函数

int getCurrentNProbe_(const SearchParameters *params) const

从当前设置的 nprobe 或可用的 SearchParameters 中,返回应该用于当前搜索的 nprobe

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

addImpl_ 是否需要 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 调用以进行搜索。

void copyFrom(const faiss::Index *index)

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

void copyTo(faiss::Index *index) const

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

受保护的属性

const GpuIndexIVFConfig ivfConfig_

我们的配置选项。

std::shared_ptr<IVFBase> baseIndex_

对于训练/初始化后的索引,这是对基类的引用。

std::shared_ptr<GpuResources> resources_

管理设备上的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

将副本从 CPU 分页到 GPU 的大小阈值。

私有函数

void init_()

共享初始化函数。

struct GpuIndexIVFFlatConfig : public faiss::gpu::GpuIndexIVFConfig

公共成员

bool interleavedLayout = true

对 IVF 列表使用替代内存布局(当前为默认布局)

IndicesOptions indicesOptions = INDICES_64_BIT

GPU的Index存储选项。

GpuIndexFlatConfig flatConfig

粗量化器对象的配置。

bool allowCpuCoarseQuantizer = false

此标志控制索引的粗量化器组件的 CPU 回退逻辑。如果设置为 false (默认),克隆器将为 GPU 上未实现的索引引发异常。 如果设置为 true,它将回退到 CPU 实现。

int device = 0

索引所在的 GPU 设备。

MemorySpace memorySpace = MemorySpace::Device

用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。

bool use_cuvs = false

索引是否应调度到 cuVS?

class GpuIndexIVFFlat : public faiss::gpu::GpuIndexIVF
#include <GpuIndexIVFFlat.h>

Wrapper around the GPU implementation that looks like faiss::IndexIVFFlat

公共类型

using component_t = float
using distance_t = float

公共函数

GpuIndexIVFFlat(GpuResourcesProvider *provider, const faiss::IndexIVFFlat *index, GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig())

从预先存在的 faiss::IndexIVFFlat 实例构建,如果输入索引已训练,则将数据复制到给定的 GPU。

GpuIndexIVFFlat(GpuResourcesProvider *provider, int dims, idx_t nlist, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig())

构造一个带有空 flat 量化器的新实例;用户提供所需的 IVF 列表的数量。

GpuIndexIVFFlat(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, idx_t nlist, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig())

构造一个带有提供的 CPU 或 GPU 粗量化器的新实例;用户提供所需的 IVF 列表的数量。

~GpuIndexIVFFlat() override
void reserveMemory(size_t numVecs)

为这个向量的数量在我们的倒排列表中预留 GPU 内存。

void copyFrom(const faiss::IndexIVFFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexIVFFlat *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t reclaimMemory()

添加向量后,可以调用此函数来回收设备内存,使其正好达到所需的数量。返回以字节为单位回收的空间

virtual void reset() override

清除所有倒排列表,但保留粗量化的质心信息

virtual void updateQuantizer() override

如果用户手动更改 IVF 粗量化器的状态(例如,替换新实例或在训练范围之外更改粗量化器中的向量),则应调用此函数。

virtual void train(idx_t n, const float *x) override

基于给定的向量数据训练粗量化器。

virtual void reconstruct_n(idx_t i0, idx_t n, float *out) const override

重建向量 i0 到 i0 + ni - 1

对于某些索引,可能未定义此函数

参数:
  • i0 – 序列中第一个向量的索引

  • ni – 序列中向量的数量

  • recons – 重建的向量(大小为 ni * d)

void copyFrom(const faiss::IndexIVF *index)

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

void copyTo(faiss::IndexIVF *index) const

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

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 – 将倒排列表索引 + 倒排列表偏移量存储在结果的 upper/lower 32 位中,而不是 ID(用于重排序)。

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

  • stats – 要更新的搜索统计信息(可以为 null)

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 – 将倒排列表索引 + 倒排列表偏移量存储在结果的 upper/lower 32 位中,而不是 ID(用于重排序)。

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

  • stats – 要更新的搜索统计信息(可以为 null)

int getDevice() const

返回此索引所在的设备。

std::shared_ptr<GpuResources> getResources()

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

void setMinPagingSize(size_t size)

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

size_t getMinPagingSize() const

返回分页搜索的当前最小数据大小。

virtual void add(idx_t, const float *x) override

x 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

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

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

virtual size_t remove_ids(const IDSelector &sel)

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

virtual void reconstruct(idx_t key, float *recons) const

重建存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • key – 要重建的向量的 id

  • recons – 重建的向量(大小为 d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

重建多个存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • n – 要重建的向量数

  • keys – 要重建的向量的 id(大小为 n)

  • recons – 重建的向量(大小为 n * d)

virtual DistanceComputer *get_distance_computer() const

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

virtual size_t sa_code_size() const

生成的代码的大小,以字节为单位

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

virtual void add_sa_codes(idx_t n, const uint8_t *codes, const idx_t *xids)

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

void train_q1(size_t n, const float *x, bool verbose, MetricType metric_type)

训练量化器并调用 train_residual 来训练子量化器。

size_t coarse_code_size() const

计算存储列表 ID 所需的字节数。

void encode_listno(idx_t list_no, uint8_t *code) const
idx_t decode_listno(const uint8_t *code) const

公共成员

int d

向量维度

idx_t ntotal

索引向量的总数

bool verbose

详细级别

bool is_trained

如果 Index 不需要训练,或者如果已经完成训练,则设置

MetricType metric_type

此索引用于搜索的度量类型

float metric_arg

度量类型的参数

size_t nprobe = 1

查询时探针的数量

size_t max_codes = 0

查询时要访问的最大代码数量

Index *quantizer = nullptr

将向量映射到倒排列表的量化器

size_t nlist = 0

倒排列表的数量

char quantizer_trains_alone = 0

= 0: use the quantizer as index in a kmeans training = 1: just pass on the training set to the train() of the quantizer = 2: kmeans training on a flat index + add the centroids to the quantizer

bool own_fields = false

对象是否拥有量化器

ClusteringParameters cp

用于覆盖默认聚类参数

Index *clustering_index = nullptr

覆盖聚类期间使用的索引

受保护函数

void setIndex_(GpuResources *resources, int dim, int nlist, faiss::MetricType metric, float metricArg, bool useResidual, faiss::ScalarQuantizer *scalarQ, bool interleavedLayout, IndicesOptions indicesOptions, MemorySpace space)

初始化合适的索引。

参数:

scalarQ – 可选的 ScalarQuantizer

void copyFrom(const faiss::Index *index)

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

void copyTo(faiss::Index *index) const

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

int getCurrentNProbe_(const SearchParameters *params) const

从当前设置的 nprobe 或可用的 SearchParameters 中,返回应该用于当前搜索的 nprobe

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

addImpl_ 是否需要 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 GpuIndexIVFFlatConfig ivfFlatConfig_

我们的配置选项。

size_t reserveMemoryVecs_

所需的倒排列表内存预留。

std::shared_ptr<IVFFlat> index_

我们拥有的实例;包含倒排列表。

const GpuIndexIVFConfig ivfConfig_

我们的配置选项。

std::shared_ptr<IVFBase> baseIndex_

对于训练/初始化后的索引,这是对基类的引用。

std::shared_ptr<GpuResources> resources_

管理设备上的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

将副本从 CPU 分页到 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代码距离。

IndicesOptions indicesOptions = INDICES_64_BIT

GPU的Index存储选项。

GpuIndexFlatConfig flatConfig

粗量化器对象的配置。

bool allowCpuCoarseQuantizer = false

此标志控制索引的粗量化器组件的 CPU 回退逻辑。如果设置为 false (默认),克隆器将为 GPU 上未实现的索引引发异常。 如果设置为 true,它将回退到 CPU 实现。

int device = 0

索引所在的 GPU 设备。

MemorySpace memorySpace = MemorySpace::Device

用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。

bool use_cuvs = false

索引是否应调度到 cuVS?

class GpuIndexIVFPQ : public faiss::gpu::GpuIndexIVF
#include <GpuIndexIVFPQ.h>

GPU的IVFPQ索引。

公共类型

using component_t = float
using distance_t = float

公共函数

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

Copy ourselves to the given CPU index; will overwrite all data in the index instance

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^bits)。

size_t reclaimMemory()

添加向量后,可以调用此函数来回收设备内存,使其正好达到所需的数量。返回以字节为单位回收的空间

virtual void reset() override

清除所有倒排列表,但保留粗量化器和乘积质心信息。

virtual void updateQuantizer() override

如果用户手动更改 IVF 粗量化器的状态(例如,替换新实例或在训练范围之外更改粗量化器中的向量),则应调用此函数。

virtual void train(idx_t n, const float *x) override

基于给定的向量数据训练粗量化器和乘积量化器。

void copyFrom(const faiss::IndexIVF *index)

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

void copyTo(faiss::IndexIVF *index) const

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

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 – 将倒排列表索引 + 倒排列表偏移量存储在结果的 upper/lower 32 位中,而不是 ID(用于重排序)。

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

  • stats – 要更新的搜索统计信息(可以为 null)

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 – 将倒排列表索引 + 倒排列表偏移量存储在结果的 upper/lower 32 位中,而不是 ID(用于重排序)。

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

  • stats – 要更新的搜索统计信息(可以为 null)

int getDevice() const

返回此索引所在的设备。

std::shared_ptr<GpuResources> getResources()

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

void setMinPagingSize(size_t size)

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

size_t getMinPagingSize() const

返回分页搜索的当前最小数据大小。

virtual void add(idx_t, const float *x) override

x 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

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

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

virtual size_t remove_ids(const IDSelector &sel)

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

virtual void reconstruct(idx_t key, float *recons) const

重建存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • key – 要重建的向量的 id

  • recons – 重建的向量(大小为 d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

重建多个存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • n – 要重建的向量数

  • keys – 要重建的向量的 id(大小为 n)

  • recons – 重建的向量(大小为 n * d)

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

重建向量 i0 到 i0 + ni - 1

对于某些索引,可能未定义此函数

参数:
  • i0 – 序列中第一个向量的索引

  • ni – 序列中向量的数量

  • recons – 重建的向量(大小为 ni * d)

virtual DistanceComputer *get_distance_computer() const

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

virtual size_t sa_code_size() const

生成的代码的大小,以字节为单位

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

virtual void add_sa_codes(idx_t n, const uint8_t *codes, const idx_t *xids)

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

void train_q1(size_t n, const float *x, bool verbose, MetricType metric_type)

训练量化器并调用 train_residual 来训练子量化器。

size_t coarse_code_size() const

计算存储列表 ID 所需的字节数。

void encode_listno(idx_t list_no, uint8_t *code) const
idx_t decode_listno(const uint8_t *code) const

公共成员

ProductQuantizer pq

与CPU版本类似,我们公开一个可公开访问的 ProductQuantizer 用于操作

int d

向量维度

idx_t ntotal

索引向量的总数

bool verbose

详细级别

bool is_trained

如果 Index 不需要训练,或者如果已经完成训练,则设置

MetricType metric_type

此索引用于搜索的度量类型

float metric_arg

度量类型的参数

size_t nprobe = 1

查询时探针的数量

size_t max_codes = 0

查询时要访问的最大代码数量

Index *quantizer = nullptr

将向量映射到倒排列表的量化器

size_t nlist = 0

倒排列表的数量

char quantizer_trains_alone = 0

= 0: use the quantizer as index in a kmeans training = 1: just pass on the training set to the train() of the quantizer = 2: kmeans training on a flat index + add the centroids to the quantizer

bool own_fields = false

对象是否拥有量化器

ClusteringParameters cp

用于覆盖默认聚类参数

Index *clustering_index = nullptr

覆盖聚类期间使用的索引

受保护函数

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 量化器。

void copyFrom(const faiss::Index *index)

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

void copyTo(faiss::Index *index) const

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

int getCurrentNProbe_(const SearchParameters *params) const

从当前设置的 nprobe 或可用的 SearchParameters 中,返回应该用于当前搜索的 nprobe

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

addImpl_ 是否需要 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 GpuIndexIVFPQConfig ivfpqConfig_

我们的配置选项,我们在初始化时使用。

bool usePrecomputedTables_

运行时覆盖:是否使用预计算表。

int subQuantizers_

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

int bitsPerCode_

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

size_t reserveMemoryVecs_

所需的倒排列表内存预留。

std::shared_ptr<IVFPQ> index_

我们拥有的乘积量化器实例;包含倒排列表

const GpuIndexIVFConfig ivfConfig_

我们的配置选项。

std::shared_ptr<IVFBase> baseIndex_

对于训练/初始化后的索引,这是对基类的引用。

std::shared_ptr<GpuResources> resources_

管理设备上的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

将副本从 CPU 分页到 GPU 的大小阈值。

struct GpuIndexIVFScalarQuantizerConfig : public faiss::gpu::GpuIndexIVFConfig

公共成员

bool interleavedLayout = true

对 IVF 列表使用替代内存布局(当前为默认布局)

IndicesOptions indicesOptions = INDICES_64_BIT

GPU的Index存储选项。

GpuIndexFlatConfig flatConfig

粗量化器对象的配置。

bool allowCpuCoarseQuantizer = false

此标志控制索引的粗量化器组件的 CPU 回退逻辑。如果设置为 false (默认),克隆器将为 GPU 上未实现的索引引发异常。 如果设置为 true,它将回退到 CPU 实现。

int device = 0

索引所在的 GPU 设备。

MemorySpace memorySpace = MemorySpace::Device

用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。

bool use_cuvs = false

索引是否应调度到 cuVS?

class GpuIndexIVFScalarQuantizer : public faiss::gpu::GpuIndexIVF
#include <GpuIndexIVFScalarQuantizer.h>

GPU 实现的包装器,看起来像 faiss::IndexIVFScalarQuantizer

公共类型

using component_t = float
using distance_t = float

公共函数

GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, const faiss::IndexIVFScalarQuantizer *index, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())

从预先存在的 faiss::IndexIVFScalarQuantizer 实例构造,如果输入索引已训练,则将数据复制到给定的 GPU。

GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, int dims, idx_t nlist, faiss::ScalarQuantizer::QuantizerType qtype, faiss::MetricType metric = MetricType::METRIC_L2, bool encodeResidual = true, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())

构造一个带有空 flat 量化器的新实例;用户提供所需的 IVF 列表的数量。

GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, idx_t nlist, faiss::ScalarQuantizer::QuantizerType qtype, faiss::MetricType metric = MetricType::METRIC_L2, bool encodeResidual = true, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())

构造一个带有提供的 CPU 或 GPU 粗量化器的新实例;用户提供所需的 IVF 列表的数量。

~GpuIndexIVFScalarQuantizer() override
void reserveMemory(size_t numVecs)

为这个向量的数量在我们的倒排列表中预留 GPU 内存。

void copyFrom(const faiss::IndexIVFScalarQuantizer *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexIVFScalarQuantizer *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t reclaimMemory()

添加向量后,可以调用此函数来回收设备内存,使其正好达到所需的数量。返回以字节为单位回收的空间

virtual void reset() override

清除所有倒排列表,但保留粗量化器和标量量化器信息。

virtual void updateQuantizer() override

如果用户手动更改 IVF 粗量化器的状态(例如,替换新实例或在训练范围之外更改粗量化器中的向量),则应调用此函数。

virtual void train(idx_t n, const float *x) override

基于给定的向量数据训练粗量化器和标量量化器。

void copyFrom(const faiss::IndexIVF *index)

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

void copyTo(faiss::IndexIVF *index) const

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

virtual idx_t getNumLists() const

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

获取列表长度。

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

获取列表的向量数据。

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

获取列表的索引。

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

预分配的搜索。

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

参数:
  • n – 要查询的向量数

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

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

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

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

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

  • store_pairs – 将倒排列表索引 + 倒排列表偏移量存储在结果的 upper/lower 32 位中,而不是 ID(用于重排序)。

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

  • stats – 要更新的搜索统计信息(可以为 null)

预分配的范围搜索。

范围搜索一组向量,这些向量已经通过 IVF 量化器进行了预量化。 填充 RangeSearchResults 结果。 默认实现使用 InvertedListScanners 进行搜索。

参数:
  • n – 要查询的向量数

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

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

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

  • result – 输出结果

  • store_pairs – 将倒排列表索引 + 倒排列表偏移量存储在结果的 upper/lower 32 位中,而不是 ID(用于重排序)。

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

  • stats – 要更新的搜索统计信息(可以为 null)

获取设备。

返回此索引所在的设备。

获取资源。

返回对我们的 GpuResources 对象的引用,该对象管理 GPU 上的内存、流和句柄资源

设置最小分页大小。

设置用于搜索的最小数据大小(以 MiB 为单位),我们将使用 CPU -> GPU 分页

获取最小分页大小。

返回分页搜索的当前最小数据大小。

virtual void add(idx_t, const float *x) override

x 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

添加向量:virtual void add(idx_t, const float *x) override

xids 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_

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

xlabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

带ID添加向量:virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

xdistanceslabels 可以驻留在 CPU 或任何 GPU 上;根据需要执行复制

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

xdistanceslabels 以及 recons 可以位于 CPU 或任何 GPU 上; 将根据需要执行复制

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

重写以强制 GPU 索引提供他们自己的 GPU 友好型实现

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

将 n 个维度为 d 的向量查询到索引。

返回所有距离 < radius 的向量。 请注意,许多索引不实现 range_search(只有 k-NN 搜索是强制性的)。

参数:
  • n – 向量数量

  • x – 要搜索的输入向量,大小为 n * d

  • radius – 搜索半径

  • result – 结果表

virtual size_t remove_ids(const IDSelector &sel)

从索引中删除 ID。 并非所有索引都支持。 返回删除的元素数。

根据给定的键 `key` 重建向量,并将结果存储在 `recons` 中。

重建存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • key – 要重建的向量的 id

  • recons – 重建的向量(大小为 d)

根据给定的键列表 `keys` 重建 `n` 个向量,并将结果存储在 `recons` 中。

重建多个存储的向量(如果是有损编码,则重建近似值)

对于某些索引,可能未定义此函数

参数:
  • n – 要重建的向量数

  • keys – 要重建的向量的 id(大小为 n)

  • recons – 重建的向量(大小为 n * d)

重建从索引 `i0` 开始的 `ni` 个向量,并将结果存储在 `recons` 中。

重建向量 i0 到 i0 + ni - 1

对于某些索引,可能未定义此函数

参数:
  • i0 – 序列中第一个向量的索引

  • ni – 序列中向量的数量

  • recons – 重建的向量(大小为 ni * d)

获取距离计算器对象。

获取一个 DistanceComputer (在 AuxIndexStructures 中定义) 对象,用于此类型的索引。

DistanceComputer 针对支持对其向量进行随机访问的索引实现。

获取辅助数据 (SA) 代码的大小。

生成的代码的大小,以字节为单位

将 `n` 个向量 `x` 编码为字节数组 `bytes`。

编码一组向量

参数:
  • n – 向量数量

  • x – 输入向量,大小为 n * d

  • bytes – 输出编码向量,大小为 n * sa_code_size()

将 `n` 个字节数组 `bytes` 解码为向量 `x`。

解码一组向量

参数:
  • n – 向量数量

  • bytes – 输入编码向量,大小为 n * sa_code_size()

  • x – 输出向量,大小为 n * d

从另一个索引 `otherIndex` 合并数据。`add_id` 用于调整合并后的 ID。

将条目从另一个数据集移动到自身。输出时,other 为空。 add_id 被添加到所有移动的 ID(对于顺序 ID,这将是 this->ntotal)

检查另一个索引 `otherIndex` 是否与当前索引兼容,可以进行合并操作。

检查两个索引是否兼容(即,它们以相同的方式训练并具有相同的参数)。否则抛出异常。

添加 `n` 个辅助数据代码 `codes`,以及它们对应的 ID `xids`。

添加使用独立编解码器计算的向量

参数:
  • codes – 要添加的代码,大小为 n * sa_code_size()

  • xids – 相应的 ID,大小为 n

训练第一级量化器 (q1)。`n` 是训练向量的数量,`x` 是训练数据,`verbose` 控制训练过程的详细程度,`metric_type` 指定度量类型。

训练量化器并调用 train_residual 来训练子量化器。

获取粗量化 (coarse quantization) 代码的大小。

计算存储列表 ID 所需的字节数。

将列表编号 `list_no` 编码为字节数组 `code`。
将字节数组 `code` 解码为列表编号。

公共成员

标量量化器 `sq` 对象。

与 CPU 版本类似公开。

bool by_residual

与 CPU 版本类似公开。

int d

向量维度

idx_t ntotal

索引向量的总数

bool verbose

详细级别

bool is_trained

如果 Index 不需要训练,或者如果已经完成训练,则设置

MetricType metric_type

此索引用于搜索的度量类型

float metric_arg

度量类型的参数

size_t nprobe = 1

查询时探针的数量

size_t max_codes = 0

查询时要访问的最大代码数量

Index *quantizer = nullptr

将向量映射到倒排列表的量化器

size_t nlist = 0

倒排列表的数量

char quantizer_trains_alone = 0

= 0: use the quantizer as index in a kmeans training = 1: just pass on the training set to the train() of the quantizer = 2: kmeans training on a flat index + add the centroids to the quantizer

bool own_fields = false

对象是否拥有量化器

ClusteringParameters cp

用于覆盖默认聚类参数

Index *clustering_index = nullptr

覆盖聚类期间使用的索引

受保护函数

void verifySQSettings_() const

验证索引 SQ 参数。

void trainResiduals_(idx_t n, const float *x)

从 train 调用以处理 SQ 残差训练。

void copyFrom(const faiss::Index *index)

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

void copyTo(faiss::Index *index) const

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

int getCurrentNProbe_(const SearchParameters *params) const

从当前设置的 nprobe 或可用的 SearchParameters 中,返回应该用于当前搜索的 nprobe

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

addImpl_ 是否需要 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 GpuIndexIVFScalarQuantizerConfig ivfSQConfig_

我们的配置选项。

size_t reserveMemoryVecs_

所需的倒排列表内存预留。

std::shared_ptr<IVFFlat> index_

我们拥有的实例;包含倒排列表。

const GpuIndexIVFConfig ivfConfig_

我们的配置选项。

std::shared_ptr<IVFBase> baseIndex_

对于训练/初始化后的索引,这是对基类的引用。

std::shared_ptr<GpuResources> resources_

管理设备上的流、cuBLAS 句柄和暂存内存。

const GpuIndexConfig config_

我们的配置选项。

size_t minPagedSize_

将副本从 CPU 分页到 GPU 的大小阈值。

struct AllocInfo
#include <GpuResources.h>

关于分配是什么/在哪里的信息。

faiss::gpu::AllocRequest 继承

公共函数

inline AllocInfo()
inline AllocInfo(AllocType at, int dev, MemorySpace sp, cudaStream_t st)
std::string toString() const

返回此信息的字符串表示形式。

公共成员

AllocType type = AllocType::Other

分配的内部类别。

int device = 0

分配正在发生的设备。

MemorySpace space = MemorySpace::Device

分配的内存空间。

cudaStream_t stream = nullptr

将在其上对内存上的新工作进行排序的流(例如,如果缓存的一段内存将被返回给此调用,上次在流 3 上使用,并且新的内存请求是针对流 4 的,则内存管理器将同步流 4 以等待通过事件或其他流同步完成流 3)。

内存管理器保证返回的内存可以自由使用,而不会发生指定流上的数据竞争。

struct AllocRequest : public faiss::gpu::AllocInfo
#include <GpuResources.h>

关于分配是什么/在哪里,以及它应该有多大的信息。

公共函数

inline AllocRequest()
inline AllocRequest(const AllocInfo &info, size_t sz)
inline AllocRequest(AllocType at, int dev, MemorySpace sp, cudaStream_t st, size_t sz)
std::string toString() const

返回此请求的字符串表示形式。

公共成员

size_t size = 0

分配的大小,以字节为单位。

AllocType type = AllocType::Other

分配的内部类别。

int device = 0

分配正在发生的设备。

MemorySpace space = MemorySpace::Device

分配的内存空间。

cudaStream_t stream = nullptr

将在其上对内存上的新工作进行排序的流(例如,如果缓存的一段内存将被返回给此调用,上次在流 3 上使用,并且新的内存请求是针对流 4 的,则内存管理器将同步流 4 以等待通过事件或其他流同步完成流 3)。

内存管理器保证返回的内存可以自由使用,而不会发生指定流上的数据竞争。

struct GpuMemoryReservation
#include <GpuResources.h>

一个管理临时内存请求的 RAII 对象。

公共函数

GpuMemoryReservation()
GpuMemoryReservation(GpuResources *r, int dev, cudaStream_t str, void *p, size_t sz)
GpuMemoryReservation(GpuMemoryReservation &&m) noexcept
~GpuMemoryReservation()
GpuMemoryReservation &operator=(GpuMemoryReservation &&m)
inline void *get()
void release()

公共成员

GpuResources *res
int device
cudaStream_t stream
void *data
size_t size
class GpuResources
#include <GpuResources.h>

GPU端资源提供程序的基类;隐藏了cuBLAS句柄、CUDA流以及执行的所有设备内存分配的provision

faiss::gpu::StandardGpuResourcesImpl 子类化

公共函数

virtual ~GpuResources()
virtual void initializeForDevice(int device) = 0

调用以为特定设备预分配资源。 如果未调用此函数,则将在首次需求时分配资源

virtual bool supportsBFloat16(int device) = 0

给定的GPU是否支持bfloat16?

virtual cublasHandle_t getBlasHandle(int device) = 0

返回我们用于给定设备的cuBLAS句柄。

virtual cudaStream_t getDefaultStream(int device) = 0

返回我们为给定设备上的所有计算排序的流

virtual void setDefaultStream(int device, cudaStream_t stream) = 0

将设备的默认流覆盖为用户提供的流。 资源对象不拥有此流(即,它不会销毁它)。

virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0

返回我们用于给定设备的一组替代流。

virtual void *allocMemory(const AllocRequest &req) = 0

内存管理 从给定的内存空间返回分配,该分配相对于给定的流(即,第一个用户将是此流中的内核)进行排序。 所有分配在内部调整大小以成为下一个最高的16字节倍数,并且保证所有返回的分配均为16字节对齐。

virtual void deallocMemory(int device, void *in) = 0

返回先前的分配。

virtual size_t getTempMemoryAvailable(int device) const = 0

对于 MemorySpace::Temporary,在没有 cudaMalloc 分配的情况下,有多少空间可以立即使用?

virtual std::pair<void*, size_t> getPinnedMemory() = 0

返回可用的CPU固定内存缓冲区。

virtual cudaStream_t getAsyncCopyStream(int device) = 0

返回我们在其上执行异步CPU <-> GPU复制的流。

bool supportsBFloat16CurrentDevice()

当前GPU是否支持bfloat16?

默认提供的函数

cublasHandle_t getBlasHandleCurrentDevice()

使用当前设备调用 getBlasHandle。

cudaStream_t getDefaultStreamCurrentDevice()

使用当前设备调用 getDefaultStream。

size_t getTempMemoryAvailableCurrentDevice() const

调用 getTempMemoryAvailable,使用当前设备。

GpuMemoryReservation allocMemoryHandle(const AllocRequest &req)

通过 RAII 对象返回临时内存分配。

void syncDefaultStream(int device)

将 CPU 与给定设备的默认流同步

void syncDefaultStreamCurrentDevice()

调用 syncDefaultStream,使用当前设备。

std::vector<cudaStream_t> getAlternateStreamsCurrentDevice()

调用 getAlternateStreams,使用当前设备。

cudaStream_t getAsyncCopyStreamCurrentDevice()

调用 getAsyncCopyStream,使用当前设备。

class GpuResourcesProvider
#include <GpuResources.h>

共享资源对象的提供程序接口。 这是为了避免将 std::shared_ptr 与 Python 接口

faiss::gpu::GpuResourcesProviderFromInstance, faiss::gpu::StandardGpuResources 继承

公共函数

virtual ~GpuResourcesProvider()
virtual std::shared_ptr<GpuResources> getResources() = 0

返回共享资源对象。

class GpuResourcesProviderFromInstance : public faiss::gpu::GpuResourcesProvider
#include <GpuResources.h>

一个简单的 GpuResources 对象包装器,用于再次从中创建一个 GpuResourcesProvider

公共函数

explicit GpuResourcesProviderFromInstance(std::shared_ptr<GpuResources> p)
~GpuResourcesProviderFromInstance() override
virtual std::shared_ptr<GpuResources> getResources() override

返回共享资源对象。

私有成员

std::shared_ptr<GpuResources> res_
template<typename GpuIndex>
struct IndexWrapper

公共函数

IndexWrapper(int numGpus, std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)> init)
faiss::Index *getIndex()
void runOnIndices(std::function<void(GpuIndex*)> f)
void setNumProbes(size_t nprobe)

公共成员

std::vector<std::unique_ptr<faiss::gpu::StandardGpuResources>> resources
std::vector<std::unique_ptr<GpuIndex>> subIndex
std::unique_ptr<faiss::IndexReplicas> replicaIndex
class StandardGpuResourcesImpl : public faiss::gpu::GpuResources
#include <StandardGpuResources.h>

GpuResources 对象的标准实现,它提供了临时内存管理器

公共函数

StandardGpuResourcesImpl()
~StandardGpuResourcesImpl() override
virtual bool supportsBFloat16(int device) override

给定的GPU是否支持bfloat16?

void noTempMemory()

禁用临时内存的分配;所有临时内存请求都将在使用时调用 cudaMalloc / cudaFree

void setTempMemory(size_t size)

指定我们希望在所有设备上使用某个固定大小的内存作为临时内存。 这是我们将保留的 GPU 内存的上限。 我们永远不会超过任何 GPU 上的 1.5 GiB; 较小的 GPU(具有 <= 4 GiB 或 <= 8 GiB)将使用比这更少的内存。 要避免任何临时内存分配,请传递 0。

void setPinnedMemory(size_t size)

设置要分配的固定内存量,用于异步 GPU <-> CPU 传输

virtual void setDefaultStream(int device, cudaStream_t stream) override

调用以更改工作排序的流。 我们不拥有 stream; 也就是说,当 GpuResources 对象被清理时,它不会被销毁。 我们保证所有 Faiss GPU 工作都按照此流进行排序,以便在退出索引或其他 Faiss GPU 调用时执行。

void revertDefaultStream(int device)

如果有人调用了 setDefaultStream,则将默认流恢复为由此资源对象管理的原始流。

virtual cudaStream_t getDefaultStream(int device) override

返回给定设备上的流,所有 Faiss GPU 工作都在该流上排序。 我们保证所有 Faiss GPU 工作都按照此流进行排序,以便在退出索引或其他 Faiss GPU 调用时执行。

void setDefaultNullStreamAllDevices()

调用以将工作排序流更改为所有设备的空流

void setLogMemoryAllocations(bool enable)

如果启用,会将每个 GPU 内存分配和取消分配打印到标准输出

virtual void initializeForDevice(int device) override

内部系统调用。

初始化此设备的资源

virtual cublasHandle_t getBlasHandle(int device) override

返回我们用于给定设备的cuBLAS句柄。

virtual std::vector<cudaStream_t> getAlternateStreams(int device) override

返回我们用于给定设备的一组替代流。

virtual void *allocMemory(const AllocRequest &req) override

分配非临时GPU内存。

virtual void deallocMemory(int device, void *in) override

返回先前的分配。

virtual size_t getTempMemoryAvailable(int device) const override

对于 MemorySpace::Temporary,在没有 cudaMalloc 分配的情况下,有多少空间可以立即使用?

std::map<int, std::map<std::string, std::pair<int, size_t>>> getMemoryInfo() const

导出用于Python的内存使用情况描述。

virtual std::pair<void*, size_t> getPinnedMemory() override

返回可用的CPU固定内存缓冲区。

virtual cudaStream_t getAsyncCopyStream(int device) override

返回我们在其上执行异步CPU <-> GPU复制的流。

bool supportsBFloat16CurrentDevice()

当前GPU是否支持bfloat16?

默认提供的函数

cublasHandle_t getBlasHandleCurrentDevice()

使用当前设备调用 getBlasHandle。

cudaStream_t getDefaultStreamCurrentDevice()

使用当前设备调用 getDefaultStream。

size_t getTempMemoryAvailableCurrentDevice() const

调用 getTempMemoryAvailable,使用当前设备。

GpuMemoryReservation allocMemoryHandle(const AllocRequest &req)

通过 RAII 对象返回临时内存分配。

void syncDefaultStream(int device)

将 CPU 与给定设备的默认流同步

void syncDefaultStreamCurrentDevice()

调用 syncDefaultStream,使用当前设备。

std::vector<cudaStream_t> getAlternateStreamsCurrentDevice()

调用 getAlternateStreams,使用当前设备。

cudaStream_t getAsyncCopyStreamCurrentDevice()

调用 getAsyncCopyStream,使用当前设备。

受保护函数

bool isInitialized(int device) const

GPU资源是否已经为此设备初始化?

受保护的属性

std::unordered_map<int, std::unordered_map<void*, AllocRequest>> allocs_

每个设备的当前未完成内存分配集合 device -> (分配请求, 已分配的指针)

std::unordered_map<int, std::unique_ptr<StackDeviceMemory>> tempMemory_

临时内存提供器,每个设备一个。

std::unordered_map<int, cudaStream_t> defaultStreams_

我们的默认流,工作在其上排序,每个设备一个。

std::unordered_map<int, cudaStream_t> userDefaultStreams_

如果用户设置了特定的流用于排序,则包含在此处

std::unordered_map<int, std::vector<cudaStream_t>> alternateStreams_

我们可以使用的其他流,每个设备一个。

std::unordered_map<int, cudaStream_t> asyncCopyStreams_

用于 GPU <-> CPU 锁页内存复制的异步复制流。

std::unordered_map<int, cublasHandle_t> blasHandles_

每个设备的 cuBLAS 句柄

void *pinnedMemAlloc_

用于此 GPU 的锁页内存分配。

size_t pinnedMemAllocSize_
size_t tempMemSize_

另一个选项是在所有设备上使用指定数量的内存

size_t pinnedMemSize_

我们应该分配的锁页内存量。

bool allocLogging_

我们是否记录每个 GPU 内存分配和释放。

受保护的静态函数

static size_t getDefaultTempMemForGPU(int device, size_t requested)

根据总 GPU 内存大小调整默认的临时内存分配

class StandardGpuResources : public faiss::gpu::GpuResourcesProvider
#include <StandardGpuResources.h>

GpuResources 的默认实现,它分配一个 cuBLAS 流和 2 个流以供使用,以及临时内存。 在内部,Faiss GPU 代码使用 getResources 管理的实例,但这是面向用户的对象,在内部进行引用计数。

公共函数

StandardGpuResources()
~StandardGpuResources() override
virtual std::shared_ptr<GpuResources> getResources() override

返回共享资源对象。

bool supportsBFloat16(int device)

给定的设备是否支持原生 bfloat16 算法。

bool supportsBFloat16CurrentDevice()

当前设备是否支持原生 bfloat16 算法。

void noTempMemory()

禁用临时内存的分配;所有临时内存请求都将在使用时调用 cudaMalloc / cudaFree

void setTempMemory(size_t size)

指定我们希望在所有设备上使用某个固定大小的内存作为临时内存。 这是我们将保留的 GPU 内存的上限。 我们永远不会超过任何 GPU 上的 1.5 GiB; 较小的 GPU(具有 <= 4 GiB 或 <= 8 GiB)将使用比这更少的内存。 要避免任何临时内存分配,请传递 0。

void setPinnedMemory(size_t size)

设置要分配的固定内存量,用于异步 GPU <-> CPU 传输

void setDefaultStream(int device, cudaStream_t stream)

调用以更改工作排序的流。 我们不拥有 stream; 也就是说,当 GpuResources 对象被清理时,它不会被销毁。 我们保证所有 Faiss GPU 工作都按照此流进行排序,以便在退出索引或其他 Faiss GPU 调用时执行。

void revertDefaultStream(int device)

如果有人调用了 setDefaultStream,则将默认流恢复为由此资源对象管理的原始流。

void setDefaultNullStreamAllDevices()

调用以将工作排序流更改为所有设备的空流

std::map<int, std::map<std::string, std::pair<int, size_t>>> getMemoryInfo() const

导出用于Python的内存使用情况描述。

cudaStream_t getDefaultStream(int device)

返回当前默认流。

size_t getTempMemoryAvailable(int device) const

返回当前可用的临时内存量。

void syncDefaultStreamCurrentDevice()

将我们的默认流与 CPU 同步。

void setLogMemoryAllocations(bool enable)

如果启用,会将每个 GPU 内存分配和取消分配打印到标准输出

私有成员

std::shared_ptr<StandardGpuResourcesImpl> res_
class DeviceScope
#include <DeviceUtils.h>

RAII 对象用于设置当前设备,并在销毁时恢复之前的设备

公共函数

explicit DeviceScope(int device)
~DeviceScope()

私有成员

int prevDevice_
class CublasHandleScope
#include <DeviceUtils.h>

RAII 对象用于管理 cublasHandle_t。

公共函数

CublasHandleScope()
~CublasHandleScope()
inline cublasHandle_t get()

私有成员

cublasHandle_t blasHandle_
class CudaEvent

公共函数

explicit CudaEvent(cudaStream_t stream, bool timer = false)

创建一个事件并将其记录在此流中。

CudaEvent(const CudaEvent &event) = delete
CudaEvent(CudaEvent &&event) noexcept
~CudaEvent()
inline cudaEvent_t get()
void streamWaitOnEvent(cudaStream_t stream)

在此流中等待此事件。

void cpuWaitOnEvent()

让 CPU 等待此事件完成。

CudaEvent &operator=(CudaEvent &&event) noexcept
CudaEvent &operator=(CudaEvent &event) = delete

私有成员

cudaEvent_t event_
class StackDeviceMemory
#include <StackDeviceMemory.h>

设备内存管理器,用于从内存区域中为单个设备提供临时内存分配

公共函数

StackDeviceMemory(GpuResources *res, int device, size_t allocPerDevice)

分配一个新的我们管理的内存区域。

StackDeviceMemory(int device, void *p, size_t size, bool isOwner)

为特定设备管理内存区域,无论是否具有所有权

~StackDeviceMemory()
int getDevice() const
void *allocMemory(cudaStream_t stream, size_t size)

所有请求的分配都应该是 16 字节的倍数。

void deallocMemory(int device, cudaStream_t, size_t size, void *p)
size_t getSizeAvailable() const
std::string toString() const

受保护的属性

int device_

我们的设备。

Stack stack_

内存栈。

struct Range
#include <StackDeviceMemory.h>

先前的分配范围以及需要同步的流

公共函数

inline Range(char *s, char *e, cudaStream_t str)

公共成员

char *start_
char *end_
cudaStream_t stream_
struct Stack

公共函数

Stack(GpuResources *res, int device, size_t size)

通过 cudaMalloc 分配内存的构造函数。

~Stack()
size_t getSizeAvailable() const

返回在不调用 cudaMalloc 的情况下可用于分配的大小

char *getAlloc(size_t size, cudaStream_t stream)

获取分配;保证所有分配都是 16 字节对齐的

void returnAlloc(char *p, size_t size, cudaStream_t stream)

返回分配。

std::string toString() const

返回堆栈状态。

公共成员

GpuResources *res_

我们的 GpuResources 对象。

int device_

分配所在的设备。

char *alloc_

临时内存缓冲区分配的位置;我们从这里开始分配 16 个字节

size_t allocSize_

分配的总大小。

char *start_

我们的临时内存区域; [start_, end_) 有效。

char *end_
char *head_

Stack [start, end) 中的堆栈头

std::list<Range> lastUsers_

我们堆栈上分配的先前最后用户的列表,用于可能的同步目的

size_t highWaterMemoryUsed_

从临时缓冲区使用的内存方面来看,高水位线是什么?

class KernelTimer
#include <Timer.h>

用于定时内核执行的实用程序类。

公共函数

KernelTimer(cudaStream_t stream = nullptr)

构造函数启动计时器,并将事件添加到当前设备流中

~KernelTimer()

析构函数释放事件资源。

float elapsedMilliseconds()

添加一个停止事件,然后同步停止事件,以获取在当前流中启动的任何内核的实际 GPU 端内核计时。 返回经过的毫秒数。 只能调用一次。

私有成员

cudaEvent_t startEvent_
cudaEvent_t stopEvent_
cudaStream_t stream_
bool valid_
class CpuTimer
#include <Timer.h>

CPU 挂钟经过时间计时器。

公共函数

CpuTimer()

创建并启动一个新的计时器。

float elapsedMilliseconds()

返回经过的时间,以毫秒为单位。

私有成员

std::chrono::time_point<std::chrono::steady_clock> start_
namespace utils

函数

template<typename U, typename V> constexpr __host__ __device__ auto divDown (U a, V b) -> decltype(a+b)
template<typename U, typename V> constexpr __host__ __device__ auto divUp (U a, V b) -> decltype(a+b)
template<typename U, typename V> constexpr __host__ __device__ auto roundDown (U a, V b) -> decltype(a+b)
template<typename U, typename V> constexpr __host__ __device__ auto roundUp (U a, V b) -> decltype(a+b)
template<class T> constexpr __host__ __device__ T pow (T n, T power)
template<class T> constexpr __host__ __device__ T pow2 (T n)
template<typename T> constexpr __host__ __device__ int log2 (T n, int p=0)
template<typename T> constexpr __host__ __device__ bool isPowerOf2 (T v)
template<typename T> constexpr __host__ __device__ T nextHighestPowerOf2 (T v)