diff --git a/WORKSPACE b/WORKSPACE index 094d3e0c99..9218ca3581 100644 --- 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,7 +71,7 @@ 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" @@ -79,7 +79,7 @@ http_archive( 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 index 1a7fdd0931..4107dff63d 100644 --- 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"], )