experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #43949 from sanjoy:cudnn-version-check

Browse Source 
PiperOrigin-RevId: 339899425
Change-Id: If9bbb9c68787dc484129102c716da2203c209dab
This commit is contained in:
TensorFlower Gardener 2020-10-30 10:51:29 -07:00
commit 6bc291d831
1 changed files with 17 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
tensorflow/stream_executor/cuda/cuda_dnn.cc
View File

@ -291,6 +291,17 @@ port::Status GetLoadedCudnnVersion(CudnnVersion* version) {
return port::Status::OK(); return port::Status::OK();
} }
#if CUDNN_MAJOR >= 8 && (CUDNN_MINOR > 0 || CUDNN_PATCHLEVEL >= 4)
void PreloadCudnnLibrary(cudnnStatus_t (*version_check_fn)(),
absl::string_view sub_library) {
cudnnStatus_t status = version_check_fn();
if (status != CUDNN_STATUS_SUCCESS) {
VLOG(1) << "Could not pre-initialize cuDNN sub-library " << sub_library
<< ". Error: " << cudnnGetErrorString(status) << ".";
}
}
#endif
} // namespace } // namespace
CudnnSupport::CudnnSupport(GpuExecutor* parent) : parent_(parent) {} CudnnSupport::CudnnSupport(GpuExecutor* parent) : parent_(parent) {}
@ -320,12 +331,12 @@ port::Status CudnnSupport::Init() {
// Preload sub libs for cudnn 8.0.4+ // Preload sub libs for cudnn 8.0.4+
#if CUDNN_MAJOR >= 8 && (CUDNN_MINOR > 0 || CUDNN_PATCHLEVEL >= 4) #if CUDNN_MAJOR >= 8 && (CUDNN_MINOR > 0 || CUDNN_PATCHLEVEL >= 4)
cudnnOpsInferVersionCheck(); PreloadCudnnLibrary(cudnnOpsInferVersionCheck, "cudnn_ops_infer");
cudnnOpsTrainVersionCheck(); PreloadCudnnLibrary(cudnnOpsTrainVersionCheck, "cudnn_ops_train");
cudnnCnnInferVersionCheck(); PreloadCudnnLibrary(cudnnCnnInferVersionCheck, "cudnn_cnn_infer");
cudnnCnnTrainVersionCheck(); PreloadCudnnLibrary(cudnnCnnTrainVersionCheck, "cudnn_cnn_train");
cudnnAdvInferVersionCheck(); PreloadCudnnLibrary(cudnnAdvInferVersionCheck, "cudnn_adv_infer");
cudnnAdvTrainVersionCheck(); PreloadCudnnLibrary(cudnnAdvTrainVersionCheck, "cudnn_adv_train");
#endif #endif
cudnn_.reset(new CudnnAccess(cudnn_handle)); cudnn_.reset(new CudnnAccess(cudnn_handle));