文件 StandardGpuResources.h

namespace faiss

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

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

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

PQ4 SIMD 打包和累加函数

基本内核使用 bbs = nb * 2 * 16 个向量累加 nq 个查询向量，并生成一个输出矩阵。对于 nq * nb <= 4 来说很有意思，否则寄存器溢出变得太大。

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

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

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

x[ i * d + j ]

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

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

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

反向列表的定义 + 几个实现该接口的常见类。

由于 IVF（反向文件）索引对于大规模用例非常有用，因此我们将与它们相关的几个函数分组到这个小库中。大多数函数都可以处理 IndexIVF 和嵌入在 IndexPreTransform 中的 IndexIVF。

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

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

定义了一些对向量集应用转换的对象。通常这些是预处理步骤。

namespace gpu

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 isInitialized(int device) const: 是否已为此设备初始化 GPU 资源？

保护属性

std::unordered_map<int, std::unordered_map<void*, AllocRequest>> allocs_: 每个设备的当前未完成内存分配集合。设备 ->（分配请求，已分配指针）

std::unordered_map<int, std::unique_ptr<StackDeviceMemory>> tempMemory_: 每个设备的临时内存提供程序。

std::unordered_map<int, cudaStream_t> defaultStreams_: 我们的默认流，工作基于此排序，每个设备一个。

std::unordered_map<int, cudaStream_t> userDefaultStreams_: 如果用户设置了用于排序的特定流，则包含在此处

std::unordered_map<int, std::vector<cudaStream_t>> alternateStreams_: 每个设备我们可以使用的其他流。

std::unordered_map<int, cudaStream_t> asyncCopyStreams_: 用于 GPU <-> CPU 锁页内存复制的异步复制流。

std::unordered_map<int, cublasHandle_t> blasHandles_: 每个设备的 cuBLAS 句柄

void *pinnedMemAlloc_: 用于此 GPU 的锁页内存分配。

size_t pinnedMemAllocSize_

size_t tempMemSize_: 另一种选择是在所有设备上使用指定的内存量

size_t pinnedMemSize_: 我们应该分配的锁页内存量。

bool allocLogging_: 是否记录每次 GPU 内存分配和释放。

保护静态函数

static size_t getDefaultTempMemForGPU(int device, size_t requested): 根据 GPU 总内存大小调整默认临时内存分配

class StandardGpuResources : public faiss::gpu::GpuResourcesProvider

#include <StandardGpuResources.h>

GpuResources 的默认实现，它分配一个 cuBLAS 流和 2 个流供使用，以及临时内存。在内部，Faiss GPU 代码使用 getResources 管理的实例，但这是面向用户的对象，它在内部被引用计数。

公共函数

StandardGpuResources()

~StandardGpuResources() override

virtual std::shared_ptr<GpuResources> getResources() override: 返回共享资源对象。

bool supportsBFloat16(int device): 判断给定的设备是否支持原生 bfloat16 运算。

bool supportsBFloat16CurrentDevice(): 判断当前设备是否支持原生 bfloat16 运算。

void noTempMemory(): 禁用临时内存的分配；所有临时内存请求都将在使用时调用 cudaMalloc / cudaFree

void setTempMemory(size_t size): 指定我们希望在所有设备上使用特定大小的固定内存作为临时内存。这是我们将保留的GPU内存的上限。我们在任何GPU上都不会超过1.5 GiB；较小的GPU（具有<= 4 GiB或<= 8 GiB）将使用比这更少的内存。要避免任何临时内存分配，请传递0。

void setPinnedMemory(size_t size): 设置要分配的固定内存量，用于异步 GPU <-> CPU 传输

void setDefaultStream(int device, cudaStream_t stream): 调用以更改工作排序的流。我们不拥有 stream; 也就是说，当 GpuResources 对象被清理时，它不会被销毁。我们保证在退出索引或其他 Faiss GPU 调用时，所有 Faiss GPU 工作都与此流有关。

void revertDefaultStream(int device): 如果有人调用了 setDefaultStream，则将默认流恢复为此资源对象管理的原始流。

void setDefaultNullStreamAllDevices(): 调用以将所有设备的工作排序流更改为空流

std::map<int, std::map<std::string, std::pair<int, size_t>>> getMemoryInfo() const: 导出用于 Python 的内存使用情况描述。

cudaStream_t getDefaultStream(int device): 返回当前默认流。

size_t getTempMemoryAvailable(int device) const: 返回当前可用临时内存量。

void syncDefaultStreamCurrentDevice(): 将我们的默认流与 CPU 同步。

void setLogMemoryAllocations(bool enable): 如果启用，会将每个 GPU 内存分配和释放打印到标准输出

私有成员

std::shared_ptr<StandardGpuResourcesImpl> res_