Struct caffe2_context::ThreadLocalCUDAObjects

pub struct ThreadLocalCUDAObjects { /* private fields */ }

| A struct to host thread-local cuda objects. | | In Caffe2, each thread has its own non-default | cuda stream as well as related objects | such as cublas and curand handles. | | This is achieved by having the ThreadLocalCUDAObjects | wrapper that takes care of allocating | and deallocating these objects at the | thread scope. | | This class is solely used inside CUDAContext | and should not be used externally. | | This class manages the mapping from | logical stream ID (int stream_id passed | around in Caffe2) and CudaStream objects. | | We intend to eventually deprecate the | logical stream ID interface, but not | for now. |

Implementations

impl ThreadLocalCUDAObjects

pub fn set_current_stream_id(&mut self, gpu: DeviceIndex, stream_id: StreamId)

| Record current stream id for the current | thread. | | This is the new API we’re trying to | migrate use cases to and get rid of | explicit stream id passing. | | For now it’s invoked in | CUDAContext::SwitchToDevice

pub fn get_stream(&mut self, gpu: DeviceIndex) -> CudaStream

| Uses the logical stream id from the thread | local to pick the stream | | We’re going to migrate all usages to this | case API instead of passing the stream id | directly

pub fn get_stream_with_stream_id( &mut self, gpu: DeviceIndex, stream_id: StreamId ) -> CudaStream

pub fn get_handle_with_device_index(&mut self, gpu: DeviceIndex) -> CublasHandle

| Uses the logical stream id from the thread | local to pick the stream | | We’re going to migrate all usages to this | case API instead of passing the stream id | directly

pub fn get_cudastream( &mut self, gpu: DeviceIndex, stream_id: StreamId ) -> CudaStream

| Retrieves the CudaStream corresponding to | a logical stream ID, ensuring that it | exists in cuda_streams_ if it has not been | allocated yet.

pub fn get_handle_with_cuda_stream( &mut self, cuda_stream: CudaStream ) -> CublasHandle

pub fn get_cudnn_handle(&mut self, gpu: DeviceIndex)

| Uses the logical stream id from the thread | local to pick the stream | | We’re going to migrate all usages to this | case API instead of passing the stream id | directly

Trait Implementations

impl Default for ThreadLocalCUDAObjects

fn default() -> Self

Returns the “default value” for a type.

impl Drop for ThreadLocalCUDAObjects

fn drop(&mut self)

Executes the destructor for this type.