For review only: Build torch_tensorrt for jetpack6.0, do not merge

.github/scripts/build_jetson_6.0.sh

@@ -0,0 +1,89 @@
+#!/usr/bin/env bash
+
+set -euxo pipefail
+
+# get jetpack version: eg: Version: 6.0+b106 ---> 6.0
+jetpack_version=$(apt show nvidia-jetpack 2>/dev/null | grep Version: | cut -d ' ' -f 2 | cut -d '+' -f 1)
+python_version=$(python --version)
+cuda_version=$(nvcc --version | grep Cuda | grep release | cut -d ',' -f 2 | sed -e 's/ release //g')
+echo "Current jetpack_version: ${jetpack_version} cuda_version: ${cuda_version} python_version: ${python_version} "
+
+export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/lib/aarch64-linux-gnu:/usr/include/aarch64-linux-gnu:/usr/local/cuda-${cuda_version}/lib64
+
+# make sure nvidia-jetpack dev package is installed:
+# go to /usr/include/aarch64-linux-gnu/ if you can see NvInfer.h(tensorrt related header files) which means dev package is installed
+# if not installed, install via the below cmd:
+# sudo apt update
+# sudo apt install nvidia-jetpack
+
+# make sure cuda is installed:
+# nvcc --version or go to /usr/local/cuda/bin to see whether it is installed
+# the install nvidia-jetpack dev package step will automatically install the cuda tool
+# if not installed, install via the below cmd:
+# sudo apt update
+# sudo apt install cuda-toolkit-12-2
+
+# make sure bazel is installed via the below cmd:
+# wget -v https://github.com/bazelbuild/bazelisk/releases/download/v1.20.0/bazelisk-linux-arm64
+# sudo mv bazelisk-linux-arm64 /usr/bin/bazel
+# chmod +x /usr/bin/bazel
+
+# make sure setuptools is installed
+# sudo apt install python3-pip
+# make sure setuptools is upgraded via the below cmd:
+# pip install -U pip setuptools
+
+# make sure torch is installed via the below cmd:
+# wget https://developer.download.nvidia.cn/compute/redist/jp/v60/pytorch/torch-2.4.0a0+3bcc3cddb5.nv24.07.16234504-cp310-cp310-linux_aarch64.whl
+# python -m pip install torch-2.4.0a0+3bcc3cddb5.nv24.07.16234504-cp310-cp310-linux_aarch64.whl
+
+# make sure libcusparseLt.so exists if not download and copy via the below cmd:
+# wget https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-sbsa/libcusparse_lt-linux-sbsa-0.5.2.1-archive.tar.xz
+# tar xf libcusparse_lt-linux-sbsa-0.5.2.1-archive.tar.xz
+# sudo cp -a libcusparse_lt-linux-sbsa-0.5.2.1-archive/include/* /usr/local/cuda/include/
+# sudo cp -a libcusparse_lt-linux-sbsa-0.5.2.1-archive/lib/* /usr/local/cuda/lib64/
+
+# make sure tensorrt is upgraded from 8.6.2 to tensorrt10.1.0
+# wget http://cuda-repo/release-candidates/Libraries/TensorRT/v10.1/10.1.0.16-e64cb73a/12.4-r550/l4t-aarch64/tar/TensorRT-10.1.0.16.Ubuntu-22.04.aarch64-gnu.cuda-12.4.tar.gz
+# gunzip TensorRT-10.1.0.16.Ubuntu-22.04.aarch64-gnu.cuda-12.4.tar.gz
+# tar xvf TensorRT-10.1.0.16.Ubuntu-22.04.aarch64-gnu.cuda-12.4.tar
+
+# copy tensorrt 10.1.0 header files to /usr/include/aarch64-linux-gnu/:
+# cd ~/Desktop/lan/Downloads/TensorRT-10.1.0.16/include
+# sudo cp * /usr/include/aarch64-linux-gnu/
+
+# copy tensorrt 10.1.0 .so files to /usr/lib/aarch64-linux-gnu/:
+# cd ~/Desktop/lan/Downloads/TensorRT-10.1.0.16/lib
+# sudo cp libnvinfer.so.10.1.0 /usr/lib/aarch64-linux-gnu/
+# sudo cp libnvinfer_plugin.so.10.1.0 /usr/lib/aarch64-linux-gnu/
+# sudo cp libnvinfer_vc_plugin.so.10.1.0 /usr/lib/aarch64-linux-gnu/
+# sudo cp libnvonnxparser.so.10.1.0 /usr/lib/aarch64-linux-gnu/
+# sudo cp libnvinfer_lean.so.10.1.0  /usr/lib/aarch64-linux-gnu/
+# sudo cp libnvinfer_dispatch.so.10.1.0 /usr/lib/aarch64-linux-gnu/
+# sudo cp libnvinfer_builder_resource.so.10.1.0 /usr/lib/aarch64-linux-gnu/
+
+# copy tensorrt 10.1.0 *_static.a and stubs to /usr/lib/aarch64-linux-gnu/:
+# cd ~/Desktop/lan/Downloads/TensorRT-10.1.0.16/lib
+# sudo cp stub/* /usr/lib/aarch64-linux-gnu/stub/
+# sudo cp *_static.a /usr/lib/aarch64-linux-gnu/
+# sudo cp libonnx_proto.a /usr/lib/aarch64-linux-gnu/
+
+# create symbolic link under /usr/lib/aarch64-linux-gnu/:
+# cd /usr/lib/aarch64-linux-gnu/
+# sudo ln -s libnvinfer.so.10.1.0 libnvinfer.so
+# sudo ln -s libnvinfer_plugin.so.10.1.0 libnvinfer_plugin.so
+# sudo ln -s libnvinfer_vc_plugin.so.10.1.0 libnvinfer_vc_plugin.so
+# sudo ln -s libnvonnxparser.so.10.1.0 libnvonnxparser.so
+# sudo ln -s libnvinfer_lean.so.10.1.0 libnvinfer_lean.so
+# sudo ln -s libnvinfer_dispatch.so.10.1.0 libnvinfer_dispatch.so
+
+
+export TORCH_INSTALL_PATH=$(python -c "import torch, os; print(os.path.dirname(torch.__file__))")
+export SITE_PACKAGE_PATH=${TORCH_INSTALL_PATH::-6}
+# replace the WORKSPACE file with jetpack one
+cat WORKSPACE > WORKSPACE.orig
+cat toolchains/jp_workspaces/WORKSPACE.jp60 | envsubst > WORKSPACE
+
+# build on jetpack
+python setup.py  --use-cxx11-abi  install --user
+

setup.py

@@ -152,12 +152,14 @@ def load_dep_info():
             JETPACK_VERSION = "4.6"
         elif version == "5.0":
             JETPACK_VERSION = "5.0"
+        elif version == "6.0":
+            JETPACK_VERSION = "6.0"
 
     if not JETPACK_VERSION:
         warnings.warn(
-            "Assuming jetpack version to be 5.0, if not use the --jetpack-version option"
+            "Assuming jetpack version to be 6.0, if not use the --jetpack-version option"
         )
-        JETPACK_VERSION = "5.0"
+        JETPACK_VERSION = "6.0"
 
     if not CXX11_ABI:
         warnings.warn(
@@ -209,12 +211,15 @@ def build_libtorchtrt_pre_cxx11_abi(
     elif JETPACK_VERSION == "5.0":
         cmd.append("--platforms=//toolchains:jetpack_5.0")
         print("Jetpack version: 5.0")
+    elif JETPACK_VERSION == "6.0":
+        cmd.append("--platforms=//toolchains:jetpack_6.0")
+        print("Jetpack version: 6.0")
 
     if CI_BUILD:
         cmd.append("--platforms=//toolchains:ci_rhel_x86_64_linux")
         print("CI based build")
 
-    print("building libtorchtrt")
+    print(f"building libtorchtrt {cmd=}")
     status_code = subprocess.run(cmd).returncode
 
     if status_code != 0:

toolchains/BUILD

@@ -35,6 +35,15 @@ platform(
     ],
 )
 
+platform(
+    name = "jetpack_6.0",
+    constraint_values = [
+        "@platforms//os:linux",
+        "@platforms//cpu:aarch64",
+        "@//toolchains/jetpack:4.6",
+    ],
+)
+
 platform(
     name = "ci_rhel_x86_64_linux",
     constraint_values = [

toolchains/jp_workspaces/WORKSPACE.jp60

@@ -0,0 +1,92 @@
+workspace(name = "Torch-TensorRT")
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+http_archive(
+    name = "rules_python",
+    sha256 = "863ba0fa944319f7e3d695711427d9ad80ba92c6edd0b7c7443b84e904689539",
+    strip_prefix = "rules_python-0.22.0",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.22.0/rules_python-0.22.0.tar.gz",
+)
+
+load("@rules_python//python:repositories.bzl", "py_repositories")
+
+py_repositories()
+
+http_archive(
+    name = "rules_pkg",
+    sha256 = "8f9ee2dc10c1ae514ee599a8b42ed99fa262b757058f65ad3c384289ff70c4b8",
+    urls = [
+        "https://mirror.bazel.build/github.com/bazelbuild/rules_pkg/releases/download/0.9.1/rules_pkg-0.9.1.tar.gz",
+        "https://github.com/bazelbuild/rules_pkg/releases/download/0.9.1/rules_pkg-0.9.1.tar.gz",
+    ],
+)
+
+load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
+
+rules_pkg_dependencies()
+
+http_archive(
+    name = "googletest",
+    sha256 = "755f9a39bc7205f5a0c428e920ddad092c33c8a1b46997def3f1d4a82aded6e1",
+    strip_prefix = "googletest-5ab508a01f9eb089207ee87fd547d290da39d015",
+    urls = ["https://github.com/google/googletest/archive/5ab508a01f9eb089207ee87fd547d290da39d015.zip"],
+)
+
+# External dependency for torch_tensorrt if you already have precompiled binaries.
+local_repository(
+    name = "torch_tensorrt",
+    path = "${SITE_PACKAGE_PATH}/torch_tensorrt",
+)
+
+# CUDA should be installed on the system locally
+new_local_repository(
+    name = "cuda",
+    build_file = "@//third_party/cuda:BUILD",
+    path = "/usr/local/cuda-12.2/",
+)
+
+####################################################################################
+# 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 = "${TORCH_INSTALL_PATH}",
+    build_file = "third_party/libtorch/BUILD"
+)
+
+# NOTE: Unused on aarch64-jetson with NVIDIA provided PyTorch distribu†ion
+new_local_repository(
+    name = "libtorch_pre_cxx11_abi",
+    path = "${TORCH_INSTALL_PATH}",
+    build_file = "third_party/libtorch/BUILD"
+)
+
+new_local_repository(
+   name = "tensorrt",
+   path = "/usr/",
+   build_file = "@//third_party/tensorrt/local:BUILD"
+)
+
+#########################################################################
+# Development Dependencies (optional - comment out on aarch64)
+#########################################################################
+
+load("@rules_python//python:pip.bzl", "pip_parse")
+
+pip_parse(
+    name = "devtools_deps",
+    requirements = "//:requirements-dev.txt",
+)
+
+load("@devtools_deps//:requirements.bzl", "install_deps")
+
+install_deps()