From 0e01ac5091b4f32fd92f2175c20625697ec56315 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Fri, 11 Apr 2025 14:28:47 -0600 Subject: [PATCH 1/8] infra: Enabling aarch64 builds --- .github/scripts/filter-matrix.py | 7 +- .../scripts/generate_binary_build_matrix.py | 82 ++-- .../workflows/build-test-linux-aarch64.yml | 350 ++++++++++++++++++ ...-linux.yml => build-test-linux-x86_64.yml} | 6 +- .../workflows/build-test-tensorrt-linux.yml | 4 +- .github/workflows/build-test-windows.yml | 2 +- .../generate_binary_build_matrix.yml | 9 +- MODULE.bazel | 45 +-- packaging/pre_build_script.sh | 14 +- pyproject.toml | 8 +- version.txt | 2 +- 11 files changed, 444 insertions(+), 85 deletions(-) create mode 100644 .github/workflows/build-test-linux-aarch64.yml rename .github/workflows/{build-test-linux.yml => build-test-linux-x86_64.yml} (98%) diff --git a/.github/scripts/filter-matrix.py b/.github/scripts/filter-matrix.py index 3710539f59..685ea98985 100644 --- a/.github/scripts/filter-matrix.py +++ b/.github/scripts/filter-matrix.py @@ -1,13 +1,15 @@ #!/usr/bin/env python3 +from typing import List import argparse import json import sys disabled_python_versions = "3.13" +cpu_tag = "cpu" -def main(args: list[str]) -> None: +def main(args: List[str]) -> None: parser = argparse.ArgumentParser() parser.add_argument( "--matrix", @@ -24,7 +26,8 @@ def main(args: list[str]) -> None: includes = matrix_dict["include"] filtered_includes = [] for item in includes: - if item["python_version"] not in disabled_python_versions: + if all([item["python_version"] not in disabled_python_versions, + item["desired_cuda"] != cpu_tag]): filtered_includes.append(item) filtered_matrix_dict = {} filtered_matrix_dict["include"] = filtered_includes diff --git a/.github/scripts/generate_binary_build_matrix.py b/.github/scripts/generate_binary_build_matrix.py index 52063ce9fc..b7e038f98f 100644 --- a/.github/scripts/generate_binary_build_matrix.py +++ b/.github/scripts/generate_binary_build_matrix.py @@ -2,6 +2,9 @@ """Generates a matrix to be utilized through github actions +Important. After making changes to this file please run following command: +python -m tools.tests.test_generate_binary_build_matrix --update-reference-files + Will output a condensed version of the matrix if on a pull request that only includes the latest version of python we support built on four different architectures: @@ -11,13 +14,13 @@ * Latest XPU """ - import argparse import json import os import sys from typing import Any, Callable, Dict, List, Optional, Tuple + PYTHON_ARCHES_DICT = { "nightly": ["3.9", "3.10", "3.11", "3.12"], "test": ["3.9", "3.10", "3.11", "3.12"], @@ -26,16 +29,30 @@ CUDA_ARCHES_DICT = { "nightly": ["11.8", "12.6", "12.8"], "test": ["11.8", "12.6", "12.8"], - "release": ["11.8", "12.6", "12.8"], + "release": ["11.8", "12.4", "12.6"], } ROCM_ARCHES_DICT = { - "nightly": ["6.1", "6.2"], - "test": ["6.1", "6.2"], - "release": ["6.1", "6.2"], + "nightly": ["6.2.4", "6.3"], + "test": ["6.2.4", "6.3"], + "release": ["6.1", "6.2.4"], +} + +CUDA_CUDNN_VERSIONS = { + "11.8": {"cuda": "11.8.0", "cudnn": "9"}, + "12.4": {"cuda": "12.4.1", "cudnn": "9"}, + "12.6": {"cuda": "12.6.3", "cudnn": "9"}, + "12.8": {"cuda": "12.8.0", "cudnn": "9"}, +} + +STABLE_CUDA_VERSIONS = { + "nightly": "12.6", + "test": "12.6", + "release": "12.4", } +CUDA_AARCH64_ARCHES = ["12.8-aarch64", "12.6-aarch64", "11.8-aarch64"] + PACKAGE_TYPES = ["wheel", "conda", "libtorch"] -PRE_CXX11_ABI = "pre-cxx11" CXX11_ABI = "cxx11-abi" RELEASE = "release" DEBUG = "debug" @@ -59,7 +76,7 @@ CURRENT_NIGHTLY_VERSION = "2.8.0" CURRENT_CANDIDATE_VERSION = "2.7.0" -CURRENT_STABLE_VERSION = "2.7.0" +CURRENT_STABLE_VERSION = "2.6.0" CURRENT_VERSION = CURRENT_STABLE_VERSION # By default use Nightly for CUDA arches @@ -94,7 +111,7 @@ def arch_type(arch_version: str) -> str: return ROCM elif arch_version == CPU_AARCH64: return CPU_AARCH64 - elif arch_version == CUDA_AARCH64: + elif arch_version in CUDA_AARCH64_ARCHES: return CUDA_AARCH64 elif arch_version == XPU: return XPU @@ -140,11 +157,14 @@ def initialize_globals(channel: str, build_python_only: bool) -> None: else: PYTHON_ARCHES = PYTHON_ARCHES_DICT[channel] WHEEL_CONTAINER_IMAGES = { - "11.8": "pytorch/manylinux2_28-builder:cuda11.8", - "12.1": "pytorch/manylinux2_28-builder:cuda12.1", - "12.4": "pytorch/manylinux2_28-builder:cuda12.4", - "12.6": "pytorch/manylinux2_28-builder:cuda12.6", - "12.8": "pytorch/manylinux2_28-builder:cuda12.8", + **{ + gpu_arch: f"pytorch/manylinux2_28-builder:cuda{gpu_arch}" + for gpu_arch in CUDA_ARCHES + }, + **{ + gpu_arch: f"pytorch/manylinuxaarch64-builder:cuda{gpu_arch.replace('-aarch64', '')}" + for gpu_arch in CUDA_AARCH64_ARCHES + }, **{ gpu_arch: f"pytorch/manylinux2_28-builder:rocm{gpu_arch}" for gpu_arch in ROCM_ARCHES @@ -153,26 +173,17 @@ def initialize_globals(channel: str, build_python_only: bool) -> None: XPU: "pytorch/manylinux2_28-builder:xpu", # TODO: Migrate CUDA_AARCH64 image to manylinux2_28_aarch64-builder:cuda12.4 CPU_AARCH64: "pytorch/manylinux2_28_aarch64-builder:cpu-aarch64", - CUDA_AARCH64: "pytorch/manylinuxaarch64-builder:cuda12.4", + CUDA_AARCH64: "pytorch/manylinuxaarch64-builder:cuda12.6", } LIBTORCH_CONTAINER_IMAGES = { - **{ - (gpu_arch, PRE_CXX11_ABI): f"pytorch/manylinux2_28-builder:cuda{gpu_arch}" - for gpu_arch in CUDA_ARCHES - }, **{ (gpu_arch, CXX11_ABI): f"pytorch/libtorch-cxx11-builder:cuda{gpu_arch}" for gpu_arch in CUDA_ARCHES }, - **{ - (gpu_arch, PRE_CXX11_ABI): f"pytorch/manylinux2_28-builder:rocm{gpu_arch}" - for gpu_arch in ROCM_ARCHES - }, **{ (gpu_arch, CXX11_ABI): f"pytorch/libtorch-cxx11-builder:rocm{gpu_arch}" for gpu_arch in ROCM_ARCHES }, - (CPU, PRE_CXX11_ABI): "pytorch/manylinux2_28-builder:cpu", (CPU, CXX11_ABI): "pytorch/libtorch-cxx11-builder:cpu", } @@ -181,7 +192,7 @@ def translate_desired_cuda(gpu_arch_type: str, gpu_arch_version: str) -> str: return { CPU: "cpu", CPU_AARCH64: CPU, - CUDA_AARCH64: "cu124", + CUDA_AARCH64: f"cu{gpu_arch_version.replace('-aarch64', '').replace('.', '')}", CUDA: f"cu{gpu_arch_version.replace('.', '')}", ROCM: f"rocm{gpu_arch_version}", XPU: "xpu", @@ -272,7 +283,7 @@ def get_wheel_install_command( return f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL} --index-url {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}_pypi_pkg" # noqa: E501 else: raise ValueError( - "Split build is not supported for this configuration. It is only supported for CUDA 11.8, 12.4, 12.6, 12.8 on Linux nightly builds." # noqa: E501 + "Split build is not supported for this configuration. It is only supported for CUDA 11.8, 12.4, 12.6 on Linux nightly builds." # noqa: E501 ) if ( channel == RELEASE @@ -343,7 +354,7 @@ def generate_libtorch_matrix( if os == WINDOWS: abi_versions = [RELEASE, DEBUG] elif os == LINUX: - abi_versions = [PRE_CXX11_ABI, CXX11_ABI] + abi_versions = [CXX11_ABI] elif os in [MACOS_ARM64]: abi_versions = [CXX11_ABI] else: @@ -422,11 +433,6 @@ def generate_wheels_matrix( # Define default python version python_versions = list(PYTHON_ARCHES) - # If the list of python versions is set explicitly by the caller, stick with it instead - # of trying to add more versions behind the scene - if channel == NIGHTLY and (os in (LINUX, MACOS_ARM64, LINUX_AARCH64)): - python_versions += ["3.13"] - if os == LINUX: # NOTE: We only build manywheel packages for linux package_type = "manywheel" @@ -442,7 +448,11 @@ def generate_wheels_matrix( if os == LINUX_AARCH64: # Only want the one arch as the CPU type is different and # uses different build/test scripts - arches = [CPU_AARCH64, CUDA_AARCH64] + arches = [] + if with_cpu == ENABLE: + arches += [CPU_AARCH64] + elif with_cuda == ENABLE: + arches += CUDA_AARCH64_ARCHES if with_cuda == ENABLE: upload_to_base_bucket = "no" @@ -463,15 +473,15 @@ def generate_wheels_matrix( ret: List[Dict[str, Any]] = [] for python_version in python_versions: for arch_version in arches: - # TODO: Enable Python 3.13 support for ROCM - if arch_version in ROCM_ARCHES and python_version == "3.13": - continue - gpu_arch_type = arch_type(arch_version) gpu_arch_version = ( "" if arch_version in [CPU, CPU_AARCH64, XPU] else arch_version ) + # TODO: Enable python 3.13t on cpu-s390x or Windows + if (gpu_arch_type == "cpu-s390x") and python_version == "3.13t": + continue + desired_cuda = translate_desired_cuda(gpu_arch_type, gpu_arch_version) entry = { "python_version": python_version, diff --git a/.github/workflows/build-test-linux-aarch64.yml b/.github/workflows/build-test-linux-aarch64.yml new file mode 100644 index 0000000000..966ce31d04 --- /dev/null +++ b/.github/workflows/build-test-linux-aarch64.yml @@ -0,0 +1,350 @@ +name: Build and test Linux aarch64 wheels + +on: + pull_request: + push: + branches: + - main + - nightly + - release/* + tags: + # NOTE: Binary build pipelines should only get triggered on release candidate builds + # Release candidate tags look like: v1.11.0-rc1 + - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ + workflow_dispatch: + +jobs: + generate-matrix: + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: wheel + os: linux-aarch64 + test-infra-repository: pytorch/test-infra + test-infra-ref: main + with-rocm: false + with-cpu: false + python-versions: '["3.11", "3.12", "3.10", "3.9"]' + + filter-matrix: + needs: [generate-matrix] + outputs: + matrix: ${{ steps.generate.outputs.matrix }} + runs-on: ubuntu-latest + steps: + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + - uses: actions/checkout@v4 + with: + repository: pytorch/tensorrt + - name: Generate release matrix + id: generate + run: | + set -eou pipefail + MATRIX_BLOB=${{ toJSON(needs.generate-matrix.outputs.matrix) }} + MATRIX_BLOB="$(python3 .github/scripts/filter-matrix.py --matrix "${MATRIX_BLOB}")" + echo "${MATRIX_BLOB}" + echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}" + + build: + needs: filter-matrix + permissions: + id-token: write + contents: read + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + pre-script: packaging/pre_build_script.sh + env-var-script: packaging/env_vars.txt + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + package-name: torch_tensorrt + name: Build torch-tensorrt whl package + uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main + with: + repository: ${{ matrix.repository }} + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + env-var-script: ${{ matrix.env-var-script }} + post-script: ${{ matrix.post-script }} + package-name: ${{ matrix.package-name }} + smoke-test-script: ${{ matrix.smoke-test-script }} + trigger-event: ${{ github.event_name }} + architecture: "aarch64" + + tests-py-torchscript-fe: + name: Test torchscript frontend [Python] + needs: [filter-matrix, build] + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + package-name: torch_tensorrt + pre-script: packaging/pre_build_script.sh + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + uses: ./.github/workflows/linux-test.yml + with: + job-name: tests-py-torchscript-fe + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + architecture: "aarch64" + script: | + export USE_HOST_DEPS=1 + export CI_BUILD=1 + export LD_LIBRARY_PATH=/usr/lib64:$LD_LIBRARY_PATH + pushd . + cd tests/modules + python hub.py + popd + pushd . + cd tests/py/ts + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_api_test_results.xml api/ + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_models_test_results.xml models/ + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/ts_integrations_test_results.xml integrations/ + popd + + tests-py-dynamo-converters: + name: Test dynamo converters [Python] + needs: [filter-matrix, build] + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + package-name: torch_tensorrt + pre-script: packaging/pre_build_script.sh + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + uses: ./.github/workflows/linux-test.yml + with: + job-name: tests-py-dynamo-converters + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + architecture: "aarch64" + script: | + export USE_HOST_DEPS=1 + export CI_BUILD=1 + pushd . + cd tests/py + python -m pip install -r requirements.txt + cd dynamo + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml -n 4 conversion/ + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml automatic_plugin/test_automatic_plugin.py + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml automatic_plugin/test_automatic_plugin_with_attrs.py + popd + + tests-py-dynamo-fe: + name: Test dynamo frontend [Python] + needs: [filter-matrix, build] + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + package-name: torch_tensorrt + pre-script: packaging/pre_build_script.sh + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + uses: ./.github/workflows/linux-test.yml + with: + job-name: tests-py-dynamo-fe + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + architecture: "aarch64" + script: | + export USE_HOST_DEPS=1 + export CI_BUILD=1 + pushd . + cd tests/py + python -m pip install -r requirements.txt + cd dynamo + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dyn_models_export.xml --ir dynamo models/ + popd + + tests-py-dynamo-serde: + name: Test dynamo export serde [Python] + needs: [filter-matrix, build] + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + package-name: torch_tensorrt + pre-script: packaging/pre_build_script.sh + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + uses: ./.github/workflows/linux-test.yml + with: + job-name: tests-py-dynamo-serde + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + architecture: "aarch64" + script: | + export USE_HOST_DEPS=1 + export CI_BUILD=1 + pushd . + cd tests/py + python -m pip install -r requirements.txt + cd dynamo + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/export_serde_test_results.xml --ir dynamo models/test_export_serde.py + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/reexport_test_results.xml --ir dynamo models/test_reexport.py + popd + + tests-py-torch-compile-be: + name: Test torch compile backend [Python] + needs: [filter-matrix, build] + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + package-name: torch_tensorrt + pre-script: packaging/pre_build_script.sh + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + uses: ./.github/workflows/linux-test.yml + with: + job-name: tests-py-torch-compile-be + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + architecture: "aarch64" + script: | + export USE_HOST_DEPS=1 + export CI_BUILD=1 + pushd . + cd tests/py + python -m pip install -r requirements.txt + cd dynamo + python -m pytest -ra -n 10 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_be_test_results.xml backend/ + python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_complete_be_e2e_test_results.xml --ir torch_compile models/test_models.py + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/torch_compile_dyn_models_export.xml --ir torch_compile models/test_dyn_models.py + popd + + tests-py-dynamo-core: + name: Test dynamo core [Python] + needs: [filter-matrix, build] + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + package-name: torch_tensorrt + pre-script: packaging/pre_build_script.sh + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + uses: ./.github/workflows/linux-test.yml + with: + job-name: tests-py-dynamo-core + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + architecture: "aarch64" + script: | + export USE_HOST_DEPS=1 + export CI_BUILD=1 + pushd . + cd tests/py + python -m pip install -r requirements.txt + cd dynamo + python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_test_results.xml --ignore runtime/test_002_cudagraphs_py.py --ignore runtime/test_002_cudagraphs_cpp.py runtime/ + python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_partitioning_test_results.xml partitioning/ + python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_lowering_test_results.xml lowering/ + popd + + tests-py-dynamo-cudagraphs: + name: Test dynamo cudagraphs [Python] + needs: [filter-matrix, build] + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + package-name: torch_tensorrt + pre-script: packaging/pre_build_script.sh + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + uses: ./.github/workflows/linux-test.yml + with: + job-name: tests-py-dynamo-cudagraphs + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + architecture: "aarch64" + script: | + export USE_HOST_DEPS=1 + export CI_BUILD=1 + pushd . + cd tests/py + python -m pip install -r requirements.txt + cd dynamo + nvidia-smi + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_cudagraphs_cpp_test_results.xml runtime/test_002_cudagraphs_cpp.py || true + python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_dynamo_core_runtime_cudagraphs_py_test_results.xml runtime/test_002_cudagraphs_py.py || true + popd + + tests-py-core: + name: Test core [Python] + needs: [filter-matrix, build] + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/tensorrt + package-name: torch_tensorrt + pre-script: packaging/pre_build_script.sh + post-script: packaging/post_build_script.sh + smoke-test-script: packaging/smoke_test_script.sh + uses: ./.github/workflows/linux-test.yml + with: + job-name: tests-py-core + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + architecture: "aarch64" + script: | + export USE_HOST_DEPS=1 + export CI_BUILD=1 + pushd . + cd tests/py/core + python -m pytest -ra -n 4 --junitxml=${RUNNER_TEST_RESULTS_DIR}/tests_py_core_test_results.xml . + popd + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }}-${{ inputs.job-name }} + cancel-in-progress: true diff --git a/.github/workflows/build-test-linux.yml b/.github/workflows/build-test-linux-x86_64.yml similarity index 98% rename from .github/workflows/build-test-linux.yml rename to .github/workflows/build-test-linux-x86_64.yml index 024afd8c62..66aae7ea1e 100644 --- a/.github/workflows/build-test-linux.yml +++ b/.github/workflows/build-test-linux-x86_64.yml @@ -1,4 +1,4 @@ -name: Build and test Linux wheels +name: Build and test Linux x86-64 wheels on: pull_request: @@ -15,7 +15,7 @@ on: jobs: generate-matrix: - uses: ./.github/workflows/generate_binary_build_matrix.yml + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main with: package-type: wheel os: linux @@ -33,7 +33,7 @@ jobs: steps: - uses: actions/setup-python@v5 with: - python-version: '3.11' + python-version: "3.11" - uses: actions/checkout@v4 with: repository: pytorch/tensorrt diff --git a/.github/workflows/build-test-tensorrt-linux.yml b/.github/workflows/build-test-tensorrt-linux.yml index 6ddd0e336d..2237ecb50d 100644 --- a/.github/workflows/build-test-tensorrt-linux.yml +++ b/.github/workflows/build-test-tensorrt-linux.yml @@ -3,7 +3,7 @@ name: Build and Test Torch-TensorRT on Linux with Future TensorRT Versions on: workflow_dispatch: schedule: - - cron: '0 0 * * 0' # Runs at 00:00 UTC every Sunday (minute hour day-of-month month-of-year day-of-week) + - cron: "0 0 * * 0" # Runs at 00:00 UTC every Sunday (minute hour day-of-month month-of-year day-of-week) permissions: id-token: write @@ -30,7 +30,7 @@ jobs: steps: - uses: actions/setup-python@v5 with: - python-version: '3.11' + python-version: "3.11" - uses: actions/checkout@v4 with: repository: pytorch/tensorrt diff --git a/.github/workflows/build-test-windows.yml b/.github/workflows/build-test-windows.yml index f78218e75d..00bea34649 100644 --- a/.github/workflows/build-test-windows.yml +++ b/.github/workflows/build-test-windows.yml @@ -15,7 +15,7 @@ on: jobs: generate-matrix: - uses: ./.github/workflows/generate_binary_build_matrix.yml + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main with: package-type: wheel os: windows diff --git a/.github/workflows/generate_binary_build_matrix.yml b/.github/workflows/generate_binary_build_matrix.yml index fceb8cf1ee..6bbbc90852 100644 --- a/.github/workflows/generate_binary_build_matrix.yml +++ b/.github/workflows/generate_binary_build_matrix.yml @@ -42,6 +42,7 @@ on: use-only-dl-pytorch-org: description: "Use only download.pytorch.org when generating wheel install command?" default: "false" + required: false type: string build-python-only: description: "Generate binary build matrix for a python only package (i.e. only one python version)" @@ -72,16 +73,13 @@ jobs: steps: - uses: actions/setup-python@v5 with: - python-version: '3.11' + python-version: "3.11" - name: Checkout test-infra repository uses: actions/checkout@v4 with: repository: ${{ inputs.test-infra-repository }} ref: ${{ inputs.test-infra-ref }} - uses: ./.github/actions/set-channel - - uses: actions/checkout@v4 - with: - repository: pytorch/tensorrt - name: Generate test matrix id: generate env: @@ -103,10 +101,11 @@ jobs: PYTHON_VERSIONS: ${{ inputs.python-versions }} run: | set -eou pipefail + ls -la . MATRIX_BLOB="$(python3 .github/scripts/generate_binary_build_matrix.py)" echo "${MATRIX_BLOB}" echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}" concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.package-type }}-${{ inputs.os }}-${{ inputs.test-infra-repository }}-${{ inputs.test-infra-ref }} - cancel-in-progress: true \ No newline at end of file + cancel-in-progress: true diff --git a/MODULE.bazel b/MODULE.bazel index 58de67800e..89a57cbc4e 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -1,7 +1,7 @@ module( name = "torch_tensorrt", - version = "2.8.0a0", repo_name = "org_pytorch_tensorrt", + version = "" ) bazel_dep(name = "googletest", version = "1.14.0") @@ -30,21 +30,23 @@ local_repository( path = "/opt/conda/lib/python3.8/site-packages/torch_tensorrt", ) + new_local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "new_local_repository") # CUDA should be installed on the system locally new_local_repository( name = "cuda", build_file = "@//third_party/cuda:BUILD", - path = "/usr/local/cuda-12.8/", + path = "", ) new_local_repository( name = "cuda_win", build_file = "@//third_party/cuda:BUILD", - path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.8/", + path = "", ) + http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") ############################################################################################################# @@ -55,21 +57,7 @@ http_archive( name = "libtorch", build_file = "@//third_party/libtorch:BUILD", strip_prefix = "libtorch", - urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-cxx11-abi-shared-with-deps-latest.zip"], -) - -http_archive( - name = "libtorch_pre_cxx11_abi", - build_file = "@//third_party/libtorch:BUILD", - strip_prefix = "libtorch", - urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-shared-with-deps-latest.zip"], -) - -http_archive( - name = "libtorch_win", - build_file = "@//third_party/libtorch:BUILD", - strip_prefix = "libtorch", - urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-win-shared-with-deps-latest.zip"], + urls = ["https://download.pytorch.org/libtorch///libtorch-cxx11-abi-shared-with-deps-latest.zip"], ) # Download these tarballs manually from the NVIDIA website @@ -94,6 +82,7 @@ http_archive( ], ) + #################################################################################### # Locally installed dependencies (use in cases of custom dependencies or aarch64) #################################################################################### @@ -105,17 +94,17 @@ http_archive( # x86_64 python distribution. If using NVIDIA's version just point to the root of the package # for both versions here and do not use --config=pre-cxx11-abi -#new_local_repository( -# name = "libtorch", -# path = "/usr/local/lib/python3.6/dist-packages/torch", -# build_file = "third_party/libtorch/BUILD" -#) +new_local_repository( + name = "libtorch_win", + path = "", + build_file = "third_party/libtorch/BUILD" +) -#new_local_repository( -# name = "libtorch_pre_cxx11_abi", -# path = "/usr/local/lib/python3.6/dist-packages/torch", -# build_file = "third_party/libtorch/BUILD" -#) +new_local_repository( + name = "libtorch_pre_cxx11_abi", + path = "", + build_file = "third_party/libtorch/BUILD" +) #new_local_repository( # name = "tensorrt", diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh index 173498201f..661b38764c 100755 --- a/packaging/pre_build_script.sh +++ b/packaging/pre_build_script.sh @@ -5,10 +5,18 @@ set -x # Install dependencies python3 -m pip install pyyaml -yum install -y ninja-build gettext +PLATFORM="amd64" -wget https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-amd64 \ - && mv bazelisk-linux-amd64 /usr/bin/bazel \ +if [[ $(uname -m) == "aarch64" ]]; then + PLATFORM="arm64" + rm -rf /opt/openssl # probs unsafe but not sure why there is a different openssl version +fi + +dnf update -y openssl-libs krb5-libs +dnf install -y ninja-build gettext + +wget https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelisk-linux-${PLATFORM} \ + && mv bazelisk-linux-${PLATFORM} /usr/bin/bazel \ && chmod +x /usr/bin/bazel TORCH_TORCHVISION=$(grep "^torch" py/requirements.txt) diff --git a/pyproject.toml b/pyproject.toml index 3feacec679..120a767822 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ requires = [ "cffi>=1.15.1", "typing-extensions>=4.7.0", "future>=0.18.3", - "tensorrt-cu12>=10.9.0,<10.10.0", + "tensorrt->=10.9.0,<10.10.0", "torch>=2.8.0.dev,<2.9.0", "pybind11==2.6.2", "numpy", @@ -57,9 +57,9 @@ keywords = [ dependencies = [ "torch>=2.8.0.dev,<2.9.0", "tensorrt>=10.9.0,<10.10.0", - "tensorrt-cu12>=10.9.0,<10.10.0", - "tensorrt-cu12-bindings>=10.9.0,<10.10.0", - "tensorrt-cu12-libs>=10.9.0,<10.10.0", + "tensorrt->=10.9.0,<10.10.0", + "tensorrt--bindings>=10.9.0,<10.10.0", + "tensorrt--libs>=10.9.0,<10.10.0", "packaging>=23", "numpy", "typing-extensions>=4.7.0", diff --git a/version.txt b/version.txt index 787e6e4ab7..11922a5ce1 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -2.7.0a0 +2.8.0a0 From eb946cd03e4db7b959b4e1a1956d646e8b264fca Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Mon, 14 Apr 2025 18:16:15 -0600 Subject: [PATCH 2/8] infra: bazel aarch64 colocation --- MODULE.bazel | 39 +++++++++++- core/BUILD | 43 +++++++++++--- core/conversion/BUILD | 37 ++++++++++-- core/conversion/conversionctx/BUILD | 37 ++++++++++-- core/conversion/converters/BUILD | 59 +++++++++++++++---- core/conversion/evaluators/BUILD | 33 +++++++++-- core/conversion/tensorcontainer/BUILD | 37 ++++++++++-- core/conversion/var/BUILD | 37 ++++++++++-- core/ir/BUILD | 37 ++++++++++-- core/lowering/BUILD | 33 +++++++++-- core/lowering/passes/BUILD | 33 +++++++++-- core/partitioning/BUILD | 33 +++++++++-- core/partitioning/partitioningctx/BUILD | 33 +++++++++-- core/partitioning/partitioninginfo/BUILD | 33 +++++++++-- core/partitioning/segmentedblock/BUILD | 33 +++++++++-- core/plugins/BUILD | 49 +++++++++++++-- core/runtime/BUILD | 46 ++++++++++----- core/util/BUILD | 45 ++++++++++---- core/util/logging/BUILD | 37 ++++++++++-- cpp/bin/torchtrtc/BUILD | 14 ++--- examples/int8/benchmark/BUILD | 24 +++++++- examples/int8/ptq/BUILD | 49 +++++++++++++++ examples/int8/qat/BUILD | 50 ++++++++++++++++ examples/torchtrt_runtime_example/BUILD | 30 ++++++++++ packaging/pre_build_script.sh | 16 ++--- tests/core/BUILD | 28 +++++++-- tests/core/conversion/converters/BUILD | 4 +- tests/core/conversion/evaluators/BUILD | 4 +- tests/core/lowering/BUILD | 6 +- tests/core/lowering/lowering_test.bzl | 4 +- tests/core/partitioning/BUILD | 15 ++--- tests/core/partitioning/partitioning_test.bzl | 4 +- tests/core/runtime/BUILD | 7 ++- tests/core/runtime/runtime_test.bzl | 4 +- tests/cpp/BUILD | 18 +++--- tests/util/BUILD | 34 +++++++++-- toolchains/build_srcs/BUILD | 7 +++ toolchains/build_srcs/defs.bzl | 18 ++++++ toolchains/dep_collection/BUILD | 7 +++ toolchains/dep_collection/defs.bzl | 18 ++++++ 40 files changed, 917 insertions(+), 178 deletions(-) create mode 100644 toolchains/build_srcs/BUILD create mode 100644 toolchains/build_srcs/defs.bzl create mode 100644 toolchains/dep_collection/BUILD create mode 100644 toolchains/dep_collection/defs.bzl diff --git a/MODULE.bazel b/MODULE.bazel index 89a57cbc4e..10283be49c 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -27,7 +27,7 @@ local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl" # External dependency for torch_tensorrt if you already have precompiled binaries. local_repository( name = "torch_tensorrt", - path = "/opt/conda/lib/python3.8/site-packages/torch_tensorrt", + path = "/opt/conda/lib/python3.10/site-packages/torch_tensorrt", ) @@ -37,7 +37,7 @@ new_local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local. new_local_repository( name = "cuda", build_file = "@//third_party/cuda:BUILD", - path = "", + path = "/usr/local/cuda-12.8", ) new_local_repository( @@ -57,7 +57,19 @@ http_archive( name = "libtorch", build_file = "@//third_party/libtorch:BUILD", strip_prefix = "libtorch", +<<<<<<< HEAD urls = ["https://download.pytorch.org/libtorch///libtorch-cxx11-abi-shared-with-deps-latest.zip"], +======= + urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-cxx11-abi-shared-with-deps-latest.zip"], +) + +http_archive( + name = "torch_whl", + build_file = "@//third_party/libtorch:BUILD", + strip_prefix = "torch", + type = "zip", + urls = ["https://download.pytorch.org/whl/nightly/cu128/torch-2.8.0.dev20250414%2Bcu128-cp39-cp39-manylinux_2_28_aarch64.whl"], +>>>>>>> 75173f897 (infra: bazel aarch64 colocation) ) # Download these tarballs manually from the NVIDIA website @@ -73,6 +85,15 @@ http_archive( ], ) +http_archive( + name = "tensorrt_sbsa", + build_file = "@//third_party/tensorrt/archive:BUILD", + strip_prefix = "TensorRT-10.9.0.34", + urls = [ + "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.aarch64-gnu.cuda-12.8.tar.gz", + ], +) + http_archive( name = "tensorrt_win", build_file = "@//third_party/tensorrt/archive:BUILD", @@ -94,6 +115,7 @@ http_archive( # x86_64 python distribution. If using NVIDIA's version just point to the root of the package # for both versions here and do not use --config=pre-cxx11-abi +<<<<<<< HEAD new_local_repository( name = "libtorch_win", path = "", @@ -105,6 +127,19 @@ new_local_repository( path = "", build_file = "third_party/libtorch/BUILD" ) +======= +# new_local_repository( +# name = "libtorch", +# path = "/workspace/tensorrt/.venv/lib/python3.9/site-packages/torch", +# build_file = "third_party/libtorch/BUILD" +# ) + +# new_local_repository( +# name = "libtorch_pre_cxx11_abi", +# path = "", +# build_file = "third_party/libtorch/BUILD" +# ) +>>>>>>> 75173f897 (infra: bazel aarch64 colocation) #new_local_repository( # name = "tensorrt", diff --git a/core/BUILD b/core/BUILD index 4fd1a7cf64..c03f6036f3 100644 --- a/core/BUILD +++ b/core/BUILD @@ -4,16 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" }, ) config_setting( - name = "python_core", - values = { - "define": "target_lang=python", + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -24,6 +37,13 @@ config_setting( ], ) +config_setting( + name = "python_core", + values = { + "define": "target_lang=python", + }, +) + cc_library( name = "core", srcs = [ @@ -39,9 +59,14 @@ cc_library( "//core/runtime", "//core/util/logging", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/conversion/BUILD b/core/conversion/BUILD index d00defcbd8..af3371fe07 100644 --- a/core/conversion/BUILD +++ b/core/conversion/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -34,9 +54,14 @@ cc_library( "//core/ir", "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/conversion/conversionctx/BUILD b/core/conversion/conversionctx/BUILD index 0e910d7127..209e8a4c65 100644 --- a/core/conversion/conversionctx/BUILD +++ b/core/conversion/conversionctx/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -29,9 +49,14 @@ cc_library( "//core/ir", "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index 062f700cab..251f63108b 100644 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -29,9 +49,14 @@ cc_library( "//core/conversion/conversionctx", "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) @@ -49,9 +74,14 @@ cc_library( "//core/conversion/conversionctx", "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) @@ -106,9 +136,14 @@ cc_library( "//core/plugins:torch_tensorrt_plugins", "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/conversion/evaluators/BUILD b/core/conversion/evaluators/BUILD index 5116f3f76f..024945cdbe 100644 --- a/core/conversion/evaluators/BUILD +++ b/core/conversion/evaluators/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -35,8 +55,13 @@ cc_library( "//core/conversion/var", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), alwayslink = True, diff --git a/core/conversion/tensorcontainer/BUILD b/core/conversion/tensorcontainer/BUILD index c07e8f5516..e6851eee47 100644 --- a/core/conversion/tensorcontainer/BUILD +++ b/core/conversion/tensorcontainer/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -28,9 +48,14 @@ cc_library( deps = [ "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/conversion/var/BUILD b/core/conversion/var/BUILD index 7042a1f402..12ee259f7d 100644 --- a/core/conversion/var/BUILD +++ b/core/conversion/var/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -31,9 +51,14 @@ cc_library( "//core/conversion/tensorcontainer", "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/ir/BUILD b/core/ir/BUILD index 1e4640f08f..5272a2af0d 100644 --- a/core/ir/BUILD +++ b/core/ir/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -31,9 +51,14 @@ cc_library( deps = [ "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/lowering/BUILD b/core/lowering/BUILD index 2e9d236b05..568ad23414 100644 --- a/core/lowering/BUILD +++ b/core/lowering/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -33,8 +53,13 @@ cc_library( "//core/lowering/passes", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), alwayslink = True, diff --git a/core/lowering/passes/BUILD b/core/lowering/passes/BUILD index 5a99139db3..2e016ff79b 100644 --- a/core/lowering/passes/BUILD +++ b/core/lowering/passes/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -56,8 +76,13 @@ cc_library( deps = [ "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), alwayslink = True, diff --git a/core/partitioning/BUILD b/core/partitioning/BUILD index 784f20c719..e871df658a 100644 --- a/core/partitioning/BUILD +++ b/core/partitioning/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -36,8 +56,13 @@ cc_library( "//core/partitioning/segmentedblock", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), alwayslink = True, diff --git a/core/partitioning/partitioningctx/BUILD b/core/partitioning/partitioningctx/BUILD index c595ca001d..3916b8f260 100644 --- a/core/partitioning/partitioningctx/BUILD +++ b/core/partitioning/partitioningctx/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -32,8 +52,13 @@ cc_library( "//core/partitioning/segmentedblock", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), alwayslink = True, diff --git a/core/partitioning/partitioninginfo/BUILD b/core/partitioning/partitioninginfo/BUILD index 96194bf629..138e7c7f1b 100644 --- a/core/partitioning/partitioninginfo/BUILD +++ b/core/partitioning/partitioninginfo/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -31,8 +51,13 @@ cc_library( "//core/lowering", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), alwayslink = True, diff --git a/core/partitioning/segmentedblock/BUILD b/core/partitioning/segmentedblock/BUILD index 44cd1da98d..c6c7db7ccc 100644 --- a/core/partitioning/segmentedblock/BUILD +++ b/core/partitioning/segmentedblock/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -31,8 +51,13 @@ cc_library( "//core/lowering", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), alwayslink = True, diff --git a/core/plugins/BUILD b/core/plugins/BUILD index 2b7a28848f..19aecf4a7d 100644 --- a/core/plugins/BUILD +++ b/core/plugins/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -38,9 +58,26 @@ cc_library( deps = [ "//core/util:prelude", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@tensorrt_win//:nvinferplugin", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@tensorrt//:nvinferplugin", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@tensorrt//:nvinferplugin", "@libtorch"], + ":windows": [ + "@tensorrt_win//:nvinfer", + "@tensorrt_win//:nvinferplugin", + ], + ":sbsa": [ + "@tensorrt_sbsa//:nvinfer", + "@tensorrt_sbsa//:nvinferplugin", + ], + ":jetpack": [ + "@tensorrt_jetpack//:nvinfer", + "@tensorrt_jetpack//:nvinferplugin", + ], + "//conditions:default": [ + "@tensorrt//:nvinfer", + "@tensorrt//:nvinferplugin", + ], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/runtime/BUILD b/core/runtime/BUILD index 5a8825b235..aa206a5a6d 100644 --- a/core/runtime/BUILD +++ b/core/runtime/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -43,18 +63,14 @@ cc_library( "//core/plugins:torch_tensorrt_plugins", "//core/util:prelude", ] + select({ - ":use_pre_cxx11_abi": [ - "@libtorch_pre_cxx11_abi//:libtorch", - "@tensorrt//:nvinfer", - ], - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/util/BUILD b/core/util/BUILD index 56e5182160..60aa6cdf08 100644 --- a/core/util/BUILD +++ b/core/util/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -41,7 +61,7 @@ cc_library( ":macros", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -74,9 +94,9 @@ cc_library( "build_info.h", ], deps = select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), ) @@ -92,9 +112,14 @@ cc_library( ":macros", "//core/util/logging", ] + select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/core/util/logging/BUILD b/core/util/logging/BUILD index 4796d5def1..be3b80cf02 100644 --- a/core/util/logging/BUILD +++ b/core/util/logging/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -26,9 +46,14 @@ cc_library( "TorchTRTLogger.h", ], deps = select({ - ":windows": ["@tensorrt_win//:nvinfer", "@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@tensorrt//:nvinfer", "@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@tensorrt//:nvinfer", "@libtorch"], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) diff --git a/cpp/bin/torchtrtc/BUILD b/cpp/bin/torchtrtc/BUILD index 00e2490daa..d3a2b72093 100644 --- a/cpp/bin/torchtrtc/BUILD +++ b/cpp/bin/torchtrtc/BUILD @@ -3,9 +3,9 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" }, ) @@ -37,11 +37,11 @@ cc_binary( ] + select({ ":windows": [ "@libtorch_win//:caffe2", - "@libtorch_win//:libtorch", + "@libtorch_win//:libtorch" ], - ":use_pre_cxx11_abi": [ - "@libtorch_pre_cxx11_abi//:caffe2", - "@libtorch_pre_cxx11_abi//:libtorch", + ":use_torch_whl": [ + "@torch_whl//:caffe2", + "@torch_whl//:libtorch" ], "//conditions:default": [ "@libtorch", diff --git a/examples/int8/benchmark/BUILD b/examples/int8/benchmark/BUILD index 23d147c4c3..115dd050c0 100644 --- a/examples/int8/benchmark/BUILD +++ b/examples/int8/benchmark/BUILD @@ -2,6 +2,13 @@ load("@rules_cc//cc:defs.bzl", "cc_library") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + cc_library( name = "benchmark", srcs = [ @@ -13,7 +20,18 @@ cc_library( ], deps = [ "//cpp:torch_tensorrt", - "@libtorch", - "@libtorch//:caffe2", - ], + ] + select({ + ":windows": [ + "@libtorch_win//:libtorch", + "@libtorch_win//:caffe2", + ], + ":use_torch_whl": [ + "@torch_whl//:libtorch", + "@torch_whl//:caffe2", + ], + "//conditions:default": [ + "@libtorch//:libtorch", + "@libtorch//:caffe2", + ], + }), ) diff --git a/examples/int8/ptq/BUILD b/examples/int8/ptq/BUILD index c9fa200220..3f98253af0 100644 --- a/examples/int8/ptq/BUILD +++ b/examples/int8/ptq/BUILD @@ -2,6 +2,40 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" + }, +) + +config_setting( + name = "windows", + constraint_values = [ + "@platforms//os:windows", + ], +) + cc_binary( name = "ptq", srcs = [ @@ -20,7 +54,22 @@ cc_binary( "@libtorch", "@libtorch//:caffe2", ] + select({ + ":windows": [ + "@libtorch_win//:libtorch", + "@libtorch_win//:caffe2", + ], + ":use_torch_whl": [ + "@torch_whl//:libtorch", + "@torch_whl//:caffe2", + ], + "//conditions:default": [ + "@libtorch//:libtorch", + "@libtorch//:caffe2", + ], + }) + select({ ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/examples/int8/qat/BUILD b/examples/int8/qat/BUILD index e97398d932..d8407acee9 100644 --- a/examples/int8/qat/BUILD +++ b/examples/int8/qat/BUILD @@ -2,6 +2,41 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" + }, +) + +config_setting( + name = "windows", + constraint_values = [ + "@platforms//os:windows", + ], +) + + cc_binary( name = "qat", srcs = [ @@ -20,7 +55,22 @@ cc_binary( "@libtorch", "@libtorch//:caffe2", ] + select({ + ":windows": [ + "@libtorch_win//:libtorch", + "@libtorch_win//:caffe2", + ], + ":use_torch_whl": [ + "@torch_whl//:libtorch", + "@torch_whl//:caffe2", + ], + "//conditions:default": [ + "@libtorch//:libtorch", + "@libtorch//:caffe2", + ], + }) + select({ ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/examples/torchtrt_runtime_example/BUILD b/examples/torchtrt_runtime_example/BUILD index 957caedd23..5e86b89f14 100644 --- a/examples/torchtrt_runtime_example/BUILD +++ b/examples/torchtrt_runtime_example/BUILD @@ -2,6 +2,34 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" + }, +) + +config_setting( + name = "windows", + constraint_values = [ + "@platforms//os:windows", + ], +) + + cc_binary( name = "torchtrt_runtime_example", srcs = [ @@ -13,6 +41,8 @@ cc_binary( "@libtorch//:caffe2", ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh index 661b38764c..247cf2105b 100755 --- a/packaging/pre_build_script.sh +++ b/packaging/pre_build_script.sh @@ -6,17 +6,17 @@ set -x python3 -m pip install pyyaml PLATFORM="amd64" - +PLATFORM=x86_64 +BAZEL_PLATFORM=amd64 if [[ $(uname -m) == "aarch64" ]]; then - PLATFORM="arm64" - rm -rf /opt/openssl # probs unsafe but not sure why there is a different openssl version -fi + PLATFORM=aarch64 + BAZEL_PLATFORM=arm64 -dnf update -y openssl-libs krb5-libs -dnf install -y ninja-build gettext + rm -rf /opt/openssl # Not sure whats up with the openssl mismatch +fi -wget https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelisk-linux-${PLATFORM} \ - && mv bazelisk-linux-${PLATFORM} /usr/bin/bazel \ +wget https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-linux-${BAZEL_PLATFORM} \ + && mv bazelisk-linux-${BAZEL_PLATFORM} /usr/bin/bazel \ && chmod +x /usr/bin/bazel TORCH_TORCHVISION=$(grep "^torch" py/requirements.txt) diff --git a/tests/core/BUILD b/tests/core/BUILD index cc18d03ea7..ee17e64a4b 100644 --- a/tests/core/BUILD +++ b/tests/core/BUILD @@ -1,9 +1,29 @@ load("@rules_cc//cc:defs.bzl", "cc_test") config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" }, ) @@ -30,7 +50,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) diff --git a/tests/core/conversion/converters/BUILD b/tests/core/conversion/converters/BUILD index 05ad6e356c..dfab90a978 100644 --- a/tests/core/conversion/converters/BUILD +++ b/tests/core/conversion/converters/BUILD @@ -1,9 +1,9 @@ load("//tests/core/conversion/converters:converter_test.bzl", "converter_test") config_setting( - name = "use_pre_cxx11_abi", + name = "use_torch_whl", values = { - "define": "abi=pre_cxx11_abi", + "define": "torch_src=whl", }, ) diff --git a/tests/core/conversion/evaluators/BUILD b/tests/core/conversion/evaluators/BUILD index 4059b97df3..9b2a708297 100644 --- a/tests/core/conversion/evaluators/BUILD +++ b/tests/core/conversion/evaluators/BUILD @@ -1,9 +1,9 @@ load("//tests/core/conversion/evaluators:evaluator_test.bzl", "evaluator_test") config_setting( - name = "use_pre_cxx11_abi", + name = "use_torch_whl", values = { - "define": "abi=pre_cxx11_abi", + "define": "torch_src=whl", }, ) diff --git a/tests/core/lowering/BUILD b/tests/core/lowering/BUILD index b2d3609ccf..3bc26285ed 100644 --- a/tests/core/lowering/BUILD +++ b/tests/core/lowering/BUILD @@ -2,9 +2,9 @@ load("@rules_cc//cc:defs.bzl", "cc_test") load("//tests/core/lowering:lowering_test.bzl", "lowering_test") config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" }, ) diff --git a/tests/core/lowering/lowering_test.bzl b/tests/core/lowering/lowering_test.bzl index ecdac73778..9d2c38a969 100644 --- a/tests/core/lowering/lowering_test.bzl +++ b/tests/core/lowering/lowering_test.bzl @@ -20,8 +20,8 @@ def lowering_test(name, visibility = None): "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@libtorch//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), timeout = "short", ) diff --git a/tests/core/partitioning/BUILD b/tests/core/partitioning/BUILD index ed2328897c..935f085e1e 100644 --- a/tests/core/partitioning/BUILD +++ b/tests/core/partitioning/BUILD @@ -2,9 +2,9 @@ load("@rules_cc//cc:defs.bzl", "cc_test") load("//tests/core/partitioning:partitioning_test.bzl", "partitioning_test") config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" }, ) @@ -15,6 +15,7 @@ config_setting( ], ) + filegroup( name = "jit_models", srcs = [ @@ -62,7 +63,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -78,7 +79,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -94,7 +95,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -110,7 +111,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) diff --git a/tests/core/partitioning/partitioning_test.bzl b/tests/core/partitioning/partitioning_test.bzl index 322721aa93..f89dd22f40 100644 --- a/tests/core/partitioning/partitioning_test.bzl +++ b/tests/core/partitioning/partitioning_test.bzl @@ -20,8 +20,8 @@ def partitioning_test(name, visibility = None): "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@libtorch//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), #timeout = "short", ) diff --git a/tests/core/runtime/BUILD b/tests/core/runtime/BUILD index d25379b44d..212ce65679 100644 --- a/tests/core/runtime/BUILD +++ b/tests/core/runtime/BUILD @@ -3,9 +3,9 @@ load("//tests/core/runtime:runtime_test.bzl", "runtime_test") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" }, ) @@ -16,6 +16,7 @@ config_setting( ], ) + runtime_test( name = "test_multi_device_safe_mode", ) diff --git a/tests/core/runtime/runtime_test.bzl b/tests/core/runtime/runtime_test.bzl index 021321c40b..3363edb627 100644 --- a/tests/core/runtime/runtime_test.bzl +++ b/tests/core/runtime/runtime_test.bzl @@ -20,7 +20,7 @@ def runtime_test(name, visibility = None): "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], - "//conditions:default": ["@libtorch//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), ) diff --git a/tests/cpp/BUILD b/tests/cpp/BUILD index 2917f57130..9b93d4bb56 100644 --- a/tests/cpp/BUILD +++ b/tests/cpp/BUILD @@ -3,9 +3,9 @@ load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" }, ) @@ -94,7 +94,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -133,7 +133,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -149,7 +149,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -162,7 +162,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -178,7 +178,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -214,7 +214,7 @@ cc_library( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], - ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) diff --git a/tests/util/BUILD b/tests/util/BUILD index c49b88acf3..64750fa794 100644 --- a/tests/util/BUILD +++ b/tests/util/BUILD @@ -3,9 +3,9 @@ load("@rules_cc//cc:defs.bzl", "cc_library") package(default_visibility = ["//visibility:public"]) config_setting( - name = "use_pre_cxx11_abi", - values = { - "define": "abi=pre_cxx11_abi", + name = "use_torch_whl", + flag_values = { + "//toolchains/build_srcs:torch": "whl" }, ) @@ -16,6 +16,26 @@ config_setting( ], ) +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "datacenter" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_lib_collection": "jetpack" + }, +) + config_setting( name = "ci_build_testing", values = { @@ -39,15 +59,17 @@ cc_library( "@googletest//:gtest_main", ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_jetpack//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": [ "@libtorch_win//:caffe2", "@libtorch_win//:libtorch", ], - ":use_pre_cxx11_abi": [ - "@libtorch_pre_cxx11_abi//:caffe2", - "@libtorch_pre_cxx11_abi//:libtorch", + ":use_torch_whl": [ + "@torch_whl//:caffe2", + "@torch_whl//:libtorch", ], "//conditions:default": [ "@libtorch", diff --git a/toolchains/build_srcs/BUILD b/toolchains/build_srcs/BUILD new file mode 100644 index 0000000000..555b13272d --- /dev/null +++ b/toolchains/build_srcs/BUILD @@ -0,0 +1,7 @@ +package(default_visibility = ["//visibility:public"]) + +load(":defs.bzl", "build_src") +build_src( + name = "torch", + build_setting_default = "archive", +) diff --git a/toolchains/build_srcs/defs.bzl b/toolchains/build_srcs/defs.bzl new file mode 100644 index 0000000000..55da439436 --- /dev/null +++ b/toolchains/build_srcs/defs.bzl @@ -0,0 +1,18 @@ +# buildifier: disable=module-docstring +BuildSrcInfo = provider(doc = "", fields = ["type"]) + +src_types = ["archive", "whl", "local"] + +def _impl(ctx): + src = ctx.build_setting_value + if src not in src_types: + fail(str(ctx.label) + " build setting allowed to take values {" + + ", ".join(src_types) + "} but was set to unallowed value " + + src) + + return BuildSrcInfo(type = src) + +build_src = rule( + implementation = _impl, + build_setting = config.string(flag = True), +) diff --git a/toolchains/dep_collection/BUILD b/toolchains/dep_collection/BUILD new file mode 100644 index 0000000000..711a3732fc --- /dev/null +++ b/toolchains/dep_collection/BUILD @@ -0,0 +1,7 @@ +package(default_visibility = ["//visibility:public"]) + +load(":defs.bzl", "dep_collection") +dep_collection( + name = "compute_lib_collection", + build_setting_default = "datacenter", +) diff --git a/toolchains/dep_collection/defs.bzl b/toolchains/dep_collection/defs.bzl new file mode 100644 index 0000000000..35f3bc654a --- /dev/null +++ b/toolchains/dep_collection/defs.bzl @@ -0,0 +1,18 @@ +# buildifier: disable=module-docstring +DependencyCollectionInfo = provider(doc = "", fields = ["type"]) + +collection_types = ["datacenter", "jetpack"] + +def _impl(ctx): + _type = ctx.build_setting_value + if _type not in collection_types: + fail(str(ctx.label) + " build setting allowed to take values {" + + ", ".join(collection_types) + "} but was set to unallowed value " + + _type) + + return DependencyCollectionInfo(type = _type) + +dep_collection = rule( + implementation = _impl, + build_setting = config.string(flag = True), +) From ebef14ebede12a0bd0a88bdb3b9215d955aecf8c Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Mon, 14 Apr 2025 23:16:51 -0600 Subject: [PATCH 3/8] L4T collocated toolchain --- .bazelrc | 4 +- .bazelversion | 2 +- MODULE.bazel | 62 +++++++++---------- core/BUILD | 6 +- core/conversion/BUILD | 6 +- core/conversion/conversionctx/BUILD | 6 +- core/conversion/converters/BUILD | 10 +-- core/conversion/evaluators/BUILD | 6 +- core/conversion/tensorcontainer/BUILD | 6 +- core/conversion/var/BUILD | 6 +- core/ir/BUILD | 6 +- core/lowering/BUILD | 6 +- core/lowering/passes/BUILD | 6 +- core/partitioning/BUILD | 6 +- core/partitioning/partitioningctx/BUILD | 6 +- core/partitioning/partitioninginfo/BUILD | 6 +- core/partitioning/segmentedblock/BUILD | 6 +- core/plugins/BUILD | 8 +-- core/runtime/BUILD | 6 +- core/util/BUILD | 6 +- core/util/logging/BUILD | 6 +- cpp/bin/torchtrtc/BUILD | 2 +- docker/README.md | 23 ++++--- docs/_sources/getting_started/jetpack.rst.txt | 3 +- docsrc/getting_started/jetpack.rst | 3 +- examples/int8/benchmark/BUILD | 2 +- examples/int8/ptq/BUILD | 6 +- examples/int8/qat/BUILD | 6 +- examples/torchtrt_runtime_example/BUILD | 4 +- tests/core/BUILD | 4 +- tests/core/lowering/BUILD | 2 +- tests/core/partitioning/BUILD | 2 +- tests/core/runtime/BUILD | 2 +- tests/cpp/BUILD | 2 +- tests/util/BUILD | 6 +- third_party/tensorrt/archive/BUILD | 8 ++- toolchains/dep_collection/BUILD | 2 +- toolchains/dep_collection/defs.bzl | 2 +- toolchains/{build_srcs => dep_src}/BUILD | 4 +- toolchains/{build_srcs => dep_src}/defs.bzl | 6 +- toolchains/jetpack/BUILD | 2 +- 41 files changed, 138 insertions(+), 135 deletions(-) rename toolchains/{build_srcs => dep_src}/BUILD (73%) rename toolchains/{build_srcs => dep_src}/defs.bzl (79%) diff --git a/.bazelrc b/.bazelrc index 04ac0ba143..542608f849 100644 --- a/.bazelrc +++ b/.bazelrc @@ -36,9 +36,7 @@ build:cxx11_abi --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=1" build:cxx11_abi --linkopt="-D_GLIBCXX_USE_CXX11_ABI=1" build:cxx11_abi --define=abi=cxx11_abi -build:pre_cxx11_abi --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" -build:pre_cxx11_abi --linkopt="-D_GLIBCXX_USE_CXX11_ABI=0" -build:pre_cxx11_abi --define=abi=pre_cxx11_abi +build:jetpack --//toolchains/dep_collection:compute_lib_collection=jetpack build:ci_testing --define=torchtrt_src=prebuilt --cxxopt="-DDISABLE_TEST_IN_CI" --action_env "NVIDIA_TF32_OVERRIDE=0" build:use_precompiled_torchtrt --define=torchtrt_src=prebuilt diff --git a/.bazelversion b/.bazelversion index b26a34e470..0e79152459 100644 --- a/.bazelversion +++ b/.bazelversion @@ -1 +1 @@ -7.2.1 +8.1.1 diff --git a/MODULE.bazel b/MODULE.bazel index 10283be49c..9b555971d8 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -1,13 +1,13 @@ module( name = "torch_tensorrt", repo_name = "org_pytorch_tensorrt", - version = "" + version = "2.8.0" ) -bazel_dep(name = "googletest", version = "1.14.0") -bazel_dep(name = "platforms", version = "0.0.10") -bazel_dep(name = "rules_cc", version = "0.0.9") -bazel_dep(name = "rules_python", version = "0.34.0") +bazel_dep(name = "googletest", version = "1.16.0") +bazel_dep(name = "platforms", version = "0.0.11") +bazel_dep(name = "rules_cc", version = "0.1.1") +bazel_dep(name = "rules_python", version = "1.3.0") python = use_extension("@rules_python//python/extensions:python.bzl", "python") python.toolchain( @@ -40,10 +40,16 @@ new_local_repository( path = "/usr/local/cuda-12.8", ) +new_local_repository( + name = "cuda_l4t", + build_file = "@//third_party/cuda:BUILD", + path = "/usr/local/cuda-12.8", +) + new_local_repository( name = "cuda_win", build_file = "@//third_party/cuda:BUILD", - path = "", + path = "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.8/", ) @@ -57,19 +63,24 @@ http_archive( name = "libtorch", build_file = "@//third_party/libtorch:BUILD", strip_prefix = "libtorch", -<<<<<<< HEAD - urls = ["https://download.pytorch.org/libtorch///libtorch-cxx11-abi-shared-with-deps-latest.zip"], -======= urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-cxx11-abi-shared-with-deps-latest.zip"], ) + +http_archive( + name = "libtorch_win", + build_file = "@//third_party/libtorch:BUILD", + strip_prefix = "libtorch", + urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-win-shared-with-deps-latest.zip"], +) + + http_archive( name = "torch_whl", build_file = "@//third_party/libtorch:BUILD", strip_prefix = "torch", type = "zip", urls = ["https://download.pytorch.org/whl/nightly/cu128/torch-2.8.0.dev20250414%2Bcu128-cp39-cp39-manylinux_2_28_aarch64.whl"], ->>>>>>> 75173f897 (infra: bazel aarch64 colocation) ) # Download these tarballs manually from the NVIDIA website @@ -94,6 +105,15 @@ http_archive( ], ) +http_archive( + name = "tensorrt_l4t", + build_file = "@//third_party/tensorrt/archive:BUILD", + strip_prefix = "TensorRT-10.3.0.26", + urls = [ + "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.3.0/tars/TensorRT-10.3.0.26.l4t.aarch64-gnu.cuda-12.6.tar.gz", + ], +) + http_archive( name = "tensorrt_win", build_file = "@//third_party/tensorrt/archive:BUILD", @@ -115,32 +135,12 @@ http_archive( # x86_64 python distribution. If using NVIDIA's version just point to the root of the package # for both versions here and do not use --config=pre-cxx11-abi -<<<<<<< HEAD -new_local_repository( - name = "libtorch_win", - path = "", - build_file = "third_party/libtorch/BUILD" -) - -new_local_repository( - name = "libtorch_pre_cxx11_abi", - path = "", - build_file = "third_party/libtorch/BUILD" -) -======= # new_local_repository( # name = "libtorch", -# path = "/workspace/tensorrt/.venv/lib/python3.9/site-packages/torch", +# path = "/workspace/tensorrt/.venv/lib/python3.10/site-packages/torch", # build_file = "third_party/libtorch/BUILD" # ) -# new_local_repository( -# name = "libtorch_pre_cxx11_abi", -# path = "", -# build_file = "third_party/libtorch/BUILD" -# ) ->>>>>>> 75173f897 (infra: bazel aarch64 colocation) - #new_local_repository( # name = "tensorrt", # path = "/usr/", diff --git a/core/BUILD b/core/BUILD index c03f6036f3..f5c9b3e7ee 100644 --- a/core/BUILD +++ b/core/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -61,7 +61,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/conversion/BUILD b/core/conversion/BUILD index af3371fe07..f5ee98df74 100644 --- a/core/conversion/BUILD +++ b/core/conversion/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -56,7 +56,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/conversion/conversionctx/BUILD b/core/conversion/conversionctx/BUILD index 209e8a4c65..02d78118bd 100644 --- a/core/conversion/conversionctx/BUILD +++ b/core/conversion/conversionctx/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -51,7 +51,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index 251f63108b..61a5b0231f 100644 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -51,7 +51,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], @@ -76,7 +76,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], @@ -138,7 +138,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/conversion/evaluators/BUILD b/core/conversion/evaluators/BUILD index 024945cdbe..6251499cfc 100644 --- a/core/conversion/evaluators/BUILD +++ b/core/conversion/evaluators/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -57,7 +57,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/conversion/tensorcontainer/BUILD b/core/conversion/tensorcontainer/BUILD index e6851eee47..a570e9f6cf 100644 --- a/core/conversion/tensorcontainer/BUILD +++ b/core/conversion/tensorcontainer/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -50,7 +50,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/conversion/var/BUILD b/core/conversion/var/BUILD index 12ee259f7d..fa3f6ad451 100644 --- a/core/conversion/var/BUILD +++ b/core/conversion/var/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -53,7 +53,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/ir/BUILD b/core/ir/BUILD index 5272a2af0d..9a7cdf33fd 100644 --- a/core/ir/BUILD +++ b/core/ir/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -53,7 +53,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/lowering/BUILD b/core/lowering/BUILD index 568ad23414..d08a52463c 100644 --- a/core/lowering/BUILD +++ b/core/lowering/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -55,7 +55,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/lowering/passes/BUILD b/core/lowering/passes/BUILD index 2e016ff79b..961263d289 100644 --- a/core/lowering/passes/BUILD +++ b/core/lowering/passes/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -78,7 +78,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/partitioning/BUILD b/core/partitioning/BUILD index e871df658a..295e84d6cf 100644 --- a/core/partitioning/BUILD +++ b/core/partitioning/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -58,7 +58,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/partitioning/partitioningctx/BUILD b/core/partitioning/partitioningctx/BUILD index 3916b8f260..210086d9ec 100644 --- a/core/partitioning/partitioningctx/BUILD +++ b/core/partitioning/partitioningctx/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -54,7 +54,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/partitioning/partitioninginfo/BUILD b/core/partitioning/partitioninginfo/BUILD index 138e7c7f1b..e0b287a6fd 100644 --- a/core/partitioning/partitioninginfo/BUILD +++ b/core/partitioning/partitioninginfo/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -53,7 +53,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/partitioning/segmentedblock/BUILD b/core/partitioning/segmentedblock/BUILD index c6c7db7ccc..f00b3b2bb6 100644 --- a/core/partitioning/segmentedblock/BUILD +++ b/core/partitioning/segmentedblock/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -53,7 +53,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/plugins/BUILD b/core/plugins/BUILD index 19aecf4a7d..df1428b885 100644 --- a/core/plugins/BUILD +++ b/core/plugins/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -67,8 +67,8 @@ cc_library( "@tensorrt_sbsa//:nvinferplugin", ], ":jetpack": [ - "@tensorrt_jetpack//:nvinfer", - "@tensorrt_jetpack//:nvinferplugin", + "@tensorrt_l4t//:nvinfer", + "@tensorrt_l4t//:nvinferplugin", ], "//conditions:default": [ "@tensorrt//:nvinfer", diff --git a/core/runtime/BUILD b/core/runtime/BUILD index aa206a5a6d..623e56ab04 100644 --- a/core/runtime/BUILD +++ b/core/runtime/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -65,7 +65,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/util/BUILD b/core/util/BUILD index 60aa6cdf08..1532c2f52a 100644 --- a/core/util/BUILD +++ b/core/util/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -114,7 +114,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/core/util/logging/BUILD b/core/util/logging/BUILD index be3b80cf02..6d8b19e065 100644 --- a/core/util/logging/BUILD +++ b/core/util/logging/BUILD @@ -6,7 +6,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -48,7 +48,7 @@ cc_library( deps = select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": ["@libtorch_win//:libtorch"], diff --git a/cpp/bin/torchtrtc/BUILD b/cpp/bin/torchtrtc/BUILD index d3a2b72093..51ee4ca2ab 100644 --- a/cpp/bin/torchtrtc/BUILD +++ b/cpp/bin/torchtrtc/BUILD @@ -5,7 +5,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) diff --git a/docker/README.md b/docker/README.md index 3effebbfbf..3852e0e22a 100644 --- a/docker/README.md +++ b/docker/README.md @@ -1,19 +1,20 @@ # Building a Torch-TensorRT container -* Use `Dockerfile` to build a container which provides the exact development environment that our main branch is usually tested against. +- Use `Dockerfile` to build a container which provides the exact development environment that our main branch is usually tested against. -* The `Dockerfile` currently uses Bazelisk to select the Bazel version, and uses the exact library versions of Torch and CUDA listed in dependencies. - * The desired versions of TensorRT must be specified as build-args, with major and minor versions as in: `--build-arg TENSORRT_VERSION=a.b` - * [**Optional**] The desired base image be changed by explicitly setting a base image, as in `--build-arg BASE_IMG=nvidia/cuda:11.8.0-devel-ubuntu22.04`, though this is optional. - * [**Optional**] Additionally, the desired Python version can be changed by explicitly setting a version, as in `--build-arg PYTHON_VERSION=3.11`, though this is optional as well. +- The `Dockerfile` currently uses Bazelisk to select the Bazel version, and uses the exact library versions of Torch and CUDA listed in dependencies. -* This `Dockerfile` installs `cxx11-abi` versions of Pytorch and builds Torch-TRT using `cxx11-abi` libtorch as well. As of torch 2.7, torch requires `cxx11-abi` for all CUDA 11.8, 12.4, 12.6, and later versions. + - The desired versions of TensorRT must be specified as build-args, with major and minor versions as in: `--build-arg TENSORRT_VERSION=a.b` + - [**Optional**] The desired base image be changed by explicitly setting a base image, as in `--build-arg BASE_IMG=nvidia/cuda:11.8.0-devel-ubuntu22.04`, though this is optional. + - [**Optional**] Additionally, the desired Python version can be changed by explicitly setting a version, as in `--build-arg PYTHON_VERSION=3.11`, though this is optional as well. + +- This `Dockerfile` installs `cxx11-abi` versions of Pytorch and builds Torch-TRT using `cxx11-abi` libtorch as well. As of torch 2.7, torch requires `cxx11-abi` for all CUDA 11.8, 12.4, 12.6, and later versions. Note: By default the container uses the `cxx11-abi` version of Torch + Torch-TRT. If you are using a workflow that requires a build of PyTorch on the PRE CXX11 ABI, please add the Docker build argument: `--build-arg USE_PRE_CXX11_ABI=1` ### Dependencies -* Install nvidia-docker by following https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker +- Install nvidia-docker by following https://docs.nvidia.com/default/cloud-native/container-toolkit/install-guide.html#docker ### Instructions @@ -23,25 +24,27 @@ Note: By default the container uses the `cxx11-abi` version of Torch + Torch-TRT > From root of Torch-TensorRT repo Build: + ``` DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=10.9.0 -f docker/Dockerfile -t torch_tensorrt:latest . ``` Run: + ``` nvidia-docker run --gpus all -it --shm-size=8gb --env="DISPLAY" --volume="/tmp/.X11-unix:/tmp/.X11-unix:rw" --name=torch_tensorrt --ipc=host --net=host torch_tensorrt:latest ``` Test: - You can run any converter test to verify if Torch-TRT built successfully inside the container. Once you launch the container, you can run + ``` bazel test //tests/core/conversion/converters:test_activation --compilation_mode=opt --test_output=summary --config use_precompiled_torchtrt --config pre_cxx11_abi ``` -* `--config use_precompiled_torchtrt` : Indicates bazel to use pre-installed Torch-TRT library to test an application. -* `--config pre_cxx11_abi` : This flag ensures `bazel test` uses `pre_cxx11_abi` version of `libtorch`. Use this flag corresponding to the ABI format of your Torch-TensorRT installation. +- `--config use_precompiled_torchtrt` : Indicates bazel to use pre-installed Torch-TRT library to test an application. +- `--config pre_cxx11_abi` : This flag ensures `bazel test` uses `pre_cxx11_abi` version of `libtorch`. Use this flag corresponding to the ABI format of your Torch-TensorRT installation. ### Pytorch NGC containers diff --git a/docs/_sources/getting_started/jetpack.rst.txt b/docs/_sources/getting_started/jetpack.rst.txt index ddbf89dc63..b99c45addd 100644 --- a/docs/_sources/getting_started/jetpack.rst.txt +++ b/docs/_sources/getting_started/jetpack.rst.txt @@ -1,4 +1,4 @@ -.. _Torch_TensorRT_in_JetPack_6.1 +.. _Torch_TensorRT_in_l4t_6.1 Overview ################## @@ -116,4 +116,3 @@ Please make sure to build torch_tensorrt wheel file from source release/2.5 bran cat toolchains/jp_workspaces/MODULE.bazel.tmpl | envsubst > MODULE.bazel # build and install torch_tensorrt wheel file python setup.py install --user - diff --git a/docsrc/getting_started/jetpack.rst b/docsrc/getting_started/jetpack.rst index ddbf89dc63..b99c45addd 100644 --- a/docsrc/getting_started/jetpack.rst +++ b/docsrc/getting_started/jetpack.rst @@ -1,4 +1,4 @@ -.. _Torch_TensorRT_in_JetPack_6.1 +.. _Torch_TensorRT_in_l4t_6.1 Overview ################## @@ -116,4 +116,3 @@ Please make sure to build torch_tensorrt wheel file from source release/2.5 bran cat toolchains/jp_workspaces/MODULE.bazel.tmpl | envsubst > MODULE.bazel # build and install torch_tensorrt wheel file python setup.py install --user - diff --git a/examples/int8/benchmark/BUILD b/examples/int8/benchmark/BUILD index 115dd050c0..3b464b3a99 100644 --- a/examples/int8/benchmark/BUILD +++ b/examples/int8/benchmark/BUILD @@ -5,7 +5,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) diff --git a/examples/int8/ptq/BUILD b/examples/int8/ptq/BUILD index 3f98253af0..26f84a40d9 100644 --- a/examples/int8/ptq/BUILD +++ b/examples/int8/ptq/BUILD @@ -5,7 +5,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -15,7 +15,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -69,7 +69,7 @@ cc_binary( }) + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/examples/int8/qat/BUILD b/examples/int8/qat/BUILD index d8407acee9..9fe1ab1c96 100644 --- a/examples/int8/qat/BUILD +++ b/examples/int8/qat/BUILD @@ -5,7 +5,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -15,7 +15,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -70,7 +70,7 @@ cc_binary( }) + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/examples/torchtrt_runtime_example/BUILD b/examples/torchtrt_runtime_example/BUILD index 5e86b89f14..3c58d960ec 100644 --- a/examples/torchtrt_runtime_example/BUILD +++ b/examples/torchtrt_runtime_example/BUILD @@ -8,7 +8,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -42,7 +42,7 @@ cc_binary( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/tests/core/BUILD b/tests/core/BUILD index ee17e64a4b..2489715993 100644 --- a/tests/core/BUILD +++ b/tests/core/BUILD @@ -3,7 +3,7 @@ load("@rules_cc//cc:defs.bzl", "cc_test") config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -13,7 +13,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) diff --git a/tests/core/lowering/BUILD b/tests/core/lowering/BUILD index 3bc26285ed..2de6b6a7f7 100644 --- a/tests/core/lowering/BUILD +++ b/tests/core/lowering/BUILD @@ -4,7 +4,7 @@ load("//tests/core/lowering:lowering_test.bzl", "lowering_test") config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) diff --git a/tests/core/partitioning/BUILD b/tests/core/partitioning/BUILD index 935f085e1e..b2f99dbd34 100644 --- a/tests/core/partitioning/BUILD +++ b/tests/core/partitioning/BUILD @@ -4,7 +4,7 @@ load("//tests/core/partitioning:partitioning_test.bzl", "partitioning_test") config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) diff --git a/tests/core/runtime/BUILD b/tests/core/runtime/BUILD index 212ce65679..cd02a54b40 100644 --- a/tests/core/runtime/BUILD +++ b/tests/core/runtime/BUILD @@ -5,7 +5,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) diff --git a/tests/cpp/BUILD b/tests/cpp/BUILD index 9b93d4bb56..b50a3c6783 100644 --- a/tests/cpp/BUILD +++ b/tests/cpp/BUILD @@ -5,7 +5,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) diff --git a/tests/util/BUILD b/tests/util/BUILD index 64750fa794..3c9730d8c4 100644 --- a/tests/util/BUILD +++ b/tests/util/BUILD @@ -5,7 +5,7 @@ package(default_visibility = ["//visibility:public"]) config_setting( name = "use_torch_whl", flag_values = { - "//toolchains/build_srcs:torch": "whl" + "//toolchains/dep_src:torch": "whl" }, ) @@ -22,7 +22,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "datacenter" + "//toolchains/dep_collection:compute_lib_collection": "default" }, ) @@ -60,7 +60,7 @@ cc_library( ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], ":sbsa": ["@tensorrt_sbsa//:nvinfer"], - ":jetpack": ["@tensorrt_jetpack//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": [ diff --git a/third_party/tensorrt/archive/BUILD b/third_party/tensorrt/archive/BUILD index c9b32dc23a..80eae0c4b1 100644 --- a/third_party/tensorrt/archive/BUILD +++ b/third_party/tensorrt/archive/BUILD @@ -3,11 +3,13 @@ load("@rules_cc//cc:defs.bzl", "cc_import", "cc_library") package(default_visibility = ["//visibility:public"]) config_setting( - name = "aarch64_linux", + name = "jetpack", constraint_values = [ "@platforms//cpu:aarch64", - "@platforms//os:linux", ], + flag_values = { + "@//toolchains/dep_collection:compute_lib_collection": "jetpack" + }, ) config_setting( @@ -58,6 +60,7 @@ cc_library( "nvinfer_lib", ] + select({ ":windows": ["@cuda_win//:cudart", "nvinfer_static_lib"], + ":jetpack": ["@cuda_l4t//:cudart", "nvinfer_static_lib"], "//conditions:default": ["@cuda//:cudart"], }), ) @@ -200,6 +203,7 @@ cc_library( "nvinfer", ] + select({ ":windows": ["@cuda_win//:cudart"], + ":jetpack": ["@cuda_l4t//:cudart"], "//conditions:default": ["@cuda//:cudart"], }), alwayslink = True, diff --git a/toolchains/dep_collection/BUILD b/toolchains/dep_collection/BUILD index 711a3732fc..f7fe8b33ff 100644 --- a/toolchains/dep_collection/BUILD +++ b/toolchains/dep_collection/BUILD @@ -3,5 +3,5 @@ package(default_visibility = ["//visibility:public"]) load(":defs.bzl", "dep_collection") dep_collection( name = "compute_lib_collection", - build_setting_default = "datacenter", + build_setting_default = "default", ) diff --git a/toolchains/dep_collection/defs.bzl b/toolchains/dep_collection/defs.bzl index 35f3bc654a..6eaa710261 100644 --- a/toolchains/dep_collection/defs.bzl +++ b/toolchains/dep_collection/defs.bzl @@ -1,7 +1,7 @@ # buildifier: disable=module-docstring DependencyCollectionInfo = provider(doc = "", fields = ["type"]) -collection_types = ["datacenter", "jetpack"] +collection_types = ["default", "jetpack"] def _impl(ctx): _type = ctx.build_setting_value diff --git a/toolchains/build_srcs/BUILD b/toolchains/dep_src/BUILD similarity index 73% rename from toolchains/build_srcs/BUILD rename to toolchains/dep_src/BUILD index 555b13272d..b51ef75b78 100644 --- a/toolchains/build_srcs/BUILD +++ b/toolchains/dep_src/BUILD @@ -1,7 +1,7 @@ package(default_visibility = ["//visibility:public"]) -load(":defs.bzl", "build_src") -build_src( +load(":defs.bzl", "dep_src") +dep_src( name = "torch", build_setting_default = "archive", ) diff --git a/toolchains/build_srcs/defs.bzl b/toolchains/dep_src/defs.bzl similarity index 79% rename from toolchains/build_srcs/defs.bzl rename to toolchains/dep_src/defs.bzl index 55da439436..bec0efb8cf 100644 --- a/toolchains/build_srcs/defs.bzl +++ b/toolchains/dep_src/defs.bzl @@ -1,5 +1,5 @@ # buildifier: disable=module-docstring -BuildSrcInfo = provider(doc = "", fields = ["type"]) +DepSrcInfo = provider(doc = "", fields = ["type"]) src_types = ["archive", "whl", "local"] @@ -10,9 +10,9 @@ def _impl(ctx): ", ".join(src_types) + "} but was set to unallowed value " + src) - return BuildSrcInfo(type = src) + return DepSrcInfo(type = src) -build_src = rule( +dep_src = rule( implementation = _impl, build_setting = config.string(flag = True), ) diff --git a/toolchains/jetpack/BUILD b/toolchains/jetpack/BUILD index ed37a864b6..b51796bb3d 100644 --- a/toolchains/jetpack/BUILD +++ b/toolchains/jetpack/BUILD @@ -15,4 +15,4 @@ constraint_value( constraint_value( name = "6.1", constraint_setting = ":jetpack", -) \ No newline at end of file +) From 925913b7693b28db0b08588d477c1bbd17660922 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Mon, 14 Apr 2025 23:40:15 -0600 Subject: [PATCH 4/8] some naming stuff --- .bazelrc | 2 +- MODULE.bazel | 10 ++++++++-- core/BUILD | 4 ++-- core/conversion/BUILD | 4 ++-- core/conversion/conversionctx/BUILD | 4 ++-- core/conversion/converters/BUILD | 4 ++-- core/conversion/evaluators/BUILD | 4 ++-- core/conversion/tensorcontainer/BUILD | 4 ++-- core/conversion/var/BUILD | 4 ++-- core/ir/BUILD | 4 ++-- core/lowering/BUILD | 4 ++-- core/lowering/passes/BUILD | 4 ++-- core/partitioning/BUILD | 4 ++-- core/partitioning/partitioningctx/BUILD | 4 ++-- core/partitioning/partitioninginfo/BUILD | 4 ++-- core/partitioning/segmentedblock/BUILD | 4 ++-- core/plugins/BUILD | 4 ++-- core/runtime/BUILD | 4 ++-- core/util/BUILD | 4 ++-- core/util/logging/BUILD | 4 ++-- examples/int8/ptq/BUILD | 4 ++-- examples/int8/qat/BUILD | 4 ++-- examples/torchtrt_runtime_example/BUILD | 4 ++-- tests/core/BUILD | 4 ++-- tests/util/BUILD | 4 ++-- third_party/tensorrt/archive/BUILD | 2 +- toolchains/dep_collection/BUILD | 2 +- 27 files changed, 57 insertions(+), 51 deletions(-) diff --git a/.bazelrc b/.bazelrc index 542608f849..801b7193d4 100644 --- a/.bazelrc +++ b/.bazelrc @@ -36,7 +36,7 @@ build:cxx11_abi --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=1" build:cxx11_abi --linkopt="-D_GLIBCXX_USE_CXX11_ABI=1" build:cxx11_abi --define=abi=cxx11_abi -build:jetpack --//toolchains/dep_collection:compute_lib_collection=jetpack +build:jetpack --//toolchains/dep_collection:compute_libs=jetpack build:ci_testing --define=torchtrt_src=prebuilt --cxxopt="-DDISABLE_TEST_IN_CI" --action_env "NVIDIA_TF32_OVERRIDE=0" build:use_precompiled_torchtrt --define=torchtrt_src=prebuilt diff --git a/MODULE.bazel b/MODULE.bazel index 9b555971d8..aad30c837c 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -40,10 +40,13 @@ new_local_repository( path = "/usr/local/cuda-12.8", ) +# Server Arm (SBSA) and Jetson Jetpack (L4T) use different versions of CUDA and TensorRT +# These versions can be selected using the flag `--//toolchains/dep_collection:compute_libs="jetpack"` + new_local_repository( name = "cuda_l4t", build_file = "@//third_party/cuda:BUILD", - path = "/usr/local/cuda-12.8", + path = "/usr/local/cuda-12.6", ) new_local_repository( @@ -75,12 +78,15 @@ http_archive( ) +# It is possible to specify a wheel file to use as the libtorch source by providing the URL below and +# using the build flag `--//toolchains/dep_src:torch="whl"` + http_archive( name = "torch_whl", build_file = "@//third_party/libtorch:BUILD", strip_prefix = "torch", type = "zip", - urls = ["https://download.pytorch.org/whl/nightly/cu128/torch-2.8.0.dev20250414%2Bcu128-cp39-cp39-manylinux_2_28_aarch64.whl"], + urls = ["https://pypi.jetson-ai-lab.dev/jp6/cu126/+f/52c/2cbdd62b78f32/torch-2.7.0-cp310-cp310-linux_aarch64.whl#sha256=52c2cbdd62b78f32c51fa178212e4721241a2ba9e0c4d7d690dd808bd890d51b"], ) # Download these tarballs manually from the NVIDIA website diff --git a/core/BUILD b/core/BUILD index f5c9b3e7ee..28cd460690 100644 --- a/core/BUILD +++ b/core/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/conversion/BUILD b/core/conversion/BUILD index f5ee98df74..13696550e6 100644 --- a/core/conversion/BUILD +++ b/core/conversion/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/conversion/conversionctx/BUILD b/core/conversion/conversionctx/BUILD index 02d78118bd..89ff7f613c 100644 --- a/core/conversion/conversionctx/BUILD +++ b/core/conversion/conversionctx/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index 61a5b0231f..9571d91604 100644 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/conversion/evaluators/BUILD b/core/conversion/evaluators/BUILD index 6251499cfc..172e8f6670 100644 --- a/core/conversion/evaluators/BUILD +++ b/core/conversion/evaluators/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/conversion/tensorcontainer/BUILD b/core/conversion/tensorcontainer/BUILD index a570e9f6cf..c812b74a9f 100644 --- a/core/conversion/tensorcontainer/BUILD +++ b/core/conversion/tensorcontainer/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/conversion/var/BUILD b/core/conversion/var/BUILD index fa3f6ad451..65d9583474 100644 --- a/core/conversion/var/BUILD +++ b/core/conversion/var/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/ir/BUILD b/core/ir/BUILD index 9a7cdf33fd..d522a6a018 100644 --- a/core/ir/BUILD +++ b/core/ir/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/lowering/BUILD b/core/lowering/BUILD index d08a52463c..e9b1e1ae07 100644 --- a/core/lowering/BUILD +++ b/core/lowering/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/lowering/passes/BUILD b/core/lowering/passes/BUILD index 961263d289..459eaf25ea 100644 --- a/core/lowering/passes/BUILD +++ b/core/lowering/passes/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/partitioning/BUILD b/core/partitioning/BUILD index 295e84d6cf..2cbcec34b1 100644 --- a/core/partitioning/BUILD +++ b/core/partitioning/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/partitioning/partitioningctx/BUILD b/core/partitioning/partitioningctx/BUILD index 210086d9ec..011a48c6be 100644 --- a/core/partitioning/partitioningctx/BUILD +++ b/core/partitioning/partitioningctx/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/partitioning/partitioninginfo/BUILD b/core/partitioning/partitioninginfo/BUILD index e0b287a6fd..d2a86a2688 100644 --- a/core/partitioning/partitioninginfo/BUILD +++ b/core/partitioning/partitioninginfo/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/partitioning/segmentedblock/BUILD b/core/partitioning/segmentedblock/BUILD index f00b3b2bb6..c463d17b92 100644 --- a/core/partitioning/segmentedblock/BUILD +++ b/core/partitioning/segmentedblock/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/plugins/BUILD b/core/plugins/BUILD index df1428b885..ba167d5f2a 100644 --- a/core/plugins/BUILD +++ b/core/plugins/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/runtime/BUILD b/core/runtime/BUILD index 623e56ab04..f30519619b 100644 --- a/core/runtime/BUILD +++ b/core/runtime/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/util/BUILD b/core/util/BUILD index 1532c2f52a..bc9b53ec8d 100644 --- a/core/util/BUILD +++ b/core/util/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/core/util/logging/BUILD b/core/util/logging/BUILD index 6d8b19e065..d29568cf97 100644 --- a/core/util/logging/BUILD +++ b/core/util/logging/BUILD @@ -16,7 +16,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -26,7 +26,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/examples/int8/ptq/BUILD b/examples/int8/ptq/BUILD index 26f84a40d9..d30c7d3c03 100644 --- a/examples/int8/ptq/BUILD +++ b/examples/int8/ptq/BUILD @@ -15,7 +15,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -25,7 +25,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/examples/int8/qat/BUILD b/examples/int8/qat/BUILD index 9fe1ab1c96..0aab56a02a 100644 --- a/examples/int8/qat/BUILD +++ b/examples/int8/qat/BUILD @@ -15,7 +15,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -25,7 +25,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/examples/torchtrt_runtime_example/BUILD b/examples/torchtrt_runtime_example/BUILD index 3c58d960ec..7102a7f3db 100644 --- a/examples/torchtrt_runtime_example/BUILD +++ b/examples/torchtrt_runtime_example/BUILD @@ -8,7 +8,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -18,7 +18,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/tests/core/BUILD b/tests/core/BUILD index 2489715993..a0e19fa232 100644 --- a/tests/core/BUILD +++ b/tests/core/BUILD @@ -13,7 +13,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -23,7 +23,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/tests/util/BUILD b/tests/util/BUILD index 3c9730d8c4..38e4b5ff6d 100644 --- a/tests/util/BUILD +++ b/tests/util/BUILD @@ -22,7 +22,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "default" + "//toolchains/dep_collection:compute_libs": "default" }, ) @@ -32,7 +32,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "//toolchains/dep_collection:compute_lib_collection": "jetpack" + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/third_party/tensorrt/archive/BUILD b/third_party/tensorrt/archive/BUILD index 80eae0c4b1..4aef027d12 100644 --- a/third_party/tensorrt/archive/BUILD +++ b/third_party/tensorrt/archive/BUILD @@ -8,7 +8,7 @@ config_setting( "@platforms//cpu:aarch64", ], flag_values = { - "@//toolchains/dep_collection:compute_lib_collection": "jetpack" + "@//toolchains/dep_collection:compute_libs": "jetpack" }, ) diff --git a/toolchains/dep_collection/BUILD b/toolchains/dep_collection/BUILD index f7fe8b33ff..4c0dff784c 100644 --- a/toolchains/dep_collection/BUILD +++ b/toolchains/dep_collection/BUILD @@ -2,6 +2,6 @@ package(default_visibility = ["//visibility:public"]) load(":defs.bzl", "dep_collection") dep_collection( - name = "compute_lib_collection", + name = "compute_libs", build_setting_default = "default", ) From 41e8ed055fae237801cde45d3a6ac02350ac98c9 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Tue, 15 Apr 2025 00:07:43 -0600 Subject: [PATCH 5/8] infra: CI updates --- .../workflows/build-test-linux-aarch64.yml | 2 +- .../generate_binary_build_matrix.yml | 6 +- .gitignore | 3 +- MODULE.bazel | 15 ++- docs/_sources/getting_started/jetpack.rst.txt | 3 +- pyproject.toml | 8 +- setup.py | 104 +++++++++--------- toolchains/ci_workspaces/MODULE.bazel.tmpl | 72 +++++++++--- 8 files changed, 128 insertions(+), 85 deletions(-) diff --git a/.github/workflows/build-test-linux-aarch64.yml b/.github/workflows/build-test-linux-aarch64.yml index 966ce31d04..dce15a5cc0 100644 --- a/.github/workflows/build-test-linux-aarch64.yml +++ b/.github/workflows/build-test-linux-aarch64.yml @@ -15,7 +15,7 @@ on: jobs: generate-matrix: - uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + uses: ./.github/workflows/generate_binary_build_matrix.yml with: package-type: wheel os: linux-aarch64 diff --git a/.github/workflows/generate_binary_build_matrix.yml b/.github/workflows/generate_binary_build_matrix.yml index 6bbbc90852..3352b6cb2d 100644 --- a/.github/workflows/generate_binary_build_matrix.yml +++ b/.github/workflows/generate_binary_build_matrix.yml @@ -79,7 +79,10 @@ jobs: with: repository: ${{ inputs.test-infra-repository }} ref: ${{ inputs.test-infra-ref }} - - uses: ./.github/actions/set-channel + - uses: actions/checkout@v4 + with: + repository: pytorch/tensorrt + - uses: pytorch/test-infra/.github/actions/set-channel@main - name: Generate test matrix id: generate env: @@ -101,7 +104,6 @@ jobs: PYTHON_VERSIONS: ${{ inputs.python-versions }} run: | set -eou pipefail - ls -la . MATRIX_BLOB="$(python3 .github/scripts/generate_binary_build_matrix.py)" echo "${MATRIX_BLOB}" echo "matrix=${MATRIX_BLOB}" >> "${GITHUB_OUTPUT}" diff --git a/.gitignore b/.gitignore index 16e4f4f838..786b9542d6 100644 --- a/.gitignore +++ b/.gitignore @@ -74,4 +74,5 @@ tests/py/dynamo/models/*.ts tests/py/dynamo/models/*.ep *.deb *.tar.xz -MODULE.bazel.lock \ No newline at end of file +MODULE.bazel.lock +*.whl diff --git a/MODULE.bazel b/MODULE.bazel index aad30c837c..b752abf10b 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -69,7 +69,6 @@ http_archive( urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-cxx11-abi-shared-with-deps-latest.zip"], ) - http_archive( name = "libtorch_win", build_file = "@//third_party/libtorch:BUILD", @@ -81,13 +80,13 @@ http_archive( # It is possible to specify a wheel file to use as the libtorch source by providing the URL below and # using the build flag `--//toolchains/dep_src:torch="whl"` -http_archive( - name = "torch_whl", - build_file = "@//third_party/libtorch:BUILD", - strip_prefix = "torch", - type = "zip", - urls = ["https://pypi.jetson-ai-lab.dev/jp6/cu126/+f/52c/2cbdd62b78f32/torch-2.7.0-cp310-cp310-linux_aarch64.whl#sha256=52c2cbdd62b78f32c51fa178212e4721241a2ba9e0c4d7d690dd808bd890d51b"], -) +# http_archive( +# name = "torch_whl", +# build_file = "@//third_party/libtorch:BUILD", +# strip_prefix = "torch", +# type = "zip", +# urls = ["https://pypi.jetson-ai-lab.dev/jp6/cu126/+f/52c/2cbdd62b78f32/torch-2.7.0-cp310-cp310-linux_aarch64.whl#sha256=52c2cbdd62b78f32c51fa178212e4721241a2ba9e0c4d7d690dd808bd890d51b"], +# ) # Download these tarballs manually from the NVIDIA website # Either place them in the distdir directory in third_party and use the --distdir flag diff --git a/docs/_sources/getting_started/jetpack.rst.txt b/docs/_sources/getting_started/jetpack.rst.txt index b99c45addd..ddbf89dc63 100644 --- a/docs/_sources/getting_started/jetpack.rst.txt +++ b/docs/_sources/getting_started/jetpack.rst.txt @@ -1,4 +1,4 @@ -.. _Torch_TensorRT_in_l4t_6.1 +.. _Torch_TensorRT_in_JetPack_6.1 Overview ################## @@ -116,3 +116,4 @@ Please make sure to build torch_tensorrt wheel file from source release/2.5 bran cat toolchains/jp_workspaces/MODULE.bazel.tmpl | envsubst > MODULE.bazel # build and install torch_tensorrt wheel file python setup.py install --user + diff --git a/pyproject.toml b/pyproject.toml index 120a767822..3feacec679 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ requires = [ "cffi>=1.15.1", "typing-extensions>=4.7.0", "future>=0.18.3", - "tensorrt->=10.9.0,<10.10.0", + "tensorrt-cu12>=10.9.0,<10.10.0", "torch>=2.8.0.dev,<2.9.0", "pybind11==2.6.2", "numpy", @@ -57,9 +57,9 @@ keywords = [ dependencies = [ "torch>=2.8.0.dev,<2.9.0", "tensorrt>=10.9.0,<10.10.0", - "tensorrt->=10.9.0,<10.10.0", - "tensorrt--bindings>=10.9.0,<10.10.0", - "tensorrt--libs>=10.9.0,<10.10.0", + "tensorrt-cu12>=10.9.0,<10.10.0", + "tensorrt-cu12-bindings>=10.9.0,<10.10.0", + "tensorrt-cu12-libs>=10.9.0,<10.10.0", "packaging>=23", "numpy", "typing-extensions>=4.7.0", diff --git a/setup.py b/setup.py index 09933307c8..18b7756167 100644 --- a/setup.py +++ b/setup.py @@ -79,12 +79,13 @@ def load_dep_info(): dir_path = os.path.join(str(get_root_dir()), "py") PRE_CXX11_ABI = False -JETPACK_VERSION = None +IS_JETPACK = False PY_ONLY = False NO_TS = False LEGACY = False RELEASE = False CI_BUILD = False +IS_SBSA = False if "--fx-only" in sys.argv: PY_ONLY = True @@ -121,6 +122,9 @@ def load_dep_info(): if (gpu_arch_version := os.environ.get("CU_VERSION")) is None: gpu_arch_version = f"cu{__cuda_version__.replace('.','')}" +if (jetpack := os.environ.get("JETPACK_BUILD")) is None: + if jetpack == "1": + IS_JETPACK = True if RELEASE: __version__ = os.environ.get("BUILD_VERSION") @@ -136,40 +140,12 @@ def load_dep_info(): if ci_env_var == "1": CI_BUILD = True -if "--use-pre-cxx11-abi" in sys.argv: - sys.argv.remove("--use-pre-cxx11-abi") - PRE_CXX11_ABI = True - -if (pre_cxx11_abi_env_var := os.environ.get("USE_PRE_CXX11_ABI")) is not None: - if pre_cxx11_abi_env_var == "1": - PRE_CXX11_ABI = True - if platform.uname().processor == "aarch64": - if "--jetpack-version" in sys.argv: - version_idx = sys.argv.index("--jetpack-version") + 1 - version = sys.argv[version_idx] - sys.argv.remove(version) - sys.argv.remove("--jetpack-version") - if version == "4.5": - JETPACK_VERSION = "4.5" - elif version == "4.6": - JETPACK_VERSION = "4.6" - elif version == "5.0": - JETPACK_VERSION = "5.0" - elif version == "6.1": - JETPACK_VERSION = "6.1" - - if not JETPACK_VERSION: - warnings.warn( - "Assuming jetpack version to be 6.1, if not use the --jetpack-version option" - ) - JETPACK_VERSION = "6.1" - - if PRE_CXX11_ABI: - warnings.warn( - "Jetson platform detected. Please remove --use-pre-cxx11-abi flag if you are using it." - ) - + if "--jetpack" in sys.argv: + sys.argv.remove("--jetpack") + IS_JETPACK = True + else: + IS_SBSA = True BAZEL_EXE = None if not PY_ONLY: @@ -204,30 +180,13 @@ def build_libtorchtrt_cxx11_abi( if target_python: cmd.append("--config=python") - if pre_cxx11_abi: - cmd.append("--config=pre_cxx11_abi") - print("using PRE CXX11 ABI build") - else: - cmd.append("--config=cxx11_abi") - print("using CXX11 ABI build") - if IS_WINDOWS: cmd.append("--config=windows") else: cmd.append("--config=linux") - if JETPACK_VERSION == "4.5": - cmd.append("--platforms=//toolchains:jetpack_4.5") - print("Jetpack version: 4.5") - elif JETPACK_VERSION == "4.6": - cmd.append("--platforms=//toolchains:jetpack_4.6") - print("Jetpack version: 4.6") - elif JETPACK_VERSION == "5.0": - cmd.append("--platforms=//toolchains:jetpack_5.0") - print("Jetpack version: 5.0") - elif JETPACK_VERSION == "6.1": - cmd.append("--platforms=//toolchains:jetpack_6.1") - print("Jetpack version: 6.1") + if IS_JETPACK: + cmd.append("--config=jetpack") if CI_BUILD: cmd.append("--platforms=//toolchains:ci_rhel_x86_64_linux") @@ -497,7 +456,16 @@ def run(self): package_data = {} if not (PY_ONLY or NO_TS): - tensorrt_linux_external_dir = ( + tensorrt_windows_external_dir = ( + lambda: subprocess.check_output( + [BAZEL_EXE, "query", "@tensorrt_win//:nvinfer", "--output", "location"] + ) + .decode("ascii") + .strip() + .split("/BUILD.bazel")[0] + ) + + tensorrt_x86_64_external_dir = ( lambda: subprocess.check_output( [BAZEL_EXE, "query", "@tensorrt//:nvinfer", "--output", "location"] ) @@ -505,9 +473,35 @@ def run(self): .strip() .split("/BUILD.bazel")[0] ) + + tensorrt_sbsa_external_dir = ( + lambda: subprocess.check_output( + [BAZEL_EXE, "query", "@tensorrt_sbsa//:nvinfer", "--output", "location"] + ) + .decode("ascii") + .strip() + .split("/BUILD.bazel")[0] + ) + + tensorrt_jetpack_external_dir = ( + lambda: subprocess.check_output( + [BAZEL_EXE, "query", "@tensorrt_l4t//:nvinfer", "--output", "location"] + ) + .decode("ascii") + .strip() + .split("/BUILD.bazel")[0] + ) + + if IS_SBSA: + tensorrt_linux_external_dir = tensorrt_sbsa_external_dir() + elif IS_JETPACK: + tensorrt_linux_external_dir = tensorrt_jetpack_external_dir() + else: + tensorrt_linux_external_dir = tensorrt_x86_64_external_dir() + tensorrt_windows_external_dir = ( lambda: subprocess.check_output( - [BAZEL_EXE, "query", "@tensorrt_win//:nvinfer", "--output", "location"] + [BAZEL_EXE, "query", "@tensorrt_windows//:nvinfer", "--output", "location"] ) .decode("ascii") .strip() diff --git a/toolchains/ci_workspaces/MODULE.bazel.tmpl b/toolchains/ci_workspaces/MODULE.bazel.tmpl index e7a651c92d..dde6edc651 100644 --- a/toolchains/ci_workspaces/MODULE.bazel.tmpl +++ b/toolchains/ci_workspaces/MODULE.bazel.tmpl @@ -4,10 +4,10 @@ module( version = "${BUILD_VERSION}" ) -bazel_dep(name = "googletest", version = "1.14.0") -bazel_dep(name = "platforms", version = "0.0.10") -bazel_dep(name = "rules_cc", version = "0.0.9") -bazel_dep(name = "rules_python", version = "0.34.0") +bazel_dep(name = "googletest", version = "1.16.0") +bazel_dep(name = "platforms", version = "0.0.11") +bazel_dep(name = "rules_cc", version = "0.1.1") +bazel_dep(name = "rules_python", version = "1.3.0") python = use_extension("@rules_python//python/extensions:python.bzl", "python") python.toolchain( @@ -27,7 +27,7 @@ local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl" # External dependency for torch_tensorrt if you already have precompiled binaries. local_repository( name = "torch_tensorrt", - path = "/opt/conda/lib/python3.8/site-packages/torch_tensorrt", + path = "/opt/conda/lib/python3.10/site-packages/torch_tensorrt", ) @@ -40,6 +40,15 @@ new_local_repository( path = "${CUDA_HOME}", ) +# Server Arm (SBSA) and Jetson Jetpack (L4T) use different versions of CUDA and TensorRT +# These versions can be selected using the flag `--//toolchains/dep_collection:compute_libs="jetpack"` + +new_local_repository( + name = "cuda_l4t", + build_file = "@//third_party/cuda:BUILD", + path = "/usr/local/cuda-12.6", +) + new_local_repository( name = "cuda_win", build_file = "@//third_party/cuda:BUILD", @@ -53,12 +62,31 @@ http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "ht # Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs) ############################################################################################################# -http_archive( - name = "libtorch", - build_file = "@//third_party/libtorch:BUILD", - strip_prefix = "libtorch", - urls = ["https://download.pytorch.org/libtorch/${CHANNEL}/${CU_VERSION}/libtorch-cxx11-abi-shared-with-deps-latest.zip"], -) +# http_archive( +# name = "libtorch", +# build_file = "@//third_party/libtorch:BUILD", +# strip_prefix = "libtorch", +# urls = ["https://download.pytorch.org/libtorch/${CHANNEL}/${CU_VERSION}/libtorch-cxx11-abi-shared-with-deps-latest.zip"], +# ) + +# http_archive( +# name = "libtorch_win", +# build_file = "@//third_party/libtorch:BUILD", +# strip_prefix = "libtorch", +# urls = ["https://download.pytorch.org/libtorch//${CHANNEL}/${CU_VERSION}/libtorch-win-shared-with-deps-latest.zip"], +# ) + + +# It is possible to specify a wheel file to use as the libtorch source by providing the URL below and +# using the build flag `--//toolchains/dep_src:torch="whl"` + +# http_archive( +# name = "torch_whl", +# build_file = "@//third_party/libtorch:BUILD", +# strip_prefix = "torch", +# type = "zip", +# urls = ["file:///${TORCH_WHL_PATH}"], +# ) # Download these tarballs manually from the NVIDIA website # Either place them in the distdir directory in third_party and use the --distdir flag @@ -73,6 +101,24 @@ http_archive( ], ) +http_archive( + name = "tensorrt_sbsa", + build_file = "@//third_party/tensorrt/archive:BUILD", + strip_prefix = "TensorRT-10.9.0.34", + urls = [ + "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.9.0/tars/TensorRT-10.9.0.34.Linux.aarch64-gnu.cuda-12.8.tar.gz", + ], +) + +http_archive( + name = "tensorrt_l4t", + build_file = "@//third_party/tensorrt/archive:BUILD", + strip_prefix = "TensorRT-10.3.0.26", + urls = [ + "https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.3.0/tars/TensorRT-10.3.0.26.l4t.aarch64-gnu.cuda-12.6.tar.gz", + ], +) + http_archive( name = "tensorrt_win", build_file = "@//third_party/tensorrt/archive:BUILD", @@ -95,13 +141,13 @@ http_archive( # for both versions here and do not use --config=pre-cxx11-abi new_local_repository( - name = "libtorch_win", + name = "libtorch", path = "${TORCH_INSTALL_PATH}", build_file = "third_party/libtorch/BUILD" ) new_local_repository( - name = "libtorch_pre_cxx11_abi", + name = "libtorch_win", path = "${TORCH_INSTALL_PATH}", build_file = "third_party/libtorch/BUILD" ) From d59d84d554ee3046ea0e4b7f81bf5fd57a6d0086 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Tue, 15 Apr 2025 09:36:18 -0600 Subject: [PATCH 6/8] aarch64 python build --- .github/scripts/filter-matrix.py | 8 ++++++-- .../scripts/generate_binary_build_matrix.py | 2 +- MODULE.bazel | 14 ++++++------- setup.py | 20 ++++++++----------- 4 files changed, 22 insertions(+), 22 deletions(-) diff --git a/.github/scripts/filter-matrix.py b/.github/scripts/filter-matrix.py index 685ea98985..401e1cd68d 100644 --- a/.github/scripts/filter-matrix.py +++ b/.github/scripts/filter-matrix.py @@ -26,8 +26,12 @@ def main(args: List[str]) -> None: includes = matrix_dict["include"] filtered_includes = [] for item in includes: - if all([item["python_version"] not in disabled_python_versions, - item["desired_cuda"] != cpu_tag]): + if all( + [ + item["python_version"] not in disabled_python_versions, + item["desired_cuda"] != cpu_tag + ] + ): filtered_includes.append(item) filtered_matrix_dict = {} filtered_matrix_dict["include"] = filtered_includes diff --git a/.github/scripts/generate_binary_build_matrix.py b/.github/scripts/generate_binary_build_matrix.py index b7e038f98f..a159a25f8d 100644 --- a/.github/scripts/generate_binary_build_matrix.py +++ b/.github/scripts/generate_binary_build_matrix.py @@ -50,7 +50,7 @@ "release": "12.4", } -CUDA_AARCH64_ARCHES = ["12.8-aarch64", "12.6-aarch64", "11.8-aarch64"] +CUDA_AARCH64_ARCHES = ["12.8-aarch64", "12.6-aarch64"] PACKAGE_TYPES = ["wheel", "conda", "libtorch"] CXX11_ABI = "cxx11-abi" diff --git a/MODULE.bazel b/MODULE.bazel index b752abf10b..17f1c05bbc 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -80,13 +80,13 @@ http_archive( # It is possible to specify a wheel file to use as the libtorch source by providing the URL below and # using the build flag `--//toolchains/dep_src:torch="whl"` -# http_archive( -# name = "torch_whl", -# build_file = "@//third_party/libtorch:BUILD", -# strip_prefix = "torch", -# type = "zip", -# urls = ["https://pypi.jetson-ai-lab.dev/jp6/cu126/+f/52c/2cbdd62b78f32/torch-2.7.0-cp310-cp310-linux_aarch64.whl#sha256=52c2cbdd62b78f32c51fa178212e4721241a2ba9e0c4d7d690dd808bd890d51b"], -# ) +http_archive( + name = "torch_whl", + build_file = "@//third_party/libtorch:BUILD", + strip_prefix = "torch", + type = "zip", + urls = ["https://download.pytorch.org/whl/nightly/cu128/torch-2.8.0.dev20250415%2Bcu128-cp310-cp310-manylinux_2_28_aarch64.whl"], +) # Download these tarballs manually from the NVIDIA website # Either place them in the distdir directory in third_party and use the --distdir flag diff --git a/setup.py b/setup.py index 18b7756167..3798a20fcd 100644 --- a/setup.py +++ b/setup.py @@ -188,6 +188,10 @@ def build_libtorchtrt_cxx11_abi( if IS_JETPACK: cmd.append("--config=jetpack") + if IS_SBSA: + if CI_BUILD: + cmd.append("--//toolchains/dep_src:torch=whl") + if CI_BUILD: cmd.append("--platforms=//toolchains:ci_rhel_x86_64_linux") print("CI based build") @@ -456,14 +460,6 @@ def run(self): package_data = {} if not (PY_ONLY or NO_TS): - tensorrt_windows_external_dir = ( - lambda: subprocess.check_output( - [BAZEL_EXE, "query", "@tensorrt_win//:nvinfer", "--output", "location"] - ) - .decode("ascii") - .strip() - .split("/BUILD.bazel")[0] - ) tensorrt_x86_64_external_dir = ( lambda: subprocess.check_output( @@ -493,15 +489,15 @@ def run(self): ) if IS_SBSA: - tensorrt_linux_external_dir = tensorrt_sbsa_external_dir() + tensorrt_linux_external_dir = tensorrt_sbsa_external_dir elif IS_JETPACK: - tensorrt_linux_external_dir = tensorrt_jetpack_external_dir() + tensorrt_linux_external_dir = tensorrt_jetpack_external_dir else: - tensorrt_linux_external_dir = tensorrt_x86_64_external_dir() + tensorrt_linux_external_dir = tensorrt_x86_64_external_dir tensorrt_windows_external_dir = ( lambda: subprocess.check_output( - [BAZEL_EXE, "query", "@tensorrt_windows//:nvinfer", "--output", "location"] + [BAZEL_EXE, "query", "@tensorrt_win//:nvinfer", "--output", "location"] ) .decode("ascii") .strip() From 5b7a0349712d0bca0f13312a08717296b853b0d5 Mon Sep 17 00:00:00 2001 From: Naren Dasan <1790613+narendasan@users.noreply.github.com> Date: Tue, 15 Apr 2025 09:55:13 -0600 Subject: [PATCH 7/8] Update build-test-windows.yml --- .github/workflows/build-test-windows.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build-test-windows.yml b/.github/workflows/build-test-windows.yml index 00bea34649..f78218e75d 100644 --- a/.github/workflows/build-test-windows.yml +++ b/.github/workflows/build-test-windows.yml @@ -15,7 +15,7 @@ on: jobs: generate-matrix: - uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + uses: ./.github/workflows/generate_binary_build_matrix.yml with: package-type: wheel os: windows From 23ae053eba595f04c2f8115191f6e6f594b80502 Mon Sep 17 00:00:00 2001 From: Naren Dasan <1790613+narendasan@users.noreply.github.com> Date: Tue, 15 Apr 2025 09:55:45 -0600 Subject: [PATCH 8/8] Update build-test-linux-x86_64.yml --- .github/workflows/build-test-linux-x86_64.yml | 2 +- packaging/pre_build_script.sh | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build-test-linux-x86_64.yml b/.github/workflows/build-test-linux-x86_64.yml index 66aae7ea1e..9951757ef3 100644 --- a/.github/workflows/build-test-linux-x86_64.yml +++ b/.github/workflows/build-test-linux-x86_64.yml @@ -15,7 +15,7 @@ on: jobs: generate-matrix: - uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + uses: ./.github/workflows/generate_binary_build_matrix.yml with: package-type: wheel os: linux diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh index 247cf2105b..edc4ff8b76 100755 --- a/packaging/pre_build_script.sh +++ b/packaging/pre_build_script.sh @@ -5,6 +5,8 @@ set -x # Install dependencies python3 -m pip install pyyaml +install -y ninja-build gettext + PLATFORM="amd64" PLATFORM=x86_64 BAZEL_PLATFORM=amd64