文件 GpuResources.h

namespace faiss

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

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

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

PQ4 SIMD 打包和累加函数

基本内核将 nq 查询向量与 bbs = nb * 2 * 16 个向量累加，并为此生成一个输出矩阵。这对于 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

枚举

enum AllocType

值

enumerator Other: 未知分配类型或杂项（当前未分类）

枚举值 FlatData: GpuIndexFlat 的主要数据存储（向量的原始矩阵和所需的向量范数）。

枚举值 IVFLists: GpuIndexIVF* 的主要数据存储（每个单独的 IVF 列表的存储）。

枚举值 Quantizer: Quantizer (PQ, SQ) 字典信息。

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

枚举值 TemporaryMemoryBuffer: StandardGpuResources 实现特定类型。使用 StandardGpuResources 时，临时内存分配 (MemorySpace::Temporary) 来自于为每个 gpu 预先分配的内存堆栈区域（例如，初始化时为 1.5 GiB）。 StandardGpuResources 的此分配用此 AllocType 标记。

枚举值 TemporaryMemoryOverflow: 使用 StandardGpuResources 时，任何无法在 TemporaryMemoryBuffer 区域内满足的 MemorySpace::Temporary 分配都会回退到调用 cudaMalloc，该调用的大小正好适合手头的请求。这些“溢出”临时分配用此 AllocType 标记。

枚举 MemorySpace

GPU 可访问的内存区域。

值

枚举值 Temporary: 临时设备内存（保证在顶级索引调用退出后不再使用，并且使用它的流已完成 GPU 工作）。通常由设备内存（cudaMalloc/cudaFree）支持。

枚举值 Device: 使用 cudaMalloc/cudaFree 进行管理（典型的 GPU 设备内存）

枚举值 Unified: 使用 cudaMallocManaged/cudaFree 进行管理（典型的统一 CPU/GPU 内存）

函数

std::string allocTypeToString(AllocType t): 将 AllocType 转换为字符串。

std::string memorySpaceToString(MemorySpace s): 将 MemorySpace 转换为字符串。

AllocInfo makeDevAlloc(AllocType at, cudaStream_t st): 使用 MemorySpace::Device 为当前设备创建 AllocInfo。

AllocInfo makeTempAlloc(AllocType at, cudaStream_t st): 使用 MemorySpace::Temporary 为当前设备创建 AllocInfo。

AllocInfo makeSpaceAlloc(AllocType at, MemorySpace sp, cudaStream_t st): 为当前设备创建 AllocInfo。

结构体 AllocInfo

#include <GpuResources.h>

有关分配内容/位置的信息。

由 faiss::gpu::AllocRequest 继承

公共函数

inline AllocInfo()

inline AllocInfo(AllocType at, int dev, MemorySpace sp, cudaStream_t st)

std::string toString() const: 返回此信息的字符串表示形式。

公共成员

AllocType type = AllocType::Other : 分配的内部类别。

int device = 0: 发生分配的设备。

MemorySpace space = MemorySpace::Device : 分配的内存空间。

cudaStream_t stream = nullptr

将在内存上进行新工作的流（例如，如果缓存的一块内存并返回以供此调用最后在流3上使用，而新的内存请求是针对流4，则内存管理器将同步流4以等待流3通过事件或其他流同步完成）。

内存管理器保证返回的内存可以自由使用，而不会在此指定的流上发生数据竞争。

struct AllocRequest : public faiss::gpu::AllocInfo 

#include <GpuResources.h>

关于分配的内容/位置以及它应该有多大的信息。

公共函数

inline AllocRequest()

inline AllocRequest(const AllocInfo &info, size_t sz)

inline AllocRequest(AllocType at, int dev, MemorySpace sp, cudaStream_t st, size_t sz)

std::string toString() const: 返回此请求的字符串表示形式。

公共成员

size_t size = 0: 分配的大小（以字节为单位）。

AllocType type = AllocType::Other : 分配的内部类别。

int device = 0: 发生分配的设备。

MemorySpace space = MemorySpace::Device : 分配的内存空间。

cudaStream_t stream = nullptr

将在内存上进行新工作的流（例如，如果缓存的一块内存并返回以供此调用最后在流3上使用，而新的内存请求是针对流4，则内存管理器将同步流4以等待流3通过事件或其他流同步完成）。

内存管理器保证返回的内存可以自由使用，而不会在此指定的流上发生数据竞争。

struct GpuMemoryReservation

#include <GpuResources.h>

一个管理临时内存请求的 RAII 对象。

公共函数

GpuMemoryReservation()

GpuMemoryReservation(GpuResources *r, int dev, cudaStream_t str, void *p, size_t sz)

GpuMemoryReservation(GpuMemoryReservation &&m) noexcept

~GpuMemoryReservation()

GpuMemoryReservation &operator=(GpuMemoryReservation &&m)

inline void *get()

void release()

公共成员

GpuResources *res

int device

cudaStream_t stream

void *data

size_t size

class GpuResources

#include <GpuResources.h>

GPU端资源提供器的基类；隐藏了cuBLAS句柄、CUDA流和所有执行的设备内存分配的供应

faiss::gpu::StandardGpuResourcesImpl 的子类

公共函数

virtual ~GpuResources()

virtual void initializeForDevice(int device) = 0: 调用以为特定设备预先分配资源。如果未调用此函数，则将在首次需求时分配资源

virtual bool supportsBFloat16(int device) = 0: 给定的GPU是否支持bfloat16？

virtual cublasHandle_t getBlasHandle(int device) = 0: 返回我们用于给定设备的cuBLAS句柄。

virtual cudaStream_t getDefaultStream(int device) = 0: 返回我们为给定设备排序所有计算的流

virtual void setDefaultStream(int device, cudaStream_t stream) = 0: 将设备的默认流覆盖为用户提供的流。资源对象不拥有此流（即，它不会销毁它）。

virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0: 返回我们用于给定设备的一组备用流。

virtual void *allocMemory(const AllocRequest &req) = 0: 内存管理：从给定的内存空间返回一个分配，并根据给定的流进行排序（即，第一个用户将是此流中的内核）。所有分配在内部都被调整大小为最接近的16字节的倍数，并且保证所有返回的分配都是16字节对齐的。

virtual void deallocMemory(int device, void *in) = 0: 返回先前的分配。

virtual size_t getTempMemoryAvailable(int device) const = 0: 对于 MemorySpace::Temporary，无需 cudaMalloc 分配即可立即使用多少空间？

virtual std::pair<void*, size_t> getPinnedMemory() = 0: 返回可用的 CPU 锁页内存缓冲区。

virtual cudaStream_t getAsyncCopyStream(int device) = 0: 返回执行异步 CPU <-> GPU 复制的流。

bool supportsBFloat16CurrentDevice()

当前 GPU 是否支持 bfloat16？

默认提供的函数

cublasHandle_t getBlasHandleCurrentDevice(): 使用当前设备调用 getBlasHandle。

cudaStream_t getDefaultStreamCurrentDevice(): 使用当前设备调用 getDefaultStream。

size_t getTempMemoryAvailableCurrentDevice() const: 使用当前设备调用 getTempMemoryAvailable。

GpuMemoryReservation allocMemoryHandle(const AllocRequest &req): 通过 RAII 对象返回临时内存分配。

void syncDefaultStream(int device): 使 CPU 与给定设备的默认流同步

void syncDefaultStreamCurrentDevice(): 为当前设备调用 syncDefaultStream。

std::vector<cudaStream_t> getAlternateStreamsCurrentDevice(): 为当前设备调用 getAlternateStreams。

cudaStream_t getAsyncCopyStreamCurrentDevice(): 为当前设备调用 getAsyncCopyStream。

class GpuResourcesProvider

#include <GpuResources.h>

共享资源对象的提供程序接口。这是为了避免将 std::shared_ptr 连接到 Python

由 faiss::gpu::GpuResourcesProviderFromInstance、faiss::gpu::StandardGpuResources 继承

公共函数

virtual ~GpuResourcesProvider()

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

class GpuResourcesProviderFromInstance : public faiss::gpu::GpuResourcesProvider

#include <GpuResources.h>

一个简单的 GpuResources 对象包装器，用于再次从中创建 GpuResourcesProvider

公共函数

explicit GpuResourcesProviderFromInstance(std::shared_ptr<GpuResources> p)

~GpuResourcesProviderFromInstance() override

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

私有成员

std::shared_ptr<GpuResources> res_