Creates a CudnnHandle instance for stream. cuDNN API calls using the same handle instance need to be serialized across threads. This is guaranteed by CudnnHandle instances locking the mutex owned by this class. Most cuDNN APIs taking a handle perform work on a CUDA stream. The CudnnHandle instance acquires the executor's CUDA context and sets cuDNN to use the provided stream. The stream argumen
| 221 | // therefore a bad idea (performance wise) to call any cuDNN APIs that |
| 222 | // enqueue work in the stream. |
| 223 | CudnnHandle GetHandle(GpuExecutor* executor, Stream* stream) { |
| 224 | auto lock = absl::make_unique<absl::MutexLock>(&mutex_); |
| 225 | mutex_.AssertHeld(); |
| 226 | gpu::ScopedActivateExecutorContext context(executor); |
| 227 | CUstream cu_stream = stream ? AsGpuStreamValue(stream) : cudaStreamLegacy; |
| 228 | const auto status = cudnnSetStream(handle_, cu_stream); |
| 229 | CHECK_EQ(status, CUDNN_STATUS_SUCCESS) << "Failed to set cuDNN stream."; |
| 230 | return CudnnHandle(std::move(context), std::move(lock), handle_); |
| 231 | } |
| 232 | |
| 233 | private: |
| 234 | // Guards the enqueueing of cuDNN operations via the handle_ below. |
no test coverage detected