(//bazel): Native compilation support for Jetson AGX platform

.bazelrc

@@ -32,4 +32,4 @@ build:python --define=abi=pre_cxx11_abi
 
 build:pre_cxx11_abi --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0"
 build:pre_cxx11_abi --linkopt="-D_GLIBCXX_USE_CXX11_ABI=0"
-build:pre_cxx11_abi --define=abi=pre_cxx11_abi
+build:pre_cxx11_abi --define=abi=pre_cxx11_abi

BUILD

@@ -8,57 +8,56 @@ pkg_tar(
         "//core/conversion:include",
         "//core/conversion/conversionctx:include",
         "//core/conversion/converters:include",
-        "//core/conversion/var:include",
-        "//core/conversion/tensorcontainer:include",
+        "//core/conversion/converters/impl/plugins:include",
         "//core/conversion/evaluators:include",
+        "//core/conversion/tensorcontainer:include",
+        "//core/conversion/var:include",
         "//core/conversion/converters/impl/plugins:include",
         "//core/execution:include",
         "//core/lowering:include",
         "//core/lowering/passes:include",
         "//core/util:include",
-        "//core/util/logging:include"
+        "//core/util/logging:include",
     ],
 )
 
 pkg_tar(
     name = "include",
-    package_dir = "include/trtorch/",
     srcs = [
         "//cpp/api:api_headers",
     ],
+    package_dir = "include/trtorch/",
 )
 
 pkg_tar(
     name = "lib",
-    package_dir = "lib/",
     srcs = [
         "//cpp/api/lib:libtrtorch.so",
     ],
     mode = "0755",
+    package_dir = "lib/",
 )
 
-
 pkg_tar(
     name = "bin",
-    package_dir = "bin/",
     srcs = [
-        "//cpp/trtorchc:trtorchc",
+        "//cpp/trtorchc",
     ],
     mode = "0755",
+    package_dir = "bin/",
 )
 
-
 pkg_tar(
     name = "libtrtorch",
-    extension = "tar.gz",
-    package_dir = "trtorch",
     srcs = [
-        "//:LICENSE"
+        "//:LICENSE",
     ],
+    extension = "tar.gz",
+    package_dir = "trtorch",
     deps = [
-        ":lib",
         ":bin",
         ":include",
         ":include_core",
+        ":lib",
     ],
 )

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 --cpu=aarch64 --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

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,61 @@ 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",
+#    # Use below for aarch64
+#    #path = "/home/nvidia/.local/lib/python3.6/site-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",
+#    # Use below for aarch64
+#    #path = "/home/nvidia/.local/lib/python3.6/site-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()

core/BUILD

@@ -4,35 +4,34 @@ config_setting(
     name = "use_pre_cxx11_abi",
     values = {
         "define": "abi=pre_cxx11_abi",
-    }
+    },
 )
 
 cc_library(
     name = "core",
-    hdrs = [
-        "compiler.h",
-    ],
     srcs = [
         "compiler.cpp",
     ],
+    hdrs = [
+        "compiler.h",
+    ],
     deps = [
         "//core/conversion",
         "//core/execution",
         "//core/lowering",
         "//core/util/logging",
-        "@tensorrt//:nvinfer"
+        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default":  ["@libtorch//:libtorch"],
+        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@libtorch//:libtorch"],
     }),
-    alwayslink=True,
+    alwayslink = True,
 )
 
-
 load("@rules_pkg//:pkg.bzl", "pkg_tar")
 
 pkg_tar(
     name = "include",
-    package_dir = "core/",
     srcs = ["compiler.h"],
+    package_dir = "core/",
 )

core/conversion/BUILD

@@ -4,36 +4,36 @@ config_setting(
     name = "use_pre_cxx11_abi",
     values = {
         "define": "abi=pre_cxx11_abi",
-    }
+    },
 )
 
 cc_library(
     name = "conversion",
-    hdrs = [
-        "conversion.h",
-    ],
     srcs = [
+        "InterfaceTypes.cpp",
         "conversion.cpp",
         "conversion_blacklist.cpp",
-        "InterfaceTypes.cpp"
+    ],
+    hdrs = [
+        "conversion.h",
     ],
     deps = [
-        "@tensorrt//:nvinfer",
         "//core/conversion/var",
         "//core/conversion/conversionctx",
         "//core/conversion/converters",
         "//core/conversion/evaluators",
-        "//core/util:prelude"
+        "//core/util:prelude",
+        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default":  ["@libtorch//:libtorch"],
+        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@libtorch//:libtorch"],
     }),
 )
 
 load("@rules_pkg//:pkg.bzl", "pkg_tar")
 
 pkg_tar(
     name = "include",
-    package_dir = "core/conversion/",
     srcs = ["conversion.h"],
+    package_dir = "core/conversion/",
 )

core/conversion/conversionctx/BUILD

@@ -4,30 +4,30 @@ config_setting(
     name = "use_pre_cxx11_abi",
     values = {
         "define": "abi=pre_cxx11_abi",
-    }
+    },
 )
 
 cc_library(
     name = "conversionctx",
-    hdrs = [
-        "ConversionCtx.h",
-    ],
     srcs = [
         "ConversionCtx.cpp",
     ],
+    hdrs = [
+        "ConversionCtx.h",
+    ],
     deps = [
-        "@tensorrt//:nvinfer",
         "//core/util:prelude",
+        "@tensorrt//:nvinfer",
     ] + select({
-        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
-        "//conditions:default":  ["@libtorch//:libtorch"],
+        ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default": ["@libtorch//:libtorch"],
     }),
 )
 
 load("@rules_pkg//:pkg.bzl", "pkg_tar")
 
 pkg_tar(
     name = "include",
-    package_dir = "core/conversion/conversionctx/",
     srcs = ["ConversionCtx.h"],
+    package_dir = "core/conversion/conversionctx/",
 )