diff --git a/.github/workflows/build_windows.yml b/.github/workflows/build_windows.yml index 13a70d4ee6..b16d68b900 100644 --- a/.github/workflows/build_windows.yml +++ b/.github/workflows/build_windows.yml @@ -105,6 +105,14 @@ on: required: false default: false type: boolean + cross-compile: + description: 'Build ARM64 from an x64 Windows host' + default: false + type: boolean + skip-binary-upload: + description: 'Skip the test-infra binary upload job' + default: false + type: boolean secrets: R2_ACCOUNT_ID: description: Cloudflare R2 account ID @@ -186,7 +194,7 @@ jobs: setup-miniconda: false python-version: ${{ env.PYTHON_VERSION }} cuda-version: ${{ env.CU_VERSION }} - arch: ${{ inputs.architecture }} + arch: ${{ inputs.cross-compile && 'arm64' || inputs.architecture }} - name: Shorten Conda environment path if: inputs.architecture == 'x64' run: | @@ -233,6 +241,22 @@ jobs: source "${BUILD_ENV_FILE}" # shellcheck disable=SC2086 ${CONDA_RUN} ${PIP_INSTALL_TORCH} + - name: Install CUDA 13.4 toolkit with ARM64 target libraries + if: inputs.cross-compile + shell: powershell + env: + CUDA_INSTALLER_URL: https://packages.nvidia.com/prerelease/cuda/13.4.0/local_installers/cuda_13.4.0_windows_x86_64.exe + CUDA_ROOT: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.4 + run: | + $installer = Join-Path $env:RUNNER_TEMP "cuda_13.4.0_windows_x86_64.exe" + Invoke-WebRequest -Uri $env:CUDA_INSTALLER_URL -OutFile $installer + $process = Start-Process -FilePath $installer -ArgumentList "-s" -Wait -PassThru + if ($process.ExitCode -ne 0) { throw "CUDA 13.4 installer failed with exit code $($process.ExitCode)" } + $cudaArm64Lib = Join-Path $env:CUDA_ROOT "lib\arm64\cudart.lib" + if (-not (Test-Path -LiteralPath $cudaArm64Lib -PathType Leaf)) { throw "CUDA 13.4 ARM64 target library was not installed: $cudaArm64Lib" } + "CUDA_PATH=$env:CUDA_ROOT" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append + "CUDA_HOME=$env:CUDA_ROOT" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append + (Join-Path $env:CUDA_ROOT "bin") | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append - name: Run Pre-Script with Caching if: ${{ inputs.pre-script != '' && inputs.architecture == 'x64' }} uses: ./test-infra/.github/actions/run-script-with-cache @@ -242,6 +266,58 @@ jobs: repository: ${{ inputs.repository }} script: ${{ inputs.pre-script }} is_windows: 'enabled' + - name: Prepare ARM64 target sysroot + if: inputs.cross-compile + working-directory: ${{ inputs.repository }} + shell: powershell + run: | + ./packaging/prepare_windows_arm64_cross.ps1 ` + -TargetRoot '${{ runner.temp }}\torchtrt-windows-arm64' ` + -CudaRoot $env:CUDA_PATH ` + -PythonExecutable (Join-Path $env:CONDA_ENV 'python.exe') + - name: Ensure MSVC ARM64 cross tools + if: inputs.cross-compile + shell: powershell + run: | + $vswhere = Join-Path ${env:ProgramFiles(x86)} "Microsoft Visual Studio/Installer/vswhere.exe" + if (-not (Test-Path -LiteralPath $vswhere -PathType Leaf)) { throw "Visual Studio locator not found: $vswhere" } + $vsInstall = (& $vswhere -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -property installationPath | Select-Object -First 1).Trim() + if ([string]::IsNullOrWhiteSpace($vsInstall)) { throw "No Visual Studio instance with x64 MSVC tools was found" } + $vcvarsall = [IO.Path]::Combine($vsInstall, "VC", "Auxiliary", "Build", "vcvarsall.bat") + if (-not (Test-Path -LiteralPath $vcvarsall -PathType Leaf)) { throw "vcvarsall.bat not found: $vcvarsall" } + $arm64Compiler = Get-ChildItem -Path (Join-Path $vsInstall "VC/Tools/MSVC/*/bin/Hostx64/arm64/cl.exe") -File -ErrorAction SilentlyContinue | Select-Object -First 1 + if (-not $arm64Compiler) { + $installStopwatch = [Diagnostics.Stopwatch]::StartNew() + $setup = Join-Path ${env:ProgramFiles(x86)} "Microsoft Visual Studio/Installer/setup.exe" + if (-not (Test-Path -LiteralPath $setup -PathType Leaf)) { throw "Visual Studio Installer not found: $setup" } + & $setup modify --installPath $vsInstall --add Microsoft.VisualStudio.Component.VC.Tools.ARM64 --passive --norestart + if ($LASTEXITCODE -ne 0) { throw "Installing the MSVC ARM64 component failed with exit code $LASTEXITCODE" } + $deadline = (Get-Date).AddMinutes(20) + while (-not $arm64Compiler -and (Get-Date) -lt $deadline) { + Write-Host "Waiting for the MSVC ARM64 compiler to be installed ($([Math]::Floor($installStopwatch.Elapsed.TotalSeconds)) seconds elapsed; next check in 30 seconds)..." + Start-Sleep -Seconds 30 + $arm64Compiler = Get-ChildItem -Path (Join-Path $vsInstall "VC/Tools/MSVC/*/bin/Hostx64/arm64/cl.exe") -File -ErrorAction SilentlyContinue | Select-Object -First 1 + } + $installStopwatch.Stop() + if ($arm64Compiler) { Write-Host "MSVC ARM64 compiler became available after $([Math]::Round($installStopwatch.Elapsed.TotalSeconds, 1)) seconds" } + $arm64Compiler = Get-ChildItem -Path (Join-Path $vsInstall "VC/Tools/MSVC/*/bin/Hostx64/arm64/cl.exe") -File -ErrorAction SilentlyContinue | Select-Object -First 1 + } + if (-not $arm64Compiler) { throw "MSVC ARM64 compiler was not installed after $([Math]::Round($installStopwatch.Elapsed.TotalSeconds, 1)) seconds" } + "TORCHTRT_VCVARSALL=$vcvarsall" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append + - name: Cross-compile Windows ARM64 wheel + if: inputs.cross-compile + working-directory: ${{ inputs.repository }} + shell: cmd + run: | + call "%TORCHTRT_VCVARSALL%" amd64_arm64 + %CONDA_RUN% cmd /c packaging\build_windows_arm64_cross.bat + - name: Validate Windows ARM64 wheel + if: inputs.cross-compile + working-directory: ${{ inputs.repository }} + shell: cmd + run: | + call "%TORCHTRT_VCVARSALL%" amd64_arm64 + powershell -NoProfile -ExecutionPolicy Bypass -File packaging\validate_windows_arm64_wheel.ps1 -DistDirectory dist - name: Run Pre-Script Arm64 if: ${{ inputs.pre-script != '' && inputs.architecture == 'arm64' }} env: @@ -278,7 +354,7 @@ jobs: - name: Build clean # Only relevant for setup-py builds; python-build-package projects may # not have a setup.py (e.g. scikit-build-core based builds). - if: inputs.build-platform == 'setup-py' && inputs.architecture == 'x64' + if: inputs.build-platform == 'setup-py' && inputs.architecture == 'x64' && inputs.cross-compile == false working-directory: ${{ inputs.repository }} env: ENV_SCRIPT: ${{ inputs.env-script }} @@ -295,7 +371,7 @@ jobs: fi fi - name: Set PYTORCH_VERSION on x64 - if: inputs.architecture == 'x64' + if: inputs.architecture == 'x64' && inputs.cross-compile == false working-directory: ${{ inputs.repository }} run: | source "${BUILD_ENV_FILE}" @@ -307,7 +383,7 @@ jobs: export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//')" fi - name: Build the wheel (python-build-package) X64 - if: ${{ inputs.build-platform == 'python-build-package' && inputs.architecture == 'x64' }} + if: ${{ inputs.build-platform == 'python-build-package' && inputs.architecture == 'x64' && inputs.cross-compile == false }} working-directory: ${{ inputs.repository }} env: ENV_SCRIPT: ${{ inputs.env-script }} @@ -322,7 +398,7 @@ jobs: ${CONDA_RUN} ${ENV_SCRIPT} ${{ inputs.build-command }} fi - name: Build the wheel (setup-py) X64 - if: ${{ inputs.build-platform == 'setup-py' && inputs.architecture == 'x64' }} + if: ${{ inputs.build-platform == 'setup-py' && inputs.architecture == 'x64' && inputs.cross-compile == false }} working-directory: ${{ inputs.repository }} env: ENV_SCRIPT: ${{ inputs.env-script }} @@ -398,13 +474,13 @@ jobs: env: POST_SCRIPT: ${{ inputs.post-script }} ENV_SCRIPT: ${{ inputs.env-script }} - if: ${{ inputs.post-script != '' && inputs.architecture == 'x64'}} + if: ${{ inputs.post-script != '' && inputs.architecture == 'x64' && inputs.cross-compile == false }} run: | set -euxo pipefail source "${BUILD_ENV_FILE}" ${CONDA_RUN} ${ENV_SCRIPT} ${POST_SCRIPT} - name: Smoke Test X64 - if: ${{ inputs.architecture == 'x64' && inputs.is-release-tarball == false }} + if: ${{ inputs.architecture == 'x64' && inputs.cross-compile == false && inputs.is-release-tarball == false }} env: ENV_SCRIPT: ${{ inputs.env-script }} PACKAGE_NAME: ${{ inputs.package-name }} @@ -454,15 +530,20 @@ jobs: fi # NB: Only upload to GitHub after passing smoke tests - name: Get Artifact name - if: inputs.architecture == 'arm64' + if: inputs.architecture == 'arm64' || inputs.cross-compile env: REPOSITORY: ${{ inputs.repository }} REF: ${{ inputs.ref }} PYTHON_VERSION: ${{ matrix.python_version }} CU_VERSION: ${{ env.CU_VERSION }} - ARCH: ${{ inputs.architecture }} + ARCH: ${{ inputs.cross-compile && 'arm64' || inputs.architecture }} run: | - echo "ARTIFACT_NAME=${REPOSITORY//\//_}_${REF//\//_}_${PYTHON_VERSION}_${CU_VERSION}_${ARCH}" >> "${GITHUB_ENV}" + artifact_cuda_version="${CU_VERSION}" + # The cross-compiled ARM64 sysroot is CUDA 13.4; retain the host CUDA label otherwise. + if [[ "${{ inputs.cross-compile }}" == "true" ]]; then + artifact_cuda_version="cu134" + fi + echo "ARTIFACT_NAME=${REPOSITORY//\//_}_${REF//\//_}_${PYTHON_VERSION}_${artifact_cuda_version}_${ARCH}" >> "${GITHUB_ENV}" - name: Upload wheel to GitHub if: ${{ inputs.is-release-tarball == false }} continue-on-error: true @@ -478,14 +559,14 @@ jobs: name: cxx11-zip-${{ env.PYTHON_VERSION }}-${{ env.CU_VERSION }} path: ${{ inputs.repository }}/release/tarball/ - uses: ./test-infra/.github/actions/teardown-windows - if: inputs.architecture == 'x64' + if: inputs.architecture == 'x64' || inputs.cross-compile name: Teardown Windows upload: needs: build uses: pytorch/test-infra/.github/workflows/_binary_upload.yml@main # if it is the release wheel or tarball, then skip upload to pytorch index(release wheel are uploaded to pypi) - if: ${{ inputs.is-release-wheel == false && inputs.is-release-tarball == false }} + if: ${{ inputs.skip-binary-upload == false && inputs.is-release-wheel == false && inputs.is-release-tarball == false }} with: repository: ${{ inputs.repository }} ref: ${{ inputs.ref }} diff --git a/.github/workflows/ci-linux-x86_64.yml b/.github/workflows/ci-linux-x86_64.yml index 11c1d03723..8359314daa 100644 --- a/.github/workflows/ci-linux-x86_64.yml +++ b/.github/workflows/ci-linux-x86_64.yml @@ -5,8 +5,8 @@ name: CI Linux x86_64 # come from the shared _decide reusable. on: - pull_request: - types: [opened, synchronize, reopened, labeled] + #pull_request: + #types: [opened, synchronize, reopened, labeled] pull_request_review: types: [submitted] push: diff --git a/.github/workflows/ci-sbsa.yml b/.github/workflows/ci-sbsa.yml index 30d242f50c..0917afc64b 100644 --- a/.github/workflows/ci-sbsa.yml +++ b/.github/workflows/ci-sbsa.yml @@ -7,8 +7,8 @@ name: CI SBSA # RTX channels build @tensorrt_rtx_sbsa via USE_TRT_RTX=1.) on: - pull_request: - types: [opened, synchronize, reopened, labeled] + #pull_request: + #types: [opened, synchronize, reopened, labeled] pull_request_review: types: [submitted] push: diff --git a/.github/workflows/ci-windows-arm64.yml b/.github/workflows/ci-windows-arm64.yml new file mode 100644 index 0000000000..e48aa660b9 --- /dev/null +++ b/.github/workflows/ci-windows-arm64.yml @@ -0,0 +1,53 @@ +name: CI Windows ARM64 + +on: + pull_request: + types: [opened, synchronize, reopened, labeled] + pull_request_review: + types: [submitted] + push: + branches: [main] + schedule: + - cron: "0 7 * * *" + workflow_dispatch: +concurrency: + # Keep one active ARM64 run per PR/ref; CI-control labels supersede the prior run, + # while unrelated label churn remains isolated. + group: "${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}${{ (github.event.action == 'labeled' || github.event.action == 'unlabeled') && !contains(fromJSON('[\"ci: full\", \"ci: nightly\", \"backend: TensorRT\", \"backend: TensorRT-RTX\"]'), github.event.label.name) && format('-label-{0}', github.event.label.name) || '' }}" + cancel-in-progress: true + + +permissions: + id-token: write + contents: read + +jobs: + + build: + uses: ./.github/workflows/build_windows.yml + with: + repository: ${{ github.repository }} + ref: ${{ github.sha }} + build-matrix: >- + {"include":[{"build_name":"windows-arm64-cross","python_version":"3.13","desired_cuda":"cu132","validation_runner":"windows.g5.4xlarge.nvidia.gpu"}]} + architecture: x64 + cross-compile: true + pre-script: packaging/pre_build_script_windows.sh + use-rtx: true + timeout: 180 + skip-binary-upload: true + + gate: + needs: [build] + if: always() + runs-on: ubuntu-latest + steps: + - name: Gate Windows ARM64 builds + env: + BUILD_RESULT: ${{ needs.build.result }} + run: | + set -euo pipefail + if [[ "${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..e4e1a99a3a 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,16 @@ http_archive( urls = ["https://download.pytorch.org/libtorch/nightly/cu130/libtorch-win-shared-with-deps-latest.zip"], ) +# Pinned Windows ARM64 PyTorch wheel used by the RTX cross-build. +http_archive( + name = "libtorch_win_arm64", + build_file = "@//third_party/libtorch:BUILD", + sha256 = "23862a93476cb038ffd26f7141cd476717d97b69b50074f2ab14036eb6093200", + strip_prefix = "torch", + type = "zip", + urls = ["https://pypi.nvidia.com/nvtorch_oot_nightly/torch/torch-2.14.0.dev20260728%2Bcu134-cp313-cp313-win_arm64.whl"], +) + http_archive( name = "torch_l4t", build_file = "@//third_party/libtorch:BUILD", @@ -241,3 +261,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..b7b490d06f --- /dev/null +++ b/docsrc/getting_started/windows_arm64.rst @@ -0,0 +1,198 @@ +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 2.14 package containing the C++ headers + and libraries under ``torch\include`` and ``torch\lib``. Nightly wheels are + available from NVIDIA's ``nvtorch_oot_nightly`` index at + ``https://pypi.nvidia.com/nvtorch_oot_nightly/torch/``; +* 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 a command prompt. You can + also initialize one from ``cmd.exe``: + + .. code-block:: bat + # confirm your visual studio version and path + call "%ProgramFiles%\Microsoft Visual Studio\18\Community\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 + + # Install the Windows ARM64 PyTorch 2.14 nightly. + python -m pip install --pre "torch>=2.14.0.dev,<2.15.0" --index-url https://pypi.nvidia.com/nvtorch_oot_nightly/torch/ + + # 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 2.14 installation. Download the ARM64 + wheel from ``https://pypi.nvidia.com/nvtorch_oot_nightly/torch/`` and + install or extract it to create this target root. + * ``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 + # this is the target root for the pytorch windows on arm64 artifacts + set "TORCHTRT_TARGET_TORCH_ROOT=C:/torchtrt-arm64-artifacts/pytorch/torch" + + # this is the cuda root for the cuda toolkit installed on windows x86 host machine + set "TORCHTRT_TARGET_CUDA_ROOT=C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v13.4/" + + # this is the python root for the python 3.13 on windows for arm64 + set "TORCHTRT_TARGET_PYTHON_ROOT=C:\torchtrt-arm64-artifacts\python313-arm64" + + # set the cuda path and home path + 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..9c8170ee00 --- /dev/null +++ b/packaging/build_windows_arm64_cross.bat @@ -0,0 +1,29 @@ +@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 "%VCToolsInstallDir%"=="" ( + echo VCToolsInstallDir is unset; initialize MSVC with vcvarsall.bat amd64_arm64 first + exit /b 1 +) +rem Bazel's VS discovery can select a different, incomplete installation. Pin it +rem to the VC root selected by the active amd64_arm64 vcvarsall environment. +for %%I in ("%VCToolsInstallDir%\..\..\..") do set "BAZEL_VC=%%~fI" +echo Using BAZEL_VC=%BAZEL_VC% +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..8ead757957 --- /dev/null +++ b/packaging/prepare_windows_arm64_cross.ps1 @@ -0,0 +1,96 @@ +param( + [string] $PyTorchArtifact = "https://pypi.nvidia.com/nvtorch_oot_nightly/torch/torch-2.14.0.dev20260728%2Bcu134-cp313-cp313-win_arm64.whl", + [string] $PythonArtifact = "https://api.nuget.org/v3-flatcontainer/pythonarm64/3.13.0/pythonarm64.3.13.0.nupkg", + [Parameter(Mandatory = $true)] [string] $TargetRoot, + [Parameter(Mandatory = $true)] [string] $CudaRoot, + [string] $PythonExecutable = "python" +) + +$ErrorActionPreference = "Stop" + +$PyTorchArtifactSha256 = "23862a93476cb038ffd26f7141cd476717d97b69b50074f2ab14036eb6093200" +$PythonArtifactSha256 = "f44428dc94e6f9c72cd69ad6436280784e6f9eed46a149641fee71866d3081f3" + +$ArtifactTempRoot = $env:RUNNER_TEMP +if (-not $ArtifactTempRoot) { + $ArtifactTempRoot = [System.IO.Path]::GetTempPath() +} +function Expand-ZipArtifact { + param([string] $Artifact, [string] $Destination, [string] $Name, [string] $ExpectedSha256 = "") + $source = $Artifact + if ($Artifact -match '^https?://') { + $source = Join-Path $ArtifactTempRoot "$Name.zip" + Invoke-WebRequest -Uri $Artifact -OutFile $source + } + if (-not (Test-Path -LiteralPath $source -PathType Leaf)) { + throw "$Name artifact was not found: $source" + } + if ($ExpectedSha256) { + $actualSha256 = (Get-FileHash -LiteralPath $source -Algorithm SHA256).Hash.ToLowerInvariant() + if ($actualSha256 -ne $ExpectedSha256) { + throw "$Name artifact SHA256 mismatch: expected $ExpectedSha256, got $actualSha256" + } + } + New-Item -ItemType Directory -Force -Path $Destination | Out-Null + & $PythonExecutable -m zipfile -e $source $Destination + if ($LASTEXITCODE -ne 0) { + throw "Failed to extract $Name artifact" + } +} + +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" +} + +New-Item -ItemType Directory -Force -Path $TargetRoot | Out-Null + +$torchExtract = Join-Path $TargetRoot "pytorch" +if (-not (Test-Path -LiteralPath $torchExtract -PathType Container)) { + Expand-ZipArtifact -Artifact $PyTorchArtifact -Destination $torchExtract -Name "pytorch-arm64" -ExpectedSha256 $PyTorchArtifactSha256 +} + +$torchRoot = Find-Root -SearchRoot $torchExtract -IncludePath "include" -LibraryPath "lib" -Name "PyTorch" +$cudaArm64Lib = Join-Path $CudaRoot "lib\arm64\cudart.lib" +$pythonExtract = Join-Path $TargetRoot "pythonarm64" +if (-not (Test-Path -LiteralPath $pythonExtract -PathType Container)) { + Expand-ZipArtifact -Artifact $PythonArtifact -Destination $pythonExtract -Name "pythonarm64" -ExpectedSha256 $PythonArtifactSha256 +} +$pythonRoot = Find-Root -SearchRoot $pythonExtract -IncludePath "include\Python.h" -LibraryPath "libs\python313.lib" -Name "Python ARM64" +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" + } +} +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_CUDA_ROOT" $CudaRoot +Export-GitHubEnvironment "TORCHTRT_TARGET_PYTHON_ROOT" $pythonRoot +Write-Host "Prepared Windows ARM64 target sysroot: TORCHTRT_TARGET_TORCH_ROOT=$torchRoot TORCHTRT_TARGET_CUDA_ROOT=$CudaRoot TORCHTRT_TARGET_PYTHON_ROOT=$pythonRoot" diff --git a/packaging/validate_windows_arm64_wheel.ps1 b/packaging/validate_windows_arm64_wheel.ps1 new file mode 100644 index 0000000000..b29db58a85 --- /dev/null +++ b/packaging/validate_windows_arm64_wheel.ps1 @@ -0,0 +1,36 @@ +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) +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..6777a5858e 100644 --- a/toolchains/ci_workspaces/MODULE.bazel.tmpl +++ b/toolchains/ci_workspaces/MODULE.bazel.tmpl @@ -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, @@ -73,6 +76,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 +169,15 @@ http_archive( ], ) +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", + ], +) + #################################################################################### # Locally installed dependencies (use in cases of custom dependencies or aarch64) #################################################################################### @@ -176,6 +196,15 @@ new_local_repository( build_file = "third_party/libtorch/BUILD" ) +http_archive( + name = "libtorch_win_arm64", + build_file = "third_party/libtorch/BUILD", + sha256 = "23862a93476cb038ffd26f7141cd476717d97b69b50074f2ab14036eb6093200", + strip_prefix = "torch", + type = "zip", + urls = ["https://pypi.nvidia.com/nvtorch_oot_nightly/torch/torch-2.14.0.dev20260728%2Bcu134-cp313-cp313-win_arm64.whl"], +) + new_local_repository( name = "torch_l4t", path = "${TORCH_INSTALL_PATH}",