diff --git a/README.md b/README.md index 3eefd76320..fd65ff86a1 100644 --- 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 index 555d599dbd..26a24dd0e5 100644 --- 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:////.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 index 6b53bf826f..3e0ed92dc9 100644 --- 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 index 5aeacef5c9..78e897f878 100644 --- 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 new file mode 100644 index 0000000000..e7d3126dbb --- /dev/null +++ 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 index ba4e0a2935..5f9216cf47 100644 --- 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 index 17bdfe6bbc..96e66506aa 100644 --- 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"], ) diff --git a/third_party/tensorrt/local/BUILD b/third_party/tensorrt/local/BUILD index 3c0e1c33a1..1a7fdd0931 100644 --- a/third_party/tensorrt/local/BUILD +++ b/third_party/tensorrt/local/BUILD @@ -1,20 +1,52 @@ package(default_visibility = ["//visibility:public"]) +config_setting( + name = "aarch64_linux", + constraint_values = [ + "@platforms//cpu:aarch64", + "@platforms//os:linux", + ], +) + cc_library( name = "nvinfer_headers", - hdrs = glob([ - "include/NvInfer*.h", - ], exclude=[ - "include/NvInferPlugin.h", - "include/NvInferPluginUtils.h" - ]), - includes = ["include/"], + hdrs = select({ + ":aarch64_linux": [ + "include/aarch64-linux-gnu/NvUtils.h", + ] + glob( + [ + "include/aarch64-linux-gnu/NvInfer*.h", + ], + exclude = [ + "include/aarch64-linux-gnu/NvInferPlugin.h", + "include/aarch64-linux-gnu/NvInferPluginUtils.h", + ], + ), + "//conditions:default": [ + "include/NvUtils.h", + ] + glob( + [ + "include/NvInfer*.h", + ], + exclude = [ + "include/NvInferPlugin.h", + "include/NvInferPluginUtils.h", + ], + ), + }), + includes = select({ + ":aarch64_linux": ["include/aarch64-linux-gnu"], + "//conditions:default": ["include/"], + }), visibility = ["//visibility:private"], ) cc_import( name = "nvinfer_lib", - shared_library = "lib/x86_64-linux-gnu/libnvinfer.so", + shared_library = select({ + ":aarch64_linux": "lib/aarch64-linux-gnu/libnvinfer.so", + "//conditions:default": "lib/x86_64-linux-gnu/libnvinfer.so", + }), visibility = ["//visibility:private"], ) @@ -34,7 +66,10 @@ cc_library( cc_import( name = "nvparsers_lib", - shared_library = "lib/x86_64-linux-gnu/libnvparsers.so", + shared_library = select({ + ":aarch64_linux": "lib/aarch64-linux-gnu/libnvparsers.so", + "//conditions:default": "lib/x86_64-linux-gnu/libnvparsers.so", + }), visibility = ["//visibility:private"], ) @@ -66,7 +101,10 @@ cc_library( cc_import( name = "nvonnxparser_lib", - shared_library = "lib/x86_64-linux-gnu/libnvonnxparser.so", + shared_library = select({ + ":aarch64_linux": "lib/aarch64-linux-gnu/libnvonnxparser.so", + "//conditions:default": "lib/x86_64-linux-gnu/libnvonnxparser.so", + }), visibility = ["//visibility:private"], ) @@ -95,7 +133,10 @@ cc_library( cc_import( name = "nvonnxparser_runtime_lib", - shared_library = "lib/x86_64-linux-gnu/libnvonnxparser_runtime.so", + shared_library = select({ + ":aarch64_linux": "lib/x86_64-linux-gnu/libnvonnxparser_runtime.so", + "//conditions:default": "lib/x86_64-linux-gnu/libnvonnxparser_runtime.so", + }), visibility = ["//visibility:public"], ) @@ -120,7 +161,10 @@ cc_library( cc_import( name = "nvcaffeparser_lib", - shared_library = "lib/x86_64-linux-gnu/libnvcaffe_parsers.so", + shared_library = select({ + ":aarch64_linux": "lib/aarch64-linux-gnu/libnvcaffe_parsers.so", + "//conditions:default": "lib/x86_64-linux-gnu/libnvcaffe_parsers.so", + }), visibility = ["//visibility:private"], ) @@ -145,7 +189,10 @@ cc_library( cc_import( name = "nvinferplugin_lib", - shared_library = "lib/x86_64-linux-gnu/libnvinfer_plugin.so", + shared_library = select({ + ":aarch64_linux": "lib/x86_64-linux-gnu/libnvinfer_plugin.so", + "//conditions:default": "lib/x86_64-linux-gnu/libnvinfer_plugin.so", + }), visibility = ["//visibility:private"], ) diff --git a/toolchains/BUILD b/toolchains/BUILD new file mode 100644 index 0000000000..1be750e2d1 --- /dev/null +++ b/toolchains/BUILD @@ -0,0 +1,9 @@ +package(default_visibility = ["//visibility:public"]) + +platform( + name = "aarch64_linux", + constraint_values = [ + "@platforms//os:linux", + "@platforms//cpu:aarch64", + ], +) \ No newline at end of file