From b8f5440b3823289689a9e0739f207eeda84e350e Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Thu, 30 Jul 2026 01:20:06 -0700 Subject: [PATCH 1/7] initial check in for windows on arm --- .github/workflows/build_windows_arm64.yml | 88 +++++++++++++++++++ .github/workflows/ci-windows-arm64.yml | 99 ++++++++++++++++++++++ BUILD.bazel | 35 ++++++-- MODULE.bazel | 21 +++++ core/BUILD | 25 ++++-- core/conversion/BUILD | 22 ++++- core/conversion/conversionctx/BUILD | 22 ++++- core/conversion/converters/BUILD | 38 ++++++--- core/conversion/evaluators/BUILD | 22 ++++- core/conversion/tensorcontainer/BUILD | 22 ++++- core/conversion/var/BUILD | 22 ++++- core/ir/BUILD | 22 ++++- core/lowering/BUILD | 22 ++++- core/lowering/passes/BUILD | 17 +++- core/partitioning/BUILD | 22 ++++- core/partitioning/partitioningctx/BUILD | 22 ++++- core/partitioning/partitioninginfo/BUILD | 22 ++++- core/partitioning/segmentedblock/BUILD | 22 ++++- core/plugins/BUILD | 32 +++++-- core/runtime/BUILD | 27 ++++-- core/util/BUILD | 31 +++++-- core/util/logging/BUILD | 25 ++++-- cpp/BUILD | 20 ++++- cpp/bin/torchtrtc/BUILD | 20 ++++- packaging/build_windows_arm64_cross.bat | 18 ++++ packaging/prepare_windows_arm64_cross.ps1 | 71 ++++++++++++++++ packaging/validate_windows_arm64_wheel.ps1 | 37 ++++++++ setup.py | 98 ++++++++++++++++++--- tests/util/BUILD | 19 ++++- third_party/cuda/BUILD | 21 +++++ third_party/tensorrt_rtx/archive/BUILD | 26 +++++- third_party/tensorrt_rtx/local/BUILD | 30 +++++-- toolchains/BUILD | 8 ++ toolchains/ci_workspaces/MODULE.bazel.tmpl | 18 ++++ 34 files changed, 942 insertions(+), 124 deletions(-) create mode 100644 .github/workflows/build_windows_arm64.yml create mode 100644 .github/workflows/ci-windows-arm64.yml create mode 100644 packaging/build_windows_arm64_cross.bat create mode 100644 packaging/prepare_windows_arm64_cross.ps1 create mode 100644 packaging/validate_windows_arm64_wheel.ps1 diff --git a/.github/workflows/build_windows_arm64.yml b/.github/workflows/build_windows_arm64.yml new file mode 100644 index 0000000000..a6702261ba --- /dev/null +++ b/.github/workflows/build_windows_arm64.yml @@ -0,0 +1,88 @@ +name: Build Windows ARM64 RTX wheel + +on: + workflow_call: + inputs: + runner: + description: Windows x86_64 runner with the local ARM64 PyTorch artifact + required: true + type: string + python-version: + required: true + type: string + desired-cuda: + description: cu130 or cu132 + required: true + type: string + pytorch-arm64-artifact: + description: Absolute runner-local path to the Windows ARM64 PyTorch wheel or zip + required: true + type: string + tensorrt-rtx-arm64-artifact: + description: Runner-local path or URL for the Windows ARM64 TensorRT-RTX SDK zip + default: https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-arm64-cuda-13.4-Release-external.zip + type: string + +permissions: + contents: read + +jobs: + build: + runs-on: ${{ inputs.runner }} + timeout-minutes: 120 + env: + CU_VERSION: ${{ inputs.desired-cuda }} + CHANNEL: nightly + USE_TRT_RTX: "1" + PYTHON_VERSION: ${{ inputs.python-version }} + defaults: + run: + shell: bash -l {0} + steps: + - uses: actions/checkout@v6 + with: + submodules: recursive + - uses: actions/checkout@v6 + with: + repository: pytorch/test-infra + path: test-infra + - uses: ./test-infra/.github/actions/setup-windows + with: + cuda-version: ${{ inputs.desired-cuda }} + - uses: actions/setup-python@v6 + with: + python-version: ${{ inputs.python-version }} + architecture: x64 + - name: Install x64 host build dependencies + shell: pwsh + run: | + python -m pip install --upgrade pip + python -m pip install --pre torch --index-url "https://download.pytorch.org/whl/nightly/${{ inputs.desired-cuda }}" + python -m pip install numpy packaging pyyaml setuptools wheel fmt + - name: Prepare ARM64 target sysroot + shell: pwsh + run: | + ./packaging/prepare_windows_arm64_cross.ps1 ` + -PyTorchArtifact '${{ inputs.pytorch-arm64-artifact }}' ` + -TensorRTArtifact '${{ inputs.tensorrt-rtx-arm64-artifact }}' ` + -TargetRoot '${{ runner.temp }}\torchtrt-windows-arm64' ` + -CudaRoot "$env:CUDA_PATH" + - name: Prepare Torch-TensorRT dependencies + run: bash packaging/pre_build_script_windows.sh + - name: Cross-compile Windows ARM64 wheel + shell: cmd + run: | + call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64_arm64 + call packaging\build_windows_arm64_cross.bat + - name: Validate Windows ARM64 wheel + shell: cmd + run: | + call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64_arm64 + powershell -NoProfile -ExecutionPolicy Bypass -File packaging\validate_windows_arm64_wheel.ps1 -DistDirectory dist + - uses: actions/upload-artifact@v6 + with: + name: torch-tensorrt-rtx-${{ inputs.python-version }}-${{ inputs.desired-cuda }}-win-arm64 + path: dist/*-win_arm64.whl + if-no-files-found: error + - uses: ./test-infra/.github/actions/teardown-windows + if: always() diff --git a/.github/workflows/ci-windows-arm64.yml b/.github/workflows/ci-windows-arm64.yml new file mode 100644 index 0000000000..161065868b --- /dev/null +++ b/.github/workflows/ci-windows-arm64.yml @@ -0,0 +1,99 @@ +name: CI Windows ARM64 + +# RTX-only build channel. Until Windows ARM64 PyTorch is published remotely, +# this is manually dispatched onto an x86_64 Windows runner that can access the +# two runner-local PyTorch artifacts. Runtime tests begin when a native ARM64 +# runner is available; cross-builds perform PE/header validation only. +on: + workflow_dispatch: + inputs: + runner: + description: Windows x86_64 runner label with access to the local artifacts + required: true + type: string + python-version: + description: Target CPython version matching both ARM64 PyTorch wheels + default: "3.12" + required: true + type: string + cuda: + description: CUDA variants to build + default: both + type: choice + options: [both, cu130, cu132] + pytorch-cu130-artifact: + description: Absolute runner-local path to the cu130 Windows ARM64 PyTorch artifact + required: false + type: string + pytorch-cu132-artifact: + description: Absolute runner-local path to the cu132 Windows ARM64 PyTorch artifact + required: false + type: string + tensorrt-rtx-artifact: + description: Optional local TensorRT-RTX ARM64 zip; blank uses the NVIDIA URL + required: false + type: string + +permissions: + contents: read + +jobs: + validate-inputs: + runs-on: ubuntu-latest + steps: + - name: Validate selected artifacts + env: + CUDA_SELECTION: ${{ inputs.cuda }} + PYTORCH_CU130: ${{ inputs.pytorch-cu130-artifact }} + PYTORCH_CU132: ${{ inputs.pytorch-cu132-artifact }} + run: | + set -euo pipefail + if [[ "${CUDA_SELECTION}" != "cu132" && -z "${PYTORCH_CU130}" ]]; then + echo "::error::pytorch-cu130-artifact is required for ${CUDA_SELECTION}" + exit 1 + fi + if [[ "${CUDA_SELECTION}" != "cu130" && -z "${PYTORCH_CU132}" ]]; then + echo "::error::pytorch-cu132-artifact is required for ${CUDA_SELECTION}" + exit 1 + fi + + cu130: + needs: validate-inputs + if: inputs.cuda == 'both' || inputs.cuda == 'cu130' + uses: ./.github/workflows/build_windows_arm64.yml + with: + runner: ${{ inputs.runner }} + python-version: ${{ inputs.python-version }} + desired-cuda: cu130 + pytorch-arm64-artifact: ${{ inputs.pytorch-cu130-artifact }} + tensorrt-rtx-arm64-artifact: ${{ inputs.tensorrt-rtx-artifact || 'https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-arm64-cuda-13.4-Release-external.zip' }} + + cu132: + needs: validate-inputs + if: inputs.cuda == 'both' || inputs.cuda == 'cu132' + uses: ./.github/workflows/build_windows_arm64.yml + with: + runner: ${{ inputs.runner }} + python-version: ${{ inputs.python-version }} + desired-cuda: cu132 + pytorch-arm64-artifact: ${{ inputs.pytorch-cu132-artifact }} + tensorrt-rtx-arm64-artifact: ${{ inputs.tensorrt-rtx-artifact || 'https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-arm64-cuda-13.4-Release-external.zip' }} + + gate: + needs: [validate-inputs, cu130, cu132] + if: always() + runs-on: ubuntu-latest + steps: + - name: Gate Windows ARM64 builds + env: + VALIDATE_RESULT: ${{ needs.validate-inputs.result }} + CU130_RESULT: ${{ needs.cu130.result }} + CU132_RESULT: ${{ needs.cu132.result }} + run: | + set -euo pipefail + for result in "${VALIDATE_RESULT}" "${CU130_RESULT}" "${CU132_RESULT}"; do + if [[ "${result}" == "failure" || "${result}" == "cancelled" ]]; then + exit 1 + fi + done + echo "Windows ARM64 RTX cross-build gate passed." diff --git a/BUILD.bazel b/BUILD.bazel index 4dbc88f847..ec3c7fa730 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1,5 +1,5 @@ -load("@hedron_compile_commands//:refresh_compile_commands.bzl", "refresh_compile_commands") load("@bazel_skylib//rules:common_settings.bzl", "bool_flag") +load("@hedron_compile_commands//:refresh_compile_commands.bzl", "refresh_compile_commands") load("@rules_pkg//:pkg.bzl", "pkg_tar", "pkg_zip") load("@rules_pkg//pkg:mappings.bzl", "pkg_files") @@ -16,8 +16,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -226,9 +238,10 @@ pkg_tar( srcs = ["//cpp/bin/torchtrtc"] + select({ ":jetpack": [], ":no_package_executorch": [], - ":rtx_win": [], - ":rtx_x86_64": [], ":rtx_sbsa": [], + ":rtx_win_arm64": [], + ":rtx_win_x86_64": [], + ":rtx_x86_64": [], ":windows": [], "//conditions:default": [ "//examples/executorch_reference_runner:example_executorch_runner", @@ -270,16 +283,18 @@ pkg_tar( ] + select({ ":jetpack": [], ":no_package_executorch": [], - ":rtx_win": [], - ":rtx_x86_64": [], ":rtx_sbsa": [], + ":rtx_win_arm64": [], + ":rtx_win_x86_64": [], + ":rtx_x86_64": [], ":windows": [], "//conditions:default": [ ":executorch_source_package", ":include_executorch", ], }) + select({ - ":rtx_win": [], + ":rtx_win_arm64": [], + ":rtx_win_x86_64": [], ":windows": [], "//conditions:default": [":bin"], }), @@ -300,7 +315,8 @@ pkg_tar( alias( name = "libtorchtrt", actual = select({ - ":rtx_win": ":libtorchtrt_zip", + ":rtx_win_arm64": ":libtorchtrt_zip", + ":rtx_win_x86_64": ":libtorchtrt_zip", ":windows": ":libtorchtrt_zip", "//conditions:default": ":libtorchtrt_tar", }), @@ -326,7 +342,8 @@ pkg_tar( alias( name = "libtorchtrt_runtime", actual = select({ - ":rtx_win": ":libtorchtrt_runtime_zip", + ":rtx_win_arm64": ":libtorchtrt_runtime_zip", + ":rtx_win_x86_64": ":libtorchtrt_runtime_zip", ":windows": ":libtorchtrt_runtime_zip", "//conditions:default": ":libtorchtrt_runtime_tar", }), diff --git a/MODULE.bazel b/MODULE.bazel index 7409f40f5d..bb51173959 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -94,6 +94,12 @@ new_local_repository( path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.2/", ) +new_local_repository( + name = "cuda_win_arm64", + build_file = "@//third_party/cuda:BUILD", + path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.2/", +) + http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") #################################################################################### @@ -163,6 +169,12 @@ http_archive( urls = ["https://download.pytorch.org/libtorch/nightly/cu130/libtorch-win-shared-with-deps-latest.zip"], ) +new_local_repository( + name = "libtorch_win_arm64", + build_file = "@//third_party/libtorch:BUILD", + path = "C:/torchtrt-arm64-artifacts/pytorch", +) + http_archive( name = "torch_l4t", build_file = "@//third_party/libtorch:BUILD", @@ -241,3 +253,12 @@ http_archive( "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-amd64-cuda-13.4-Release-external.zip", ], ) + +http_archive( + name = "tensorrt_rtx_win_arm64", + build_file = "@//third_party/tensorrt_rtx/archive:BUILD", + strip_prefix = "TensorRT-RTX-1.6.1.120", + urls = [ + "https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.6/TensorRT-RTX-1.6.1.120-Windows-arm64-cuda-13.4-Release-external.zip", + ], +) diff --git a/core/BUILD b/core/BUILD index 51ddecb75b..cb9ca5720a 100644 --- a/core/BUILD +++ b/core/BUILD @@ -1,6 +1,7 @@ load("@rules_cc//cc:defs.bzl", "cc_library") load("@rules_pkg//:pkg.bzl", "pkg_tar") load("@rules_pkg//pkg:mappings.bzl", "pkg_files") + package(default_visibility = ["//visibility:public"]) config_setting( @@ -33,8 +34,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", + constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -92,15 +105,17 @@ cc_library( "//core/util/logging", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], @@ -121,6 +136,6 @@ pkg_files( srcs = [ "compiler.h", ], - visibility = ["//visibility:public"], prefix = "include/torch_tensorrt/core/", + visibility = ["//visibility:public"], ) diff --git a/core/conversion/BUILD b/core/conversion/BUILD index 4196f2eb1d..8853db0cc3 100644 --- a/core/conversion/BUILD +++ b/core/conversion/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -87,15 +99,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/conversion/conversionctx/BUILD b/core/conversion/conversionctx/BUILD index b9f9b53d62..33cd2c0db2 100644 --- a/core/conversion/conversionctx/BUILD +++ b/core/conversion/conversionctx/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -82,15 +94,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index 6df31e1ff9..e1aad9035e 100644 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -82,15 +94,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], @@ -112,15 +126,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], @@ -180,15 +196,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/conversion/evaluators/BUILD b/core/conversion/evaluators/BUILD index f71df9f9e0..963b0db166 100644 --- a/core/conversion/evaluators/BUILD +++ b/core/conversion/evaluators/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -88,15 +100,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/conversion/tensorcontainer/BUILD b/core/conversion/tensorcontainer/BUILD index 629afce89f..dbfee3fa44 100644 --- a/core/conversion/tensorcontainer/BUILD +++ b/core/conversion/tensorcontainer/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -81,15 +93,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/conversion/var/BUILD b/core/conversion/var/BUILD index e6e66441c6..b02e7d5798 100644 --- a/core/conversion/var/BUILD +++ b/core/conversion/var/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -84,15 +96,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/ir/BUILD b/core/ir/BUILD index 7c3919bea9..25eb5bc7bb 100644 --- a/core/ir/BUILD +++ b/core/ir/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -84,15 +96,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/lowering/BUILD b/core/lowering/BUILD index 8f463cd40d..05f9ffb4d5 100644 --- a/core/lowering/BUILD +++ b/core/lowering/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -86,15 +98,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/lowering/passes/BUILD b/core/lowering/passes/BUILD index bd8462eed1..6ebdf3ae5d 100644 --- a/core/lowering/passes/BUILD +++ b/core/lowering/passes/BUILD @@ -31,8 +31,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -87,7 +99,8 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/partitioning/BUILD b/core/partitioning/BUILD index 9e84dbff21..52caab81d4 100644 --- a/core/partitioning/BUILD +++ b/core/partitioning/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -89,15 +101,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/partitioning/partitioningctx/BUILD b/core/partitioning/partitioningctx/BUILD index 9b2b9bcd5c..32205fd650 100644 --- a/core/partitioning/partitioningctx/BUILD +++ b/core/partitioning/partitioningctx/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -85,15 +97,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/partitioning/partitioninginfo/BUILD b/core/partitioning/partitioninginfo/BUILD index f2b28e4f5a..9a94de51dd 100644 --- a/core/partitioning/partitioninginfo/BUILD +++ b/core/partitioning/partitioninginfo/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -84,15 +96,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/partitioning/segmentedblock/BUILD b/core/partitioning/segmentedblock/BUILD index 0aafc0b581..622f062f29 100644 --- a/core/partitioning/segmentedblock/BUILD +++ b/core/partitioning/segmentedblock/BUILD @@ -33,8 +33,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -84,15 +96,17 @@ cc_library( "//core/util:prelude", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/plugins/BUILD b/core/plugins/BUILD index 1b9fc164bc..7b7eb05e50 100644 --- a/core/plugins/BUILD +++ b/core/plugins/BUILD @@ -34,8 +34,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -78,9 +90,10 @@ cc_library( "impl/normalize_plugin.cpp", "register_plugins.cpp", ], - ":rtx_win": [], - ":rtx_x86_64": [], ":rtx_sbsa": [], + ":rtx_win_arm64": [], + ":rtx_win_x86_64": [], + ":rtx_x86_64": [], "//conditions:default": [ "impl/interpolate_plugin.cpp", "impl/normalize_plugin.cpp", @@ -94,9 +107,10 @@ cc_library( "impl/normalize_plugin.h", "plugins.h", ], - ":rtx_win": [], - ":rtx_x86_64": [], ":rtx_sbsa": [], + ":rtx_win_arm64": [], + ":rtx_win_x86_64": [], + ":rtx_x86_64": [], "//conditions:default": [ "impl/interpolate_plugin.h", "impl/normalize_plugin.h", @@ -126,9 +140,10 @@ cc_library( "@tensorrt_l4t//:nvinfer", "@tensorrt_l4t//:nvinferplugin", ], - ":rtx_win": [], - ":rtx_x86_64": [], ":rtx_sbsa": [], + ":rtx_win_arm64": [], + ":rtx_win_x86_64": [], + ":rtx_x86_64": [], ":sbsa": [ "@tensorrt_sbsa//:nvinfer", "@tensorrt_sbsa//:nvinferplugin", @@ -143,7 +158,8 @@ cc_library( ], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/runtime/BUILD b/core/runtime/BUILD index 51e77a8663..254e0a9425 100644 --- a/core/runtime/BUILD +++ b/core/runtime/BUILD @@ -35,8 +35,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -83,9 +95,10 @@ cc_library( ], deps = select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], @@ -139,15 +152,17 @@ cc_library( "//core/util:prelude", ] + if_torch_nccl(["@torch_nccl//:nccl_headers"]) + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], diff --git a/core/util/BUILD b/core/util/BUILD index 22609d97ee..917d99854b 100644 --- a/core/util/BUILD +++ b/core/util/BUILD @@ -1,6 +1,7 @@ load("@rules_cc//cc:defs.bzl", "cc_library") load("@rules_pkg//:pkg.bzl", "pkg_tar") load("@rules_pkg//pkg:mappings.bzl", "pkg_files") + package(default_visibility = ["//visibility:public"]) config_setting( @@ -33,8 +34,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", + constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -93,7 +106,8 @@ cc_library( ":macros", ] + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], @@ -129,7 +143,8 @@ cc_library( ], deps = select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], @@ -149,15 +164,17 @@ cc_library( "//core/util/logging", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], @@ -187,6 +204,6 @@ pkg_tar( pkg_files( name = "include_pkg_files", srcs = ["//core/util:include_files"], - visibility = ["//visibility:public"], prefix = "include/torch_tensorrt/core/util/", + visibility = ["//visibility:public"], ) diff --git a/core/util/logging/BUILD b/core/util/logging/BUILD index ff9c542d8a..75008b51ab 100644 --- a/core/util/logging/BUILD +++ b/core/util/logging/BUILD @@ -1,6 +1,7 @@ load("@rules_cc//cc:defs.bzl", "cc_library") load("@rules_pkg//:pkg.bzl", "pkg_tar") load("@rules_pkg//pkg:mappings.bzl", "pkg_files") + package(default_visibility = ["//visibility:public"]) config_setting( @@ -29,13 +30,25 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", "@platforms//os:windows", ], flag_values = {"//toolchains/dep_collection:compute_libs": "rtx"}, ) +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + config_setting( name = "sbsa", constraint_values = [ @@ -74,15 +87,17 @@ cc_library( ], deps = select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":jetpack": ["@torch_l4t//:libtorch"], - ":rtx_win": ["@libtorch_win//:libtorch"], + ":rtx_win_arm64": ["@libtorch_win_arm64//:libtorch"], + ":rtx_win_x86_64": ["@libtorch_win//:libtorch"], ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], @@ -99,6 +114,6 @@ pkg_tar( pkg_files( name = "include_pkg_files", srcs = ["TorchTRTLogger.h"], - visibility = ["//visibility:public"], prefix = "include/torch_tensorrt/core/util/logging/", + visibility = ["//visibility:public"], ) diff --git a/cpp/BUILD b/cpp/BUILD index b54b2d24e5..2bd6745cf9 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -64,8 +64,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -126,6 +138,7 @@ cc_library( hdrs = [ "include/torch_tensorrt/executorch/TensorRTBackend.h", ], + strip_include_prefix = "include", # Build the TensorRT backend as a static library. The final application # links this target together with the ExecuTorch runtime it was compiled # against, avoiding any runtime plugin/dlopen dependency. @@ -134,19 +147,18 @@ cc_library( ":sbsa": [], "//conditions:default": ["@platforms//:incompatible"], }), - strip_include_prefix = "include", deps = [ ":tensorrt_executorch_binding_names", ":tensorrt_executorch_blob_header", ] + select({ ":linux_x86_64": [ - "@executorch//:executorch_headers", "@cuda//:cudart", + "@executorch//:executorch_headers", "@tensorrt//:nvinfer", ], ":sbsa": [ - "@executorch//:executorch_headers", "@cuda//:cudart", + "@executorch//:executorch_headers", "@tensorrt_sbsa//:nvinfer", ], "//conditions:default": [], diff --git a/cpp/bin/torchtrtc/BUILD b/cpp/bin/torchtrtc/BUILD index 2c87eddae2..926b56d915 100644 --- a/cpp/bin/torchtrtc/BUILD +++ b/cpp/bin/torchtrtc/BUILD @@ -27,8 +27,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -59,7 +71,11 @@ cc_binary( "@torch_l4t//:caffe2", "@torch_l4t//:libtorch", ], - ":rtx_win": [ + ":rtx_win_arm64": [ + "@libtorch_win_arm64//:caffe2", + "@libtorch_win_arm64//:libtorch", + ], + ":rtx_win_x86_64": [ "@libtorch_win//:caffe2", "@libtorch_win//:libtorch", ], diff --git a/packaging/build_windows_arm64_cross.bat b/packaging/build_windows_arm64_cross.bat new file mode 100644 index 0000000000..7a4f392f32 --- /dev/null +++ b/packaging/build_windows_arm64_cross.bat @@ -0,0 +1,18 @@ +@echo off +setlocal enabledelayedexpansion +if not "%VSCMD_ARG_HOST_ARCH%"=="x64" ( + echo Expected the x64-hosted MSVC toolchain, got %VSCMD_ARG_HOST_ARCH% + exit /b 1 +) +if not "%VSCMD_ARG_TGT_ARCH%"=="arm64" ( + echo Expected the ARM64 MSVC target, got %VSCMD_ARG_TGT_ARCH% + exit /b 1 +) +if not "%TORCHTRT_TARGET_PLATFORM%"=="windows-arm64" ( + echo TORCHTRT_TARGET_PLATFORM is not windows-arm64 + exit /b 1 +) +python -m pip install --upgrade setuptools==72.1.0 wheel +if errorlevel 1 exit /b %errorlevel% +python setup.py bdist_wheel --use-rtx +exit /b %errorlevel% diff --git a/packaging/prepare_windows_arm64_cross.ps1 b/packaging/prepare_windows_arm64_cross.ps1 new file mode 100644 index 0000000000..c55f91c431 --- /dev/null +++ b/packaging/prepare_windows_arm64_cross.ps1 @@ -0,0 +1,71 @@ +param( + [Parameter(Mandatory = $true)] [string] $PyTorchArtifact, + [Parameter(Mandatory = $true)] [string] $TensorRTArtifact, + [Parameter(Mandatory = $true)] [string] $TargetRoot, + [Parameter(Mandatory = $true)] [string] $CudaRoot +) + +$ErrorActionPreference = "Stop" + +function Expand-ZipArtifact { + param([string] $Artifact, [string] $Destination, [string] $Name) + $source = $Artifact + if ($Artifact -match '^https?://') { + $source = Join-Path $env:RUNNER_TEMP "$Name.zip" + Invoke-WebRequest -Uri $Artifact -OutFile $source + } + if (-not (Test-Path -LiteralPath $source -PathType Leaf)) { + throw "$Name artifact was not found: $source" + } + New-Item -ItemType Directory -Force -Path $Destination | Out-Null + python -m zipfile -e $source $Destination +} + +function Find-Root { + param([string] $SearchRoot, [string] $IncludePath, [string] $LibraryPath, [string] $Name) + $candidates = @($SearchRoot) + @(Get-ChildItem -LiteralPath $SearchRoot -Directory -Recurse | Select-Object -ExpandProperty FullName) + foreach ($candidate in $candidates) { + if ((Test-Path (Join-Path $candidate $IncludePath)) -and (Test-Path (Join-Path $candidate $LibraryPath))) { + return $candidate + } + } + throw "Could not locate the normalized $Name root under $SearchRoot" +} + +if ($env:CU_VERSION -notin @('cu130', 'cu132')) { + throw "Windows ARM64 RTX supports only cu130 and cu132; got '$env:CU_VERSION'" +} +if (Test-Path -LiteralPath $TargetRoot) { + Remove-Item -LiteralPath $TargetRoot -Recurse -Force +} +New-Item -ItemType Directory -Force -Path $TargetRoot | Out-Null + +$torchExtract = Join-Path $TargetRoot "pytorch" +$trtExtract = Join-Path $TargetRoot "tensorrt-rtx" +Expand-ZipArtifact -Artifact $PyTorchArtifact -Destination $torchExtract -Name "pytorch-arm64" +Expand-ZipArtifact -Artifact $TensorRTArtifact -Destination $trtExtract -Name "tensorrt-rtx-arm64" + +$torchRoot = Find-Root -SearchRoot $torchExtract -IncludePath "include" -LibraryPath "lib" -Name "PyTorch" +$trtRoot = Find-Root -SearchRoot $trtExtract -IncludePath "include" -LibraryPath "lib" -Name "TensorRT-RTX" +$cudaArm64Lib = Join-Path $CudaRoot "lib\arm64\cudart.lib" +if (-not (Test-Path -LiteralPath $cudaArm64Lib)) { + throw "CUDA target library is missing: $cudaArm64Lib" +} +foreach ($library in @('c10.lib', 'torch.lib', 'torch_cpu.lib', 'torch_python.lib')) { + if (-not (Test-Path -LiteralPath (Join-Path $torchRoot "lib\$library"))) { + throw "ARM64 PyTorch artifact is missing lib\$library" + } +} +if (-not (Get-ChildItem -LiteralPath (Join-Path $trtRoot "lib") -Filter "tensorrt_rtx*.lib" | Select-Object -First 1)) { + throw "ARM64 TensorRT-RTX artifact contains no tensorrt_rtx import library" +} + +function Export-GitHubEnvironment([string] $Name, [string] $Value) { + "$Name=$($Value -replace '\\', '/')" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append +} +Export-GitHubEnvironment "TORCHTRT_TARGET_PLATFORM" "windows-arm64" +Export-GitHubEnvironment "TORCHTRT_BUILD_MODE" "cross" +Export-GitHubEnvironment "TORCHTRT_TARGET_TORCH_ROOT" $torchRoot +Export-GitHubEnvironment "TORCHTRT_TARGET_TRT_ROOT" $trtRoot +Export-GitHubEnvironment "TORCHTRT_TARGET_CUDA_ROOT" $CudaRoot +Write-Host "Prepared Windows ARM64 target sysroot: PyTorch=$torchRoot TensorRT-RTX=$trtRoot CUDA=$CudaRoot" diff --git a/packaging/validate_windows_arm64_wheel.ps1 b/packaging/validate_windows_arm64_wheel.ps1 new file mode 100644 index 0000000000..282999b38f --- /dev/null +++ b/packaging/validate_windows_arm64_wheel.ps1 @@ -0,0 +1,37 @@ +param([Parameter(Mandatory = $true)] [string] $DistDirectory) + +$ErrorActionPreference = "Stop" +$wheel = Get-ChildItem -LiteralPath $DistDirectory -Filter "*-win_arm64.whl" | Select-Object -First 1 +if (-not $wheel) { + throw "No win_arm64 wheel found in $DistDirectory" +} +if (Get-ChildItem -LiteralPath $DistDirectory -Filter "*-win_amd64.whl") { + throw "An unexpected win_amd64 wheel was produced" +} + +$extractRoot = Join-Path $env:RUNNER_TEMP "torchtrt-win-arm64-validation" +if (Test-Path -LiteralPath $extractRoot) { + Remove-Item -LiteralPath $extractRoot -Recurse -Force +} +New-Item -ItemType Directory -Force -Path $extractRoot | Out-Null +python -m zipfile -e $wheel.FullName $extractRoot + +$nativeFiles = @(Get-ChildItem -LiteralPath $extractRoot -Recurse -File | Where-Object { $_.Extension -in @('.dll', '.pyd') }) +$targetLibraries = @( + Join-Path $env:TORCHTRT_TARGET_TORCH_ROOT "lib\c10.lib" + Join-Path $env:TORCHTRT_TARGET_TORCH_ROOT "lib\torch_cpu.lib" + Join-Path $env:TORCHTRT_TARGET_CUDA_ROOT "lib\arm64\cudart.lib" +) +$nativeFiles += @($targetLibraries | Get-Item) +$nativeFiles += @(Get-ChildItem -LiteralPath (Join-Path $env:TORCHTRT_TARGET_TRT_ROOT "lib") -Filter "tensorrt_rtx*.lib" | Select-Object -First 1) + +if (-not $nativeFiles) { + throw "No native outputs or target import libraries were found" +} +foreach ($file in $nativeFiles) { + $headers = & dumpbin.exe /headers $file.FullName 2>&1 | Out-String + if ($LASTEXITCODE -ne 0 -or $headers -notmatch 'AA64 machine \(ARM64\)') { + throw "Native file is not ARM64: $($file.FullName)" + } +} +Write-Host "Validated $($nativeFiles.Count) ARM64 native files in $($wheel.Name) and its target sysroot" diff --git a/setup.py b/setup.py index 152d370b37..5aedd36c3f 100644 --- a/setup.py +++ b/setup.py @@ -90,6 +90,13 @@ def load_dep_info(): dir_path = os.path.join(str(get_root_dir()), "py") IS_AARCH64 = platform.machine() == "aarch64" +TARGET_WINDOWS_ARM64 = ( + os.environ.get("TORCHTRT_TARGET_PLATFORM", "").lower() == "windows-arm64" +) +WINDOWS_CROSS_COMPILE = ( + TARGET_WINDOWS_ARM64 + and os.environ.get("TORCHTRT_BUILD_MODE", "cross").lower() == "cross" +) IS_JETPACK = False PY_ONLY = False @@ -245,6 +252,8 @@ def build_libtorchtrt_cxx11_abi( if IS_WINDOWS: cmd.append("--config=windows") + if TARGET_WINDOWS_ARM64: + cmd.append("--platforms=//toolchains:windows_arm64") else: cmd.append("--config=linux") @@ -415,6 +424,14 @@ def run(self): install.run(self) +class TorchTRTBuildExtension(BuildExtension): + def get_ext_filename(self, ext_name): + filename = super().get_ext_filename(ext_name) + if TARGET_WINDOWS_ARM64: + filename = filename.replace("win_amd64", "win_arm64") + return filename + + class BdistCommand(bdist_wheel): description = "Builds the package" @@ -423,9 +440,19 @@ def initialize_options(self): def finalize_options(self): bdist_wheel.finalize_options(self) + if TARGET_WINDOWS_ARM64: + # Cross-builds run under x64 Python, so the target platform cannot be + # inferred from the build interpreter. + self.plat_name = "win_arm64" if NO_TS or PY_ONLY: self.root_is_pure = False + def get_tag(self): + python_tag, abi_tag, platform_tag = bdist_wheel.get_tag(self) + if TARGET_WINDOWS_ARM64: + platform_tag = "win_arm64" + return python_tag, abi_tag, platform_tag + def run(self): # Ensure wheel metadata/project name reflects RTX vs standard build, # even when pyproject.toml provides static [project].name. @@ -684,12 +711,15 @@ def run(self): tensorrt_linux_external_dir = tensorrt_x86_64_external_dir if USE_TRT_RTX: - tensorrt_windows_external_dir = ( - lambda: subprocess.check_output( + tensorrt_windows_repo = ( + "@tensorrt_rtx_win_arm64" if TARGET_WINDOWS_ARM64 else "@tensorrt_rtx_win" + ) + tensorrt_windows_external_dir = lambda: ( + subprocess.check_output( [ BAZEL_EXE, "query", - "@tensorrt_rtx_win//:nvinfer", + f"{tensorrt_windows_repo}//:nvinfer", "--output", "location", ] @@ -708,8 +738,37 @@ def run(self): .split("/BUILD.bazel")[0] ) + extension_type = setuptools.Extension if WINDOWS_CROSS_COMPILE else CUDAExtension + target_torch_root = os.environ.get("TORCHTRT_TARGET_TORCH_ROOT") + target_cuda_root = os.environ.get("TORCHTRT_TARGET_CUDA_ROOT") + if WINDOWS_CROSS_COMPILE and not (target_torch_root and target_cuda_root): + raise RuntimeError( + "Windows ARM64 cross-compilation requires " + "TORCHTRT_TARGET_TORCH_ROOT and TORCHTRT_TARGET_CUDA_ROOT" + ) + + extension_kwargs = {} + if WINDOWS_CROSS_COMPILE: + extension_kwargs = { + "library_dirs": [ + dir_path + "/torch_tensorrt/lib/", + os.path.join(target_torch_root, "lib"), + os.path.join(target_cuda_root, "lib", "arm64"), + ], + "libraries": [ + "torchtrt", + "c10", + "torch", + "torch_cpu", + "torch_python", + "c10_cuda", + "torch_cuda", + "cudart", + ], + } + ext_modules += [ - CUDAExtension( + extension_type( "torch_tensorrt._C", [ "py/" + f @@ -720,11 +779,14 @@ def run(self): "torch_tensorrt/csrc/register_tensorrt_classes.cpp", ] ], - library_dirs=[ - (dir_path + "/torch_tensorrt/lib/"), - "/opt/conda/lib/python3.6/config-3.6m-x86_64-linux-gnu", - ], - libraries=["torchtrt"], + library_dirs=extension_kwargs.get( + "library_dirs", + [ + (dir_path + "/torch_tensorrt/lib/"), + "/opt/conda/lib/python3.6/config-3.6m-x86_64-linux-gnu", + ], + ), + libraries=extension_kwargs.get("libraries", ["torchtrt"]), include_dirs=( [ dir_path + "torch_tensorrt/csrc", @@ -733,6 +795,22 @@ def run(self): dir_path + "/../cpp/include", "/usr/local/cuda", ] + + ( + [ + os.path.join(target_torch_root, "include"), + os.path.join( + target_torch_root, + "include", + "torch", + "csrc", + "api", + "include", + ), + os.path.join(target_cuda_root, "include"), + ] + if WINDOWS_CROSS_COMPILE + else [] + ) + ( [ dir_path + "/../bazel-TRTorch/external/tensorrt_win/include", @@ -939,7 +1017,7 @@ def get_requirements(): "install": InstallCommand, "clean": CleanCommand, "develop": DevelopCommand, - "build_ext": BuildExtension, + "build_ext": TorchTRTBuildExtension, "bdist_wheel": BdistCommand, "editable_wheel": EditableWheelCommand, }, diff --git a/tests/util/BUILD b/tests/util/BUILD index 6a4a18f51c..b1577cf93c 100644 --- a/tests/util/BUILD +++ b/tests/util/BUILD @@ -66,8 +66,20 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "rtx", + }, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = { @@ -91,9 +103,10 @@ cc_library( "@googletest//:gtest_main", ] + select({ ":jetpack": ["@tensorrt_l4t//:nvinfer"], - ":rtx_win": ["@tensorrt_rtx_win//:nvinfer"], - ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":rtx_sbsa": ["@tensorrt_rtx_sbsa//:nvinfer"], + ":rtx_win_arm64": ["@tensorrt_rtx_win_arm64//:nvinfer"], + ":rtx_win_x86_64": ["@tensorrt_rtx_win//:nvinfer"], + ":rtx_x86_64": ["@tensorrt_rtx//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], ":windows": ["@tensorrt_win//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], diff --git a/third_party/cuda/BUILD b/third_party/cuda/BUILD index b5359d42b3..b887b7e785 100644 --- a/third_party/cuda/BUILD +++ b/third_party/cuda/BUILD @@ -9,12 +9,23 @@ config_setting( ], ) +config_setting( + name = "windows_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", + "@platforms//os:windows", + ], +) + cc_library( name = "cudart", srcs = select({ ":windows": [ "lib/x64/cudart.lib", ], + ":windows_arm64": [ + "lib/arm64/cudart.lib", + ], "//conditions:default": [ "lib64/libcudart.so", ], @@ -34,6 +45,9 @@ cc_library( ":windows": [ "bin/nvToolsExt64_1.dll", ], + ":windows_arm64": [ + "bin/arm64/nvToolsExt64_1.dll", + ], "//conditions:default": [ "lib64/libnvToolsExt.so.1", ], @@ -46,6 +60,10 @@ cc_library( ":windows": [ "bin/*.dll", ], + ":windows_arm64": glob( + ["bin/arm64/*.dll"], + allow_empty = True, + ), "//conditions:default": glob( [ "lib64/**/lib*.so", @@ -69,6 +87,9 @@ cc_library( ":windows": [ "lib/x64/cublas.lib", ], + ":windows_arm64": [ + "lib/arm64/cublas.lib", + ], "//conditions:default": glob( [ "lib64/**/*libcublas.so", diff --git a/third_party/tensorrt_rtx/archive/BUILD b/third_party/tensorrt_rtx/archive/BUILD index 6a558a4b69..83331e35e3 100644 --- a/third_party/tensorrt_rtx/archive/BUILD +++ b/third_party/tensorrt_rtx/archive/BUILD @@ -21,13 +21,25 @@ config_setting( ) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", "@platforms//os:windows", ], flag_values = {"@//toolchains/dep_collection:compute_libs": "rtx"}, ) +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", + "@platforms//os:windows", + ], + flag_values = { + "@//toolchains/dep_collection:compute_libs": "rtx", + }, +) + cc_library( name = "nvinfer_headers", hdrs = glob( @@ -48,7 +60,8 @@ cc_import( name = "nvinfer_lib", shared_library = select({ ":rtx_sbsa": "lib/libtensorrt_rtx.so", - ":rtx_win": "bin/tensorrt_rtx_1_6.dll", + ":rtx_win_arm64": "bin/tensorrt_rtx_1_6.dll", + ":rtx_win_x86_64": "bin/tensorrt_rtx_1_6.dll", ":rtx_x86_64": "lib/libtensorrt_rtx.so", }), visibility = ["//visibility:private"], @@ -57,7 +70,8 @@ cc_import( cc_import( name = "nvinfer_static_lib", static_library = select({ - ":rtx_win": "lib/tensorrt_rtx_1_6.lib", + ":rtx_win_arm64": "lib/tensorrt_rtx_1_6.lib", + ":rtx_win_x86_64": "lib/tensorrt_rtx_1_6.lib", }), visibility = ["//visibility:private"], ) @@ -70,7 +84,11 @@ cc_library( "nvinfer_lib", ] + select({ ":rtx_sbsa": ["@cuda//:cudart"], - ":rtx_win": [ + ":rtx_win_arm64": [ + "nvinfer_static_lib", + "@cuda_win_arm64//:cudart", + ], + ":rtx_win_x86_64": [ "nvinfer_static_lib", "@cuda_win//:cudart", ], diff --git a/third_party/tensorrt_rtx/local/BUILD b/third_party/tensorrt_rtx/local/BUILD index 737f5551e6..a1141ac2c7 100644 --- a/third_party/tensorrt_rtx/local/BUILD +++ b/third_party/tensorrt_rtx/local/BUILD @@ -3,8 +3,18 @@ load("@rules_cc//cc:defs.bzl", "cc_import", "cc_library") package(default_visibility = ["//visibility:public"]) config_setting( - name = "rtx_win", + name = "rtx_win_x86_64", constraint_values = [ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], + flag_values = {"@//toolchains/dep_collection:compute_libs": "rtx"}, +) + +config_setting( + name = "rtx_win_arm64", + constraint_values = [ + "@platforms//cpu:aarch64", "@platforms//os:windows", ], flag_values = {"@//toolchains/dep_collection:compute_libs": "rtx"}, @@ -41,7 +51,11 @@ cc_library( # "include/NvInferPluginUtils.h", # ], ), - ":rtx_win": glob( + ":rtx_win_arm64": glob( + ["include/NvInfer*.h"], + allow_empty = True, + ), + ":rtx_win_x86_64": glob( [ "include/NvInfer*.h", ], @@ -69,7 +83,8 @@ cc_library( cc_import( name = "nvinfer_static_lib", static_library = select({ - ":rtx_win": "lib/tensorrt_rtx_1_6.lib", + ":rtx_win_arm64": "lib/tensorrt_rtx_1_6.lib", + ":rtx_win_x86_64": "lib/tensorrt_rtx_1_6.lib", }), visibility = ["//visibility:private"], ) @@ -78,7 +93,8 @@ cc_import( name = "nvinfer_lib", shared_library = select({ ":rtx_sbsa": "lib/libtensorrt_rtx.so", - ":rtx_win": "bin/tensorrt_rtx_1_6.dll", + ":rtx_win_arm64": "bin/tensorrt_rtx_1_6.dll", + ":rtx_win_x86_64": "bin/tensorrt_rtx_1_6.dll", ":rtx_x86_64": "lib/libtensorrt_rtx.so", }), visibility = ["//visibility:private"], @@ -91,7 +107,11 @@ cc_library( "nvinfer_headers", "nvinfer_lib", ] + select({ - ":rtx_win": [ + ":rtx_win_arm64": [ + "nvinfer_static_lib", + "@cuda_win_arm64//:cudart", + ], + ":rtx_win_x86_64": [ "nvinfer_static_lib", "@cuda_win//:cudart", ], diff --git a/toolchains/BUILD b/toolchains/BUILD index a3cd275373..73090f3e34 100644 --- a/toolchains/BUILD +++ b/toolchains/BUILD @@ -8,6 +8,14 @@ platform( ], ) +platform( + name = "windows_arm64", + constraint_values = [ + "@platforms//os:windows", + "@platforms//cpu:aarch64", + ], +) + platform( name = "jetpack_4.5", constraint_values = [ diff --git a/toolchains/ci_workspaces/MODULE.bazel.tmpl b/toolchains/ci_workspaces/MODULE.bazel.tmpl index 8998010090..720c1f4498 100644 --- a/toolchains/ci_workspaces/MODULE.bazel.tmpl +++ b/toolchains/ci_workspaces/MODULE.bazel.tmpl @@ -73,6 +73,12 @@ new_local_repository( path = "${CUDA_HOME}", ) +new_local_repository( + name = "cuda_win_arm64", + build_file = "@//third_party/cuda:BUILD", + path = "${TORCHTRT_TARGET_CUDA_ROOT}", +) + http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") @@ -158,6 +164,12 @@ http_archive( ], ) +new_local_repository( + name = "tensorrt_rtx_win_arm64", + build_file = "@//third_party/tensorrt_rtx/archive:BUILD", + path = "${TORCHTRT_TARGET_TRT_ROOT}", +) + #################################################################################### # Locally installed dependencies (use in cases of custom dependencies or aarch64) #################################################################################### @@ -176,6 +188,12 @@ new_local_repository( build_file = "third_party/libtorch/BUILD" ) +new_local_repository( + name = "libtorch_win_arm64", + path = "${TORCHTRT_TARGET_TORCH_ROOT}", + build_file = "third_party/libtorch/BUILD" +) + new_local_repository( name = "torch_l4t", path = "${TORCH_INSTALL_PATH}", From 8f136728fe348542cb491964028cbf0a21458f2c Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Thu, 30 Jul 2026 10:57:07 -0700 Subject: [PATCH 2/7] update the cuda 13.2 woa artifact location --- MODULE.bazel | 4 +++- toolchains/ci_workspaces/MODULE.bazel.tmpl | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/MODULE.bazel b/MODULE.bazel index bb51173959..717e857e6d 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -94,10 +94,12 @@ new_local_repository( path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.2/", ) +# ARM64 target sysroot copied from the Windows ARM64 machine. Keep this +# separate from the x86_64 host CUDA installation above. new_local_repository( name = "cuda_win_arm64", build_file = "@//third_party/cuda:BUILD", - path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.2/", + path = "C:/torchtrt-arm64-artifacts/cuda-v13.2-arm64/", ) http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") diff --git a/toolchains/ci_workspaces/MODULE.bazel.tmpl b/toolchains/ci_workspaces/MODULE.bazel.tmpl index 720c1f4498..eff946660e 100644 --- a/toolchains/ci_workspaces/MODULE.bazel.tmpl +++ b/toolchains/ci_workspaces/MODULE.bazel.tmpl @@ -73,6 +73,8 @@ new_local_repository( path = "${CUDA_HOME}", ) +# ARM64 target sysroot; this must contain lib/arm64/cudart.lib and must not +# point at the x86_64 host CUDA installation above. new_local_repository( name = "cuda_win_arm64", build_file = "@//third_party/cuda:BUILD", From 567e6d108dd7fcb3b501b4d0e95bcbbf52ca9258 Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Thu, 30 Jul 2026 14:21:27 -0700 Subject: [PATCH 3/7] add --windows-on-arm build flag --- setup.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 5aedd36c3f..1b502a026f 100644 --- a/setup.py +++ b/setup.py @@ -90,8 +90,15 @@ def load_dep_info(): dir_path = os.path.join(str(get_root_dir()), "py") IS_AARCH64 = platform.machine() == "aarch64" +WINDOWS_ON_ARM = "--windows-on-arm" in sys.argv +if WINDOWS_ON_ARM: + # setuptools does not know this project-specific option, so consume it + # before command-line processing. The environment variable remains + # supported for CI and scripted builds. + sys.argv.remove("--windows-on-arm") TARGET_WINDOWS_ARM64 = ( - os.environ.get("TORCHTRT_TARGET_PLATFORM", "").lower() == "windows-arm64" + WINDOWS_ON_ARM + or os.environ.get("TORCHTRT_TARGET_PLATFORM", "").lower() == "windows-arm64" ) WINDOWS_CROSS_COMPILE = ( TARGET_WINDOWS_ARM64 From b629511953cf5572cebabc60bc306209d1390f9f Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Thu, 30 Jul 2026 14:25:12 -0700 Subject: [PATCH 4/7] change cpu:aarch64 to cpu:arm64 --- BUILD.bazel | 2 +- core/BUILD | 2 +- core/conversion/BUILD | 2 +- core/conversion/conversionctx/BUILD | 2 +- core/conversion/converters/BUILD | 2 +- core/conversion/evaluators/BUILD | 2 +- core/conversion/tensorcontainer/BUILD | 2 +- core/conversion/var/BUILD | 2 +- core/ir/BUILD | 2 +- core/lowering/BUILD | 2 +- core/lowering/passes/BUILD | 2 +- core/partitioning/BUILD | 2 +- core/partitioning/partitioningctx/BUILD | 2 +- core/partitioning/partitioninginfo/BUILD | 2 +- core/partitioning/segmentedblock/BUILD | 2 +- core/plugins/BUILD | 2 +- core/runtime/BUILD | 2 +- core/util/BUILD | 2 +- core/util/logging/BUILD | 2 +- cpp/BUILD | 2 +- cpp/bin/torchtrtc/BUILD | 2 +- tests/util/BUILD | 2 +- third_party/cuda/BUILD | 2 +- third_party/tensorrt_rtx/archive/BUILD | 2 +- third_party/tensorrt_rtx/local/BUILD | 2 +- toolchains/BUILD | 2 +- 26 files changed, 26 insertions(+), 26 deletions(-) diff --git a/BUILD.bazel b/BUILD.bazel index ec3c7fa730..c46a9d931b 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -29,7 +29,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/BUILD b/core/BUILD index cb9ca5720a..15b90ff1de 100644 --- a/core/BUILD +++ b/core/BUILD @@ -47,7 +47,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/conversion/BUILD b/core/conversion/BUILD index 8853db0cc3..90bd854073 100644 --- a/core/conversion/BUILD +++ b/core/conversion/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/conversion/conversionctx/BUILD b/core/conversion/conversionctx/BUILD index 33cd2c0db2..a66ba9bd24 100644 --- a/core/conversion/conversionctx/BUILD +++ b/core/conversion/conversionctx/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index e1aad9035e..ea8a3208bc 100644 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/conversion/evaluators/BUILD b/core/conversion/evaluators/BUILD index 963b0db166..968f12cdaf 100644 --- a/core/conversion/evaluators/BUILD +++ b/core/conversion/evaluators/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/conversion/tensorcontainer/BUILD b/core/conversion/tensorcontainer/BUILD index dbfee3fa44..bb5826ef38 100644 --- a/core/conversion/tensorcontainer/BUILD +++ b/core/conversion/tensorcontainer/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/conversion/var/BUILD b/core/conversion/var/BUILD index b02e7d5798..206db84324 100644 --- a/core/conversion/var/BUILD +++ b/core/conversion/var/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/ir/BUILD b/core/ir/BUILD index 25eb5bc7bb..ff0eaf5751 100644 --- a/core/ir/BUILD +++ b/core/ir/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/lowering/BUILD b/core/lowering/BUILD index 05f9ffb4d5..ff664742ac 100644 --- a/core/lowering/BUILD +++ b/core/lowering/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/lowering/passes/BUILD b/core/lowering/passes/BUILD index 6ebdf3ae5d..f44696b3eb 100644 --- a/core/lowering/passes/BUILD +++ b/core/lowering/passes/BUILD @@ -44,7 +44,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/partitioning/BUILD b/core/partitioning/BUILD index 52caab81d4..4daa975dee 100644 --- a/core/partitioning/BUILD +++ b/core/partitioning/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/partitioning/partitioningctx/BUILD b/core/partitioning/partitioningctx/BUILD index 32205fd650..f7fc7ca216 100644 --- a/core/partitioning/partitioningctx/BUILD +++ b/core/partitioning/partitioningctx/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/partitioning/partitioninginfo/BUILD b/core/partitioning/partitioninginfo/BUILD index 9a94de51dd..1ccad74527 100644 --- a/core/partitioning/partitioninginfo/BUILD +++ b/core/partitioning/partitioninginfo/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/partitioning/segmentedblock/BUILD b/core/partitioning/segmentedblock/BUILD index 622f062f29..f1673ca846 100644 --- a/core/partitioning/segmentedblock/BUILD +++ b/core/partitioning/segmentedblock/BUILD @@ -46,7 +46,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/plugins/BUILD b/core/plugins/BUILD index 7b7eb05e50..68d2146f21 100644 --- a/core/plugins/BUILD +++ b/core/plugins/BUILD @@ -47,7 +47,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/runtime/BUILD b/core/runtime/BUILD index 254e0a9425..8c3560c271 100644 --- a/core/runtime/BUILD +++ b/core/runtime/BUILD @@ -48,7 +48,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/util/BUILD b/core/util/BUILD index 917d99854b..2f1885544c 100644 --- a/core/util/BUILD +++ b/core/util/BUILD @@ -47,7 +47,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/core/util/logging/BUILD b/core/util/logging/BUILD index 75008b51ab..d11b8bfe86 100644 --- a/core/util/logging/BUILD +++ b/core/util/logging/BUILD @@ -41,7 +41,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/cpp/BUILD b/cpp/BUILD index 2bd6745cf9..3a765fc675 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -77,7 +77,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/cpp/bin/torchtrtc/BUILD b/cpp/bin/torchtrtc/BUILD index 926b56d915..161afd23c2 100644 --- a/cpp/bin/torchtrtc/BUILD +++ b/cpp/bin/torchtrtc/BUILD @@ -40,7 +40,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/tests/util/BUILD b/tests/util/BUILD index b1577cf93c..1dba512622 100644 --- a/tests/util/BUILD +++ b/tests/util/BUILD @@ -79,7 +79,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/third_party/cuda/BUILD b/third_party/cuda/BUILD index b887b7e785..204b9cee23 100644 --- a/third_party/cuda/BUILD +++ b/third_party/cuda/BUILD @@ -12,7 +12,7 @@ config_setting( config_setting( name = "windows_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], ) diff --git a/third_party/tensorrt_rtx/archive/BUILD b/third_party/tensorrt_rtx/archive/BUILD index 83331e35e3..488bf0e074 100644 --- a/third_party/tensorrt_rtx/archive/BUILD +++ b/third_party/tensorrt_rtx/archive/BUILD @@ -32,7 +32,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = { diff --git a/third_party/tensorrt_rtx/local/BUILD b/third_party/tensorrt_rtx/local/BUILD index a1141ac2c7..53715b585a 100644 --- a/third_party/tensorrt_rtx/local/BUILD +++ b/third_party/tensorrt_rtx/local/BUILD @@ -14,7 +14,7 @@ config_setting( config_setting( name = "rtx_win_arm64", constraint_values = [ - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", "@platforms//os:windows", ], flag_values = {"@//toolchains/dep_collection:compute_libs": "rtx"}, diff --git a/toolchains/BUILD b/toolchains/BUILD index 73090f3e34..73358b1863 100644 --- a/toolchains/BUILD +++ b/toolchains/BUILD @@ -12,7 +12,7 @@ platform( name = "windows_arm64", constraint_values = [ "@platforms//os:windows", - "@platforms//cpu:aarch64", + "@platforms//cpu:arm64", ], ) From e77547aebee336d0e7faed9e69c55eddf7449b6b Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Thu, 30 Jul 2026 16:47:59 -0700 Subject: [PATCH 5/7] fix the python issue --- MODULE.bazel | 10 ++++++---- setup.py | 11 +++++++++-- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/MODULE.bazel b/MODULE.bazel index 717e857e6d..0d7043cd91 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -11,6 +11,9 @@ bazel_dep(name = "rules_python", version = "1.3.0") bazel_dep(name = "bazel_skylib", version = "1.7.1") bazel_dep(name = "rules_foreign_cc", version = "0.15.1") +cc_configure = use_extension("@rules_cc//cc:extensions.bzl", "cc_configure_extension") +use_repo(cc_configure, "local_config_cc") + python = use_extension("@rules_python//python/extensions:python.bzl", "python") python.toolchain( ignore_root_user_error = True, @@ -94,12 +97,11 @@ new_local_repository( path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.2/", ) -# ARM64 target sysroot copied from the Windows ARM64 machine. Keep this -# separate from the x86_64 host CUDA installation above. +# supposedly cuda 13.4 on windows x86 preview has everything required for cross compile for windows on arm new_local_repository( name = "cuda_win_arm64", build_file = "@//third_party/cuda:BUILD", - path = "C:/torchtrt-arm64-artifacts/cuda-v13.2-arm64/", + path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.4/", ) http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") @@ -174,7 +176,7 @@ http_archive( new_local_repository( name = "libtorch_win_arm64", build_file = "@//third_party/libtorch:BUILD", - path = "C:/torchtrt-arm64-artifacts/pytorch", + path = "C:/torchtrt-arm64-artifacts/pytorch/torch", ) http_archive( diff --git a/setup.py b/setup.py index 1b502a026f..244de522e3 100644 --- a/setup.py +++ b/setup.py @@ -261,6 +261,9 @@ def build_libtorchtrt_cxx11_abi( cmd.append("--config=windows") if TARGET_WINDOWS_ARM64: cmd.append("--platforms=//toolchains:windows_arm64") + cmd.append( + "--extra_toolchains=@local_config_cc//:cc-toolchain-arm64_windows" + ) else: cmd.append("--config=linux") @@ -748,21 +751,24 @@ def run(self): extension_type = setuptools.Extension if WINDOWS_CROSS_COMPILE else CUDAExtension target_torch_root = os.environ.get("TORCHTRT_TARGET_TORCH_ROOT") target_cuda_root = os.environ.get("TORCHTRT_TARGET_CUDA_ROOT") - if WINDOWS_CROSS_COMPILE and not (target_torch_root and target_cuda_root): + target_python_root = os.environ.get("TORCHTRT_TARGET_PYTHON_ROOT") + if WINDOWS_CROSS_COMPILE and not (target_torch_root and target_cuda_root and target_python_root): raise RuntimeError( "Windows ARM64 cross-compilation requires " - "TORCHTRT_TARGET_TORCH_ROOT and TORCHTRT_TARGET_CUDA_ROOT" + "TORCHTRT_TARGET_TORCH_ROOT and TORCHTRT_TARGET_CUDA_ROOT and TORCHTRT_TARGET_PYTHON_ROOT" ) extension_kwargs = {} if WINDOWS_CROSS_COMPILE: extension_kwargs = { "library_dirs": [ + os.path.join(target_python_root, "libs"), dir_path + "/torch_tensorrt/lib/", os.path.join(target_torch_root, "lib"), os.path.join(target_cuda_root, "lib", "arm64"), ], "libraries": [ + f"python{sys.version_info.major}{sys.version_info.minor}", "torchtrt", "c10", "torch", @@ -804,6 +810,7 @@ def run(self): ] + ( [ + os.path.join(target_python_root, "include"), os.path.join(target_torch_root, "include"), os.path.join( target_torch_root, From 86f643a6f415743789f98b1d6e50541ef2a8821c Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Thu, 30 Jul 2026 16:50:17 -0700 Subject: [PATCH 6/7] add the vc variable --- packaging/build_windows_arm64_cross.bat | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packaging/build_windows_arm64_cross.bat b/packaging/build_windows_arm64_cross.bat index 7a4f392f32..03619d0f97 100644 --- a/packaging/build_windows_arm64_cross.bat +++ b/packaging/build_windows_arm64_cross.bat @@ -12,6 +12,9 @@ if not "%TORCHTRT_TARGET_PLATFORM%"=="windows-arm64" ( echo TORCHTRT_TARGET_PLATFORM is not windows-arm64 exit /b 1 ) +rem Reuse the amd64_arm64 environment initialized by vcvarsall.bat. Without +rem this, torch.utils.cpp_extension rejects the already-active VC environment. +set DISTUTILS_USE_SDK=1 python -m pip install --upgrade setuptools==72.1.0 wheel if errorlevel 1 exit /b %errorlevel% python setup.py bdist_wheel --use-rtx From 3ffc38e520ceda0cab78ebfbff890d9bf2baf8e3 Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Thu, 30 Jul 2026 22:53:26 -0700 Subject: [PATCH 7/7] add doc --- MODULE.bazel | 8 +- docsrc/getting_started/installation.rst | 1 + docsrc/getting_started/windows_arm64.rst | 187 +++++++++++++++++++++++ setup.py | 48 ++++-- 4 files changed, 225 insertions(+), 19 deletions(-) create mode 100644 docsrc/getting_started/windows_arm64.rst diff --git a/MODULE.bazel b/MODULE.bazel index 0d7043cd91..7b14e4f901 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -173,11 +173,9 @@ http_archive( urls = ["https://download.pytorch.org/libtorch/nightly/cu130/libtorch-win-shared-with-deps-latest.zip"], ) -new_local_repository( - name = "libtorch_win_arm64", - build_file = "@//third_party/libtorch:BUILD", - path = "C:/torchtrt-arm64-artifacts/pytorch/torch", -) +# Native Windows ARM64 builds use the Torch package installed in the active +# Python environment. Cross-builds override discovery through TORCH_PATH. +local_torch(name = "libtorch_win_arm64") http_archive( name = "torch_l4t", diff --git a/docsrc/getting_started/installation.rst b/docsrc/getting_started/installation.rst index d1c72d1b7c..f072e01ffe 100644 --- a/docsrc/getting_started/installation.rst +++ b/docsrc/getting_started/installation.rst @@ -259,6 +259,7 @@ Platform-specific Installation jetpack tensorrt_rtx + windows_arm64 diff --git a/docsrc/getting_started/windows_arm64.rst b/docsrc/getting_started/windows_arm64.rst new file mode 100644 index 0000000000..a2ca19c7cd --- /dev/null +++ b/docsrc/getting_started/windows_arm64.rst @@ -0,0 +1,187 @@ +Building Torch-TensorRT-RTX for Windows ARM64 +######################################### + +This guide describes local TensorRT-RTX builds for Windows ARM64. It covers: + +* a native build on a Windows ARM64 machine; and +* a cross-build on a Windows x64 machine that targets Windows ARM64. + +The commands below must be run from the root of the Torch-TensorRT repository +in a Visual Studio developer command prompt. The resulting wheel is written to +the ``dist`` directory. + +Prerequisites +------------- + +Both build types require: + +* Visual Studio 2022 Build Tools with the C++ build tools and ARM64 compiler + components; +* Bazelisk via choco; +* CUDA Toolkit 13.4 Preview. The Windows x86 installation contains the + headers and ARM64 libraries required to cross-compile for Windows ARM64. + Native Windows ARM64 builds also use CUDA 13.4; +* a compatible Windows ARM64 PyTorch package containing the C++ headers and + libraries under ``torch\include`` and ``torch\lib``; +* Python packages required by Torch-TensorRT-RTX. The currently available + Windows ARM64 PyTorch wheel targets CPython 3.13, so native and cross-builds + are currently supported only with Python 3.13. + +If the CUDA Toolkit is installed somewhere other than the path in +``MODULE.bazel``, update the ``cuda_win_arm64`` repository path. Use the same +CUDA root for ``CUDA_PATH``, ``CUDA_HOME``, and +``TORCHTRT_TARGET_CUDA_ROOT`` where those variables are required below. + +Native build on Windows ARM64 +----------------------------- + +1. Open an ARM64-native Visual Studio 2022 developer command prompt. You can + also initialize one from ``cmd.exe``: + + .. code-block:: bat + + call "%ProgramFiles%\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat" arm64 + echo %VSCMD_ARG_HOST_ARCH% + echo %VSCMD_ARG_TGT_ARCH% + where cl.exe + cl.exe + + Both printed architectures should be ``arm64``. Adjust ``BuildTools`` in + the path if Visual Studio was installed in another edition. ``where + cl.exe`` should resolve to an ARM64-native compiler path containing + ``Hostarm64\arm64\cl.exe``, and the ``cl.exe`` banner should identify + ARM64 as the target architecture. + +2. Create and activate a virtual environment using a native ARM64 Python 3.13 + interpreter: + + .. code-block:: bat + + C:\path\to\arm64-python\python.exe -m venv .venv + call .venv\Scripts\activate.bat + python -c "import platform; print(platform.machine())" + + The last command must print ``ARM64``. Install the compatible ARM64 PyTorch + wheel and the prerequisite packages listed above in this environment. + +3. Configure CUDA and tell setuptools to reuse the active Visual Studio + environment: + + .. code-block:: bat + + set "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4" + set "CUDA_HOME=%CUDA_PATH%" + set DISTUTILS_USE_SDK=1 + +4. Build the ARM64 TensorRT-RTX wheel: + + .. code-block:: bat + + python -m pip install --upgrade pip + python -m pip install numpy packaging pyyaml setuptools==72.1.0 wheel fmt build + + # only for run the setup.py, not use this torch wheelto build windows on arm + python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130 + + # build the native windows on arm TensorRT-RTX wheel + python setup.py bdist_wheel --use-rtx + + ``setup.py`` detects the native ARM64 Python interpreter, selects the + Windows ARM64 Bazel platform, and produces a wheel tagged + ``win_arm64``. Do not pass ``--windows-on-arm`` for a native build. + +5. Confirm the wheel tag and test it in the ARM64 environment: + + .. code-block:: bat + + dir dist\*win_arm64.whl + python -m pip install C:\path\to\the\generated-win_arm64.whl + python -c "import torch_tensorrt; print(torch_tensorrt.__version__)" + +Cross-build on Windows x64 +-------------------------- + +The cross-build uses windows-x64 to run the build and the x64-hosted ARM64 MSVC +compiler to create ARM64 binaries. The host and target Python installations +must have the same major and minor version. Use Python 3.13 for both the x64 +host and the ARM64 target because the available Windows ARM64 PyTorch wheel +and the resulting Torch-TensorRT wheel target CPython 3.13. + +1. Open an x64-to-ARM64 Visual Studio 2022 cross-tools prompt, or initialize + one from ``cmd.exe``: + + .. code-block:: bat + + call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvarsall.bat" amd64_arm64 + echo Host=%VSCMD_ARG_HOST_ARCH% + echo Target=%VSCMD_ARG_TGT_ARCH% + echo Tools=%VCToolsInstallDir% + + The host architecture should be ``x64`` and the target architecture should + be ``arm64``. + +2. Create and activate an environment using x64 Python 3.13. Install an x64 + PyTorch package so that ``setup.py`` can run, followed by the prerequisite + packages listed above: + + .. code-block:: bat + + C:\path\to\x64-python\python.exe -m venv .venv-cross + call .venv-cross\Scripts\activate.bat + python -c "import platform; print(platform.machine())" + + The last command should print ``AMD64``. + +3. Prepare the ARM64 target dependencies: + + * ``TORCHTRT_TARGET_TORCH_ROOT`` must contain ``include`` and ``lib`` + directories from an ARM64 PyTorch installation. + * ``TORCHTRT_TARGET_CUDA_ROOT`` must contain the CUDA headers and + ``lib\arm64\cudart.lib``. + * ``TORCHTRT_TARGET_PYTHON_ROOT`` must point to an ARM64 Python 3.13 + installation. It must contain the ARM64 Python headers and import library: + ``include\Python.h`` and ``libs\python313.lib``. Do not use headers or + libraries from the x64 host Python installation. + + Set the target roots and reuse the active cross-tools environment: + + .. code-block:: bat + + set "TORCHTRT_TARGET_TORCH_ROOT=C:/torchtrt-arm64-artifacts/pytorch/torch" + set "TORCHTRT_TARGET_CUDA_ROOT=C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.4/" + set "TORCHTRT_TARGET_PYTHON_ROOT=C:\torchtrt-arm64-artifacts\python313-arm64" + set "CUDA_PATH=%TORCHTRT_TARGET_CUDA_ROOT%" + set "CUDA_HOME=%TORCHTRT_TARGET_CUDA_ROOT%" + set DISTUTILS_USE_SDK=1 + +4. Cross-build the ARM64 TensorRT-RTX wheel: + + .. code-block:: bat + python -m pip install --upgrade pip + + python -m pip install numpy packaging pyyaml setuptools==72.1.0 wheel fmt build + + #only for run the setup.py, not use this torch wheelto build windows on arm + python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130 + + # cross-build the ARM64 TensorRT-RTX wheel + python setup.py bdist_wheel --use-rtx --windows-on-arm + + The ``--windows-on-arm`` option explicitly selects ARM64 while the build is + running under x64 Python. It also selects Bazel's x64-to-ARM64 MSVC + toolchain and emits a ``win_arm64`` wheel. + +5. Confirm that the generated wheel is tagged for ARM64: + + .. code-block:: bat + + dir dist\*win_arm64.whl + + An ARM64 wheel cannot be imported by the x64 host Python used for the + cross-build. Copy the wheel to a Windows ARM64 machine, install it into a + matching ARM64 Python environment, and run: + + .. code-block:: bat + + python -m pip install C:\path\to\the\generated-win_arm64.whl + python -c "import torch_tensorrt; print(torch_tensorrt.__version__)" diff --git a/setup.py b/setup.py index 244de522e3..7768b6b2d9 100644 --- a/setup.py +++ b/setup.py @@ -89,21 +89,25 @@ def load_dep_info(): dir_path = os.path.join(str(get_root_dir()), "py") -IS_AARCH64 = platform.machine() == "aarch64" +MACHINE_ARCH = platform.machine().lower() +NATIVE_WINDOWS_ARM64 = IS_WINDOWS and MACHINE_ARCH in ("arm64", "aarch64") +IS_AARCH64 = not IS_WINDOWS and MACHINE_ARCH == "aarch64" WINDOWS_ON_ARM = "--windows-on-arm" in sys.argv if WINDOWS_ON_ARM: # setuptools does not know this project-specific option, so consume it # before command-line processing. The environment variable remains # supported for CI and scripted builds. sys.argv.remove("--windows-on-arm") -TARGET_WINDOWS_ARM64 = ( +FORCED_WINDOWS_ARM64 = ( WINDOWS_ON_ARM or os.environ.get("TORCHTRT_TARGET_PLATFORM", "").lower() == "windows-arm64" ) -WINDOWS_CROSS_COMPILE = ( - TARGET_WINDOWS_ARM64 - and os.environ.get("TORCHTRT_BUILD_MODE", "cross").lower() == "cross" -) +if FORCED_WINDOWS_ARM64 and not IS_WINDOWS: + raise RuntimeError("Windows ARM64 builds are supported only on Windows") +TARGET_WINDOWS_ARM64 = NATIVE_WINDOWS_ARM64 or FORCED_WINDOWS_ARM64 +# An ARM64 Python process builds natively. An x64 Python process must opt in to +# the ARM64 target via --windows-on-arm or TORCHTRT_TARGET_PLATFORM. +WINDOWS_CROSS_COMPILE = TARGET_WINDOWS_ARM64 and not NATIVE_WINDOWS_ARM64 IS_JETPACK = False PY_ONLY = False @@ -248,7 +252,7 @@ def build_libtorchtrt_cxx11_abi( cmd.append("--compilation_mode=dbg") else: cmd.append("--compilation_mode=opt") - if use_dist_dir: + if use_dist_dir and not IS_WINDOWS: if IS_AARCH64: cmd.append("--distdir=third_party/dist_dir/aarch64-linux-gnu") else: @@ -261,6 +265,7 @@ def build_libtorchtrt_cxx11_abi( cmd.append("--config=windows") if TARGET_WINDOWS_ARM64: cmd.append("--platforms=//toolchains:windows_arm64") + if WINDOWS_CROSS_COMPILE: cmd.append( "--extra_toolchains=@local_config_cc//:cc-toolchain-arm64_windows" ) @@ -300,7 +305,11 @@ def build_libtorchtrt_cxx11_abi( print(f"Using bazel --disk_cache={disk_cache}") env = os.environ.copy() - if "TORCH_PATH" not in env: + target_torch_path = env.get("TORCHTRT_TARGET_TORCH_ROOT") + if TARGET_WINDOWS_ARM64 and target_torch_path: + env["TORCH_PATH"] = target_torch_path + print(f"Using target TORCH_PATH={target_torch_path}") + elif "TORCH_PATH" not in env: stable_torch_path = resolve_torch_path() if stable_torch_path is not None: env["TORCH_PATH"] = stable_torch_path @@ -748,18 +757,29 @@ def run(self): .split("/BUILD.bazel")[0] ) - extension_type = setuptools.Extension if WINDOWS_CROSS_COMPILE else CUDAExtension + extension_type = setuptools.Extension if TARGET_WINDOWS_ARM64 else CUDAExtension target_torch_root = os.environ.get("TORCHTRT_TARGET_TORCH_ROOT") target_cuda_root = os.environ.get("TORCHTRT_TARGET_CUDA_ROOT") target_python_root = os.environ.get("TORCHTRT_TARGET_PYTHON_ROOT") - if WINDOWS_CROSS_COMPILE and not (target_torch_root and target_cuda_root and target_python_root): + if NATIVE_WINDOWS_ARM64: + target_torch_root = target_torch_root or os.path.dirname(torch.__file__) + target_cuda_root = ( + target_cuda_root + or os.environ.get("CUDA_PATH") + or os.environ.get("CUDA_HOME") + ) + target_python_root = target_python_root or sys.base_prefix + if TARGET_WINDOWS_ARM64 and not ( + target_torch_root and target_cuda_root and target_python_root + ): raise RuntimeError( - "Windows ARM64 cross-compilation requires " - "TORCHTRT_TARGET_TORCH_ROOT and TORCHTRT_TARGET_CUDA_ROOT and TORCHTRT_TARGET_PYTHON_ROOT" + "Windows ARM64 builds require Torch, CUDA, and Python roots; set " + "TORCHTRT_TARGET_TORCH_ROOT, TORCHTRT_TARGET_CUDA_ROOT, and " + "TORCHTRT_TARGET_PYTHON_ROOT" ) extension_kwargs = {} - if WINDOWS_CROSS_COMPILE: + if TARGET_WINDOWS_ARM64: extension_kwargs = { "library_dirs": [ os.path.join(target_python_root, "libs"), @@ -822,7 +842,7 @@ def run(self): ), os.path.join(target_cuda_root, "include"), ] - if WINDOWS_CROSS_COMPILE + if TARGET_WINDOWS_ARM64 else [] ) + (