Repository: pytorch/TensorRT

Issue:
❓ [Question] Issue for native compilation on NVIDIA Jetson (Jetpack 4.6)
### Question
I am trying to build latest TRTorch with Torch 1.9.0 and facing this issue.
I folllowed Instructions from this repo and https://nvidia.github.io/TRTorch/tutorials/installation.html 
Please Help!

Environment:
PyTorch Version: 1.9.0
CPU Architecture: aarch64 (Jetson Nano)
OS: Linux (Ubuntu 18.04)
How you installed PyTorch (conda, pip, libtorch, source): wheel for Jetson Nano torch-1.9.0a0+gitd69c22d-cp36-cp36m-linux_aarch64.whl
Build command you used (if compiling from source): bazel build //:libtrtorch --distdir third_party/dist_dir/aarch64-linux-gnu and cd py && sudo python3 setup.py install --use-cxx11-abi
Are you using local sources or building from archives: local
Python version: 3.6
CUDA version: 10.2
CUDNN version: 8.2.1
TensorRT: 8.0.1.6
Bazel: 4.0.0
Jetpack: 4.6


## Context:

Error:
```
nvidia@nvidia-desktop:~/development/TRTorch/py$ sudo python3 setup.py install --use-cxx11-abi
running install
using CXX11 ABI build
building libtrtorch
Extracting Bazel installation...
Starting local Bazel server and connecting to it...
INFO: Analyzed target //cpp/api/lib:libtrtorch.so (36 packages loaded, 2237 targets configured).
INFO: Found 1 target...
ERROR: /home/nvidia/.cache/bazel/_bazel_root/48274610a0e93cd5aa1f0ffb3ba9c3c7/external/cublas/BUILD.bazel:20:11: Middleman _middlemen/@cublas_S_S_Ccublas_Uheaders-BazelCppSemantics_build_arch_aarch64-opt failed: missing input file 'external/cublas/include/cublas.h', owner: '@cublas//:include/cublas.h'
Target //cpp/api/lib:libtrtorch.so failed to build
Use --verbose_failures to see the command lines of failed build steps.
ERROR: /home/nvidia/.cache/bazel/_bazel_root/48274610a0e93cd5aa1f0ffb3ba9c3c7/external/cublas/BUILD.bazel:20:11 Middleman _middlemen/@cublas_S_S_Ccublas_Uheaders-BazelCppSemantics_build_arch_aarch64-opt failed: 1 input file(s) do not exist
INFO: Elapsed time: 24.183s, Critical Path: 0.29s
INFO: 5 processes: 5 internal.
FAILED: Build did NOT complete successfully
```

My WORKSPACE file:
```
workspace(name = "TRTorch")

load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")

http_archive(
    name = "rules_python",
    sha256 = "778197e26c5fbeb07ac2a2c5ae405b30f6cb7ad1f5510ea6fdac03bded96cc6f",
    url = "https://github.com/bazelbuild/rules_python/releases/download/0.2.0/rules_python-0.2.0.tar.gz",
)

load("@rules_python//python:pip.bzl", "pip_install")

http_archive(
    name = "rules_pkg",
    sha256 = "038f1caa773a7e35b3663865ffb003169c6a71dc995e39bf4815792f385d837d",
    urls = [
        "https://mirror.bazel.build/github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
        "https://github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
    ],
)

load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")

rules_pkg_dependencies()

git_repository(
    name = "googletest",
    commit = "703bd9caab50b139428cea1aaff9974ebee5742e",
    remote = "https://github.com/google/googletest",
    shallow_since = "1570114335 -0400",
)

# CUDA should be installed on the system locally
new_local_repository(
    name = "cuda",
    build_file = "@//third_party/cuda:BUILD",
    path = "/usr/local/cuda-10.2/",
)

new_local_repository(
    name = "cublas",
    build_file = "@//third_party/cublas:BUILD",
    path = "/usr",
)
#############################################################################################################
# Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs)
#############################################################################################################

#http_archive(
#    name = "libtorch",
#    build_file = "@//third_party/libtorch:BUILD",
#    sha256 = "edc12091193ba772db77a6ec14e05cef6da881288fca0dfc89a031f631601f60",
#    strip_prefix = "libtorch",
#    urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-cxx11-abi-shared-with-deps-1.9.0%2Bcu111.zip"],
#)

#http_archive(
#    name = "libtorch_pre_cxx11_abi",
#    build_file = "@//third_party/libtorch:BUILD",
#    sha256 = "af9435fa4b44bb395c1a7645391c00228a72af4305f43a61e9300c0abdbe0819",
#    strip_prefix = "libtorch",
#    urls = ["https://download.pytorch.org/libtorch/cu111/libtorch-shared-with-deps-1.9.0%2Bcu111.zip"],
#)

# Download these tarballs manually from the NVIDIA website
# Either place them in the distdir directory in third_party and use the --distdir flag
# or modify the urls to "file:///<PATH TO TARBALL>/<TARBALL NAME>.tar.gz

#http_archive(
#    name = "cudnn",
#    build_file = "@//third_party/cudnn/archive:BUILD",
#    sha256 = "39412acd9ef5dd27954b6b9f5df75bd381c5d7ceb7979af6c743a7f4521f9c77",
#    strip_prefix = "cuda",
#    urls = [
#        "https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.2.1.32/11.3_06072021/cudnn-11.3-linux-x64-v8.2.1.32.tgz",
#    ],
#)

#http_archive(
#    name = "tensorrt",
#    build_file = "@//third_party/tensorrt/archive:BUILD",
#    sha256 = "def6a5ee50bed25a68a9c9e22ec671a8f29ee5414bde47c5767bd279e5596f88",
#    strip_prefix = "TensorRT-8.0.1.6",
#    urls = [
#        "https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/8.0.1/tars/tensorrt-8.0.1.6.linux.x86_64-gnu.cuda-11.3.cudnn8.2.tar.gz",
#    ],
#)

####################################################################################
# Locally installed dependencies (use in cases of custom dependencies or aarch64)
####################################################################################

# NOTE: In the case you are using just the pre-cxx11-abi path or just the cxx11 abi path
# with your local libtorch, just point deps at the same path to satisfy bazel.

# NOTE: NVIDIA's aarch64 PyTorch (python) wheel file uses the CXX11 ABI unlike PyTorch's standard
# x86_64 python distribution. If using NVIDIA's version just point to the root of the package
# for both versions here and do not use --config=pre-cxx11-abi

new_local_repository(
    name = "libtorch",
    path = "/usr/local/lib/python3.6/dist-packages/torch",
    build_file = "third_party/libtorch/BUILD"
)

new_local_repository(
    name = "libtorch_pre_cxx11_abi",
    path = "/usr/local/lib/python3.6/dist-packages/torch",
    build_file = "third_party/libtorch/BUILD"
)

new_local_repository(
    name = "cudnn",
    path = "/usr/",
    build_file = "@//third_party/cudnn/local:BUILD"
)

new_local_repository(
   name = "tensorrt",
   path = "/usr/",
   build_file = "@//third_party/tensorrt/local:BUILD"
)

#########################################################################
# Testing Dependencies (optional - comment out on aarch64)
#########################################################################
#pip_install(
#    name = "trtorch_py_deps",
#    requirements = "//py:requirements.txt",
#)

#pip_install(
#    name = "py_test_deps",
#    requirements = "//tests/py:requirements.txt",
#)

#pip_install(
#    name = "pylinter_deps",
#    requirements = "//tools/linter:requirements.txt",
#)
```


Discussion:
What is your jetpack version?
> What is your jetpack version?

Jetpack: 4.6
Edit: Tested it on Jetson TX2 flashed with Jetpack 4.6 and facing the same issue. 

Resolving patch:
diff --git a/py/setup.py b/py/setup.py
--- a/py/setup.py
+++ b/py/setup.py
@@ -13,6 +13,8 @@
 from shutil import copyfile, rmtree
 
 import subprocess
+import platform
+import warnings
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 
@@ -20,10 +22,26 @@
 
 CXX11_ABI = False
 
+JETPACK_VERSION = None
+
 if "--use-cxx11-abi" in sys.argv:
     sys.argv.remove("--use-cxx11-abi")
     CXX11_ABI = True
 
+if platform.uname().processor == "aarch64":
+    if "--jetpack-version" in sys.argv:
+        version_idx = sys.argv.index("--jetpack-version") + 1
+        version = sys.argv[version_idx]
+        sys.argv.remove(version)
+        sys.argv.remove("--jetpack-version")
+        if version == "4.5":
+            JETPACK_VERSION = "4.5"
+        elif version == "4.6":
+            JETPACK_VERSION = "4.6"
+    if not JETPACK_VERSION:
+        warnings.warn("Assuming jetpack version to be 4.6, if not use the --jetpack-version option")
+        JETPACK_VERSION = "4.6"
+
 
 def which(program):
     import os
@@ -66,6 +84,13 @@ def build_libtrtorch_pre_cxx11_abi(develop=True, use_dist_dir=True, cxx11_abi=Fa
     else:
         print("using CXX11 ABI build")
 
+    if JETPACK_VERSION == "4.5":
+        cmd.append("--platforms=//toolchains:jetpack_4.5")
+        print("Jetpack version: 4.5")
+    elif JETPACK_VERSION == "4.6":
+        cmd.append("--platforms=//toolchains:jetpack_4.6")
+        print("Jetpack version: 4.6")
+
     print("building libtrtorch")
     status_code = subprocess.run(cmd).returncode
 


Test patch:
diff --git a/tests/core/conversion/evaluators/test_aten_evaluators.cpp b/tests/core/conversion/evaluators/test_aten_evaluators.cpp
--- a/tests/core/conversion/evaluators/test_aten_evaluators.cpp
+++ b/tests/core/conversion/evaluators/test_aten_evaluators.cpp
@@ -3,6 +3,7 @@
 #include "gtest/gtest.h"
 #include "tests/util/util.h"
 #include "torch/csrc/jit/ir/irparser.h"
+#include "torch/torch.h"
 
 TEST(Evaluators, DivIntEvaluatesCorrectly) {
   const auto graph = R"IR(
diff --git a/tests/cpp/BUILD b/tests/cpp/BUILD
--- a/tests/cpp/BUILD
+++ b/tests/cpp/BUILD
@@ -77,6 +77,7 @@ cc_test(
     deps = [
         ":cpp_api_test",
     ],
+    timeout="long"
 )
 
 cc_test(
diff --git a/tests/modules/hub.py b/tests/modules/hub.py
--- a/tests/modules/hub.py
+++ b/tests/modules/hub.py
@@ -4,6 +4,8 @@
 import torchvision.models as models
 import timm
 
+torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
+
 models = {
     "alexnet": {
         "model": models.alexnet(pretrained=True),
diff --git a/tests/modules/requirements.txt b/tests/modules/requirements.txt
new file mode 100644
--- /dev/null
+++ b/tests/modules/requirements.txt
@@ -0,0 +1 @@
+timm==v0.4.12
diff --git a/tests/py/test_api_dla.py b/tests/py/test_api_dla.py
--- a/tests/py/test_api_dla.py
+++ b/tests/py/test_api_dla.py
@@ -37,7 +37,7 @@ def test_compile_traced(self):
                 "dla_core": 0,
                 "allow_gpu_fallback": True
             },
-            "enabled_precision": {torch.float, torch.half}
+            "enabled_precisions": {torch.half}
         }
 
         trt_mod = trtorch.compile(self.traced_model, compile_spec)
@@ -53,7 +53,7 @@ def test_compile_script(self):
                 "dla_core": 0,
                 "allow_gpu_fallback": True
             },
-            "enabled_precision": {torch.float, torch.half}
+            "enabled_precisions": {torch.half}
         }
 
         trt_mod = trtorch.compile(self.scripted_model, compile_spec)