diff --git a/.bazelrc b/.bazelrc index 83df4f686b..801b7193d4 100644 --- a/.bazelrc +++ b/.bazelrc @@ -36,6 +36,8 @@ 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_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/.bazelversion b/.bazelversion index b26a34e470..0e79152459 100644 --- a/.bazelversion +++ b/.bazelversion @@ -1 +1 @@ -7.2.1 +8.1.1 diff --git a/.github/scripts/filter-matrix.py b/.github/scripts/filter-matrix.py index 69ee24080a..bfa24ebc0e 100644 --- a/.github/scripts/filter-matrix.py +++ b/.github/scripts/filter-matrix.py @@ -2,10 +2,12 @@ import argparse import json +import os import sys from typing import List -disabled_python_versions: List[str] = [] +# currently we don't support python 3.13t due to tensorrt does not support 3.13t +disabled_python_versions: List[str] = ["3.13t"] def main(args: list[str]) -> None: @@ -21,11 +23,20 @@ def main(args: list[str]) -> None: if options.matrix == "": raise Exception("--matrix needs to be provided") + matrix_dict = json.loads(options.matrix) includes = matrix_dict["include"] filtered_includes = [] for item in includes: - if item["python_version"] not in disabled_python_versions: + if item["python_version"] in disabled_python_versions: + continue + if item["gpu_arch_type"] == "cuda-aarch64": + # pytorch image:pytorch/manylinuxaarch64-builder:cuda12.8 comes with glibc2.28 + # however, TensorRT requires glibc2.31 on aarch64 platform + # TODO: in future, if pytorch supports aarch64 with glibc2.31, we should switch to use the pytorch image + item["container_image"] = "quay.io/pypa/manylinux_2_34_aarch64" + filtered_includes.append(item) + else: filtered_includes.append(item) filtered_matrix_dict = {} filtered_matrix_dict["include"] = filtered_includes diff --git a/.github/scripts/generate-tensorrt-test-matrix.py b/.github/scripts/generate-tensorrt-test-matrix.py index 02b0f746ca..17f61b9ccf 100644 --- a/.github/scripts/generate-tensorrt-test-matrix.py +++ b/.github/scripts/generate-tensorrt-test-matrix.py @@ -21,8 +21,8 @@ # channel: test if the future tensorRT version test workflow is triggered from the release branch(release/2.5 etc....) PYTHON_VERSIONS_DICT = { "nightly": ["3.11"], - "test": ["3.9", "3.10", "3.11", "3.12"], - "release": ["3.9", "3.10", "3.11", "3.12"], + "test": ["3.9", "3.10", "3.11", "3.12", "3.13"], + "release": ["3.9", "3.10", "3.11", "3.12", "3.13"], } # please update the future tensorRT version you want to test here diff --git a/.github/scripts/generate_binary_build_matrix.py b/.github/scripts/generate_binary_build_matrix.py deleted file mode 100644 index 4bdcd897e3..0000000000 --- a/.github/scripts/generate_binary_build_matrix.py +++ /dev/null @@ -1,679 +0,0 @@ -#!/usr/bin/env python3 - -"""Generates a matrix to be utilized through github actions - -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: - * CPU - * Latest CUDA - * Latest ROCM - * Latest XPU -""" - - -import argparse -import json -import os -import sys -from typing import Any, Callable, Dict, List, Optional, Tuple - -PYTHON_VERSIONS_FOR_PR_BUILD = ["3.11"] -PYTHON_ARCHES_DICT = { - "nightly": ["3.9", "3.10", "3.11", "3.12", "3.13"], - "test": ["3.9", "3.10", "3.11", "3.12", "3.13"], - "release": ["3.9", "3.10", "3.11", "3.12", "3.13"], -} -CUDA_ARCHES_DICT = { - "nightly": ["11.8", "12.6", "12.8"], - "test": ["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"], -} - -PACKAGE_TYPES = ["wheel", "conda", "libtorch"] -CXX11_ABI = "cxx11-abi" -RELEASE = "release" -DEBUG = "debug" -NIGHTLY = "nightly" -TEST = "test" - -# OS constants -LINUX = "linux" -LINUX_AARCH64 = "linux-aarch64" -MACOS_ARM64 = "macos-arm64" -WINDOWS = "windows" - -# Accelerator architectures -CPU = "cpu" -CPU_AARCH64 = "cpu-aarch64" -CUDA_AARCH64 = "cuda-aarch64" -CUDA = "cuda" -ROCM = "rocm" -XPU = "xpu" - - -CURRENT_NIGHTLY_VERSION = "2.8.0" -CURRENT_CANDIDATE_VERSION = "2.7.0" -CURRENT_STABLE_VERSION = "2.7.0" -CURRENT_VERSION = CURRENT_STABLE_VERSION - -# By default use Nightly for CUDA arches -CUDA_ARCHES = CUDA_ARCHES_DICT[NIGHTLY] -ROCM_ARCHES = ROCM_ARCHES_DICT[NIGHTLY] -PYTHON_ARCHES = PYTHON_ARCHES_DICT[NIGHTLY] - -# Container images -LIBTORCH_CONTAINER_IMAGES: Dict[Tuple[str, str], str] -WHEEL_CONTAINER_IMAGES: Dict[str, str] - -LINUX_GPU_RUNNER = "linux.g5.4xlarge.nvidia.gpu" -LINUX_CPU_RUNNER = "linux.2xlarge" -LINUX_AARCH64_RUNNER = "linux.arm64.2xlarge" -LINUX_AARCH64_GPU_RUNNER = "linux.arm64.m7g.4xlarge" -WIN_GPU_RUNNER = "windows.g4dn.xlarge" -WIN_CPU_RUNNER = "windows.4xlarge" -MACOS_M1_RUNNER = "macos-m1-stable" - -PACKAGES_TO_INSTALL_WHL = "torch torchvision torchaudio" -WHL_INSTALL_BASE = "pip3 install" -DOWNLOAD_URL_BASE = "https://download.pytorch.org" - -ENABLE = "enable" -DISABLE = "disable" - - -def arch_type(arch_version: str) -> str: - if arch_version in CUDA_ARCHES: - return CUDA - elif arch_version in ROCM_ARCHES: - return ROCM - elif arch_version == CPU_AARCH64: - return CPU_AARCH64 - elif arch_version == CUDA_AARCH64: - return CUDA_AARCH64 - elif arch_version == XPU: - return XPU - else: # arch_version should always be CPU in this case - return CPU - - -def validation_runner(arch_type: str, os: str) -> str: - if os == LINUX: - if arch_type == CUDA: - return LINUX_GPU_RUNNER - else: - return LINUX_CPU_RUNNER - elif os == LINUX_AARCH64: - if arch_type == CUDA_AARCH64: - return LINUX_AARCH64_GPU_RUNNER - else: - return LINUX_AARCH64_RUNNER - elif os == WINDOWS: - if arch_type == CUDA: - return WIN_GPU_RUNNER - else: - return WIN_CPU_RUNNER - elif os == MACOS_ARM64: - return MACOS_M1_RUNNER - else: # default to linux cpu runner - return LINUX_CPU_RUNNER - - -def initialize_globals(channel: str, build_python_only: bool) -> None: - global CURRENT_VERSION, CUDA_ARCHES, ROCM_ARCHES, PYTHON_ARCHES - global WHEEL_CONTAINER_IMAGES, LIBTORCH_CONTAINER_IMAGES - if channel == TEST: - CURRENT_VERSION = CURRENT_CANDIDATE_VERSION - else: - CURRENT_VERSION = CURRENT_STABLE_VERSION - - CUDA_ARCHES = CUDA_ARCHES_DICT[channel] - ROCM_ARCHES = ROCM_ARCHES_DICT[channel] - if build_python_only: - # Only select the oldest version of python if building a python only package - PYTHON_ARCHES = [PYTHON_ARCHES_DICT[channel][0]] - else: - PYTHON_ARCHES = PYTHON_ARCHES_DICT[channel] - WHEEL_CONTAINER_IMAGES = { - **{ - gpu_arch: f"pytorch/manylinux2_28-builder:cuda{gpu_arch}" - for gpu_arch in CUDA_ARCHES - }, - **{ - gpu_arch: f"pytorch/manylinux2_28-builder:rocm{gpu_arch}" - for gpu_arch in ROCM_ARCHES - }, - CPU: "pytorch/manylinux2_28-builder:cpu", - XPU: "pytorch/manylinux2_28-builder:xpu", - # TODO: Migrate CUDA_AARCH64 image to manylinux2_28_aarch64-builder:cuda12.6 - CPU_AARCH64: "pytorch/manylinux2_28_aarch64-builder:cpu-aarch64", - CUDA_AARCH64: "pytorch/manylinuxaarch64-builder:cuda12.6", - } - LIBTORCH_CONTAINER_IMAGES = { - **{ - (gpu_arch, CXX11_ABI): f"pytorch/libtorch-cxx11-builder:cuda{gpu_arch}" - for gpu_arch in CUDA_ARCHES - }, - **{ - (gpu_arch, CXX11_ABI): f"pytorch/libtorch-cxx11-builder:rocm{gpu_arch}" - for gpu_arch in ROCM_ARCHES - }, - (CPU, CXX11_ABI): "pytorch/libtorch-cxx11-builder:cpu", - } - - -def translate_desired_cuda(gpu_arch_type: str, gpu_arch_version: str) -> str: - return { - CPU: "cpu", - CPU_AARCH64: CPU, - CUDA_AARCH64: "cu124", - CUDA: f"cu{gpu_arch_version.replace('.', '')}", - ROCM: f"rocm{gpu_arch_version}", - XPU: "xpu", - }.get(gpu_arch_type, gpu_arch_version) - - -def list_without(in_list: List[str], without: List[str]) -> List[str]: - return [item for item in in_list if item not in without] - - -def get_base_download_url_for_repo( - repo: str, channel: str, gpu_arch_type: str, desired_cuda: str -) -> str: - base_url_for_type = f"{DOWNLOAD_URL_BASE}/{repo}" - base_url_for_type = ( - base_url_for_type if channel == RELEASE else f"{base_url_for_type}/{channel}" - ) - - if gpu_arch_type != CPU: - base_url_for_type = f"{base_url_for_type}/{desired_cuda}" - else: - base_url_for_type = f"{base_url_for_type}/{gpu_arch_type}" - - return base_url_for_type - - -def get_libtorch_install_command( - os: str, - channel: str, - gpu_arch_type: str, - libtorch_variant: str, - devtoolset: str, - desired_cuda: str, - libtorch_config: str, -) -> str: - prefix = "libtorch" if os != WINDOWS else "libtorch-win" - _libtorch_variant = ( - f"{libtorch_variant}-{libtorch_config}" - if libtorch_config == "debug" - else libtorch_variant - ) - build_name = ( - f"{prefix}-{devtoolset}-{_libtorch_variant}-latest.zip" - if devtoolset == "cxx11-abi" - else f"{prefix}-{_libtorch_variant}-latest.zip" - ) - - if os == MACOS_ARM64: - arch = "arm64" - build_name = f"libtorch-macos-{arch}-latest.zip" - if channel in [RELEASE, TEST]: - build_name = f"libtorch-macos-{arch}-{CURRENT_VERSION}.zip" - - elif os == LINUX and (channel in (RELEASE, TEST)): - build_name = ( - f"{prefix}-{devtoolset}-{_libtorch_variant}-{CURRENT_VERSION}%2B{desired_cuda}.zip" - if devtoolset == "cxx11-abi" - else f"{prefix}-{_libtorch_variant}-{CURRENT_VERSION}%2B{desired_cuda}.zip" - ) - elif os == WINDOWS and (channel in (RELEASE, TEST)): - build_name = ( - f"{prefix}-shared-with-deps-debug-{CURRENT_VERSION}%2B{desired_cuda}.zip" - if libtorch_config == "debug" - else f"{prefix}-shared-with-deps-{CURRENT_VERSION}%2B{desired_cuda}.zip" - ) - elif os == WINDOWS and channel == NIGHTLY: - build_name = ( - f"{prefix}-shared-with-deps-debug-latest.zip" - if libtorch_config == "debug" - else f"{prefix}-shared-with-deps-latest.zip" - ) - - return f"{get_base_download_url_for_repo('libtorch', channel, gpu_arch_type, desired_cuda)}/{build_name}" - - -def get_wheel_install_command( - os: str, - channel: str, - gpu_arch_type: str, - gpu_arch_version: str, - desired_cuda: str, - python_version: str, - use_only_dl_pytorch_org: bool, - use_split_build: bool = False, -) -> str: - if use_split_build: - if (gpu_arch_version in CUDA_ARCHES) and (os == LINUX) and (channel == NIGHTLY): - 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 - ) - if ( - channel == RELEASE - and (not use_only_dl_pytorch_org) - and ( - (gpu_arch_version == "12.4" and os == LINUX) - or (gpu_arch_type == CPU and os in [WINDOWS, MACOS_ARM64]) - or (os == LINUX_AARCH64) - ) - ): - return f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL}" - else: - whl_install_command = ( - f"{WHL_INSTALL_BASE} --pre {PACKAGES_TO_INSTALL_WHL}" - if channel == "nightly" - else f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL}" - ) - return f"{whl_install_command} --index-url {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}" # noqa: E501 - - -def generate_conda_matrix( - os: str, - channel: str, - with_cuda: str, - with_rocm: str, - with_cpu: str, - with_xpu: str, - limit_pr_builds: bool, - use_only_dl_pytorch_org: bool, - use_split_build: bool = False, - python_versions: Optional[List[str]] = None, -) -> List[Dict[str, str]]: - ret: List[Dict[str, str]] = [] - # return empty list. Conda builds are deprecated, see https://github.com/pytorch/pytorch/issues/138506 - return ret - - -def generate_libtorch_matrix( - os: str, - channel: str, - with_cuda: str, - with_rocm: str, - with_cpu: str, - with_xpu: str, - limit_pr_builds: bool, - use_only_dl_pytorch_org: bool, - use_split_build: bool = False, - python_versions: Optional[List[str]] = None, - abi_versions: Optional[List[str]] = None, - arches: Optional[List[str]] = None, - libtorch_variants: Optional[List[str]] = None, -) -> List[Dict[str, str]]: - ret: List[Dict[str, str]] = [] - - if arches is None: - arches = [] - - if with_cpu == ENABLE: - arches += [CPU] - - if with_cuda == ENABLE and os in (LINUX, WINDOWS): - arches += CUDA_ARCHES - - if with_rocm == ENABLE and os == LINUX: - arches += ROCM_ARCHES - - if abi_versions is None: - if os == WINDOWS: - abi_versions = [RELEASE, DEBUG] - elif os == LINUX: - abi_versions = [CXX11_ABI] - elif os in [MACOS_ARM64]: - abi_versions = [CXX11_ABI] - else: - abi_versions = [] - - if libtorch_variants is None: - libtorch_variants = [ - "shared-with-deps", - ] - - global LIBTORCH_CONTAINER_IMAGES - - for abi_version in abi_versions: - for arch_version in arches: - for libtorch_variant in libtorch_variants: - # one of the values in the following list must be exactly - # CXX11_ABI, but the precise value of the other one doesn't - # matter - gpu_arch_type = arch_type(arch_version) - gpu_arch_version = "" if arch_version == CPU else arch_version - - desired_cuda = translate_desired_cuda(gpu_arch_type, gpu_arch_version) - devtoolset = abi_version if os != WINDOWS else "" - libtorch_config = abi_version if os == WINDOWS else "" - ret.append( - { - "gpu_arch_type": gpu_arch_type, - "gpu_arch_version": gpu_arch_version, - "desired_cuda": desired_cuda, - "libtorch_variant": libtorch_variant, - "libtorch_config": libtorch_config, - "devtoolset": devtoolset, - "container_image": ( - LIBTORCH_CONTAINER_IMAGES[(arch_version, abi_version)] - if os != WINDOWS - else "" - ), - "package_type": "libtorch", - "build_name": f"libtorch-{gpu_arch_type}{gpu_arch_version}-{libtorch_variant}-{abi_version}".replace( # noqa: E501 - ".", "_" - ), - # Please noe since libtorch validations are minimal, we use CPU runners - "validation_runner": validation_runner(CPU, os), - "installation": get_libtorch_install_command( - os, - channel, - gpu_arch_type, - libtorch_variant, - devtoolset, - desired_cuda, - libtorch_config, - ), - "channel": channel, - "stable_version": CURRENT_VERSION, - } - ) - return ret - - -def generate_wheels_matrix( - os: str, - channel: str, - with_cuda: str, - with_rocm: str, - with_cpu: str, - with_xpu: str, - limit_pr_builds: bool, - use_only_dl_pytorch_org: bool, - use_split_build: bool = False, - python_versions: Optional[List[str]] = None, - arches: Optional[List[str]] = None, -) -> List[Dict[str, str]]: - package_type = "wheel" - - if not python_versions: - # Define default python version - python_versions = list(PYTHON_ARCHES) - - if os == LINUX: - # NOTE: We only build manywheel packages for linux - package_type = "manywheel" - - upload_to_base_bucket = "yes" - if arches is None: - # Define default compute architectures - arches = [] - - if with_cpu == ENABLE: - arches += [CPU] - - 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] - - if with_cuda == ENABLE: - upload_to_base_bucket = "no" - if os in (LINUX, WINDOWS): - arches += CUDA_ARCHES - - if with_rocm == ENABLE and os == LINUX: - arches += ROCM_ARCHES - - if with_xpu == ENABLE and os in (LINUX, WINDOWS): - arches += [XPU] - - if limit_pr_builds: - python_versions = PYTHON_VERSIONS_FOR_PR_BUILD - - global WHEEL_CONTAINER_IMAGES - - 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 - ) - - desired_cuda = translate_desired_cuda(gpu_arch_type, gpu_arch_version) - entry = { - "python_version": python_version, - "gpu_arch_type": gpu_arch_type, - "gpu_arch_version": gpu_arch_version, - "desired_cuda": desired_cuda, - "container_image": WHEEL_CONTAINER_IMAGES[arch_version], - "package_type": package_type, - "build_name": f"{package_type}-py{python_version}-{gpu_arch_type}{gpu_arch_version}".replace( - ".", "_" - ), - "validation_runner": validation_runner(gpu_arch_type, os), - "installation": get_wheel_install_command( - os, - channel, - gpu_arch_type, - gpu_arch_version, - desired_cuda, - python_version, - use_only_dl_pytorch_org, - ), - "channel": channel, - "upload_to_base_bucket": upload_to_base_bucket, - "stable_version": CURRENT_VERSION, - "use_split_build": False, - } - ret.append(entry) - if ( - use_split_build - and (gpu_arch_version in CUDA_ARCHES) - and (os == LINUX) - and (channel == NIGHTLY) - ): - entry = entry.copy() - entry["build_name"] = ( - f"{package_type}-py{python_version}-{gpu_arch_type}{gpu_arch_version}-split".replace( - ".", "_" - ) - ) - entry["use_split_build"] = True - ret.append(entry) - - return ret - - -GENERATING_FUNCTIONS_BY_PACKAGE_TYPE: Dict[str, Callable[..., List[Dict[str, str]]]] = { - "wheel": generate_wheels_matrix, - "conda": generate_conda_matrix, - "libtorch": generate_libtorch_matrix, -} - - -def generate_build_matrix( - package_type: str, - operating_system: str, - channel: str, - with_cuda: str, - with_rocm: str, - with_cpu: str, - with_xpu: str, - limit_pr_builds: str, - use_only_dl_pytorch_org: str, - build_python_only: str, - use_split_build: str = "false", - python_versions: Optional[List[str]] = None, -) -> Dict[str, List[Dict[str, str]]]: - includes = [] - - package_types = package_type.split(",") - if len(package_types) == 1: - package_types = PACKAGE_TYPES if package_type == "all" else [package_type] - - channels = CUDA_ARCHES_DICT.keys() if channel == "all" else [channel] - - for channel in channels: - for package in package_types: - initialize_globals(channel, build_python_only == ENABLE) - includes.extend( - GENERATING_FUNCTIONS_BY_PACKAGE_TYPE[package]( - operating_system, - channel, - with_cuda, - with_rocm, - with_cpu, - with_xpu, - limit_pr_builds == "true", - use_only_dl_pytorch_org == "true", - use_split_build == "true", - python_versions, - ) - ) - - return {"include": includes} - - -def main(args: List[str]) -> None: - parser = argparse.ArgumentParser() - parser.add_argument( - "--package-type", - help="Package type to lookup for, also supports comma separated values", - type=str, - default=os.getenv("PACKAGE_TYPE", "wheel"), - ) - parser.add_argument( - "--operating-system", - help="Operating system to generate for", - type=str, - default=os.getenv("OS", LINUX), - ) - parser.add_argument( - "--channel", - help="Channel to use, default nightly", - type=str, - choices=["nightly", "test", "release", "all"], - default=os.getenv("CHANNEL", "nightly"), - ) - parser.add_argument( - "--with-cuda", - help="Build with Cuda?", - type=str, - choices=[ENABLE, DISABLE], - default=os.getenv("WITH_CUDA", ENABLE), - ) - parser.add_argument( - "--with-rocm", - help="Build with Rocm?", - type=str, - choices=[ENABLE, DISABLE], - default=os.getenv("WITH_ROCM", ENABLE), - ) - parser.add_argument( - "--with-cpu", - help="Build with CPU?", - type=str, - choices=[ENABLE, DISABLE], - default=os.getenv("WITH_CPU", ENABLE), - ) - parser.add_argument( - "--with-xpu", - help="Build with XPU?", - type=str, - choices=[ENABLE, DISABLE], - default=os.getenv("WITH_XPU", ENABLE), - ) - # By default this is false for this script but expectation is that the caller - # workflow will default this to be true most of the time, where a pull - # request is synchronized and does not contain the label "ciflow/binaries/all" - parser.add_argument( - "--limit-pr-builds", - help="Limit PR builds to single python/cuda config", - type=str, - choices=["true", "false"], - default=os.getenv("LIMIT_PR_BUILDS", "false"), - ) - # This is used when testing release builds to test release binaries - # only from download.pytorch.org. When pipy binaries are not released yet. - parser.add_argument( - "--use-only-dl-pytorch-org", - help="Use only download.pytorch.org when gen wheel install command?", - type=str, - choices=["true", "false"], - default=os.getenv("USE_ONLY_DL_PYTORCH_ORG", "false"), - ) - # Generates a single version python for building python packages only - # This basically makes it so that we only generate a matrix including the oldest - # version of python that we support - # For packages that look similar to torchtune-0.0.1-py3-none-any.whl - parser.add_argument( - "--build-python-only", - help="Build python only", - type=str, - choices=[ENABLE, DISABLE], - default=os.getenv("BUILD_PYTHON_ONLY", ENABLE), - ) - - parser.add_argument( - "--use-split-build", - help="Use split build for wheel", - type=str, - choices=["true", "false"], - default=os.getenv("USE_SPLIT_BUILD", DISABLE), - ) - - parser.add_argument( - "--python-versions", - help="Only build the select JSON-encoded list of python versions", - type=str, - default=os.getenv("PYTHON_VERSIONS", "[]"), - ) - - options = parser.parse_args(args) - try: - python_versions = json.loads(options.python_versions) - except json.JSONDecodeError: - python_versions = None - - assert ( - options.with_cuda or options.with_rocm or options.with_xpu or options.with_cpu - ), "Must build with either CUDA, ROCM, XPU, or CPU support." - - build_matrix = generate_build_matrix( - options.package_type, - options.operating_system, - options.channel, - options.with_cuda, - options.with_rocm, - options.with_cpu, - options.with_xpu, - options.limit_pr_builds, - options.use_only_dl_pytorch_org, - options.build_python_only, - options.use_split_build, - python_versions, - ) - - print(json.dumps(build_matrix)) - - -if __name__ == "__main__": - main(sys.argv[1:]) diff --git a/.github/scripts/install-cuda-aarch64.sh b/.github/scripts/install-cuda-aarch64.sh new file mode 100755 index 0000000000..8301b333e4 --- /dev/null +++ b/.github/scripts/install-cuda-aarch64.sh @@ -0,0 +1,16 @@ + +install_cuda_aarch64() { + echo "install cuda ${CU_VERSION}" + # CU_VERSION: cu128 --> CU_VER: 12-8 + CU_VER=${CU_VERSION:2:2}-${CU_VERSION:4:1} + dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo + dnf -y install cuda-compiler-${CU_VER}.aarch64 \ + cuda-libraries-${CU_VER}.aarch64 \ + cuda-libraries-devel-${CU_VER}.aarch64 + dnf clean all + export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH + ls -lart /usr/local/ + nvcc --version + echo "cuda ${CU_VER} installed successfully" +} + diff --git a/.github/scripts/install-torch-tensorrt.sh b/.github/scripts/install-torch-tensorrt.sh index c75bfdb09d..fd88207eb4 100644 --- a/.github/scripts/install-torch-tensorrt.sh +++ b/.github/scripts/install-torch-tensorrt.sh @@ -1,9 +1,16 @@ -set -exou pipefail +#set -exou pipefail +set -x TORCH_TORCHVISION=$(grep "^torch" ${PWD}/py/requirements.txt) INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION} PLATFORM=$(python -c "import sys; print(sys.platform)") +if [[ $(uname -m) == "aarch64" ]]; then + # install cuda for aarch64 + source .github/scripts/install-cuda-aarch64.sh + install_cuda_aarch64 +fi + # Install all the dependencies required for Torch-TensorRT pip install --pre ${TORCH_TORCHVISION} --index-url ${INDEX_URL} pip install --pre -r ${PWD}/tests/py/requirements.txt diff --git a/.github/workflows/build-test-linux-aarch64.yml b/.github/workflows/build-test-linux-aarch64.yml new file mode 100644 index 0000000000..23edecd3c7 --- /dev/null +++ b/.github/workflows/build-test-linux-aarch64.yml @@ -0,0 +1,349 @@ +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 + + 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 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: ./.github/workflows/build_wheels_linux_aarch64.yml + 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 \ No newline at end of file diff --git a/.github/workflows/build-test-linux.yml b/.github/workflows/build-test-linux-x86_64.yml similarity index 99% rename from .github/workflows/build-test-linux.yml rename to .github/workflows/build-test-linux-x86_64.yml index a8e90a3222..5f24a0c7ad 100644 --- a/.github/workflows/build-test-linux.yml +++ b/.github/workflows/build-test-linux-x86_64.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: linux @@ -36,7 +36,7 @@ jobs: - uses: actions/checkout@v4 with: repository: pytorch/tensorrt - - name: Generate release matrix + - name: Generate matrix id: generate run: | set -eou pipefail diff --git a/.github/workflows/build-test-tensorrt-linux.yml b/.github/workflows/build-test-tensorrt-linux.yml index 6ddd0e336d..79adfd3980 100644 --- a/.github/workflows/build-test-tensorrt-linux.yml +++ b/.github/workflows/build-test-tensorrt-linux.yml @@ -12,7 +12,7 @@ permissions: 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 diff --git a/.github/workflows/build-test-tensorrt-windows.yml b/.github/workflows/build-test-tensorrt-windows.yml index bcb2d8865b..94f379f229 100644 --- a/.github/workflows/build-test-tensorrt-windows.yml +++ b/.github/workflows/build-test-tensorrt-windows.yml @@ -12,7 +12,7 @@ permissions: 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/build-test-windows.yml b/.github/workflows/build-test-windows.yml index 2ee31b4b74..7a974411c8 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 @@ -24,8 +24,29 @@ jobs: with-rocm: false with-cpu: false + 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 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}" + substitute-runner: - needs: generate-matrix + needs: filter-matrix outputs: matrix: ${{ steps.substitute.outputs.matrix }} runs-on: ubuntu-latest @@ -33,7 +54,7 @@ jobs: - name: Substitute runner id: substitute run: | - echo matrix="$(echo '${{ needs.generate-matrix.outputs.matrix }}' | sed -e 's/windows.g4dn.xlarge/windows.g5.4xlarge.nvidia.gpu/g')" >> ${GITHUB_OUTPUT} + echo matrix="$(echo '${{ needs.filter-matrix.outputs.matrix }}' | sed -e 's/windows.g4dn.xlarge/windows.g5.4xlarge.nvidia.gpu/g')" >> ${GITHUB_OUTPUT} build: needs: substitute-runner diff --git a/.github/workflows/build_wheels_linux_aarch64.yml b/.github/workflows/build_wheels_linux_aarch64.yml new file mode 100644 index 0000000000..ce7a0b2f98 --- /dev/null +++ b/.github/workflows/build_wheels_linux_aarch64.yml @@ -0,0 +1,334 @@ +name: Build Linux Wheels For aarch64 + +on: + workflow_call: + inputs: + repository: + description: 'Repository to checkout, defaults to ""' + default: "" + type: string + ref: + description: 'Reference to checkout, defaults to "nightly"' + default: "nightly" + type: string + test-infra-repository: + description: "Test infra repository to use" + default: "pytorch/test-infra" + type: string + test-infra-ref: + description: "Test infra reference to use" + default: "" + type: string + build-matrix: + description: "Build matrix to utilize" + default: "" + type: string + pre-script: + description: "Pre script to run prior to build" + default: "" + type: string + post-script: + description: "Post script to run prior to build" + default: "" + type: string + smoke-test-script: + description: "Script for Smoke Test for a specific domain" + default: "" + type: string + env-var-script: + description: "Script that sets Domain-Specific Environment Variables" + default: "" + type: string + package-name: + description: "Name of the actual python package that is imported in the smoke test" + default: "" + type: string + build-target: + description: "The target to build and publish (for repos that build multiple packages)" + default: "" + type: string + trigger-event: + description: "Trigger Event in caller that determines whether or not to upload" + default: "" + type: string + cache-path: + description: "The path(s) on the runner to cache or restore. The path is relative to repository." + default: "" + type: string + cache-key: + description: "The key created when saving a cache and the key used to search for a cache." + default: "" + type: string + architecture: + description: Architecture to build for x86_64 for default Linux, or aarch64 for Linux aarch64 builds + required: false + type: string + default: x86_64 + submodules: + description: Works as stated in actions/checkout, but the default value is recursive + required: false + type: string + default: recursive + setup-miniconda: + description: Set to true if setup-miniconda is needed + required: false + type: boolean + default: true + build-platform: + description: Platform to build wheels, choose from 'python-build-package' or 'setup-py' + required: false + type: string + default: 'setup-py' + upload-to-pypi: + description: The comma-separated list of CUDA arch to be uploaded to pypi + default: "" + type: string + build-command: + description: The build command to use if build-platform is python-build-package + required: false + default: "python -m build --wheel" + type: string + pip-install-torch-extra-args: + # NOTE: Why does this exist? + # Well setuptools / python packaging doesn't actually allow you to specify dependencies + # that come from other index URLs when you are building a package for "security" purposes. + # Unfortunately for us our nightlies (torch, torchvision, etc.) only exist on download.pytorch.org + # which means that if our users depend on things like torchvision then they need to have + # an ability to install these dependencies from download.pytorch.org, as part of the build process + # which currently the do not have the ability to do through normal means, hence this parameter + # Reference: https://discuss.python.org/t/specifying-extra-index-url-in-setup-cfg-option-dependencies/19377 + description: Extra arguments to pass to the command that install base torch dependency + required: false + default: "" + type: string + timeout: + description: 'Timeout for the job (in minutes)' + default: 120 + type: number + secrets: + PYPI_API_TOKEN: + description: An optional token to upload to pypi + required: false + +permissions: + id-token: write + contents: read + +jobs: + build: + strategy: + fail-fast: false + matrix: ${{ fromJSON(inputs.build-matrix) }} + env: + PYTHON_VERSION: ${{ matrix.python_version }} + PACKAGE_TYPE: wheel + REPOSITORY: ${{ inputs.repository }} + REF: ${{ inputs.ref }} + CU_VERSION: ${{ matrix.desired_cuda }} + UPLOAD_TO_BASE_BUCKET: ${{ matrix.upload_to_base_bucket }} + ARCH: ${{ inputs.architecture }} + BUILD_TARGET: ${{ inputs.build-target }} + name: build-${{ matrix.build_name }} + runs-on: ${{ matrix.validation_runner }} + environment: ${{(inputs.trigger-event == 'schedule' || (inputs.trigger-event == 'push' && (startsWith(github.event.ref, 'refs/heads/nightly') || startsWith(github.event.ref, 'refs/tags/v')))) && 'pytorchbot-env' || ''}} + container: + image: ${{ matrix.container_image }} + options: ${{ matrix.gpu_arch_type == 'cuda' && '--gpus all' || ' ' }} + timeout-minutes: ${{ inputs.timeout }} + steps: + - name: Clean workspace + shell: bash -l {0} + run: | + set -euxo pipefail + echo "::group::Cleanup debug output" + rm -rf "${GITHUB_WORKSPACE}" + mkdir -p "${GITHUB_WORKSPACE}" + + if [[ "${{ inputs.architecture }}" = "aarch64" ]]; then + rm -rf "${RUNNER_TEMP}/*" + fi + echo "::endgroup::" + + - uses: actions/checkout@v4 + with: + # Support the use case where we need to checkout someone's fork + repository: ${{ inputs.test-infra-repository }} + ref: ${{ inputs.test-infra-ref }} + path: test-infra + + - name: Install Miniforge + if: ${{ inputs.architecture == 'aarch64' }} + shell: bash -l {0} + env: + DESIRED_PYTHON: ${{ matrix.python_version }} + run: | + set -euxo pipefail + # TODO: Get rid of Conda, we already have all versions of PyThon one needs in the docker + ############################################################################### + # Install conda + # disable SSL_verify due to getting "Could not find a suitable TLS CA certificate bundle, invalid path" + # when using Python version, less than the conda latest + ############################################################################### + echo 'Installing conda-forge' + curl -L -o /mambaforge.sh https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-aarch64.sh + chmod +x /mambaforge.sh + /mambaforge.sh -b -p /opt/conda + rm /mambaforge.sh + source /opt/conda/etc/profile.d/conda.sh + conda config --set ssl_verify False + echo "/opt/conda/bin" >> $GITHUB_PATH + - uses: ./test-infra/.github/actions/set-channel + - name: Set PYTORCH_VERSION + if: ${{ env.CHANNEL == 'test' }} + run: | + # When building RC, set the version to be the current candidate version, + # otherwise, leave it alone so nightly will pick up the latest + echo "PYTORCH_VERSION=${{ matrix.stable_version }}" >> "${GITHUB_ENV}" + - uses: ./test-infra/.github/actions/setup-binary-builds + env: + PLATFORM: ${{ inputs.architecture == 'aarch64' && 'linux-aarch64' || ''}} + with: + repository: ${{ inputs.repository }} + ref: ${{ inputs.ref }} + submodules: ${{ inputs.submodules }} + setup-miniconda: false + python-version: ${{ env.PYTHON_VERSION }} + cuda-version: ${{ env.CU_VERSION }} + arch: ${{ env.ARCH }} + + - name: Combine Env Var and Build Env Files + if: ${{ inputs.env-var-script != '' }} + working-directory: ${{ inputs.repository }} + run: | + set -euxo pipefail + cat "${{ inputs.env-var-script }}" >> "${BUILD_ENV_FILE}" + - name: Add XPU Env Vars in Build Env File + if: ${{ matrix.gpu_arch_type == 'xpu' }} + run: | + { + echo "set +u" + echo "source /opt/intel/oneapi/compiler/latest/env/vars.sh" + echo "source /opt/intel/oneapi/pti/latest/env/vars.sh" + } >> "${BUILD_ENV_FILE}" + - name: Install torch dependency + run: | + set -euxo pipefail + # shellcheck disable=SC1090 + source "${BUILD_ENV_FILE}" + # shellcheck disable=SC2086 + ${CONDA_RUN} ${PIP_INSTALL_TORCH} ${{ inputs.pip-install-torch-extra-args }} + - name: Run Pre-Script with Caching + if: ${{ inputs.pre-script != '' }} + uses: ./test-infra/.github/actions/run-script-with-cache + with: + cache-path: ${{ inputs.cache-path }} + cache-key: ${{ inputs.cache-key }} + repository: ${{ inputs.repository }} + script: ${{ inputs.pre-script }} + - name: Build the wheel (python-build-package) + if: ${{ inputs.build-platform == 'python-build-package' }} + working-directory: ${{ inputs.repository }} + shell: bash -l {0} + run: | + set -euxo pipefail + source "${BUILD_ENV_FILE}" + export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')" + ${CONDA_RUN} python -m pip install build==1.2.2 + echo "Successfully installed Python build package" + ${CONDA_RUN} ${{ inputs.build-command }} + - name: Build the wheel (setup-py) + if: ${{ inputs.build-platform == 'setup-py' }} + working-directory: ${{ inputs.repository }} + shell: bash -l {0} + run: | + set -euxo pipefail + source "${BUILD_ENV_FILE}" + export PYTORCH_VERSION="$(${CONDA_RUN} pip show torch | grep ^Version: | sed 's/Version: *//' | sed 's/+.\+//')" + ${CONDA_RUN} python setup.py clean + echo "Successfully ran `python setup.py clean`" + ${CONDA_RUN} python setup.py bdist_wheel + - name: Repair Manylinux_2_28 Wheel + shell: bash -l {0} + env: + PACKAGE_NAME: ${{ inputs.package-name }} + SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }} + run: | + set -euxo pipefail + source "${BUILD_ENV_FILE}" + # for pkg in ${{ inputs.repository }}/dist/*-linux_*.whl; do + # # if the glob didn't match anything + # if [[ ! -e $pkg ]]; then + # continue + # fi + # abs_pkg=$(realpath $pkg) + # ./test-infra/.github/scripts/repair_manylinux_2_28.sh $abs_pkg + # done + echo "Repair Manylinux_2_28 Wheel is not supported for aarch64" + - name: Run Post-Script + if: ${{ inputs.post-script != '' }} + uses: ./test-infra/.github/actions/run-script-with-cache + with: + repository: ${{ inputs.repository }} + script: ${{ inputs.post-script }} + - name: Smoke Test + shell: bash -l {0} + env: + PACKAGE_NAME: ${{ inputs.package-name }} + SMOKE_TEST_SCRIPT: ${{ inputs.smoke-test-script }} + run: | + set -euxo pipefail + source "${BUILD_ENV_FILE}" + WHEEL_NAME=$(ls "${{ inputs.repository }}/dist/") + echo "$WHEEL_NAME" + + ${CONDA_RUN} pip install "${{ inputs.repository }}/dist/$WHEEL_NAME" + # Checking that we have a pinned version of torch in our dependency tree + ( + pushd "${RUNNER_TEMP}" + unzip -o "${GITHUB_WORKSPACE}/${{ inputs.repository }}/dist/$WHEEL_NAME" + # Ensure that pytorch version is pinned, should output file where it was found + grep "Requires-Dist: torch (==.*)" -r . + ) + + if [[ (! -f "${{ inputs.repository }}"/${SMOKE_TEST_SCRIPT}) ]]; then + echo "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT} not found" + if [[ "${PACKAGE_NAME}" = "torchrec" ]]; then + # Special case for torchrec temporarily since __version__ does not + # work correctly on main in torchrec. This block will be + # removed once we fix it. + ${CONDA_RUN} python -c "import ${PACKAGE_NAME}" + else + ${CONDA_RUN} python -c "import ${PACKAGE_NAME}; print('package version is ', ${PACKAGE_NAME}.__version__)" + fi + else + echo "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT} found" + ${CONDA_RUN} python "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT}" + fi + # NB: Only upload to GitHub after passing smoke tests + + - name: Upload wheel to GitHub + continue-on-error: true + uses: actions/upload-artifact@v4 + with: + name: ${{ env.ARTIFACT_NAME }} + path: ${{ inputs.repository }}/dist/ + + upload: + needs: build + uses: pytorch/test-infra/.github/workflows/_binary_upload.yml@main + if: always() + with: + repository: ${{ inputs.repository }} + ref: ${{ inputs.ref }} + test-infra-repository: ${{ inputs.test-infra-repository }} + test-infra-ref: ${{ inputs.test-infra-ref }} + build-matrix: ${{ inputs.build-matrix }} + architecture: ${{ inputs.architecture }} + trigger-event: ${{ inputs.trigger-event }} + upload-to-pypi: ${{ inputs.upload-to-pypi }} + secrets: + PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }} + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ inputs.repository }}-${{ github.event_name == 'workflow_dispatch' }} + cancel-in-progress: true diff --git a/.github/workflows/generate_binary_build_matrix.yml b/.github/workflows/generate_binary_build_matrix.yml deleted file mode 100644 index fceb8cf1ee..0000000000 --- a/.github/workflows/generate_binary_build_matrix.yml +++ /dev/null @@ -1,112 +0,0 @@ -name: Generates the binary build matrix - -on: - workflow_call: - inputs: - package-type: - description: "Package type to build from (wheel, conda, libtorch)" - default: "wheel" - type: string - os: - description: "Operating system to generate for (linux, windows, macos, macos-arm64)" - default: "linux" - type: string - channel: - description: "Channel to use (nightly, test, release, all)" - default: "" - type: string - test-infra-repository: - description: "Test infra repository to use" - default: "pytorch/test-infra" - type: string - test-infra-ref: - description: "Test infra reference to use" - default: "main" - type: string - with-cuda: - description: "Build with Cuda?" - default: "enable" - type: string - with-rocm: - description: "Build with Rocm?" - default: "enable" - type: string - with-cpu: - description: "Build with CPU?" - default: "enable" - type: string - with-xpu: - description: "Build with XPU?" - default: "disable" - type: string - use-only-dl-pytorch-org: - description: "Use only download.pytorch.org when generating wheel install command?" - default: "false" - type: string - build-python-only: - description: "Generate binary build matrix for a python only package (i.e. only one python version)" - default: "disable" - type: string - python-versions: - description: "A JSON-encoded list of python versions to build. An empty list means building all supported versions" - default: "[]" - type: string - use_split_build: - description: | - [Experimental] Build a libtorch only wheel and build pytorch such that - are built from the libtorch wheel. - required: false - type: boolean - default: false - - outputs: - matrix: - description: "Generated build matrix" - value: ${{ jobs.generate.outputs.matrix }} - -jobs: - generate: - outputs: - matrix: ${{ steps.generate.outputs.matrix }} - runs-on: ubuntu-latest - steps: - - uses: actions/setup-python@v5 - with: - 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: - PACKAGE_TYPE: ${{ inputs.package-type }} - OS: ${{ inputs.os }} - CHANNEL: ${{ inputs.channel != '' && inputs.channel || env.CHANNEL }} - WITH_CUDA: ${{ inputs.with-cuda }} - WITH_ROCM: ${{ inputs.with-rocm }} - WITH_CPU: ${{ inputs.with-cpu }} - WITH_XPU: ${{ inputs.with-xpu }} - # limit pull request builds to one version of python unless ciflow/binaries/all is applied to the workflow - # should not affect builds that are from events that are not the pull_request event - LIMIT_PR_BUILDS: ${{ github.event_name == 'pull_request' && !contains( github.event.pull_request.labels.*.name, 'ciflow/binaries/all') }} - # This is used when testing release binaries only from download.pytorch.org. - # In cases when pipy binaries are not published yet. - USE_ONLY_DL_PYTORCH_ORG: ${{ inputs.use-only-dl-pytorch-org }} - BUILD_PYTHON_ONLY: ${{ inputs.build-python-only }} - USE_SPLIT_BUILD: ${{ inputs.use_split_build }} - PYTHON_VERSIONS: ${{ inputs.python-versions }} - run: | - set -eou pipefail - 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 diff --git a/.github/workflows/release-linux.yml b/.github/workflows/release-linux.yml index e579a4f5d5..2c4be0ec1d 100644 --- a/.github/workflows/release-linux.yml +++ b/.github/workflows/release-linux.yml @@ -15,7 +15,7 @@ permissions: jobs: generate-matrix: - uses: ./.github/workflows/generate_binary_build_matrix.yml + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main if: ${{ contains(github.event.pull_request.labels.*.name, 'build-release-artifacts') || startsWith(github.event.ref, 'refs/tags/v') }} with: package-type: wheel diff --git a/.github/workflows/release-windows.yml b/.github/workflows/release-windows.yml index 2774c3ee0d..1a20eb7a71 100644 --- a/.github/workflows/release-windows.yml +++ b/.github/workflows/release-windows.yml @@ -15,7 +15,7 @@ permissions: jobs: generate-matrix: - uses: ./.github/workflows/generate_binary_build_matrix.yml + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main if: ${{ contains(github.event.pull_request.labels.*.name, 'build-release-artifacts') || startsWith(github.event.ref, 'refs/tags/v') }} with: package-type: wheel diff --git a/.gitignore b/.gitignore index 16e4f4f838..496f53ca2b 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 \ No newline at end of file diff --git a/MODULE.bazel b/MODULE.bazel index 82bb4ba79e..008c7f53fc 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -4,10 +4,10 @@ module( repo_name = "org_pytorch_tensorrt", ) -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,18 +27,28 @@ 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.11/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 +# for linux x86_64 and aarch64 new_local_repository( name = "cuda", build_file = "@//third_party/cuda:BUILD", path = "/usr/local/cuda-12.8/", ) +# for Jetson +# 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", +) + +# for windows new_local_repository( name = "cuda_win", build_file = "@//third_party/cuda:BUILD", @@ -51,6 +61,7 @@ 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) ############################################################################################################# +# this is used for linux x86_64 http_archive( name = "libtorch", build_file = "@//third_party/libtorch:BUILD", @@ -58,6 +69,17 @@ http_archive( urls = ["https://download.pytorch.org/libtorch/nightly/cu128/libtorch-cxx11-abi-shared-with-deps-latest.zip"], ) +# in aarch64 platform you can get libtorch via either local or wheel file +# 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.dev20250415%2Bcu128-cp310-cp310-manylinux_2_28_aarch64.whl"], +#) + http_archive( name = "libtorch_win", build_file = "@//third_party/libtorch:BUILD", @@ -78,6 +100,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,7 +135,8 @@ http_archive( #new_local_repository( # name = "libtorch", -# path = "/usr/local/lib/python3.6/dist-packages/torch", +# please modify this path according to your local python wheel path +# path = "/usr/local/lib/python3.11/dist-packages/torch", # build_file = "third_party/libtorch/BUILD" #) diff --git a/core/BUILD b/core/BUILD index 401534f694..28cd460690 100644 --- a/core/BUILD +++ b/core/BUILD @@ -4,9 +4,29 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) config_setting( - name = "python_core", - values = { - "define": "target_lang=python", + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" }, ) @@ -17,6 +37,13 @@ config_setting( ], ) +config_setting( + name = "python_core", + values = { + "define": "target_lang=python", + }, +) + cc_library( name = "core", srcs = [ @@ -32,14 +59,14 @@ cc_library( "//core/runtime", "//core/util/logging", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 4647610d10..13696550e6 100644 --- a/core/conversion/BUILD +++ b/core/conversion/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -27,14 +54,14 @@ cc_library( "//core/ir", "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 b8d1808857..89ff7f613c 100644 --- a/core/conversion/conversionctx/BUILD +++ b/core/conversion/conversionctx/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -22,14 +49,14 @@ cc_library( "//core/ir", "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 32ddc447c5..9571d91604 100644 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -22,14 +49,14 @@ cc_library( "//core/conversion/conversionctx", "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) @@ -47,14 +74,14 @@ cc_library( "//core/conversion/conversionctx", "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }) + select({ + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), alwayslink = True, ) @@ -109,14 +136,14 @@ cc_library( "//core/plugins:torch_tensorrt_plugins", "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 f22980c1a6..172e8f6670 100644 --- a/core/conversion/evaluators/BUILD +++ b/core/conversion/evaluators/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -28,7 +55,13 @@ cc_library( "//core/conversion/var", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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/tensorcontainer/BUILD b/core/conversion/tensorcontainer/BUILD index 04dbf22e30..c812b74a9f 100644 --- a/core/conversion/tensorcontainer/BUILD +++ b/core/conversion/tensorcontainer/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -21,14 +48,14 @@ cc_library( deps = [ "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 00f5f866d9..65d9583474 100644 --- a/core/conversion/var/BUILD +++ b/core/conversion/var/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -24,14 +51,14 @@ cc_library( "//core/conversion/tensorcontainer", "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 1377f16eea..d522a6a018 100644 --- a/core/ir/BUILD +++ b/core/ir/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -24,14 +51,14 @@ cc_library( deps = [ "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 af9fd68c71..e9b1e1ae07 100644 --- a/core/lowering/BUILD +++ b/core/lowering/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -26,7 +53,13 @@ cc_library( "//core/lowering/passes", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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/passes/BUILD b/core/lowering/passes/BUILD index 7960a312bd..42dc4faab5 100644 --- a/core/lowering/passes/BUILD +++ b/core/lowering/passes/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl", + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default", + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack", + }, +) + config_setting( name = "windows", constraint_values = [ @@ -49,6 +76,7 @@ cc_library( deps = [ "//core/util:prelude", ] + select({ + ":use_torch_whl": ["@torch_whl//:libtorch"], ":windows": ["@libtorch_win//:libtorch"], "//conditions:default": ["@libtorch"], }), diff --git a/core/partitioning/BUILD b/core/partitioning/BUILD index a1e8f9c5d7..2cbcec34b1 100644 --- a/core/partitioning/BUILD +++ b/core/partitioning/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -29,7 +56,13 @@ cc_library( "//core/partitioning/segmentedblock", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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/partitioning/partitioningctx/BUILD b/core/partitioning/partitioningctx/BUILD index 25361e2226..011a48c6be 100644 --- a/core/partitioning/partitioningctx/BUILD +++ b/core/partitioning/partitioningctx/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -25,7 +52,13 @@ cc_library( "//core/partitioning/segmentedblock", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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/partitioning/partitioninginfo/BUILD b/core/partitioning/partitioninginfo/BUILD index a048d3d1e7..d2a86a2688 100644 --- a/core/partitioning/partitioninginfo/BUILD +++ b/core/partitioning/partitioninginfo/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -24,7 +51,13 @@ cc_library( "//core/lowering", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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/partitioning/segmentedblock/BUILD b/core/partitioning/segmentedblock/BUILD index aef1adc505..c463d17b92 100644 --- a/core/partitioning/segmentedblock/BUILD +++ b/core/partitioning/segmentedblock/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -24,7 +51,13 @@ cc_library( "//core/lowering", "//core/util:prelude", ] + select({ + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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/plugins/BUILD b/core/plugins/BUILD index 3854cf7b03..ba167d5f2a 100644 --- a/core/plugins/BUILD +++ b/core/plugins/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -32,15 +59,25 @@ cc_library( "//core/util:prelude", ] + select({ ":windows": [ - "@libtorch_win//:libtorch", "@tensorrt_win//:nvinfer", "@tensorrt_win//:nvinferplugin", ], + ":sbsa": [ + "@tensorrt_sbsa//:nvinfer", + "@tensorrt_sbsa//:nvinferplugin", + ], + ":jetpack": [ + "@tensorrt_l4t//:nvinfer", + "@tensorrt_l4t//:nvinferplugin", + ], "//conditions:default": [ - "@libtorch", "@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 556d88d046..f30519619b 100644 --- a/core/runtime/BUILD +++ b/core/runtime/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -36,14 +63,14 @@ cc_library( "//core/plugins:torch_tensorrt_plugins", "//core/util:prelude", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 7042a38982..bc9b53ec8d 100644 --- a/core/util/BUILD +++ b/core/util/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -34,6 +61,7 @@ cc_library( ":macros", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -66,14 +94,9 @@ cc_library( "build_info.h", ], deps = select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], + "//conditions:default": ["@libtorch"], }), ) @@ -89,14 +112,14 @@ cc_library( ":macros", "//core/util/logging", ] + select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 262cc0e882..d29568cf97 100644 --- a/core/util/logging/BUILD +++ b/core/util/logging/BUILD @@ -3,6 +3,33 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -19,14 +46,14 @@ cc_library( "TorchTRTLogger.h", ], deps = select({ - ":windows": [ - "@libtorch_win//:libtorch", - "@tensorrt_win//:nvinfer", - ], - "//conditions:default": [ - "@libtorch", - "@tensorrt//:nvinfer", - ], + ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//: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 87d35cc57f..51ee4ca2ab 100644 --- a/cpp/bin/torchtrtc/BUILD +++ b/cpp/bin/torchtrtc/BUILD @@ -2,6 +2,13 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -30,7 +37,11 @@ cc_binary( ] + select({ ":windows": [ "@libtorch_win//:caffe2", - "@libtorch_win//:libtorch", + "@libtorch_win//: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..3b464b3a99 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/dep_src: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..d30c7d3c03 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/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "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_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/examples/int8/qat/BUILD b/examples/int8/qat/BUILD index e97398d932..0aab56a02a 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/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "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_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/examples/torchtrt_runtime_example/BUILD b/examples/torchtrt_runtime_example/BUILD index 957caedd23..7102a7f3db 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_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "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_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) ) diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh index 173498201f..d6764efe89 100755 --- a/packaging/pre_build_script.sh +++ b/packaging/pre_build_script.sh @@ -7,8 +7,22 @@ python3 -m pip install pyyaml yum install -y ninja-build gettext -wget https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-linux-amd64 \ - && mv bazelisk-linux-amd64 /usr/bin/bazel \ +BAZEL_PLATFORM="amd64" + +if [[ $(uname -m) == "aarch64" ]]; then + BAZEL_PLATFORM=arm64 + rm -rf /opt/openssl # Not sure whats up with the openssl mismatch + # aarch64 does not have envsubst pre-installed in the image, install it here + curl -L https://github.com/a8m/envsubst/releases/download/v1.4.2/envsubst-Linux-arm64 -o envsubst \ + && mv envsubst /usr/bin/envsubst && chmod +x /usr/bin/envsubst + # install cuda for aarch64 + source .github/scripts/install-cuda-aarch64.sh + install_cuda_aarch64 +fi + +curl -L https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelisk-linux-${BAZEL_PLATFORM} \ + -o 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/setup.py b/setup.py index c770ce91b6..9e386975e5 100644 --- a/setup.py +++ b/setup.py @@ -78,7 +78,9 @@ def load_dep_info(): dir_path = os.path.join(str(get_root_dir()), "py") -JETPACK_VERSION = None +IS_AARCH64 = platform.uname().processor == "aarch64" +IS_JETPACK = False +IS_SBSA = False PY_ONLY = False NO_TS = False LEGACY = False @@ -120,6 +122,13 @@ 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 not None: + if jetpack == "1": + IS_JETPACK = True + +if (sbsa := os.environ.get("SBSA_BUILD")) is not None: + if sbsa == "1": + IS_SBSA = True if RELEASE: __version__ = os.environ.get("BUILD_VERSION") @@ -135,26 +144,12 @@ def load_dep_info(): if ci_env_var == "1": CI_BUILD = 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 IS_AARCH64: + if "--jetpack" in sys.argv: + sys.argv.remove("--jetpack") + IS_JETPACK = True + else: + IS_SBSA = True BAZEL_EXE = None @@ -184,7 +179,10 @@ def build_libtorchtrt_cxx11_abi( else: cmd.append("--compilation_mode=opt") if use_dist_dir: - cmd.append("--distdir=third_party/dist_dir/x86_64-linux-gnu") + if IS_AARCH64: + cmd.append("--distdir=third_party/dist_dir/aarch64-linux-gnu") + else: + cmd.append("--distdir=third_party/dist_dir/x86_64-linux-gnu") if target_python: cmd.append("--config=python") @@ -194,22 +192,20 @@ def build_libtorchtrt_cxx11_abi( 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") + print("Jetpack build") + + if IS_SBSA: + cmd.append("--platforms=//toolchains:sbsa") + print("SBSA build") if CI_BUILD: - cmd.append("--platforms=//toolchains:ci_rhel_x86_64_linux") print("CI based build") + if IS_AARCH64: + cmd.append("--platforms=//toolchains:aarch64_linux") + else: + cmd.append("--platforms=//toolchains:ci_rhel_x86_64_linux") print(f"building libtorchtrt {cmd=}") status_code = subprocess.run(cmd).returncode @@ -463,7 +459,7 @@ def run(self): package_data = {} if not (PY_ONLY or NO_TS): - tensorrt_linux_external_dir = ( + tensorrt_x86_64_external_dir = ( lambda: subprocess.check_output( [BAZEL_EXE, "query", "@tensorrt//:nvinfer", "--output", "location"] ) @@ -471,6 +467,32 @@ 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"] diff --git a/tests/core/BUILD b/tests/core/BUILD index 96b8cc804f..a0e19fa232 100644 --- a/tests/core/BUILD +++ b/tests/core/BUILD @@ -1,5 +1,32 @@ load("@rules_cc//cc:defs.bzl", "cc_test") +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -23,6 +50,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//: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 f302683a48..dfab90a978 100644 --- a/tests/core/conversion/converters/BUILD +++ b/tests/core/conversion/converters/BUILD @@ -1,5 +1,12 @@ load("//tests/core/conversion/converters:converter_test.bzl", "converter_test") +config_setting( + name = "use_torch_whl", + values = { + "define": "torch_src=whl", + }, +) + config_setting( name = "windows", constraint_values = [ diff --git a/tests/core/conversion/evaluators/BUILD b/tests/core/conversion/evaluators/BUILD index 9c51e010a5..9b2a708297 100644 --- a/tests/core/conversion/evaluators/BUILD +++ b/tests/core/conversion/evaluators/BUILD @@ -1,5 +1,12 @@ load("//tests/core/conversion/evaluators:evaluator_test.bzl", "evaluator_test") +config_setting( + name = "use_torch_whl", + values = { + "define": "torch_src=whl", + }, +) + config_setting( name = "windows", constraint_values = [ diff --git a/tests/core/lowering/BUILD b/tests/core/lowering/BUILD index 8f942512c5..2de6b6a7f7 100644 --- a/tests/core/lowering/BUILD +++ b/tests/core/lowering/BUILD @@ -1,6 +1,13 @@ load("@rules_cc//cc:defs.bzl", "cc_test") load("//tests/core/lowering:lowering_test.bzl", "lowering_test") +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -23,6 +30,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"], "//conditions:default": ["@libtorch"], }), ) diff --git a/tests/core/lowering/lowering_test.bzl b/tests/core/lowering/lowering_test.bzl index 9e6ae63c83..9d2c38a969 100644 --- a/tests/core/lowering/lowering_test.bzl +++ b/tests/core/lowering/lowering_test.bzl @@ -20,7 +20,8 @@ def lowering_test(name, visibility = None): "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//: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 8207998365..b2f99dbd34 100644 --- a/tests/core/partitioning/BUILD +++ b/tests/core/partitioning/BUILD @@ -1,6 +1,13 @@ load("@rules_cc//cc:defs.bzl", "cc_test") load("//tests/core/partitioning:partitioning_test.bzl", "partitioning_test") +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -8,6 +15,7 @@ config_setting( ], ) + filegroup( name = "jit_models", srcs = [ @@ -55,6 +63,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -70,6 +79,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -85,6 +95,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -100,6 +111,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//: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 95fd998a14..f89dd22f40 100644 --- a/tests/core/partitioning/partitioning_test.bzl +++ b/tests/core/partitioning/partitioning_test.bzl @@ -20,7 +20,8 @@ def partitioning_test(name, visibility = None): "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//: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 82dbe008fb..cd02a54b40 100644 --- a/tests/core/runtime/BUILD +++ b/tests/core/runtime/BUILD @@ -2,6 +2,13 @@ load("//tests/core/runtime:runtime_test.bzl", "runtime_test") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -9,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 3450d0d55c..3363edb627 100644 --- a/tests/core/runtime/runtime_test.bzl +++ b/tests/core/runtime/runtime_test.bzl @@ -20,6 +20,7 @@ def runtime_test(name, visibility = None): "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//: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 581f3a9d40..b50a3c6783 100644 --- a/tests/cpp/BUILD +++ b/tests/cpp/BUILD @@ -2,6 +2,13 @@ load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") package(default_visibility = ["//visibility:public"]) +config_setting( + name = "use_torch_whl", + flag_values = { + "//toolchains/dep_src:torch": "whl" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -87,6 +94,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -125,6 +133,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -140,6 +149,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -152,6 +162,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -167,6 +178,7 @@ cc_test( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) @@ -202,6 +214,7 @@ cc_library( "@googletest//:gtest_main", ] + select({ ":windows": ["@libtorch_win//:libtorch"], + ":use_torch_whl": ["@torch_whl//:libtorch"], "//conditions:default": ["@libtorch"], }), ) diff --git a/tests/util/BUILD b/tests/util/BUILD index 8be8a86612..38e4b5ff6d 100644 --- a/tests/util/BUILD +++ b/tests/util/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/dep_src:torch": "whl" + }, +) + config_setting( name = "windows", constraint_values = [ @@ -9,6 +16,26 @@ config_setting( ], ) +config_setting( + name = "sbsa", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "default" + }, +) + +config_setting( + name = "jetpack", + constraint_values = [ + "@platforms//cpu:aarch64", + ], + flag_values = { + "//toolchains/dep_collection:compute_libs": "jetpack" + }, +) + config_setting( name = "ci_build_testing", values = { @@ -32,12 +59,18 @@ cc_library( "@googletest//:gtest_main", ] + select({ ":windows": ["@tensorrt_win//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + ":jetpack": ["@tensorrt_l4t//:nvinfer"], "//conditions:default": ["@tensorrt//:nvinfer"], }) + select({ ":windows": [ "@libtorch_win//:caffe2", "@libtorch_win//:libtorch", ], + ":use_torch_whl": [ + "@torch_whl//:caffe2", + "@torch_whl//:libtorch", + ], "//conditions:default": [ "@libtorch", "@libtorch//:caffe2", diff --git a/third_party/cuda/BUILD b/third_party/cuda/BUILD index ea2229985f..b5359d42b3 100644 --- a/third_party/cuda/BUILD +++ b/third_party/cuda/BUILD @@ -46,9 +46,12 @@ cc_library( ":windows": [ "bin/*.dll", ], - "//conditions:default": glob([ - "lib64/**/lib*.so", - ]), + "//conditions:default": glob( + [ + "lib64/**/lib*.so", + ], + allow_empty = True, + ), }), hdrs = glob([ "include/**/*.h", @@ -66,9 +69,12 @@ cc_library( ":windows": [ "lib/x64/cublas.lib", ], - "//conditions:default": glob([ - "lib64/**/*libcublas.so", - ]), + "//conditions:default": glob( + [ + "lib64/**/*libcublas.so", + ], + allow_empty = True, + ), }), hdrs = glob([ "include/**/*cublas*.h", diff --git a/third_party/tensorrt/archive/BUILD b/third_party/tensorrt/archive/BUILD index c9b32dc23a..4aef027d12 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_libs": "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/ci_workspaces/MODULE.bazel.tmpl b/toolchains/ci_workspaces/MODULE.bazel.tmpl index dc5d42207c..d4587c6a64 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,19 +27,27 @@ 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.11/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 = "${CUDA_HOME}", ) +# 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 +61,12 @@ 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"], +#) # 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 +81,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", @@ -89,6 +115,12 @@ http_archive( # NOTE: If you are using a local build of torch, just point the Libtorch dep to that path. +new_local_repository( + name = "libtorch", + path = "${TORCH_INSTALL_PATH}", + build_file = "third_party/libtorch/BUILD" +) + new_local_repository( name = "libtorch_win", path = "${TORCH_INSTALL_PATH}", diff --git a/toolchains/dep_collection/BUILD b/toolchains/dep_collection/BUILD new file mode 100644 index 0000000000..128eeb8f26 --- /dev/null +++ b/toolchains/dep_collection/BUILD @@ -0,0 +1,8 @@ +load(":defs.bzl", "dep_collection") + +package(default_visibility = ["//visibility:public"]) + +dep_collection( + name = "compute_libs", + build_setting_default = "default", +) diff --git a/toolchains/dep_collection/defs.bzl b/toolchains/dep_collection/defs.bzl new file mode 100644 index 0000000000..6eaa710261 --- /dev/null +++ b/toolchains/dep_collection/defs.bzl @@ -0,0 +1,18 @@ +# buildifier: disable=module-docstring +DependencyCollectionInfo = provider(doc = "", fields = ["type"]) + +collection_types = ["default", "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), +) diff --git a/toolchains/dep_src/BUILD b/toolchains/dep_src/BUILD new file mode 100644 index 0000000000..161ff5d1fa --- /dev/null +++ b/toolchains/dep_src/BUILD @@ -0,0 +1,8 @@ +load(":defs.bzl", "dep_src") + +package(default_visibility = ["//visibility:public"]) + +dep_src( + name = "torch", + build_setting_default = "archive", +) diff --git a/toolchains/dep_src/defs.bzl b/toolchains/dep_src/defs.bzl new file mode 100644 index 0000000000..bec0efb8cf --- /dev/null +++ b/toolchains/dep_src/defs.bzl @@ -0,0 +1,18 @@ +# buildifier: disable=module-docstring +DepSrcInfo = 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 DepSrcInfo(type = src) + +dep_src = rule( + implementation = _impl, + build_setting = config.string(flag = True), +)