命名空间 faiss::gpu
-
namespace gpu
枚举
-
enum class DistanceDataType
值
-
enumerator F32
-
enumerator F16
-
enumerator BF16
-
enumerator F32
-
enum class IndicesDataType
值
-
enumerator I64
-
enumerator I32
-
enumerator I64
-
enum class graph_build_algo
值
-
enumerator IVF_PQ
使用 IVF-PQ 构建全邻近 knn 图。
-
enumerator NN_DESCENT
使用 NN-Descent 构建全邻近 knn 图。
-
enumerator IVF_PQ
-
enum class codebook_gen
一种用于指定如何创建 PQ 码本的类型。
值
-
enumerator PER_SUBSPACE
-
enumerator PER_CLUSTER
-
enumerator PER_SUBSPACE
-
enum class search_algo
值
-
enumerator SINGLE_CTA
对于大批量大小。
-
enumerator MULTI_CTA
对于小批量大小。
-
enumerator MULTI_KERNEL
-
enumerator AUTO
-
enumerator SINGLE_CTA
-
enum class hash_mode
值
-
enumerator HASH
-
enumerator SMALL
-
enumerator AUTO
-
enumerator HASH
-
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 位整数。
-
enumerator INDICES_CPU
-
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。
-
enumerator Other
-
枚举 MemorySpace
GPU 可访问的内存区域。
值
-
枚举器 Temporary
临时设备内存(保证在顶级索引调用退出后不再使用,并且使用它的流已完成 GPU 工作)。通常由设备内存(cudaMalloc/cudaFree)支持。
-
枚举器 Device
使用 cudaMalloc/cudaFree 管理(典型的 GPU 设备内存)
-
枚举器 Unified
使用 cudaMallocManaged/cudaFree 管理(典型的统一 CPU/GPU 内存)
-
枚举器 Temporary
函数
-
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。
-
std::string memorySpaceToString(MemorySpace s)
将 MemorySpace 转换为字符串。
-
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)
生成随机位向量的集合。
-
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())
返回给定 GPU 设备可用的最大 smem。
等效于 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) 调用一个流集合等待。
-
struct GpuParameterSpace : public faiss::ParameterSpace
- #include <GpuAutoTune.h>
GPU 索引的参数空间和设置器
公共函数
-
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 模式下的抖动)
-
virtual void set_index_parameter(faiss::Index *index, const std::string &name, double val) const override
-
struct ToCPUCloner : public faiss::Cloner
- #include <GpuCloner.h>
为 GPU -> CPU 专门设计的Cloner。
公共函数
-
virtual VectorTransform *clone_VectorTransform(const VectorTransform*)
-
virtual VectorTransform *clone_VectorTransform(const VectorTransform*)
-
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 VectorTransform *clone_VectorTransform(const VectorTransform*)
公共成员
-
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 实现。
-
ToGpuCloner(GpuResourcesProvider *prov, int device, const GpuClonerOptions &options)
-
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)
-
virtual VectorTransform *clone_VectorTransform(const VectorTransform*)
公共成员
-
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 实现。
-
ToGpuClonerMultiple(std::vector<GpuResourcesProvider*> &provider, std::vector<int> &devices, const GpuMultipleClonerOptions &options)
-
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
-
explicit GpuProgressiveDimIndexFactory(int ngpu)
-
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 实现。
-
IndicesOptions indicesOptions = INDICES_64_BIT
-
`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 实现。
-
是否进行分片(sharding),默认为false。
-
`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?
-
距离度量标准,默认为 `METRIC_L2` (欧几里得距离)。
-
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::unique_ptr<IcmEncoderShards> shards
-
GpuIcmEncoder(const LocalSearchQuantizer *lsq, const std::vector<GpuResourcesProvider*> &provs, const std::vector<int> &devices)
-
struct GpuIcmEncoderFactory : public faiss::lsq::IcmEncoderFactory
公共函数
-
explicit GpuIcmEncoderFactory(int ngpus = 1)
-
virtual lsq::IcmEncoder *get(const LocalSearchQuantizer *lsq) override
-
explicit GpuIcmEncoderFactory(int ngpus = 1)
-
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?
-
int device = 0
-
class GpuIndex : public faiss::Index
被 faiss::gpu::GpuIndexCagra, faiss::gpu::GpuIndexFlat, faiss::gpu::GpuIndexIVF 继承
公共类型
-
using component_t = float
-
using distance_t = float
公共函数
-
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
x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override
x
和labels
可以驻留在 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
x
、distances
和labels
可以驻留在 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
x
、distances
和labels
以及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
度量类型的参数
受保护函数
-
virtual bool addImplRequiresIDs_() const = 0
addImpl_ 是否需要 ID? 如果是,并且未提供 ID,我们将根据添加 ID 的顺序依次生成它们
-
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 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 复制
-
using component_t = float
-
struct GpuIndexBinaryFlatConfig : public faiss::gpu::GpuIndexConfig
公共成员
-
设备 ID,默认为 0。
索引所在的 GPU 设备。
-
内存空间,默认为 Device。
用于主存储的内存空间。在 Pascal 及以上 (CC 6+) 架构上,允许 GPU 使用比 GPU 上可用内存更多的内存。
-
是否使用 cuvs,默认为 false。
索引是否应调度到 cuVS?
-
设备 ID,默认为 0。
-
GpuIndexBinaryFlat
类继承自faiss::IndexBinary
。 - 包含头文件 <GpuIndexBinaryFlat.h>。
IndexBinaryFlat
的 GPU 版本,用于通过汉明距离进行位向量的暴力比较。公共类型
-
定义别名
component_t
为uint8_t
。
-
定义别名
distance_t
为int32_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
)。但是,当dim
是pq_dim
的倍数时,此转换不是必需的(dim == rot_dim
,因此无需添加“额外”数据列/特征)。默认情况下,如果
dim == rot_dim
,则使用单位矩阵初始化旋转变换。 当force_random_rotation == true
时,无论dim
和pq_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 调整指南。
请注意,这是一个低级调整参数,如果调整不正确,可能会对搜索性能产生严重的负面影响。
-
uint32_t n_probes = 20
-
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?
-
size_t intermediate_graph_degree = 128
-
结构体 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。
-
size_t max_queries = 0
-
结构体 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
从数据库中删除所有元素。
-
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
x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override
x
和labels
可以驻留在 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
x
、distances
和labels
可以驻留在 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
x
、distances
和labels
以及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 调用以进行搜索。
受保护的属性
-
const GpuIndexCagraConfig cagraConfig_
我们的配置选项。
-
std::shared_ptr<CuvsCagra> index_
我们拥有的实例;包含倒排列表。
-
std::shared_ptr<GpuResources> resources_
管理设备上的流、cuBLAS 句柄和暂存内存。
-
const GpuIndexConfig config_
我们的配置选项。
-
size_t minPagedSize_
将副本从 CPU 分页到 GPU 的大小阈值。
-
using component_t = float
-
结构体 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 ( GpuResourcesProvider * provider , int dims , faiss :: MetricType metric , GpuIndexFlatConfig config = GpuIndexFlatConfig ())
Construct an empty instance that can be added to.
-
~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 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
x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override
x
和labels
可以驻留在 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
x
、distances
和labels
可以驻留在 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
x
、distances
和labels
以及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 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 的大小阈值。
-
using component_t = float
-
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(provider, dims, config = GpuIndexFlatConfig())
Construct an empty instance that can be added to.
-
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数组。 x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
k
: 要考虑的最近邻的数量 (默认为 1)。 x
和labels
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制
-
params
: 搜索参数 (可以为 null)。 x
、distances
和labels
可以驻留在 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
x
、distances
和labels
以及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 的大小阈值。
-
using component_t = float
-
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(GpuResourcesProvider *provider, int dims, GpuIndexFlatConfig config = GpuIndexFlatConfig())
Construct an empty instance that can be added to.
-
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的向量。
x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
分配向量到最近的聚类中心。
x
和labels
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制
-
搜索k个最近邻。
x
、distances
和labels
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制
-
搜索k个最近邻并重构它们。
x
、distances
和labels
以及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 resetIndex_(int dims)
-
virtual bool addImplRequiresIDs_() const override
Flat 索引不需要 ID,因为没有可用于存储 ID 的空间
-
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 的大小阈值。
-
组件类型定义为 float
-
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?
-
IndicesOptions indicesOptions = INDICES_64_BIT
-
class GpuIndexIVF : public faiss::gpu::GpuIndex, public faiss::IndexIVFInterface
- #include <GpuIndexIVF.h>
所有GPU IVF索引类型的基类。 目前,它有意不继承自IndexIVF,因为IndexIVF中的许多公共数据成员和功能在GPU上不支持以相同的方式使用。
被faiss::gpu::GpuIndexIVFFlat、faiss::gpu::GpuIndexIVFPQ、faiss::gpu::GpuIndexIVFScalarQuantizer继承。
公共类型
-
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
-
virtual void updateQuantizer() = 0
如果用户手动更改 IVF 粗量化器的状态(例如,替换新实例或在训练范围之外更改粗量化器中的向量),则应调用此函数。
-
virtual idx_t getNumLists() const
返回我们正在管理的倒排列表的数量。
-
virtual std::vector<uint8_t> getListVectorData(idx_t listId, bool gpuFormat = false) const
返回特定倒排列表中包含的编码向量数据,用于调试目的。如果 gpuFormat 为 true,则返回的数据采用 GPU 端表示中的编码形式。 否则,它将转换为 CPU 格式。 兼容格式,而本机 GPU 格式可能有所不同。
-
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
x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override
x
和labels
可以驻留在 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
x
、distances
和labels
可以驻留在 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
x
、distances
和labels
以及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 调用以进行搜索。
受保护的属性
-
const GpuIndexIVFConfig ivfConfig_
我们的配置选项。
-
std::shared_ptr<IVFBase> baseIndex_
对于训练/初始化后的索引,这是对基类的引用。
-
std::shared_ptr<GpuResources> resources_
管理设备上的流、cuBLAS 句柄和暂存内存。
-
const GpuIndexConfig config_
我们的配置选项。
-
size_t minPagedSize_
将副本从 CPU 分页到 GPU 的大小阈值。
私有函数
-
void init_()
共享初始化函数。
-
using component_t = float
-
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?
-
bool interleavedLayout = true
-
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)
-
virtual idx_t getNumLists() const
返回我们正在管理的倒排列表的数量。
-
virtual std::vector<uint8_t> getListVectorData(idx_t listId, bool gpuFormat = false) const
返回特定倒排列表中包含的编码向量数据,用于调试目的。如果 gpuFormat 为 true,则返回的数据采用 GPU 端表示中的编码形式。 否则,它将转换为 CPU 格式。 兼容格式,而本机 GPU 格式可能有所不同。
-
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
x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override
x
和labels
可以驻留在 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
x
、distances
和labels
可以驻留在 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
x
、distances
和labels
以及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
-
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 的大小阈值。
-
using component_t = float
-
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?
-
bool useFloat16LookupTables = false
-
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
基于给定的向量数据训练粗量化器和乘积量化器。
-
virtual idx_t getNumLists() const
返回我们正在管理的倒排列表的数量。
-
virtual std::vector<uint8_t> getListVectorData(idx_t listId, bool gpuFormat = false) const
返回特定倒排列表中包含的编码向量数据,用于调试目的。如果 gpuFormat 为 true,则返回的数据采用 GPU 端表示中的编码形式。 否则,它将转换为 CPU 格式。 兼容格式,而本机 GPU 格式可能有所不同。
-
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
x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override
x
和labels
可以驻留在 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
x
、distances
和labels
可以驻留在 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
x
、distances
和labels
以及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 量化器。
-
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 的大小阈值。
-
using component_t = float
-
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?
-
bool interleavedLayout = true
-
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
基于给定的向量数据训练粗量化器和标量量化器。
-
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
x
和ids
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制。如果添加集太大,则处理分页添加;调用 addInternal_
-
virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override
x
和labels
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制
-
带ID添加向量:virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override
x
、distances
和labels
可以驻留在 CPU 或任何 GPU 上;根据需要执行复制
-
virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override
x
、distances
和labels
以及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 残差训练。
-
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 的大小阈值。
-
using component_t = float
-
struct AllocInfo
- #include <GpuResources.h>
关于分配是什么/在哪里的信息。
公共函数
-
inline AllocInfo()
-
inline AllocInfo(AllocType at, int dev, MemorySpace sp, cudaStream_t st)
-
std::string toString() const
返回此信息的字符串表示形式。
公共成员
-
int device = 0
分配正在发生的设备。
-
MemorySpace space = MemorySpace::Device
分配的内存空间。
-
cudaStream_t stream = nullptr
将在其上对内存上的新工作进行排序的流(例如,如果缓存的一段内存将被返回给此调用,上次在流 3 上使用,并且新的内存请求是针对流 4 的,则内存管理器将同步流 4 以等待通过事件或其他流同步完成流 3)。
内存管理器保证返回的内存可以自由使用,而不会发生指定流上的数据竞争。
-
inline AllocInfo()
-
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
分配的大小,以字节为单位。
-
int device = 0
分配正在发生的设备。
-
MemorySpace space = MemorySpace::Device
分配的内存空间。
-
cudaStream_t stream = nullptr
将在其上对内存上的新工作进行排序的流(例如,如果缓存的一段内存将被返回给此调用,上次在流 3 上使用,并且新的内存请求是针对流 4 的,则内存管理器将同步流 4 以等待通过事件或其他流同步完成流 3)。
内存管理器保证返回的内存可以自由使用,而不会发生指定流上的数据竞争。
-
inline AllocRequest()
-
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()
-
GpuMemoryReservation()
-
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,使用当前设备。
-
virtual ~GpuResources()
-
class GpuResourcesProvider
- #include <GpuResources.h>
共享资源对象的提供程序接口。 这是为了避免将 std::shared_ptr 与 Python 接口
被 faiss::gpu::GpuResourcesProviderFromInstance, faiss::gpu::StandardGpuResources 继承
公共函数
-
virtual ~GpuResourcesProvider()
-
virtual std::shared_ptr<GpuResources> getResources() = 0
返回共享资源对象。
-
virtual ~GpuResourcesProvider()
-
class GpuResourcesProviderFromInstance : public faiss::gpu::GpuResourcesProvider
- #include <GpuResources.h>
一个简单的 GpuResources 对象包装器,用于再次从中创建一个 GpuResourcesProvider
公共函数
-
~GpuResourcesProviderFromInstance() override
-
virtual std::shared_ptr<GpuResources> getResources() override
返回共享资源对象。
私有成员
-
std::shared_ptr<GpuResources> res_
-
~GpuResourcesProviderFromInstance() override
-
template<typename GpuIndex>
struct IndexWrapper 公共函数
-
IndexWrapper(int numGpus, std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)> init)
-
void setNumProbes(size_t nprobe)
-
IndexWrapper(int numGpus, std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)> init)
-
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, 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 内存大小调整默认的临时内存分配
-
StandardGpuResourcesImpl()
-
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_
-
StandardGpuResources()
-
class DeviceScope
- #include <DeviceUtils.h>
RAII 对象用于设置当前设备,并在销毁时恢复之前的设备
公共函数
-
explicit DeviceScope(int device)
-
~DeviceScope()
私有成员
-
int prevDevice_
-
explicit DeviceScope(int device)
-
class CublasHandleScope
- #include <DeviceUtils.h>
RAII 对象用于管理 cublasHandle_t。
公共函数
-
CublasHandleScope()
-
~CublasHandleScope()
-
inline cublasHandle_t get()
私有成员
-
cublasHandle_t blasHandle_
-
CublasHandleScope()
-
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_t event_
-
explicit CudaEvent(cudaStream_t stream, bool timer = false)
-
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
-
struct Range
- #include <StackDeviceMemory.h>
先前的分配范围以及需要同步的流
公共函数
-
inline Range(char *s, char *e, cudaStream_t str)
公共成员
-
char *start_
-
char *end_
-
cudaStream_t stream_
-
inline Range(char *s, char *e, cudaStream_t str)
-
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) 中的堆栈头
-
size_t highWaterMemoryUsed_
从临时缓冲区使用的内存方面来看,高水位线是什么?
-
Stack(GpuResources *res, int device, size_t size)
-
StackDeviceMemory(GpuResources *res, int device, size_t allocPerDevice)
-
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_
-
KernelTimer(cudaStream_t stream = nullptr)
-
class CpuTimer
- #include <Timer.h>
CPU 挂钟经过时间计时器。
公共函数
-
CpuTimer()
创建并启动一个新的计时器。
-
float elapsedMilliseconds()
返回经过的时间,以毫秒为单位。
-
CpuTimer()
-
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)
-
enum class DistanceDataType