caffe2::CuDNNWrapper Class Reference

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

#include <common_cudnn.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 >, CAFFE2_COMPILE_TIME_MAX_GPUS >

Protected Member Functions
	DISABLE_COPY_AND_ASSIGN (CuDNNWrapper)

Static Protected Member Functions
static PerGPUCuDNNStates &	cudnn_states ()

Protected Attributes
CUDAContext *	context_

Static Protected Attributes
static thread_local CuDNNHandles	tls_cudnn_handles_
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 409 of file common_cudnn.h.

Constructor & Destructor Documentation

CuDNNWrapper()

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 415 of file common_cudnn.h.

---

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

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