类 faiss::gpu::StandardGpuResources

class StandardGpuResources : public faiss::gpu::GpuResourcesProvider

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_