tensorrt_rtx try on

.bazelrc

@@ -38,6 +38,9 @@ build:cxx11_abi --define=abi=cxx11_abi
 
 build:jetpack --//toolchains/dep_collection:compute_libs=jetpack
 
+build:rtx_x86_64 --cxxopt="-std=c++17" --cxxopt="-fdiagnostics-color=always"
+build:rtx_win --cxxopt="/GS-" --cxxopt="/std:c++17" --cxxopt="/permissive-" --cxxopt="/wd4244" --cxxopt="/wd4267" --cxxopt="/wd4819" --features=windows_export_all_symbols
+
 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-torch-tensorrt.sh

@@ -21,6 +21,24 @@ 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}
 
+# tensorrt-rtx is not publicly available, so we need to install it from the local path
+if [[ ${USE_RTX} == true ]]; then
+    echo "It is the tensorrt-rtx build, install tensorrt-rtx"
+    # python version is like 3.11, we need to convert it to cp311
+    CPYTHON_TAG="cp${PYTHON_VERSION//./}"
+    if [[ ${PLATFORM} == win32 ]]; then
+        curl -L http://cuda-repo/release-candidates/Libraries/TensorRT/v10.12/10.12.0.35-51f47a12/12.9-r575/Windows10-x64-winjit/zip/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
+        pip install TensorRT-RTX-1.0.0.21/python/tensorrt_rtx-1.0.0.21-${CPYTHON_TAG}-none-win_amd64.whl
+    else
+        curl -L http://cuda-repo/release-candidates/Libraries/TensorRT/v10.12/10.12.0.35-51f47a12/12.9-r575/Linux-x64-manylinux_2_28-winjit/tar/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
+        pip install TensorRT-RTX-1.0.0.21/python/tensorrt_rtx-1.0.0.21-${CPYTHON_TAG}-none-linux_x86_64.whl
+    fi
+else
+    echo "It is the standard tensorrt build"
+fi
+
 
 # Install Torch-TensorRT
 if [[ ${PLATFORM} == win32 ]]; then

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

@@ -66,7 +66,7 @@ jobs:
             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
+    uses: ./.github/workflows/build_wheels_linux.yml
     with:
       repository: ${{ matrix.repository }}
       ref: ""

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

@@ -63,7 +63,7 @@ jobs:
             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
+    uses: ./.github/workflows/build_wheels_linux.yml
     with:
       repository: ${{ matrix.repository }}
       ref: ""

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

@@ -1,6 +1,13 @@
 name: Build and test Linux x86_64 wheels
 
 on:
+  workflow_call:
+    inputs:
+      use-rtx:
+        description: "Use RTX TensorRT"
+        default: false
+        type: boolean
+        required: false
   pull_request:
   push:
     branches:
@@ -61,7 +68,7 @@ jobs:
             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
+    uses: ./.github/workflows/build_wheels_linux.yml
     with:
       repository: ${{ matrix.repository }}
       ref: ""
@@ -74,6 +81,8 @@ jobs:
       package-name: ${{ matrix.package-name }}
       smoke-test-script: ${{ matrix.smoke-test-script }}
       trigger-event: ${{ github.event_name }}
+      architecture: "x86_64"
+      use-rtx: ${{ inputs.use-rtx }}
 
   tests-py-torchscript-fe:
     name: Test torchscript frontend [Python]
@@ -338,5 +347,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 }}-${{ inputs.use-rtx }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}-${{ inputs.job-name }}
   cancel-in-progress: true

.github/workflows/build_wheels_linux_aarch64.yml → .github/workflows/build_wheels_linux.yml

@@ -1,4 +1,4 @@
-name: Build Linux Wheels For aarch64
+name: Build Linux Wheels
 
 on:
   workflow_call:
@@ -88,11 +88,6 @@ on:
         required: false
         default: "python -m build --wheel"
         type: string
-      is-jetpack:
-        description: Set to true if the build is for jetpack
-        required: false
-        default: false
-        type: boolean
       pip-install-torch-extra-args:
         # NOTE: Why does this exist?
         # Well setuptools / python packaging doesn't actually allow you to specify dependencies
@@ -110,11 +105,22 @@ on:
         description: 'Timeout for the job (in minutes)'
         default: 120
         type: number
+      use-rtx:
+        description: "Set to true if use TensorRT-RTX"
+        default: false
+        type: boolean
+        required: false
+      is-jetpack:
+        description: Set to true if the build is for jetpack
+        required: false
+        default: false
+        type: boolean
     secrets:
       PYPI_API_TOKEN:
         description: An optional token to upload to pypi
         required: false
 
+
 permissions:
   id-token: write
   contents: read
@@ -251,7 +257,6 @@ jobs:
         working-directory: ${{ inputs.repository }}
         shell: bash -l {0}
         run: |
-          #set -euxo pipefail
           set -x
           source "${BUILD_ENV_FILE}"
           export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')"
@@ -261,28 +266,37 @@ jobs:
             BUILD_VERSION="${BUILD_VERSION}+${CU_VERSION}"
           fi
           echo "BUILD_VERSION=$BUILD_VERSION"
-          if [[ ${{ inputs.is-jetpack }} == false ]]; then
-            ${CONDA_RUN} python setup.py bdist_wheel
+          echo "USE_RTX=$USE_RTX"
+          echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH"
+          if [[ ${{ inputs.use-rtx }} == true ]]; then
+            echo "Building tensorrt-rtx wheel"
+            ${CONDA_RUN} python setup.py bdist_wheel --use-rtx
           else
-            ${CONDA_RUN} python setup.py bdist_wheel --jetpack --plat-name=linux_tegra_aarch64
+            if [[ ${{ inputs.is-jetpack }} == true ]]; then
+              echo "Building tensorrt wheel for jetpack"
+              ${CONDA_RUN} python setup.py bdist_wheel --jetpack --plat-name=linux_tegra_aarch64
+            else
+              echo "Building standard tensorrt wheel"
+              ${CONDA_RUN} python setup.py bdist_wheel
+            fi
           fi
       - name: Repair Manylinux_2_28 Wheel
         shell: bash -l {0}
         env:
           PACKAGE_NAME: ${{ inputs.package-name }}
           SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }}
+        if: ${{ inputs.architecture == 'x86_64' }}
         run: |
           set -euxo pipefail
           source "${BUILD_ENV_FILE}"
-          # for pkg in ${{ inputs.repository }}/dist/*-linux_*.whl; do
-          #   # if the glob didn't match anything
-          #   if [[ ! -e $pkg ]]; then
-          #       continue
-          #   fi
-          #   abs_pkg=$(realpath $pkg)
-          #   ./test-infra/.github/scripts/repair_manylinux_2_28.sh $abs_pkg
-          # done
-          echo "Repair Manylinux_2_28 Wheel is not supported for aarch64"
+          for pkg in ${{ inputs.repository }}/dist/*-linux_*.whl; do
+            # if the glob didn't match anything
+            if [[ ! -e $pkg ]]; then
+                continue
+            fi
+            abs_pkg=$(realpath $pkg)
+            ./test-infra/.github/scripts/repair_manylinux_2_28.sh $abs_pkg
+          done
       - name: Run Post-Script
         if: ${{ inputs.post-script != '' }}
         uses: ./test-infra/.github/actions/run-script-with-cache
@@ -291,7 +305,6 @@ jobs:
           script: ${{ inputs.post-script }}
       - name: Smoke Test
         shell: bash -l {0}
-        if: ${{ inputs.is-jetpack == false }}
         env:
           PACKAGE_NAME: ${{ inputs.package-name }}
           SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }}
@@ -336,8 +349,10 @@ jobs:
   upload:
     needs: build
     uses: pytorch/test-infra/.github/workflows/_binary_upload.yml@main
-    # for jetpack builds, only upload to pytorch index for nightly builds
-    if: ${{ inputs.is-jetpack == false || (github.event_name == 'push' && startsWith(github.event.ref, 'refs/heads/nightly')) }}
+    # if it is not the jetpack build or tensorrt-rtx build, upload to pytorch index,
+    # if it is the jetpack build only upload to pytorch_index for nightly builds
+    # for tensorrt-rtx build, do not upload at all
+    if: ${{ (inputs.is-jetpack == false && inputs.use-rtx == false) || (inputs.is-jetpack == true && github.event_name == 'push' && startsWith(github.event.ref, 'refs/heads/nightly')) }}
     with:
       repository: ${{ inputs.repository }}
       ref: ${{ inputs.ref }}
@@ -351,5 +366,5 @@ jobs:
       PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{inputs.use-rtx}}-${{inputs.is-jetpack}}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}
   cancel-in-progress: true

.github/workflows/linux-test.yml

@@ -53,6 +53,11 @@ on:
         description: 'Name to give artifacts uploaded from ${RUNNER_ARTIFACT_DIR}'
         default: ''
         type: string
+      use-rtx:
+        description: "Use RTX TensorRT"
+        default: false
+        type: boolean
+        required: false
 
 jobs:
   test:
@@ -135,6 +140,7 @@ jobs:
         working-directory: ${{ inputs.repository }}
         env:
           ALL_SECRETS: ${{ toJSON(secrets) }}
+          USE_RTX: ${{ inputs.use-rtx }}
         run: |
           set -euxo pipefail
           # shellcheck disable=SC2086

MODULE.bazel

@@ -109,6 +109,16 @@ http_archive(
     ],
 )
 
+http_archive(
+    name = "tensorrt_rtx",
+    build_file = "@//third_party/tensorrt_rtx/archive:BUILD",
+    sha256 = "b1222e08f9d473f0bcc06c6a76bf2b1327a106dcee671415c4c46833a105a425",
+    strip_prefix = "TensorRT-RTX-1.0.0.21",
+    urls = [
+        "http://cuda-repo/release-candidates/Libraries/TensorRT/v10.12/10.12.0.35-51f47a12/12.9-r575/Linux-x64-manylinux_2_28-winjit/tar/TensorRT-RTX-1.0.0.21.Linux.x86_64-gnu.cuda-12.9.tar.gz",
+    ],
+)
+
 http_archive(
     name = "tensorrt_sbsa",
     build_file = "@//third_party/tensorrt/archive:BUILD",
@@ -136,6 +146,15 @@ http_archive(
     ],
 )
 
+http_archive(
+    name = "tensorrt_rtx_win",
+    build_file = "@//third_party/tensorrt_rtx/archive:BUILD",
+    sha256 = "49cf1247ada75faa8d538257b763b1c12b9bbb97fcd7765654c55b3ad16bd680",
+    strip_prefix = "TensorRT-RTX-1.0.0.21",
+    urls = [
+        "http://cuda-repo/release-candidates/Libraries/TensorRT/v10.12/10.12.0.35-51f47a12/12.9-r575/Windows10-x64-winjit/zip/TensorRT-RTX-1.0.0.21.Windows.win10.cuda-12.9.zip",
+    ],
+)
 ####################################################################################
 # Locally installed dependencies (use in cases of custom dependencies or aarch64)
 ####################################################################################

core/BUILD

@@ -10,6 +10,21 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "rtx_x86_64",
+    constraint_values = [
+        "@platforms//cpu:x86_64",
+        "@platforms//os:linux",
+    ],
+)
+
+config_setting(
+    name = "rtx_win",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 config_setting(
     name = "sbsa",
     constraint_values = [
@@ -59,6 +74,8 @@ cc_library(
         "//core/runtime",
         "//core/util/logging",
     ] + select({
+        ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"],
+        ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"],
         ":windows": ["@tensorrt_win//:nvinfer"],
         ":sbsa": ["@tensorrt_sbsa//:nvinfer"],
         ":jetpack": ["@tensorrt_l4t//:nvinfer"],

core/conversion/BUILD

@@ -10,6 +10,21 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "rtx_x86_64",
+    constraint_values = [
+        "@platforms//cpu:x86_64",
+        "@platforms//os:linux",
+    ],
+)
+
+config_setting(
+    name = "rtx_win",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 config_setting(
     name = "sbsa",
     constraint_values = [
@@ -54,6 +69,8 @@ cc_library(
         "//core/ir",
         "//core/util:prelude",
     ] + select({
+        ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"],
+        ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"],
         ":windows": ["@tensorrt_win//:nvinfer"],
         ":sbsa": ["@tensorrt_sbsa//:nvinfer"],
         ":jetpack": ["@tensorrt_l4t//:nvinfer"],

core/conversion/conversion.cpp

@@ -204,7 +204,7 @@ void AddInputs(ConversionCtx* ctx, c10::ArrayRef<const torch::jit::Value*> input
       "Optimization profile is invalid, please check the input range provided (conversion.AddInputs)");
 
   ctx->cfg->addOptimizationProfile(profile);
-#if NV_TENSORRT_MAJOR > 7 || (NV_TENSORRT_MAJOR == 7 && NV_TENSORRT_MINOR >= 1)
+#ifndef TRT_MAJOR_RTX && (NV_TENSORRT_MAJOR > 7 || (NV_TENSORRT_MAJOR == 7 && NV_TENSORRT_MINOR >= 1))
   if (ctx->enabled_precisions.find(nvinfer1::DataType::kINT8) != ctx->enabled_precisions.end()) {
     ctx->cfg->setCalibrationProfile(profile);
   }

core/conversion/conversionctx/BUILD

@@ -10,6 +10,21 @@ config_setting(
     },
 )
 
+config_setting(
+    name = "rtx_x86_64",
+    constraint_values = [
+        "@platforms//cpu:x86_64",
+        "@platforms//os:linux",
+    ],
+)
+
+config_setting(
+    name = "rtx_win",
+    constraint_values = [
+        "@platforms//os:windows",
+    ],
+)
+
 config_setting(
     name = "sbsa",
     constraint_values = [
@@ -49,6 +64,8 @@ cc_library(
         "//core/ir",
         "//core/util:prelude",
     ] + select({
+        ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"],
+        ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"],
         ":windows": ["@tensorrt_win//:nvinfer"],
         ":sbsa": ["@tensorrt_sbsa//:nvinfer"],
         ":jetpack": ["@tensorrt_l4t//:nvinfer"],