Class faiss::gpu::StandardGpuResourcesImpl

class StandardGpuResourcesImpl : public faiss::gpu::GpuResources

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 -> (分配请求, 已分配的 ptr)

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 总内存大小调整默认临时内存分配。