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( alias(
name = "tensorrt_lib", name = "tensorrt_lib",
actual = select({ actual = select({
"//tensorflow:oss": ":tensorrt_lib_oss", "@local_config_tensorrt//:use_static_tensorrt": "@local_config_tensorrt//:tensorrt",
"//conditions:default": "@local_config_tensorrt//:tensorrt", "//conditions:default": ":tensorrt_stub",
}), }),
visibility = ["//visibility:private"], visibility = ["//visibility:private"],
) )
@ -607,21 +591,15 @@ tf_proto_library(
protodeps = tf_additional_all_protos(), 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( cc_library(
name = "py_utils", name = "py_utils",
srcs = ["utils/py_utils.cc"], srcs = ["utils/py_utils.cc"],
hdrs = ["utils/py_utils.h"], hdrs = ["utils/py_utils.h"],
copts = tf_copts(), 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([ deps = if_tensorrt([
":common_utils", ":common_utils",
":tensorrt_lib", ":tensorrt_lib",

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

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