aboutsummaryrefslogtreecommitdiffhomepage
path: root/third_party/tensorrt
diff options
context:
space:
mode:
authorGravatar Mustafa Kasap <m@mustafakasap.com>2018-03-14 18:17:26 -0700
committerGravatar Frank Chen <frankchn@gmail.com>2018-03-14 18:17:26 -0700
commite605245fbd6ef370ac9a07dddfa733da4705017f (patch)
treeb6548acbb59ab95ee1066c2c9e11b58f81e01176 /third_party/tensorrt
parentf498991cbd36d6b745c8d0f4176079aab0e30373 (diff)
TensorRT support for ARM architectures (#17409)
i.e. without this modification Tensorflow under Nvidia Jetpack 3.2 is not compiling. Else condition in if-then-else block generates a path "%s/../include" which results in "/usr/lib/include" but it must be "/usr/include/aarch64-linux-gnu"... So this fix targets ARM architectures and with this fix, Tensorflow 1.6rc compiles fine with TensorRT support.
Diffstat (limited to 'third_party/tensorrt')
-rw-r--r--third_party/tensorrt/tensorrt_configure.bzl4
1 files changed, 4 insertions, 0 deletions
diff --git a/third_party/tensorrt/tensorrt_configure.bzl b/third_party/tensorrt/tensorrt_configure.bzl
index 8e76e5d02a..9b946505a6 100644
--- a/third_party/tensorrt/tensorrt_configure.bzl
+++ b/third_party/tensorrt/tensorrt_configure.bzl
@@ -57,6 +57,10 @@ def _find_trt_header_dir(repository_ctx, trt_install_path):
path = "/usr/include/x86_64-linux-gnu"
if _headers_exist(repository_ctx, path):
return path
+ if trt_install_path == "/usr/lib/aarch64-linux-gnu":
+ path = "/usr/include/aarch64-linux-gnu"
+ if _headers_exist(repository_ctx, path):
+ return path
path = str(repository_ctx.path("%s/../include" % trt_install_path).realpath)
if _headers_exist(repository_ctx, path):
return path