| 6439 | } // namespace gpu |
| 6440 | |
| 6441 | void initialize_cudnn() { |
| 6442 | port::Status status = |
| 6443 | PluginRegistry::Instance()->RegisterFactory<PluginRegistry::DnnFactory>( |
| 6444 | cuda::kCudaPlatformId, gpu::kCuDnnPlugin, "cuDNN", |
| 6445 | [](internal::StreamExecutorInterface* parent) -> dnn::DnnSupport* { |
| 6446 | gpu::GpuExecutor* cuda_executor = |
| 6447 | dynamic_cast<gpu::GpuExecutor*>(parent); |
| 6448 | if (cuda_executor == nullptr) { |
| 6449 | LOG(ERROR) << "Attempting to initialize an instance of the cuDNN " |
| 6450 | << "support library with a non-CUDA StreamExecutor"; |
| 6451 | return nullptr; |
| 6452 | } |
| 6453 | |
| 6454 | gpu::CudnnSupport* dnn = new gpu::CudnnSupport(cuda_executor); |
| 6455 | if (!dnn->Init().ok()) { |
| 6456 | // Note: Init() will log a more specific error. |
| 6457 | delete dnn; |
| 6458 | return nullptr; |
| 6459 | } |
| 6460 | return dnn; |
| 6461 | }); |
| 6462 | |
| 6463 | if (!status.ok()) { |
| 6464 | LOG(ERROR) << "Unable to register cuDNN factory: " |
| 6465 | << status.error_message(); |
| 6466 | } |
| 6467 | |
| 6468 | PluginRegistry::Instance()->SetDefaultFactory( |
| 6469 | cuda::kCudaPlatformId, PluginKind::kDnn, gpu::kCuDnnPlugin); |
| 6470 | } |
| 6471 | |
| 6472 | } // namespace stream_executor |
| 6473 |
no test coverage detected