fix(//third_party/tensorrt): Fix TensorRT paths for local x86 builds

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Jul 16, 2020 · 73d804b · 73d804b
1 parent 94d6a0f
commit 73d804b
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 24 deletions.
diff --git a/WORKSPACE b/WORKSPACE
@@ -54,7 +54,7 @@ http_archive(
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
     urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.5.1.zip"],
-    sha256 = "cf0691493d05062fe3239cf76773bae4c5124f4b039050dbdd291c652af3ab2a"
+    sha256 = "0efdd4e709ab11088fa75f0501c19b0e294404231442bab1d1fb953924feb6b5"
 )
 
 http_archive(
@@ -71,15 +71,15 @@ http_archive(
 
 http_archive(
     name = "cudnn",
-    urls = ["https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.0.1.13/10.2_20200626/cudnn-10.2-linux-x64-v8.0.1.13.tgz"],
+    urls = ["https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.0.1.13/10.2_20200626/cudnn-10.2-linux-x64-v8.0.1.13.tgz",],
     build_file = "@//third_party/cudnn/archive:BUILD",
     sha256 = "0c106ec84f199a0fbcf1199010166986da732f9b0907768c9ac5ea5b120772db",
     strip_prefix = "cuda"
 )
 
 http_archive(
     name = "tensorrt",
-    urls = ["https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/7.1/tars/TensorRT-7.1.3.4.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz"],
+    urls = ["https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/7.1/tars/TensorRT-7.1.3.4.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz",],
     build_file = "@//third_party/tensorrt/archive:BUILD",
     sha256 = "9205bed204e2ae7aafd2e01cce0f21309e281e18d5bfd7172ef8541771539d41",
     strip_prefix = "TensorRT-7.1.3.4"

diff --git a/third_party/tensorrt/local/BUILD b/third_party/tensorrt/local/BUILD
@@ -23,20 +23,20 @@ cc_library(
             ],
         ),
         "//conditions:default": [
-            "include/NvUtils.h",
+            "include/x86_64-linux-gnu/NvUtils.h",
         ] + glob(
             [
-                "include/NvInfer*.h",
+                "include/x86_64-linux-gnu/NvInfer*.h",
             ],
             exclude = [
-                "include/NvInferPlugin.h",
-                "include/NvInferPluginUtils.h",
+                "include/x86_64-linux-gnu/NvInferPlugin.h",
+                "include/x86_64-linux-gnu/NvInferPluginUtils.h",
             ],
         ),
     }),
     includes = select({
         ":aarch64_linux": ["include/aarch64-linux-gnu"],
-        "//conditions:default": ["include/"],
+        "//conditions:default": ["include/x86_64-linux-gnu/"],
     }),
     visibility = ["//visibility:private"],
 )
@@ -76,14 +76,26 @@ cc_import(
 
 cc_library(
     name = "nvparsers_headers",
-    hdrs = [
-        "include/NvCaffeParser.h",
-        "include/NvOnnxParser.h",
-        "include/NvOnnxParserRuntime.h",
-        "include/NvOnnxConfig.h",
-        "include/NvUffParser.h"
-    ],
-    includes = ["include/"],
+    hdrs = select({
+        ":aarch64_linux": [
+            "include/aarch64-linux-gnu/NvCaffeParser.h",
+            "include/aarch64-linux-gnu/NvOnnxParser.h",
+            "include/aarch64-linux-gnu/NvOnnxParserRuntime.h",
+            "include/aarch64-linux-gnu/NvOnnxConfig.h",
+            "include/aarch64-linux-gnu/NvUffParser.h"
+        ],
+        "//conditions:default": [
+            "include/x86_64-linux-gnu/NvCaffeParser.h",
+            "include/x86_64-linux-gnu/NvOnnxParser.h",
+            "include/x86_64-linux-gnu/NvOnnxParserRuntime.h",
+            "include/x86_64-linux-gnu/NvOnnxConfig.h",
+            "include/x86_64-linux-gnu/NvUffParser.h"
+        ]
+    }),
+    includes = select({
+        ":aarch64_linux": ["include/aarch64-linux-gnu"],
+        "//conditions:default": ["include/x86_64-linux-gnu/"],
+    }),
     visibility = ["//visibility:private"],
 )
 
@@ -110,11 +122,18 @@ cc_import(
 
 cc_library(
     name = "nvonnxparser_headers",
-    hdrs = [
-        "include/NvOnnxParser.h",
-        "include/NvOnnxParserRuntime.h",
-        "include/NvOnnxConfig.h"
-    ],
+    hdrs = select({
+        ":aarch64_linux": [
+            "include/aarch64-linux-gnu/NvOnnxParser.h",
+            "include/aarch64-linux-gnu/NvOnnxParserRuntime.h",
+            "include/aarch64-linux-gnu/NvOnnxConfig.h"
+        ],
+        "//conditions:default": [
+            "include/x86_64-linux-gnu/NvOnnxParser.h",
+            "include/x86_64-linux-gnu/NvOnnxParserRuntime.h",
+            "include/x86_64-linux-gnu/NvOnnxConfig.h"
+        ]
+    }),
     includes = ["include/"],
     visibility = ["//visibility:private"],
 )
@@ -142,7 +161,14 @@ cc_import(
 
 cc_library(
     name = "nvonnxparser_runtime_header",
-    hdrs = ["include/NvOnnxParserRuntime.h"],
+    hdrs = select({
+        ":aarch64_linux": [
+            "include/aarch64-linux-gnu/NvOnnxParserRuntime.h",
+        ],
+        "//conditions:default": [
+            "include/x86_64-linux-gnu/NvOnnxParserRuntime.h",
+        ]
+    }),
     includes = ["include/"],
     visibility = ["//visibility:private"],
 )
@@ -170,7 +196,14 @@ cc_import(
 
 cc_library(
     name = "nvcaffeparser_headers",
-    hdrs = ["include/NvCaffeParser.h"],
+    hdrs = select({
+        ":aarch64_linux": [
+            "include/aarch64-linux-gnu/NvCaffeParser.h",
+        ],
+        "//conditions:default": [
+            "include/x86_64-linux-gnu/NvCaffeParser.h",
+        ]
+    }),
     includes = ["include/"],
     visibility = ["//visibility:private"],
 )
@@ -198,7 +231,10 @@ cc_import(
 
 cc_library(
     name = "nvinferplugin_headers",
-    hdrs = glob(["include/NvInferPlugin*.h"]),
+    hdrs = select({
+        ":aarch64_linux": glob(["include/aarch64-linux-gnu/NvInferPlugin*.h"]),
+        "//conditions:default": glob(["include/x86_64-linux-gnu/NvInferPlugin*.h"])
+    }),
     includes = ["include/"],
     visibility = ["//visibility:private"],
 )