diff --git a/.github/workflows/build_windows_arm64.yml b/.github/workflows/build_windows_arm64.yml new file mode 100644 index 0000000000..2b7cf0efb6 --- /dev/null +++ b/.github/workflows/build_windows_arm64.yml @@ -0,0 +1,99 @@ +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: CUDA version for the Windows ARM64 target build + required: true + type: string + pytorch-arm64-artifact: + description: Absolute runner-local path to the Windows ARM64 PyTorch wheel or zip + required: true + type: string + python-arm64-root: + description: ARM64 Python 3.13 root containing include/Python.h and libs/python313.lib + required: true + type: string + cuda-arm64-root: + description: CUDA 13.4 Preview root containing ARM64 target libraries + 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: + # The current PyTorch WoA image only supports cu134; keep this configurable + # so future images can build against all supported CUDA versions. + 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 '${{ inputs.cuda-arm64-root }}' ` + -PythonRoot '${{ inputs.python-arm64-root }}' + - 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..5e675af6bc --- /dev/null +++ b/.github/workflows/ci-windows-arm64.yml @@ -0,0 +1,86 @@ +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 +# runner-local PyTorch artifact and ARM64 Python development files. 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: Host and target CPython version (Windows ARM64 currently supports 3.13 only) + default: "3.13" + required: true + type: choice + options: ["3.13"] + pytorch-arm64-artifact: + description: Absolute runner-local path to the Windows ARM64 PyTorch wheel or zip + required: true + type: string + python-arm64-root: + description: Absolute runner-local ARM64 Python 3.13 root containing include/Python.h and libs/python313.lib + required: true + type: string + cuda-arm64-root: + description: Absolute runner-local CUDA 13.4 Preview root containing lib/arm64/cudart.lib + default: 'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4' + required: true + 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: + PYTHON_VERSION: ${{ inputs.python-version }} + PYTORCH_ARM64: ${{ inputs.pytorch-arm64-artifact }} + PYTHON_ARM64_ROOT: ${{ inputs.python-arm64-root }} + CUDA_ARM64_ROOT: ${{ inputs.cuda-arm64-root }} + run: | + set -euo pipefail + if [[ "${PYTHON_VERSION}" != "3.13" ]]; then + echo "::error::Windows ARM64 builds currently require Python 3.13" + exit 1 + fi + for value in "${PYTORCH_ARM64}" "${PYTHON_ARM64_ROOT}" "${CUDA_ARM64_ROOT}"; do + [[ -n "${value}" ]] || { echo "::error::ARM64 PyTorch, Python, and CUDA inputs are required"; exit 1; } + done + + build: + needs: validate-inputs + uses: ./.github/workflows/build_windows_arm64.yml + with: + runner: ${{ inputs.runner }} + python-version: ${{ inputs.python-version }} + pytorch-arm64-artifact: ${{ inputs.pytorch-arm64-artifact }} + python-arm64-root: ${{ inputs.python-arm64-root }} + cuda-arm64-root: ${{ inputs.cuda-arm64-root }} + 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, build] + if: always() + runs-on: ubuntu-latest + steps: + - name: Gate Windows ARM64 builds + env: + VALIDATE_RESULT: ${{ needs.validate-inputs.result }} + BUILD_RESULT: ${{ needs.build.result }} + run: | + set -euo pipefail + if [[ "${VALIDATE_RESULT}" != "success" || "${BUILD_RESULT}" != "success" ]]; then + exit 1 + fi + echo "Windows ARM64 RTX cross-build gate passed." diff --git a/BUILD.bazel b/BUILD.bazel index 4dbc88f847..c46a9d931b 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:arm64", "@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..7b14e4f901 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,6 +97,13 @@ new_local_repository( path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.2/", ) +# 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:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.4/", +) + http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") #################################################################################### @@ -163,6 +173,10 @@ http_archive( urls = ["https://download.pytorch.org/libtorch/nightly/cu130/libtorch-win-shared-with-deps-latest.zip"], ) +# 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", build_file = "@//third_party/libtorch:BUILD", @@ -241,3 +255,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..15b90ff1de 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:arm64", "@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..90bd854073 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:arm64", "@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..a66ba9bd24 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:arm64", "@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..ea8a3208bc 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:arm64", "@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..968f12cdaf 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:arm64", "@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..bb5826ef38 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:arm64", "@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..206db84324 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:arm64", "@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..ff0eaf5751 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:arm64", "@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..ff664742ac 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:arm64", "@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..f44696b3eb 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:arm64", "@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..4daa975dee 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:arm64", "@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..f7fc7ca216 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:arm64", "@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..1ccad74527 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:arm64", "@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..f1673ca846 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:arm64", "@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..68d2146f21 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:arm64", "@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..8c3560c271 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:arm64", "@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/runtime/Platform.cpp b/core/runtime/Platform.cpp index 03d9e7580b..25b0a3c2fa 100644 --- a/core/runtime/Platform.cpp +++ b/core/runtime/Platform.cpp @@ -12,6 +12,7 @@ const std::unordered_map& get_name_to_platf {"linux_aarch64", Platform::PlatformEnum::kLINUX_AARCH64}, {"linux_x86_64", Platform::PlatformEnum::kLINUX_X86_64}, {"windows_x86_64", Platform::PlatformEnum::kWIN_X86_64}, + {"windows_arm64", Platform::PlatformEnum::kWIN_ARM64}, {"unknown", Platform::PlatformEnum::kUNKNOWN}, }; return name_to_platform_map; @@ -22,6 +23,7 @@ const std::unordered_map& _get_platform_nam {Platform::PlatformEnum::kLINUX_AARCH64, "linux_aarch64"}, {Platform::PlatformEnum::kLINUX_X86_64, "linux_x86_64"}, {Platform::PlatformEnum::kWIN_X86_64, "windows_x86_64"}, + {Platform::PlatformEnum::kWIN_ARM64, "windows_arm64"}, {Platform::PlatformEnum::kUNKNOWN, "unknown"}}; return platform_name_map; } @@ -83,6 +85,8 @@ Platform get_current_platform() { #elif defined(_WIN32) || defined(_WIN64) #if defined(_M_AMD64) || defined(_M_X64) return Platform(Platform::PlatformEnum::kWIN_X86_64); +#elif defined(_M_ARM64) + return Platform(Platform::PlatformEnum::kWIN_ARM64); #else return Platform(Platform::PlatformEnum::kUNKNOWN); #endif diff --git a/core/runtime/Platform.h b/core/runtime/Platform.h index 3f059c8b77..6c9393a445 100644 --- a/core/runtime/Platform.h +++ b/core/runtime/Platform.h @@ -12,6 +12,7 @@ struct Platform { kLINUX_AARCH64, kWIN_X86_64, kUNKNOWN, + kWIN_ARM64, } PlatformEnum; PlatformEnum _platform = Platform::kUNKNOWN; diff --git a/core/runtime/register_jit_hooks.cpp b/core/runtime/register_jit_hooks.cpp index 472cc9880d..6a44f4f10a 100644 --- a/core/runtime/register_jit_hooks.cpp +++ b/core/runtime/register_jit_hooks.cpp @@ -233,6 +233,10 @@ TORCH_LIBRARY(tensorrt, m) { auto it = get_platform_name_map().find(Platform::PlatformEnum::kWIN_X86_64); return it->second; }); + m.def("_platform_win_arm64", []() -> std::string { + auto it = get_platform_name_map().find(Platform::PlatformEnum::kWIN_ARM64); + return it->second; + }); m.def("_platform_unknown", []() -> std::string { auto it = get_platform_name_map().find(Platform::PlatformEnum::kUNKNOWN); return it->second; diff --git a/core/util/BUILD b/core/util/BUILD index 22609d97ee..2f1885544c 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:arm64", "@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..d11b8bfe86 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:arm64", + "@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..3a765fc675 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:arm64", "@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..161afd23c2 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:arm64", "@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/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/packaging/build_windows_arm64_cross.bat b/packaging/build_windows_arm64_cross.bat new file mode 100644 index 0000000000..e98df94e08 --- /dev/null +++ b/packaging/build_windows_arm64_cross.bat @@ -0,0 +1,21 @@ +@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 +) +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 --windows-on-arm +exit /b %errorlevel% diff --git a/packaging/pre_build_script_windows.sh b/packaging/pre_build_script_windows.sh index 02abf38266..6ad6712c79 100644 --- a/packaging/pre_build_script_windows.sh +++ b/packaging/pre_build_script_windows.sh @@ -18,7 +18,8 @@ if [[ ${TENSORRT_VERSION} != "" ]]; then fi TORCH=$(grep "^torch>" py/requirements.txt) -INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION} +TORCH_INDEX_CU_VERSION=${TORCH_INDEX_CU_VERSION:-${CU_VERSION}} +INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${TORCH_INDEX_CU_VERSION} # The workflow installs torch before this script runs. Avoid uninstalling and # force-reinstalling it here: with the shortened Windows conda prefix, pip can diff --git a/packaging/prepare_windows_arm64_cross.ps1 b/packaging/prepare_windows_arm64_cross.ps1 new file mode 100644 index 0000000000..bef22a0369 --- /dev/null +++ b/packaging/prepare_windows_arm64_cross.ps1 @@ -0,0 +1,87 @@ +param( + [Parameter(Mandatory = $true)] [string] $PyTorchArtifact, + [Parameter(Mandatory = $true)] [string] $TensorRTArtifact, + [Parameter(Mandatory = $true)] [string] $TargetRoot, + [Parameter(Mandatory = $true)] [string] $CudaRoot, + [Parameter(Mandatory = $true)] [string] $PythonRoot +) + +$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:PYTHON_VERSION -ne '3.13') { + throw "Windows ARM64 RTX supports only Python 3.13; got '$env:PYTHON_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" +} +$cudaVersionFile = Join-Path $CudaRoot "version.json" +if (-not (Test-Path -LiteralPath $cudaVersionFile -PathType Leaf)) { + throw "CUDA version file is missing: $cudaVersionFile" +} +$cudaVersion = (Get-Content -LiteralPath $cudaVersionFile -Raw | ConvertFrom-Json).cuda.version +if ($cudaVersion -notlike '13.4*') { + throw "Windows ARM64 builds require CUDA 13.4 Preview; got '$cudaVersion'" +} +$pythonHeader = Join-Path $PythonRoot "include\Python.h" +$pythonImportLibrary = Join-Path $PythonRoot "libs\python313.lib" +foreach ($pythonFile in @($pythonHeader, $pythonImportLibrary)) { + if (-not (Test-Path -LiteralPath $pythonFile -PathType Leaf)) { + throw "ARM64 Python 3.13 development file is missing: $pythonFile" + } +} +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_TARGET_TORCH_ROOT" $torchRoot +Export-GitHubEnvironment "TORCHTRT_TARGET_TRT_ROOT" $trtRoot +Export-GitHubEnvironment "TORCHTRT_TARGET_CUDA_ROOT" $CudaRoot +Export-GitHubEnvironment "TORCHTRT_TARGET_PYTHON_ROOT" $PythonRoot +Write-Host "Prepared Windows ARM64 target sysroot: PyTorch=$torchRoot TensorRT-RTX=$trtRoot CUDA=$CudaRoot Python=$PythonRoot" diff --git a/packaging/validate_windows_arm64_wheel.ps1 b/packaging/validate_windows_arm64_wheel.ps1 new file mode 100644 index 0000000000..dffe48df29 --- /dev/null +++ b/packaging/validate_windows_arm64_wheel.ps1 @@ -0,0 +1,38 @@ +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" + Join-Path $env:TORCHTRT_TARGET_PYTHON_ROOT "libs\python313.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/py/torch_tensorrt/_enums.py b/py/torch_tensorrt/_enums.py index c91d6cb094..46bd51987d 100644 --- a/py/torch_tensorrt/_enums.py +++ b/py/torch_tensorrt/_enums.py @@ -1377,6 +1377,13 @@ class Platform(Enum): UNKNOWN = auto() + WIN_ARM64 = auto() + """ + OS: Windows, CPU Arch: ARM64 + + :meta hide-value: + """ + @classmethod def current_platform(cls) -> Platform: """ @@ -1406,11 +1413,13 @@ def current_platform(cls) -> Platform: # Windows... machine = platform.machine().lower() python_platform = sysconfig.get_platform().lower() - if ( - machine.startswith("amd64") - or machine.startswith("x86_64") - or python_platform in ("win-amd64", "win_amd64") - ): + if python_platform in ("win-arm64", "win_arm64"): + return Platform.WIN_ARM64 + elif python_platform in ("win-amd64", "win_amd64"): + return Platform.WIN_X86_64 + elif machine.startswith(("arm64", "aarch64")): + return Platform.WIN_ARM64 + elif machine.startswith(("amd64", "x86_64")): return Platform.WIN_X86_64 return Platform.UNKNOWN @@ -1427,6 +1436,8 @@ def _to_serialized_rt_platform(self) -> str: return cast(str, torch.ops.tensorrt._platform_linux_aarch64()) elif self == Platform.WIN_X86_64: return cast(str, torch.ops.tensorrt._platform_win_x86_64()) + elif self == Platform.WIN_ARM64: + return cast(str, torch.ops.tensorrt._platform_win_arm64()) return cast(str, torch.ops.tensorrt._platform_unknown()) @@ -1439,4 +1450,6 @@ def _from_serialized_rt_platform(cls, val: str) -> Platform: return cls.LINUX_AARCH64 elif val == torch.ops.tensorrt._platform_win_x86_64(): return cls.WIN_X86_64 + elif val == torch.ops.tensorrt._platform_win_arm64(): + return cls.WIN_ARM64 return cls.UNKNOWN diff --git a/setup.py b/setup.py index 152d370b37..7768b6b2d9 100644 --- a/setup.py +++ b/setup.py @@ -89,7 +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") +FORCED_WINDOWS_ARM64 = ( + WINDOWS_ON_ARM + or os.environ.get("TORCHTRT_TARGET_PLATFORM", "").lower() == "windows-arm64" +) +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 @@ -234,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: @@ -245,6 +263,12 @@ def build_libtorchtrt_cxx11_abi( if IS_WINDOWS: 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" + ) else: cmd.append("--config=linux") @@ -281,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 @@ -415,6 +443,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 +459,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 +730,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 +757,51 @@ def run(self): .split("/BUILD.bazel")[0] ) + 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 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 builds require Torch, CUDA, and Python roots; set " + "TORCHTRT_TARGET_TORCH_ROOT, TORCHTRT_TARGET_CUDA_ROOT, and " + "TORCHTRT_TARGET_PYTHON_ROOT" + ) + + extension_kwargs = {} + if TARGET_WINDOWS_ARM64: + 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", + "torch_cpu", + "torch_python", + "c10_cuda", + "torch_cuda", + "cudart", + ], + } + ext_modules += [ - CUDAExtension( + extension_type( "torch_tensorrt._C", [ "py/" + f @@ -720,11 +812,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 +828,23 @@ def run(self): dir_path + "/../cpp/include", "/usr/local/cuda", ] + + ( + [ + os.path.join(target_python_root, "include"), + 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 TARGET_WINDOWS_ARM64 + else [] + ) + ( [ dir_path + "/../bazel-TRTorch/external/tensorrt_win/include", @@ -939,7 +1051,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/py/core/test_classes.py b/tests/py/core/test_classes.py index d87b9e3d55..8d0a1e68d2 100644 --- a/tests/py/core/test_classes.py +++ b/tests/py/core/test_classes.py @@ -67,6 +67,16 @@ def test_current_platform(self): cpp_plat_str = torch.ops.tensorrt.get_current_platform() self.assertEqual(py_plat_str, cpp_plat_str) + def test_current_platform_detects_windows_arm64_runtime(self): + with mock.patch.object( + torch.ops.tensorrt, + "get_current_platform", + return_value=torch.ops.tensorrt._platform_win_arm64(), + ): + self.assertEqual( + torchtrt.Platform.current_platform(), torchtrt.Platform.WIN_ARM64 + ) + def test_current_platform_prefers_loaded_runtime_platform(self): with mock.patch("platform.system", return_value="Windows"), mock.patch( "platform.machine", return_value="ARM64" @@ -94,6 +104,29 @@ def test_current_platform_python_fallback_uses_python_windows_arch(self): torchtrt.Platform.current_platform(), torchtrt.Platform.WIN_X86_64 ) + def test_current_platform_python_fallback_detects_windows_arm64(self): + disabled_runtime_features = torchtrt_enums.ENABLED_FEATURES._replace( + torch_tensorrt_runtime=False + ) + with mock.patch.object( + torchtrt_enums, "ENABLED_FEATURES", disabled_runtime_features + ), mock.patch("platform.system", return_value="Windows"), mock.patch( + "platform.machine", return_value="ARM64" + ), mock.patch( + "sysconfig.get_platform", return_value="win-arm64" + ): + self.assertEqual( + torchtrt.Platform.current_platform(), torchtrt.Platform.WIN_ARM64 + ) + + def test_windows_arm64_platform_serialization_round_trip(self): + serialized = torchtrt.Platform.WIN_ARM64._to_serialized_rt_platform() + self.assertEqual(serialized, torch.ops.tensorrt._platform_win_arm64()) + self.assertEqual( + torchtrt.Platform._from_serialized_rt_platform(serialized), + torchtrt.Platform.WIN_ARM64, + ) + def test_unknown_platform_serializes_to_runtime_token(self): self.assertEqual( torchtrt.Platform.UNKNOWN._to_serialized_rt_platform(), diff --git a/tests/util/BUILD b/tests/util/BUILD index 6a4a18f51c..1dba512622 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:arm64", "@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..204b9cee23 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:arm64", + "@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..488bf0e074 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:arm64", + "@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..53715b585a 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:arm64", "@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..73358b1863 100644 --- a/toolchains/BUILD +++ b/toolchains/BUILD @@ -8,6 +8,14 @@ platform( ], ) +platform( + name = "windows_arm64", + constraint_values = [ + "@platforms//os:windows", + "@platforms//cpu:arm64", + ], +) + 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..eff946660e 100644 --- a/toolchains/ci_workspaces/MODULE.bazel.tmpl +++ b/toolchains/ci_workspaces/MODULE.bazel.tmpl @@ -73,6 +73,14 @@ 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", + path = "${TORCHTRT_TARGET_CUDA_ROOT}", +) + http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") @@ -158,6 +166,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 +190,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}",