caffe2::CuDNNWrapper Class Reference

CuDNNWrapper is a class that wraps the cudnn handles and cudnn workspaces. More...

#include <cudnn_wrappers.h>

Data Structures
struct	SyncedCuDNNState

Public Member Functions
	CuDNNWrapper (CUDAContext *context)
	Creates a cudnn wrapper associated with a CUDAContext object. More...
cudnnHandle_t	inline_cudnn_handle ()
	Returns the inline cudnn handle that executes on the current thread's cuda_stream.
template<typename F >
void	with_cudnn_state (size_t state_idx, F &&f)

Protected Types
using	PerGPUCuDNNStates = std::array< std::array< SyncedCuDNNState, CAFFE2_COMPILE_TIME_MAX_CUDNN_STATES >, C10_COMPILE_TIME_MAX_GPUS >

Protected Member Functions
	C10_DISABLE_COPY_AND_ASSIGN (CuDNNWrapper)

Static Protected Member Functions
static PerGPUCuDNNStates &	cudnn_states ()

Protected Attributes
CUDAContext *	context_

Static Protected Attributes
static constexpr size_t	CAFFE2_COMPILE_TIME_MAX_CUDNN_STATES = 4

Detailed Description

CuDNNWrapper is a class that wraps the cudnn handles and cudnn workspaces.

The wrapper ensures that for each thread and each gpu, there is one identical cudnn handle, which is also associated with the thread-local per-device cuda stream. The wrapper also hosts the device-specific cudnn workspace (scratch space for some cudnn functions).

Definition at line 142 of file cudnn_wrappers.h.

Constructor & Destructor Documentation

caffe2::CuDNNWrapper::CuDNNWrapper ( CUDAContext *  context )

inlineexplicit

Creates a cudnn wrapper associated with a CUDAContext object.

Note that the CUDAContext object should outlive the CuDNNWrapper.

Definition at line 148 of file cudnn_wrappers.h.

---

The documentation for this class was generated from the following files:

• caffe2/core/cudnn_wrappers.h
• caffe2/core/common_cudnn.cc