pytorch · andi4191 · Jul 8, 2020 · Jul 8, 2020 · Jul 7, 2020 · Jul 7, 2020
diff --git a/README.md b/README.md
@@ -63,11 +63,13 @@ torch.jit.save(trt_ts_module, "trt_torchscript_module.ts")
 | Platform | Support |
 | -------- | ------- |
 | Linux AMD64 / GPU   | **Supported** |
-| Linux aarch64 / GPU | **Planned/Possible with Native Compiation but untested** |
-| Linux aarch64 / DLA | **Planned/Possible with Native Compilation but untested** |
+| Linux aarch64 / GPU | **Native Compilation Supported on JetPack-4.4** |
+| Linux aarch64 / DLA | **Native Compilation Supported on JetPack-4.4 but untested** |
 | Windows / GPU       | - |
 | Linux ppc64le / GPU | - |
 
+> Note: Refer NVIDIA NGC container(https://ngc.nvidia.com/catalog/containers/nvidia:l4t-pytorch) for PyTorch libraries on JetPack.
+
 ### Dependencies
 
 - Bazel 3.2.0
@@ -171,6 +173,11 @@ bazel build //:libtrtorch --compilation_mode opt
 bazel build //:libtrtorch --compilation_mode=dbg
 ```
 
+### Native compilation on NVIDIA Jetson AGX
+``` shell
+bazel build //:libtrtorch --distdir third_party/distdir/aarch64-linux-gnu
+```
+
 A tarball with the include files and library can then be found in bazel-bin
 
 ### Running TRTorch on a JIT Graph

diff --git a/WORKSPACE b/WORKSPACE
@@ -25,21 +25,30 @@ http_archive(
 load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
 rules_pkg_dependencies()
 
+git_repository(
+    name = "googletest",
+    remote = "https://github.com/google/googletest",
+    commit = "703bd9caab50b139428cea1aaff9974ebee5742e",
+    shallow_since = "1570114335 -0400"
+)
+
 # CUDA should be installed on the system locally
 new_local_repository(
     name = "cuda",
-    path = "/usr/local/cuda-10.2/targets/x86_64-linux/",
+    path = "/usr/local/cuda-10.2/",
     build_file = "@//third_party/cuda:BUILD",
 )
 
-http_archive(
-    name = "libtorch_pre_cxx11_abi",
-    build_file = "@//third_party/libtorch:BUILD",
-    strip_prefix = "libtorch",
-    sha256 = "ea8de17c5f70015583f3a7a43c7a5cdf91a1d4bd19a6a7bc11f074ef6cd69e27",
-    urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-shared-with-deps-1.5.0.zip"],
+new_local_repository(
+    name = "cublas",
+    path = "/usr",
+    build_file = "@//third_party/cublas:BUILD",
 )
 
+#############################################################################################################
+# 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",
@@ -48,23 +57,18 @@ http_archive(
     sha256 = "0efdd4e709ab11088fa75f0501c19b0e294404231442bab1d1fb953924feb6b5"
 )
 
-pip3_import(
-    name = "trtorch_py_deps",
-    requirements = "//py:requirements.txt"
-)
-
-load("@trtorch_py_deps//:requirements.bzl", "pip_install")
-pip_install()
-
-pip3_import(
-    name = "py_test_deps",
-    requirements = "//tests/py:requirements.txt"
+http_archive(
+    name = "libtorch_pre_cxx11_abi",
+    build_file = "@//third_party/libtorch:BUILD",
+    strip_prefix = "libtorch",
+    sha256 = "ea8de17c5f70015583f3a7a43c7a5cdf91a1d4bd19a6a7bc11f074ef6cd69e27",
+    urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-shared-with-deps-1.5.0.zip"],
 )
 
-load("@py_test_deps//:requirements.bzl", "pip_install")
-pip_install()
+# 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
 
-# Downloaded distributions to use with --distdir
 http_archive(
     name = "cudnn",
     urls = ["https://developer.nvidia.com/compute/machine-learning/cudnn/secure/7.6.5.32/Production/10.2_20191118/cudnn-10.2-linux-x64-v7.6.5.32.tgz"],
@@ -81,22 +85,57 @@ http_archive(
     strip_prefix = "TensorRT-7.0.0.11"
 )
 
-## Locally installed dependencies
-# new_local_repository(
+####################################################################################
+# 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(
+#new_local_repository(
 #   name = "tensorrt",
 #   path = "/usr/",
 #   build_file = "@//third_party/tensorrt/local:BUILD"
 #)
 
-git_repository(
-    name = "googletest",
-    remote = "https://github.com/google/googletest",
-    commit = "703bd9caab50b139428cea1aaff9974ebee5742e",
-    shallow_since = "1570114335 -0400"
+#########################################################################
+# Testing Dependencies (optional - comment out on aarch64)
+#########################################################################
+pip3_import(
+    name = "trtorch_py_deps",
+    requirements = "//py:requirements.txt"
 )
+
+load("@trtorch_py_deps//:requirements.bzl", "pip_install")
+pip_install()
+
+pip3_import(
+    name = "py_test_deps",
+    requirements = "//tests/py:requirements.txt"
+)
+
+load("@py_test_deps//:requirements.bzl", "pip_install")
+pip_install()
+
diff --git a/tests/core/converters/BUILD b/tests/core/converters/BUILD
@@ -86,5 +86,3 @@ test_suite(
     ":test_stack"
   ]
 )
-
-
diff --git a/tests/util/BUILD b/tests/util/BUILD
@@ -23,6 +23,7 @@ cc_library(
         "//core/conversion",
         "//core/util:prelude",
         "//cpp/api:trtorch",
+	"@tensorrt//:nvinfer"
     ] + select({
         ":use_pre_cxx11_abi":  [
             "@libtorch_pre_cxx11_abi//:libtorch",

diff --git a/third_party/cublas/BUILD b/third_party/cublas/BUILD
@@ -0,0 +1,34 @@
+package(default_visibility = ["//visibility:public"])
+
+config_setting(
+    name = "aarch64_linux",
+    constraint_values = [
+        "@platforms//cpu:aarch64",
+        "@platforms//os:linux",
+    ],
+)
+
+cc_library(
+    name = "cublas_headers",
+    hdrs = ["include/cublas.h"] + glob(["include/cublas+.h"]),
+    includes = ["include/"],
+    visibility = ["//visibility:private"],
+)
+
+cc_import(
+    name = "cublas_lib",
+    shared_library = select({
+        ":aarch64_linux": "lib/aarch64-linux-gnu/libcublas.so.10",
+        "//conditions:default": "lib/x86_64-linux-gnu/libcublas.so.10",
+    }),
+    visibility = ["//visibility:private"],
+)
+
+cc_library(
+    name = "cublas",
+    visibility = ["//visibility:public"],
+    deps = [
+        "cublas_headers",
+        "cublas_lib",
+    ],
+)
diff --git a/third_party/cuda/BUILD b/third_party/cuda/BUILD
@@ -1,10 +1,23 @@
 package(default_visibility = ["//visibility:public"])
 
+config_setting(
+    name = "aarch64_linux",
+    constraint_values = [
+        "@platforms//cpu:aarch64",
+        "@platforms//os:linux",
+    ],
+)
+
 cc_library(
     name = "cudart",
-    srcs = glob([
-        "lib/**/libcudart.so",
-    ]),
+    srcs = select({
+        ":aarch64_linux": [
+            "targets/aarch64-linux/lib/libcudart.so",
+        ],
+        "//conditions:default": [
+            "targets/x86_64-linux/lib/libcudart.so",
+        ],
+    }),
     hdrs = glob([
         "include/**/*.h",
         "include/**/*.hpp",
@@ -15,16 +28,26 @@ cc_library(
 
 cc_library(
     name = "nvToolsExt",
-    srcs = glob([
-        "lib/**/libnvToolsExt.so.1"
-    ])
+    srcs = select({
+        ":aarch64_linux": [
+            "targets/aarch64-linux/lib/libnvToolsExt.so.1",
+        ],
+        "//conditions:default": [
+            "targets/x86_64-linux/lib/libnvToolsExt.so.1",
+        ],
+    }),
 )
 
 cc_library(
     name = "cuda",
-    srcs = glob([
-        "lib/**/*libcuda.so",
-    ]),
+    srcs = select({
+        ":aarch64_linux": glob([
+            "targets/aarch64-linux/lib/**/lib*.so",
+        ]),
+        "//conditions:default": glob([
+            "targets/x86_64-linux/lib/**/lib*.so",
+        ]),
+    }),
     hdrs = glob([
         "include/**/*.h",
         "include/**/*.hpp",

diff --git a/third_party/cudnn/local/BUILD b/third_party/cudnn/local/BUILD
@@ -1,5 +1,13 @@
 package(default_visibility = ["//visibility:public"])
 
+config_setting(
+    name = "aarch64_linux",
+    constraint_values = [
+        "@platforms//cpu:aarch64",
+        "@platforms//os:linux",
+    ],
+)
+
 cc_library(
     name = "cudnn_headers",
     hdrs = ["include/cudnn.h"] + glob(["include/cudnn+.h"]),
@@ -9,7 +17,10 @@ cc_library(
 
 cc_import(
     name = "cudnn_lib",
-    shared_library = "lib/x86_64-linux-gnu/libcudnn.so.7.6.5",
+    shared_library = select({
+        ":aarch64_linux": "lib/aarch64-linux-gnu/libcudnn.so",
+        "//conditions:default": "lib/x86_64-linux-gnu/libcudnn.so.7.6.5",
+    }),
     visibility = ["//visibility:private"],
 )