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

Polish TensorRT static linking a little.

Browse Source 
PiperOrigin-RevId: 347545653
Change-Id: I58dd22ed1d5b3ead01c578d685660f92a55564dd
This commit is contained in:
Christian Sigg 2020-12-14 22:55:38 -08:00 committed by TensorFlower Gardener
parent 35fc7c9443
commit 6c29462797
2 changed files with 12 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
tensorflow/compiler/tf2tensorrt/BUILD
View File

@ -50,27 +50,11 @@ cc_library(
]),
)
# Copybara will replace tensorrt_lib_oss_placeholder with this one
alias(
name = "tensorrt_lib_oss",
actual = select({
"@local_config_tensorrt//:use_static_tensorrt": "@local_config_tensorrt//:tensorrt",
"//conditions:default": ":tensorrt_stub",
}),
visibility = ["//visibility:private"],
)
alias(
name = "tensorrt_lib_oss_placeholder",
actual = ":tensorrt_stub",
visibility = ["//visibility:private"],
)
alias(
name = "tensorrt_lib",
actual = select({
"//tensorflow:oss": ":tensorrt_lib_oss",
"//conditions:default": "@local_config_tensorrt//:tensorrt",
"@local_config_tensorrt//:use_static_tensorrt": "@local_config_tensorrt//:tensorrt",
"//conditions:default": ":tensorrt_stub",
}),
visibility = ["//visibility:private"],
)
@ -607,21 +591,15 @@ tf_proto_library(
protodeps = tf_additional_all_protos(),
)
tensorrt_static_define_oss_placeholder = []
# @unused
# Copybara will replace tensorrt_static_define_oss_placeholder with this one
tensorrt_static_define_oss = select({
"@local_config_tensorrt//:use_static_tensorrt": ["TF_OSS_TENSORRT_STATIC=1"],
"//conditions:default": [],
})
cc_library(
name = "py_utils",
srcs = ["utils/py_utils.cc"],
hdrs = ["utils/py_utils.h"],
copts = tf_copts(),
defines = tensorrt_static_define_oss,
local_defines = select({
"@local_config_tensorrt//:use_static_tensorrt": ["TF_USE_TENSORRT_STATIC=1"],
"//conditions:default": [],
}),
deps = if_tensorrt([
":common_utils",
":tensorrt_lib",

14
tensorflow/compiler/tf2tensorrt/utils/py_utils.cc
View File

@ -26,24 +26,22 @@ namespace tensorrt {
bool IsGoogleTensorRTEnabled() {
#if GOOGLE_CUDA && GOOGLE_TENSORRT
#if TF_OSS_TENSORRT_STATIC
#if TF_USE_TENSORRT_STATIC
LOG(INFO) << "TensorRT libraries are statically linked, skip dlopen check";
return true;
#else
#else // TF_USE_TENSORRT_STATIC
auto handle_or = se::internal::DsoLoader::TryDlopenTensorRTLibraries();
if (!handle_or.ok()) {
LOG_WARNING_WITH_PREFIX
<< "Cannot dlopen some TensorRT libraries. If you would like "
"to use Nvidia GPU with TensorRT, please make sure the "
"missing libraries mentioned above are installed properly.";
return false;
} else {
return true;
}
#endif
#else
return handle_or.ok();
#endif // TF_USE_TENSORRT_STATIC
#else // GOOGLE_CUDA && GOOGLE_TENSORRT
return false;
#endif
#endif // GOOGLE_CUDA && GOOGLE_TENSORRT
}
} // namespace tensorrt