Add support for TensorRT-RTX

.bazelrc

@@ -39,6 +39,8 @@ build:cxx11_abi --define=abi=cxx11_abi
 
 build:jetpack --//toolchains/dep_collection:compute_libs=jetpack
 
+build:rtx --//toolchains/dep_collection:compute_libs=rtx
+
 build:ci_testing --define=torchtrt_src=prebuilt --cxxopt="-DDISABLE_TEST_IN_CI" --action_env "NVIDIA_TF32_OVERRIDE=0"
 build:use_precompiled_torchtrt --define=torchtrt_src=prebuilt
 

.github/scripts/install-tensorrt-rtx.sh

@@ -0,0 +1,39 @@
+
+install_tensorrt_rtx() {
+    if [[ ${USE_RTX} == true ]]; then
+        install_wheel_or_not=${1:-false}
+        echo "It is the tensorrt-rtx build, install tensorrt-rtx with install_wheel_or_not:${install_wheel_or_not}"
+        PLATFORM=$(python -c "import sys; print(sys.platform)")
+        echo "PLATFORM: $PLATFORM"
+        # PYTHON_VERSION is always set in the CI environment, add this check for local testing
+        if [ -z "$PYTHON_VERSION" ]; then
+            echo "Error: PYTHON_VERSION environment variable is not set or empty. example format: export PYTHON_VERSION=3.11"
+            exit 1
+        fi
+
+        # python version is like 3.11, we need to convert it to cp311
+        CPYTHON_TAG="cp${PYTHON_VERSION//./}"
+        if [[ ${PLATFORM} == win32 ]]; then
+            curl -L https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.0/TensorRT-RTX-1.0.0.21.Windows.win10.cuda-12.9.zip -o TensorRT-RTX-1.0.0.21.Windows.win10.cuda-12.9.zip
+            unzip TensorRT-RTX-1.0.0.21.Windows.win10.cuda-12.9.zip
+            rtx_lib_dir=${PWD}/TensorRT-RTX-1.0.0.21/lib
+            export LD_LIBRARY_PATH=${rtx_lib_dir}:$LD_LIBRARY_PATH
+            echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH"
+            if [[ ${install_wheel_or_not} == true ]]; then
+                pip install TensorRT-RTX-1.0.0.21/python/tensorrt_rtx-1.0.0.21-${CPYTHON_TAG}-none-win_amd64.whl
+            fi
+        else
+            curl -L https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.0/TensorRT-RTX-1.0.0.21.Linux.x86_64-gnu.cuda-12.9.tar.gz -o TensorRT-RTX-1.0.0.21.Linux.x86_64-gnu.cuda-12.9.tar.gz
+            tar -xzf TensorRT-RTX-1.0.0.21.Linux.x86_64-gnu.cuda-12.9.tar.gz
+            rtx_lib_dir=${PWD}/TensorRT-RTX-1.0.0.21/lib
+            export LD_LIBRARY_PATH=${rtx_lib_dir}:$LD_LIBRARY_PATH
+            echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH"
+            if [[ ${install_wheel_or_not} == true ]]; then
+                pip install TensorRT-RTX-1.0.0.21/python/tensorrt_rtx-1.0.0.21-${CPYTHON_TAG}-none-linux_x86_64.whl
+            fi
+        fi
+    else
+        echo "It is the standard tensorrt build, skip install tensorrt-rtx"
+    fi
+
+}

.github/scripts/install-torch-tensorrt.sh

@@ -21,6 +21,12 @@ pip uninstall -y torch torchvision
 pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL}
 pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL}
 
+if [[ ${USE_RTX} == true ]]; then
+    source .github/scripts/install-tensorrt-rtx.sh
+    # tensorrt-rtx is not publicly available, so we need to install the wheel from the tar ball
+    install_wheel_or_not=true
+    install_tensorrt_rtx ${install_wheel_or_not}
+fi
 
 # Install Torch-TensorRT
 if [[ ${PLATFORM} == win32 ]]; then
@@ -29,4 +35,10 @@ else
     pip install /opt/torch-tensorrt-builds/torch_tensorrt*.whl
 fi
 
+if [[ ${USE_RTX} == true ]]; then
+    # currently tensorrt is installed automatically by install torch-tensorrt since it is a dependency of torch-tensorrt in pyproject.toml
+    # so we need to uninstall it to avoid conflict
+    pip uninstall -y tensorrt tensorrt_cu12 tensorrt_cu12_bindings tensorrt_cu12_libs
+fi
+
 echo -e "Running test script";

.github/workflows/build-test-linux-aarch64-jetpack.yml

@@ -66,8 +66,8 @@ jobs:
             post-script: packaging/post_build_script.sh
             smoke-test-script: packaging/smoke_test_script.sh
             package-name: torch_tensorrt
-    name: Build torch-tensorrt whl package
-    uses: ./.github/workflows/build_wheels_linux_aarch64.yml
+    name: Build torch-tensorrt whl package for aarch64-jetpack
+    uses: ./.github/workflows/build_wheels_linux.yml
     with:
       repository: ${{ matrix.repository }}
       ref: ""

.github/workflows/build-test-linux-aarch64.yml

@@ -62,8 +62,8 @@ jobs:
             post-script: packaging/post_build_script.sh
             smoke-test-script: packaging/smoke_test_script.sh
             package-name: torch_tensorrt
-    name: Build torch-tensorrt whl package
-    uses: ./.github/workflows/build_wheels_linux_aarch64.yml
+    name: Build torch-tensorrt whl package for aarch64
+    uses: ./.github/workflows/build_wheels_linux.yml
     with:
       repository: ${{ matrix.repository }}
       ref: ""

.github/workflows/build-test-linux-x86_64.yml

@@ -13,6 +13,7 @@ on:
       - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
   workflow_dispatch:
 
+
 jobs:
   generate-matrix:
     uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
@@ -60,8 +61,8 @@ jobs:
             post-script: packaging/post_build_script.sh
             smoke-test-script: packaging/smoke_test_script.sh
             package-name: torch_tensorrt
-    name: Build torch-tensorrt whl package
-    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    name: Build torch-tensorrt whl package for x86_64
+    uses: ./.github/workflows/build_wheels_linux.yml
     with:
       repository: ${{ matrix.repository }}
       ref: ""
@@ -74,6 +75,8 @@ jobs:
       package-name: ${{ matrix.package-name }}
       smoke-test-script: ${{ matrix.smoke-test-script }}
       trigger-event: ${{ github.event_name }}
+      architecture: "x86_64"
+      use-rtx: false
 
   tests-py-torchscript-fe:
     name: Test torchscript frontend [Python]
@@ -338,5 +341,5 @@ jobs:
         popd
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}-${{ inputs.job-name }}
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-tensorrt-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}-${{ inputs.job-name }}
   cancel-in-progress: true