Class faiss::gpu::GpuResources

class GpuResources

Base class of GPU-side resource provider; hides provision of cuBLAS handles, CUDA streams and all device memory allocation performed

Subclassed by faiss::gpu::StandardGpuResourcesImpl

Public Functions

virtual ~GpuResources()

virtual void initializeForDevice(int device) = 0

Call to pre-allocate resources for a particular device. If this is not called, then resources will be allocated at the first time of demand

virtual bool supportsBFloat16(int device) = 0

Does the given GPU support bfloat16?

virtual cublasHandle_t getBlasHandle(int device) = 0

Returns the cuBLAS handle that we use for the given device.

virtual cudaStream_t getDefaultStream(int device) = 0

Returns the stream that we order all computation on for the given device

virtual void setDefaultStream(int device, cudaStream_t stream) = 0

Overrides the default stream for a device to the user-supplied stream. The resources object does not own this stream (i.e., it will not destroy it).

virtual std::vector<cudaStream_t> getAlternateStreams(int device) = 0

Returns the set of alternative streams that we use for the given device.

virtual void *allocMemory(const AllocRequest &req) = 0

Memory management Returns an allocation from the given memory space, ordered with respect to the given stream (i.e., the first user will be a kernel in this stream). All allocations are sized internally to be the next highest multiple of 16 bytes, and all allocations returned are guaranteed to be 16 byte aligned.

virtual void deallocMemory(int device, void *in) = 0

Returns a previous allocation.

virtual size_t getTempMemoryAvailable(int device) const = 0

For MemorySpace::Temporary, how much space is immediately available without cudaMalloc allocation?

virtual std::pair<void*, size_t> getPinnedMemory() = 0

Returns the available CPU pinned memory buffer.

virtual cudaStream_t getAsyncCopyStream(int device) = 0

Returns the stream on which we perform async CPU <-> GPU copies.

bool supportsBFloat16CurrentDevice()

Does the current GPU support bfloat16?

Functions provided by default

cublasHandle_t getBlasHandleCurrentDevice()

Calls getBlasHandle with the current device.

cudaStream_t getDefaultStreamCurrentDevice()

Calls getDefaultStream with the current device.

size_t getTempMemoryAvailableCurrentDevice() const

Calls getTempMemoryAvailable with the current device.

GpuMemoryReservation allocMemoryHandle(const AllocRequest &req)

Returns a temporary memory allocation via a RAII object.

void syncDefaultStream(int device)

Synchronizes the CPU with respect to the default stream for the given device

void syncDefaultStreamCurrentDevice()

Calls syncDefaultStream for the current device.

std::vector<cudaStream_t> getAlternateStreamsCurrentDevice()

Calls getAlternateStreams for the current device.

cudaStream_t getAsyncCopyStreamCurrentDevice()

Calls getAsyncCopyStream for the current device.