diff --git a/.github/scripts/verify-executorch-reference-runner.sh b/.github/scripts/verify-executorch-reference-runner.sh index 19e00b4843..b59ed905ec 100755 --- a/.github/scripts/verify-executorch-reference-runner.sh +++ b/.github/scripts/verify-executorch-reference-runner.sh @@ -6,11 +6,12 @@ set +x # # 1. Build //:libtorchtrt first so bazel-bin/libtorchtrt.tar.gz exists. # 2. Provide an ExecuTorch source checkout with EXECUTORCH_SOURCE_DIR. -# 3. This script exports a small Torch-TensorRT ExecuTorch .pte model, -# unpacks libtorchtrt.tar.gz, configures the packaged CMake runner, -# builds example_executorch_runner, and runs one inference. +# 3. Provide a Torch-TensorRT ExecuTorch .pte model. +# 4. This script unpacks libtorchtrt.tar.gz, configures and builds the +# packaged CMake runner, and runs one inference. # # Required: +# First argument: path to an existing .pte model. # EXECUTORCH_SOURCE_DIR=/path/to/executorch # # Optional: @@ -28,6 +29,16 @@ set +x repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" cd "${repo_root}" +if [[ $# -ne 1 ]]; then + echo "Usage: $0 PATH_TO_MODEL.pte" >&2 + exit 1 +fi +model_path="$1" +if [[ ! -f "${model_path}" ]]; then + echo "ExecuTorch model not found: ${model_path}" >&2 + exit 1 +fi + python_executable="${PYTHON_EXECUTABLE:-}" if [[ -z "${python_executable}" ]]; then python_executable="$(command -v python || true)" @@ -201,14 +212,13 @@ else export LD_LIBRARY_PATH="${torch_lib_dir}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" fi -# Fail early if the Python environment cannot export a Torch-TensorRT -# ExecuTorch model or run ExecuTorch's CMake codegen. +# Fail early if the Python environment cannot run ExecuTorch's CMake codegen. if ! "${python_executable}" - <<'PY' import importlib import importlib.util missing = [] -for name in ("yaml", "torch", "torch_tensorrt", "executorch.exir"): +for name in ("yaml", "torch", "executorch.exir"): try: spec = importlib.util.find_spec(name) except ModuleNotFoundError: @@ -217,56 +227,17 @@ for name in ("yaml", "torch", "torch_tensorrt", "executorch.exir"): missing.append(name) if missing: raise SystemExit( - "Missing Python package(s) required to export the .pte and build the runner: " + "Missing Python package(s) required to build the runner: " + ", ".join(missing) ) -for name in ("yaml", "torch", "torch_tensorrt", "executorch.exir"): +for name in ("yaml", "torch", "executorch.exir"): importlib.import_module(name) PY then exit 1 fi -model_path="${verify_root}/model.pte" -"${python_executable}" - "${model_path}" <<'PY' -import importlib.util -import runpy -import sys -from pathlib import Path - -model_path = sys.argv[1] -repo_root = Path.cwd() - -# Use the installed package for native extensions and the in-tree ExecuTorch -# route for the serializer/backend under test. -import torch_tensorrt # noqa: F401 - - -def overlay_module(name: str, path: Path) -> None: - spec = importlib.util.spec_from_file_location(name, path) - if spec is None or spec.loader is None: - raise RuntimeError(f"Could not load {name} from {path}") - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - spec.loader.exec_module(module) - - -overlay_module( - "torch_tensorrt.executorch.serialization", - repo_root / "py/torch_tensorrt/executorch/serialization.py", -) -overlay_module( - "torch_tensorrt.executorch.backend", - repo_root / "py/torch_tensorrt/executorch/backend.py", -) - -export_script = repo_root / "examples/torchtrt_executorch_example/export_static_shape.py" -sys.argv = [str(export_script), "--model_path", model_path] -runpy.run_path(str(export_script), run_name="__main__") -PY -test -f "${model_path}" - if [[ -n "${TensorRT_ROOT:-}" && -d "${TensorRT_ROOT}/lib" ]]; then export LD_LIBRARY_PATH="${TensorRT_ROOT}/lib${original_ld_library_path:+:${original_ld_library_path}}" else diff --git a/.github/workflows/_linux-x86_64-core.yml b/.github/workflows/_linux-x86_64-core.yml index c0011470b2..d164a48c50 100644 --- a/.github/workflows/_linux-x86_64-core.yml +++ b/.github/workflows/_linux-x86_64-core.yml @@ -98,11 +98,22 @@ jobs: use-rtx: ${{ inputs.use-rtx }} pip-install-torch-extra-args: "--extra-index-url https://pypi.org/simple" - # Standard-TRT only: ExecuTorch static build is not part of the RTX matrix. - executorch-static-build: + # Standard-TRT only: ExecuTorch runtime build is not part of the RTX matrix. + executorch-runtime-build: needs: [filter-matrix, build] if: ${{ !inputs.use-rtx }} - uses: ./.github/workflows/executorch-static-linux.yml + uses: ./.github/workflows/executorch-build-linux.yml + with: + repository: "pytorch/tensorrt" + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.filter-matrix.outputs.matrix }} + + executorch-runtime-test: + needs: [filter-matrix, executorch-runtime-build] + if: ${{ !inputs.use-rtx && needs.executorch-runtime-build.result == 'success' }} + uses: ./.github/workflows/executorch-test-linux.yml with: repository: "pytorch/tensorrt" ref: "" @@ -639,6 +650,8 @@ jobs: needs: [ build, + executorch-runtime-build, + executorch-runtime-test, L0-dynamo-converter-tests, L0-dynamo-core-tests, L0-py-core-tests, diff --git a/.github/workflows/executorch-static-linux.yml b/.github/workflows/executorch-build-linux.yml similarity index 80% rename from .github/workflows/executorch-static-linux.yml rename to .github/workflows/executorch-build-linux.yml index fe53fdb12c..f94b60256c 100644 --- a/.github/workflows/executorch-static-linux.yml +++ b/.github/workflows/executorch-build-linux.yml @@ -1,4 +1,4 @@ -name: ExecuTorch Static Linux Build +name: ExecuTorch Delegate Linux Build on: workflow_call: @@ -60,7 +60,7 @@ jobs: needs: select-matrix uses: ./.github/workflows/linux-test.yml with: - job-name: executorch-static-build + job-name: executorch-runtime-build repository: ${{ inputs.repository }} ref: ${{ inputs.ref }} test-infra-repository: ${{ inputs.test-infra-repository }} @@ -68,6 +68,7 @@ jobs: build-matrix: ${{ needs.select-matrix.outputs.matrix }} pre-script: packaging/pre_build_script.sh fail-on-empty: false + upload-artifact: torch-tensorrt-executorch-runtime script: | set -euo pipefail BAZELISK_VERSION="1.26.0" @@ -90,11 +91,8 @@ jobs: export PATH="${RUNNER_TEMP}/bin:${PATH}" bazel --version - # this is to build the libtorchtrt.tar.gz - bazel build //:libtorchtrt --compilation_mode opt --config=linux - executorch_cmake_location="$(bazel query @executorch//:executorch/CMakeLists.txt --output=location)" - export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")" - export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}" - # this is to verify the end user's workflow - python -m pip install pyyaml "executorch>=1.3.1" - .github/scripts/verify-executorch-reference-runner.sh + python -m pip install pyyaml "executorch==1.3.1" + export TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION="$(python -c 'import importlib.metadata; print(importlib.metadata.version("torch-tensorrt"))')" + + # Build the no-compile-for-users Python runtime wheel. + python -m pip wheel --no-build-isolation --no-deps --wheel-dir dist py/torch-tensorrt-executorch-runtime diff --git a/.github/workflows/executorch-test-linux.yml b/.github/workflows/executorch-test-linux.yml new file mode 100644 index 0000000000..80ac314e34 --- /dev/null +++ b/.github/workflows/executorch-test-linux.yml @@ -0,0 +1,94 @@ +name: ExecuTorch Delegate Linux Test + +on: + workflow_call: + inputs: + build-matrix: + description: "Build matrix to utilize" + default: "" + type: string + repository: + description: 'Repository to checkout, defaults to ""' + default: "" + type: string + ref: + description: 'Reference to checkout, defaults to ""' + 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: "" + type: string + +jobs: + select-matrix: + runs-on: ubuntu-latest + outputs: + matrix: ${{ steps.pick.outputs.matrix }} + steps: + - id: pick + env: + FULL_MATRIX: ${{ inputs.build-matrix }} + run: | + set -euo pipefail + python - <<'PY' >> "${GITHUB_OUTPUT}" + import json + import os + + matrix = json.loads(os.environ["FULL_MATRIX"]) + include = matrix.get("include", []) + if not include: + raise SystemExit("build-matrix include[] is empty") + + preferred = None + for entry in include: + if entry.get("python_version") == "3.11": + preferred = entry + break + + if preferred is None: + preferred = include[0] + + print("matrix=" + json.dumps({"include": [preferred]})) + PY + + test: + needs: select-matrix + uses: ./.github/workflows/linux-test.yml + with: + job-name: executorch-runtime-test + repository: ${{ inputs.repository }} + ref: ${{ inputs.ref }} + test-infra-repository: ${{ inputs.test-infra-repository }} + test-infra-ref: ${{ inputs.test-infra-ref }} + build-matrix: ${{ needs.select-matrix.outputs.matrix }} + pre-script: packaging/pre_build_script.sh + additional-artifact: torch-tensorrt-executorch-runtime + script: | + set -euo pipefail + mkdir -p "${RUNNER_TEMP}/bin" + curl -L "https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelisk-linux-amd64" \ + -o "${RUNNER_TEMP}/bin/bazel" + chmod +x "${RUNNER_TEMP}/bin/bazel" + export PATH="${RUNNER_TEMP}/bin:${PATH}" + + python -m pip install pyyaml executorch + python -c 'from torch_tensorrt_executorch_runtime import get_runtime; assert get_runtime().backend_registry.is_available("TensorRTBackend")' + + source tests/py/utils/ci_helpers.sh + trt_tier_executorch + + bazel build //:libtorchtrt --compilation_mode opt --config=linux + executorch_cmake_location="$(bazel query @executorch//:executorch/CMakeLists.txt --output=location)" + export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")" + export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}" + python examples/torchtrt_executorch_example/export_static_shape.py \ + --model_path="${RUNNER_TEMP}/torchtrt-python.pte" + .github/scripts/verify-executorch-reference-runner.sh \ + "${RUNNER_TEMP}/torchtrt-python.pte" + python examples/executorch_reference_runner/load_model.py \ + --model_path="${RUNNER_TEMP}/torchtrt-python.pte" --num_runs=1 diff --git a/.github/workflows/linux-test.yml b/.github/workflows/linux-test.yml index d055276143..a79c869799 100644 --- a/.github/workflows/linux-test.yml +++ b/.github/workflows/linux-test.yml @@ -50,9 +50,13 @@ on: default: false type: boolean upload-artifact: - description: 'Name to give artifacts uploaded from ${RUNNER_ARTIFACT_DIR}' + description: 'Name to give wheel artifacts uploaded from the repository dist directory' default: '' type: string + additional-artifact: + description: "Optional additional artifact to install alongside the core wheel" + default: "" + type: string use-rtx: description: "Whether to use TensorRT-RTX" default: false @@ -146,6 +150,12 @@ jobs: # with: # repository: ${{ inputs.repository }} # script: .github/scripts/install-torch-tensorrt.sh + - name: Download additional artifacts + if: ${{ inputs.additional-artifact != '' }} + uses: actions/download-artifact@v7 + with: + name: ${{ inputs.additional-artifact }} + path: /opt/torch-tensorrt-builds/ - name: Pack script continue-on-error: ${{ inputs.continue-on-error }} working-directory: ${{ inputs.repository }} @@ -174,23 +184,10 @@ jobs: display-options: fEs fail-on-empty: ${{ inputs.fail-on-empty }} - - name: Prepare artifacts for upload + - name: Check documentation for upload working-directory: ${{ inputs.repository }} id: check-artifacts - env: - UPLOAD_ARTIFACT_NAME: ${{ inputs.upload-artifact }} run: | - # Only do these steps if we actually want to upload an artifact - if [[ -n "${UPLOAD_ARTIFACT_NAME}" ]]; then - # If the default execution path is followed then we should get a wheel in the dist/ folder - # attempt to just grab whatever is in there and scoop it all up - if find "dist/" -name "*.whl" >/dev/null 2>/dev/null; then - mv -v dist/*.whl "${RUNNER_ARTIFACT_DIR}/" - fi - # Set to fail upload step if there are no files for upload and expected files for upload - echo 'if-no-files-found=error' >> "${GITHUB_OUTPUT}" - fi - upload_docs=0 # Check if there are things in the documentation folder to upload if find "${RUNNER_DOCS_DIR}" -mindepth 1 -maxdepth 1 | read -r; then @@ -205,8 +202,8 @@ jobs: if: ${{ inputs.upload-artifact != '' }} with: name: ${{ inputs.upload-artifact }} - path: ${{ runner.temp }}/artifacts/ - if-no-files-found: ${{ steps.check-artifacts.outputs.if-no-files-found }} + path: ${{ inputs.repository }}/dist/ + if-no-files-found: error - name: Upload documentation to S3 (if any) uses: seemethere/upload-artifact-s3@v5 diff --git a/.gitignore b/.gitignore index 8c151f2c0d..51871e0053 100644 --- a/.gitignore +++ b/.gitignore @@ -94,3 +94,5 @@ CLAUDE.md /compile_commands.json /external/ build-executorch/ +py/torch-tensorrt-executorch-runtime/build/ +py/torch-tensorrt-executorch-runtime/dist/ \ No newline at end of file diff --git a/BUILD.bazel b/BUILD.bazel index a0123ff3ee..e234f30d1b 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -53,6 +53,14 @@ config_setting( }, ) +filegroup( + name = "torchtrt_cmake_modules", + srcs = [ + "cmake/Modules/FindTensorRT.cmake", + ], + visibility = ["//visibility:public"], +) + pkg_tar( name = "include_core", package_dir = "include/torch_tensorrt", diff --git a/MODULE.bazel b/MODULE.bazel index 2f384e8066..93d491acd8 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -32,8 +32,8 @@ git_override( bazel_dep(name = "hedron_compile_commands", dev_dependency = True) git_override( module_name = "hedron_compile_commands", - remote = "https://github.com/hedronvision/bazel-compile-commands-extractor.git", commit = "0e990032f3c5a866e72615cf67e5ce22186dcb97", + remote = "https://github.com/hedronvision/bazel-compile-commands-extractor.git", ) new_local_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:local.bzl", "new_local_repository") @@ -42,12 +42,13 @@ new_git_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:git.bzl" local_torch = use_repo_rule("//toolchains:local_torch.bzl", "local_torch") -# Pinned to the ExecuTorch release/1.3 branch head. +# Keep this pin synchronized with the ExecuTorch release installed in +# py/torch-tensorrt-executorch-runtime/README.md. new_git_repository( name = "executorch", build_file = "@//third_party/executorch:BUILD", - # latest commit in release/1.3 branch - commit = "6118688a095fd9697224f5cad72ce42db641c9cd", + # executorch==1.3.1 + commit = "e2f18eb23c45bd22ca332b0b8b49a81de304b472", patch_cmds = [ "find . -mindepth 2 \\( -name BUILD -o -name BUILD.bazel \\) -delete", "mkdir executorch && find . -mindepth 1 -maxdepth 1 ! -name executorch ! -name BUILD ! -name BUILD.bazel ! -name REPO.bazel -exec cp -a {} executorch/ \\;", diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index b2e3b08232..e335748b33 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -51,6 +51,15 @@ using ::executorch::runtime::Span; namespace { thread_local cudaStream_t g_user_stream = nullptr; thread_local bool g_user_stream_set = false; + +extern const Error kRegistrationResult; + +Error check_registration() { + if (kRegistrationResult != Error::Ok) { + ET_LOG(Error, "TensorRTBackend registration failed: %s", ::executorch::runtime::to_string(kRegistrationResult)); + } + return kRegistrationResult; +} } // namespace CudaStreamGuard::CudaStreamGuard(cudaStream_t stream) : prev_stream_(g_user_stream), prev_set_(g_user_stream_set) { @@ -209,6 +218,10 @@ bool is_cuda_accessible_ptr(const void* ptr) { // is_available // --------------------------------------------------------------------------- bool TensorRTBackend::is_available() const { + if (check_registration() != Error::Ok) { + return false; + } + TRTLogger logger; TRTUniquePtr runtime(nvinfer1::createInferRuntime(logger)); return runtime != nullptr; @@ -227,6 +240,11 @@ Result TensorRTBackend::init( ArrayRef compile_specs) const { (void)compile_specs; + const Error registration_result = check_registration(); + if (registration_result != Error::Ok) { + return registration_result; + } + TORCHTRT_ET_CHECK_NOT_NULL(processed, Error::InvalidArgument, "TensorRTBackend::init: null processed buffer"); TORCHTRT_ET_CHECK_NOT_NULL(processed->data(), Error::InvalidArgument, "TensorRTBackend::init: null processed buffer"); @@ -642,14 +660,18 @@ void TensorRTBackend::destroy(DelegateHandle* handle) const { // Static registration – links the name "TensorRTBackend" used in the .pte // file to this implementation at program startup. // --------------------------------------------------------------------------- +namespace torch_tensorrt { +namespace executorch_backend { namespace { -torch_tensorrt::executorch_backend::TensorRTBackend& get_backend() { +TensorRTBackend& get_backend() { static torch_tensorrt::executorch_backend::TensorRTBackend backend; return backend; } const ::executorch::runtime::Backend kBackendId{"TensorRTBackend", &get_backend()}; -const auto kRegistered = ::executorch::runtime::register_backend(kBackendId); +const Error kRegistrationResult = ::executorch::runtime::register_backend(kBackendId); } // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/examples/executorch_reference_runner/BUILD b/examples/executorch_reference_runner/BUILD index 4f53bc6b91..8335092725 100644 --- a/examples/executorch_reference_runner/BUILD +++ b/examples/executorch_reference_runner/BUILD @@ -7,6 +7,7 @@ filegroup( srcs = [ "CMakeLists.txt", "README.md", + "load_model.py", "main.cpp", ], ) diff --git a/examples/executorch_reference_runner/README.md b/examples/executorch_reference_runner/README.md index a518353bb6..ddf635f297 100644 --- a/examples/executorch_reference_runner/README.md +++ b/examples/executorch_reference_runner/README.md @@ -79,6 +79,28 @@ build-executorch-reference-runner/lib/libexecutorch_trt_backend.a ## Load And Run A `.pte` Model +### Python + +Install the complete prebuilt Python runtime and delegate: + +```bash +pip install "torch-tensorrt[executorch]" +``` + +Load and run the model without an ExecuTorch checkout or native build: + +```bash +python examples/executorch_reference_runner/load_model.py \ + --model_path=model.pte \ + --num_runs=1 +``` + +The extra installs `executorch` and the matching +`torch-tensorrt-executorch-runtime` wheel. That wheel contains an ExecuTorch +Python runtime with `TensorRTBackend` linked into its backend registry. + +### C++ + Run the reference runner against a Torch-TensorRT compiled ExecuTorch model: ```bash diff --git a/examples/executorch_reference_runner/load_model.py b/examples/executorch_reference_runner/load_model.py new file mode 100644 index 0000000000..1e365c6cc8 --- /dev/null +++ b/examples/executorch_reference_runner/load_model.py @@ -0,0 +1,33 @@ +import argparse +from pathlib import Path + +import torch +import torch_tensorrt + +parser = argparse.ArgumentParser() +parser.add_argument( + "--model_path", + type=Path, + required=True, + help="Path to the ExecuTorch .pte model", +) +parser.add_argument("--num_runs", type=int, default=1) +args = parser.parse_args() +if args.num_runs < 1: + raise ValueError("--num_runs must be at least 1") + +model_path = args.model_path +x = torch.ones((2, 3, 4, 4), dtype=torch.float32) + +program = torch_tensorrt.load(model_path, format="executorch") +for _ in range(args.num_runs): + outputs = program.forward(x) +y = outputs[0] + +expected = x + 1 +torch.testing.assert_close(y.cpu(), expected) + +print("methods:", sorted(program.method_names)) +print("output shape:", tuple(y.shape)) +print("output device:", y.device) +print("PASS: ExecuTorch TensorRT delegate output matches x + 1") diff --git a/py/torch-tensorrt-executorch-runtime/README.md b/py/torch-tensorrt-executorch-runtime/README.md new file mode 100644 index 0000000000..9a524d69b6 --- /dev/null +++ b/py/torch-tensorrt-executorch-runtime/README.md @@ -0,0 +1,76 @@ +# Torch-TensorRT ExecuTorch Runtime Wheel + +This directory builds `torch-tensorrt-executorch-runtime`. The Linux wheel +contains an ExecuTorch `_portable_lib` Python runtime with `TensorRTBackend` +force-linked into the same native module that owns the backend registry. + +The wheel must use the same Python, PyTorch, ExecuTorch, CUDA, TensorRT, and +C++ ABI as its matching Torch-TensorRT wheel. + +## Runtime libraries + +The wheel does not bundle PyTorch, c10, TensorRT, or CUDA shared libraries. +Its `_portable_lib.so` has origin-relative runtime search paths for the +TensorRT and CUDA library locations installed by their Python packages: + +- `tensorrt_libs` +- `nvidia/cuda_runtime/lib` (CUDA 12) +- `nvidia/cu13/lib` (CUDA 13) + +These packages are installed transitively with the matching `torch-tensorrt` +wheel. For a system TensorRT or CUDA installation outside these standard +locations, its `lib` directory must be available through the system dynamic +loader configuration or `LD_LIBRARY_PATH`. + +The CI manylinux repair step changes the wheel platform tag; it does not +bundle these external libraries. The origin-relative paths are therefore part +of the wheel runtime contract. + +## Build + +> [!IMPORTANT] +> Build this wheel with `--no-build-isolation`. Its native extension must use +> the exact PyTorch installation that the matching Torch-TensorRT artifacts +> were built against. An isolated build may download a newer, ABI-incompatible +> PyTorch version. + +```bash +export TensorRT_ROOT=/path/to/TensorRT + +python -m pip install pyyaml "executorch==1.3.1" +python -m pip wheel --no-build-isolation --no-deps \ + --wheel-dir dist py/torch-tensorrt-executorch-runtime +``` + +The native build obtains the ExecuTorch source through Bazel; no separate +source checkout or `EXECUTORCH_SOURCE_DIR` setting is required. The source +commit pinned in `MODULE.bazel` is the revision recorded by the +`executorch==1.3.1` wheel. + +The static ExecuTorch and delegate archives are intermediate build inputs; +users receive the final native Python module and do not compile anything. + +## Python tensor placement + +The ExecuTorch Python portable runtime uses CPU tensors at its API boundary. +CUDA tensor inputs passed to `Program.run()` or `Program.forward()` are copied +to CPU before dispatch. TensorRT executes the delegated graph on GPU, but the +runtime copies inputs to the device and returns outputs on CPU. + +Consequently, the Python API does not use the backend's device-resident +input/output fast path. Applications that need to keep inputs and outputs on +GPU should use the ExecuTorch C++ runner. + +## Use + +```bash +pip install "torch-tensorrt[executorch]" +``` + +```python +import torch +import torch_tensorrt + +program = torch_tensorrt.load("model.pte", format="executorch") +outputs = program.forward(torch.ones((2, 3, 4, 4))) +``` diff --git a/py/torch-tensorrt-executorch-runtime/native/BUILD.bazel b/py/torch-tensorrt-executorch-runtime/native/BUILD.bazel new file mode 100644 index 0000000000..3844ca869a --- /dev/null +++ b/py/torch-tensorrt-executorch-runtime/native/BUILD.bazel @@ -0,0 +1,56 @@ +load("@rules_foreign_cc//foreign_cc:defs.bzl", "cmake") + +package(default_visibility = ["//visibility:public"]) + +config_setting( + name = "aarch64_linux", + constraint_values = [ + "@platforms//cpu:aarch64", + "@platforms//os:linux", + ], +) + +filegroup( + name = "delegate_sources", + srcs = glob( + ["**"], + exclude = ["BUILD.bazel"], + ) + [ + "//:torchtrt_cmake_modules", + "//core/runtime:include_files", + "//cpp:executorch_api_headers", + "//cpp:executorch_backend_source_files", + "@executorch//:executorch_sources", + ], +) + +# ExecuTorch exposes its Python portable runtime through CMake. Keep that +# upstream build behind Bazel so both Torch-TensorRT wheels have one build +# entry point and share Bazel's dependency/toolchain selection. +cmake( + name = "delegate_native", + cache_entries = { + "CMAKE_BUILD_TYPE": "Release", + "EXECUTORCH_CMAKE_FILE": "$(execpath @executorch//:executorch/CMakeLists.txt)", + "PYTHON_EXECUTABLE": "$${PYTHON_BIN_PATH}", + "TORCH_HEADER_MARKER": "$(execpath @libtorch//:include/torch/headeronly/util/TypeTraits.h)", + "TORCH_TENSORRT_SOURCE_DIR": "$$EXT_BUILD_ROOT$$", + }, + data = [ + "@executorch//:executorch/CMakeLists.txt", + "@libtorch//:include/torch/headeronly/util/TypeTraits.h", + ], + lib_source = ":delegate_sources", + out_include_dir = "", + out_shared_libs = [ + "_portable_lib.so", + "data_loader.so", + ], + deps = [ + "@cuda//:cudart", + "@libtorch//:torch", + ] + select({ + ":aarch64_linux": ["@tensorrt_sbsa//:nvinfer"], + "//conditions:default": ["@tensorrt//:nvinfer"], + }), +) diff --git a/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt b/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt new file mode 100644 index 0000000000..78c95da1ea --- /dev/null +++ b/py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt @@ -0,0 +1,97 @@ +cmake_minimum_required(VERSION 3.24) +project(torch_tensorrt_executorch_runtime LANGUAGES CXX) + +if(NOT EXECUTORCH_SOURCE_DIR AND EXECUTORCH_CMAKE_FILE) + get_filename_component(EXECUTORCH_SOURCE_DIR "${EXECUTORCH_CMAKE_FILE}" DIRECTORY) +endif() +if(NOT TORCH_PRIMARY_INCLUDE_DIR AND TORCH_HEADER_MARKER) + get_filename_component(TORCH_PRIMARY_INCLUDE_DIR "${TORCH_HEADER_MARKER}/../../../.." ABSOLUTE) +endif() +if(NOT EXECUTORCH_SOURCE_DIR OR NOT TORCH_TENSORRT_SOURCE_DIR OR NOT TORCH_PRIMARY_INCLUDE_DIR) + message(FATAL_ERROR + "EXECUTORCH_SOURCE_DIR, TORCH_TENSORRT_SOURCE_DIR, and TORCH_PRIMARY_INCLUDE_DIR are required") +endif() + +list(APPEND CMAKE_MODULE_PATH "${TORCH_TENSORRT_SOURCE_DIR}/cmake/Modules") +find_package(TensorRT REQUIRED) +find_package(CUDAToolkit REQUIRED) +find_package(Threads REQUIRED) + +set(BUILD_TESTING OFF CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_PYBIND ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_MODULE ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_KERNELS_OPTIMIZED OFF CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_XNNPACK OFF CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_TESTS OFF CACHE BOOL "" FORCE) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +add_subdirectory("${EXECUTORCH_SOURCE_DIR}" executorch) + +add_library(torch_tensorrt_executorch_backend STATIC + "${TORCH_TENSORRT_SOURCE_DIR}/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp" + "${TORCH_TENSORRT_SOURCE_DIR}/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp") +target_compile_features(torch_tensorrt_executorch_backend PUBLIC cxx_std_17) +target_compile_definitions(torch_tensorrt_executorch_backend + PRIVATE C10_USING_CUSTOM_GENERATED_MACROS) +target_include_directories(torch_tensorrt_executorch_backend PRIVATE + "${TORCH_TENSORRT_SOURCE_DIR}" + "${TORCH_TENSORRT_SOURCE_DIR}/cpp/include" + "${EXECUTORCH_SOURCE_DIR}/.." + "${EXECUTORCH_SOURCE_DIR}/runtime/core/portable_type/c10") +target_link_libraries(torch_tensorrt_executorch_backend PRIVATE + CUDA::cudart TensorRT::nvinfer Threads::Threads executorch) + +if(NOT TARGET portable_lib) + message(FATAL_ERROR "ExecuTorch did not define portable_lib") +endif() + +# The pybind bridge includes both ATen and portable ExecuTorch headers. Keep the +# ATen/c10/torch header set internally consistent by resolving all of it from +# the exact PyTorch wheel that supplies the linked libraries. +if(NOT EXISTS "${TORCH_PRIMARY_INCLUDE_DIR}/torch/headeronly") + message(FATAL_ERROR + "torch/headeronly not found under TORCH_PRIMARY_INCLUDE_DIR=${TORCH_PRIMARY_INCLUDE_DIR}") +endif() +file(MAKE_DIRECTORY "${CMAKE_BINARY_DIR}/torch_header_shim/torch") +file(COPY "${TORCH_PRIMARY_INCLUDE_DIR}/c10/" + DESTINATION "${CMAKE_BINARY_DIR}/torch_header_shim/c10") +file(COPY "${TORCH_PRIMARY_INCLUDE_DIR}/torch/headeronly/" + DESTINATION "${CMAKE_BINARY_DIR}/torch_header_shim/torch/headeronly") +if(NOT EXISTS "${CMAKE_BINARY_DIR}/torch_header_shim/c10/util/complex.h" OR + NOT EXISTS "${CMAKE_BINARY_DIR}/torch_header_shim/torch/headeronly/macros/Macros.h") + message(FATAL_ERROR "Failed to stage the PyTorch header shim") +endif() +target_include_directories(util BEFORE PRIVATE "${CMAKE_BINARY_DIR}/torch_header_shim") +target_include_directories(portable_lib BEFORE PRIVATE "${CMAKE_BINARY_DIR}/torch_header_shim") + +# TensorRTBackend registers through a static initializer. Force its complete +# archive into the same pybind module that owns ExecuTorch's backend registry. +target_link_libraries(portable_lib PRIVATE + "$" + extension_threadpool + CUDA::cudart TensorRT::nvinfer Threads::Threads) + +# The runtime wheel intentionally does not bundle PyTorch, TensorRT, or CUDA. +# TensorRT and CUDA pip packages install their shared libraries in sibling +# directories under site-packages, relative to this extension package. +# Preserve these paths in both the Bazel-collected build output and installed +# extension so importing the wheel does not depend on LD_LIBRARY_PATH. +set(_torch_tensorrt_executorch_runtime_rpath + "$ORIGIN" + "$ORIGIN/../tensorrt_libs" + "$ORIGIN/../nvidia/cuda_runtime/lib" + "$ORIGIN/../nvidia/cu13/lib") +set_target_properties(portable_lib PROPERTIES + BUILD_WITH_INSTALL_RPATH ON + INSTALL_RPATH "${_torch_tensorrt_executorch_runtime_rpath}" + OUTPUT_NAME "_portable_lib" + SUFFIX ".so") +set_target_properties(data_loader PROPERTIES SUFFIX ".so") + +install(TARGETS portable_lib data_loader LIBRARY DESTINATION lib) + +add_custom_target(torch_tensorrt_executorch_portable_lib DEPENDS portable_lib data_loader) diff --git a/py/torch-tensorrt-executorch-runtime/pyproject.toml b/py/torch-tensorrt-executorch-runtime/pyproject.toml new file mode 100644 index 0000000000..a500209126 --- /dev/null +++ b/py/torch-tensorrt-executorch-runtime/pyproject.toml @@ -0,0 +1,11 @@ +[build-system] +requires = [ + "setuptools>=68", + "wheel>=0.40", + # Native ABI dependencies supplied by the active Torch-TensorRT build + # environment. + # Builds must use --no-build-isolation; see README.md. + "torch", + "executorch==1.3.1", +] +build-backend = "setuptools.build_meta" diff --git a/py/torch-tensorrt-executorch-runtime/setup.py b/py/torch-tensorrt-executorch-runtime/setup.py new file mode 100644 index 0000000000..f8d676ac80 --- /dev/null +++ b/py/torch-tensorrt-executorch-runtime/setup.py @@ -0,0 +1,155 @@ +"""Build the precompiled Torch-TensorRT backend for ExecuTorch Python.""" + +from __future__ import annotations + +import importlib.metadata +import os +import pathlib +import platform +import re +import shlex +import shutil +import subprocess +import sys +import uuid + +import torch +from setuptools import Extension, find_packages, setup +from setuptools.command.build_ext import build_ext + +HERE = pathlib.Path(__file__).resolve().parent +REPO_ROOT = HERE.parents[1] +BAZEL_TARGET = "//py/torch-tensorrt-executorch-runtime/native:delegate_native" +BUILD_NONCE = os.getenv("TORCH_TENSORRT_EXECUTORCH_BUILD_NONCE", uuid.uuid4().hex) +TENSORRT_DISTRIBUTION = "tensorrt-cu13" +CUDA_RUNTIME_DISTRIBUTION = "nvidia-cuda-runtime" + + +def torchtrt_version() -> str: + if value := os.getenv("TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION"): + return value + version_py = REPO_ROOT / "py/torch_tensorrt/_version.py" + if version_py.exists(): + if m := re.search(r'__version__\s*=\s*["\']([^"\']+)', version_py.read_text()): + return m.group(1) + return (REPO_ROOT / "version.txt").read_text().strip() + + +def public_version(version: str) -> str: + """Drop a PEP 440 local suffix that may not be present on package indexes.""" + return version.partition("+")[0] + + +def installed_version(distribution: str) -> str: + """Return the version of a dependency in the native build environment.""" + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError as error: + raise RuntimeError( + f"{distribution} must be installed to build the ExecuTorch runtime wheel" + ) from error + + +def require_cuda_13() -> None: + """Reject builds whose native dependencies do not use supported CUDA 13.""" + cuda_version = torch.version.cuda + if cuda_version is None or not cuda_version.startswith("13."): + raise RuntimeError( + "CUDA 13-enabled PyTorch is required to build this wheel " + f"(found CUDA {cuda_version or 'None'})" + ) + + +class BazelExtension(Extension): + def __init__(self, name: str) -> None: + super().__init__(name, sources=[]) + + +class BazelBuild(build_ext): + def build_extension(self, ext: Extension) -> None: + if sys.platform != "linux": + raise RuntimeError("The ExecuTorch TensorRT delegate supports Linux only") + + output = pathlib.Path(self.get_ext_fullpath(ext.name)).resolve() + output.parent.mkdir(parents=True, exist_ok=True) + + bazel = shutil.which("bazelisk") or shutil.which("bazel") + if bazel is None: + raise RuntimeError("Could not find bazelisk or bazel in PATH") + + compilation_mode = "dbg" if self.debug else "opt" + command = [ + bazel, + "build", + BAZEL_TARGET, + "--config=linux", + "--config=python", + f"--compilation_mode={compilation_mode}", + f"--action_env=PYTHON_BIN_PATH={sys.executable}", + f"--action_env=TORCH_TENSORRT_EXECUTORCH_BUILD_NONCE={BUILD_NONCE}", + ] + dist_dir_arch = ( + "aarch64-linux-gnu" + if platform.machine() in {"aarch64", "arm64"} + else "x86_64-linux-gnu" + ) + dist_dir = REPO_ROOT / "third_party/dist_dir" / dist_dir_arch + if dist_dir.is_dir(): + command.append(f"--distdir={dist_dir}") + command.extend(shlex.split(os.getenv("BAZEL_ARGS", ""))) + + env = os.environ.copy() + env.setdefault("TORCH_PATH", str(pathlib.Path(torch.__file__).resolve().parent)) + subprocess.run(command, cwd=REPO_ROOT, env=env, check=True) + + bazel_bin = pathlib.Path( + subprocess.check_output( + [ + bazel, + "info", + "bazel-bin", + f"--compilation_mode={compilation_mode}", + ], + cwd=REPO_ROOT, + env=env, + text=True, + ).strip() + ) + library_stem = ( + "_portable_lib" if ext.name.endswith("._portable_lib") else "data_loader" + ) + built = ( + bazel_bin + / "py/torch-tensorrt-executorch-runtime/native/delegate_native/lib" + / f"{library_stem}.so" + ) + if not built.is_file(): + raise RuntimeError(f"Bazel did not produce {built}") + output.unlink(missing_ok=True) + shutil.copy2(built, output) + + +require_cuda_13() +executorch_version = installed_version("executorch") +tensorrt_version = installed_version(TENSORRT_DISTRIBUTION) +cuda_runtime_version = installed_version(CUDA_RUNTIME_DISTRIBUTION) +setup( + name="torch-tensorrt-executorch-runtime", + version=torchtrt_version(), + description="Torch-TensorRT delegate for the ExecuTorch Python runtime", + packages=find_packages(), + ext_modules=[ + BazelExtension("torch_tensorrt_executorch_runtime._portable_lib"), + BazelExtension("torch_tensorrt_executorch_runtime.data_loader"), + ], + cmdclass={"build_ext": BazelBuild}, + python_requires=">=3.10", + install_requires=[ + f"torch=={public_version(torch.__version__)}", + f"executorch=={executorch_version}", + f"torch-tensorrt=={torchtrt_version()}", + f"{TENSORRT_DISTRIBUTION}=={tensorrt_version}", + f"{CUDA_RUNTIME_DISTRIBUTION}=={cuda_runtime_version}", + ], + zip_safe=False, +) diff --git a/py/torch-tensorrt-executorch-runtime/torch_tensorrt_executorch_runtime/__init__.py b/py/torch-tensorrt-executorch-runtime/torch_tensorrt_executorch_runtime/__init__.py new file mode 100644 index 0000000000..50738b8cdb --- /dev/null +++ b/py/torch-tensorrt-executorch-runtime/torch_tensorrt_executorch_runtime/__init__.py @@ -0,0 +1,75 @@ +"""Activate the TensorRT delegate-enabled ExecuTorch Python runtime.""" + +from __future__ import annotations + +import importlib +import sys +from types import ModuleType +from typing import Any, Protocol, cast + +BACKEND_NAME = "TensorRTBackend" +_NATIVE_NAME = "executorch.extension.pybindings._portable_lib" +_WRAPPER_NAME = "executorch.extension.pybindings.portable_lib" +_DATA_LOADER_NAME = "executorch.extension.pybindings.data_loader" + + +class _BackendRegistry(Protocol): + def is_available(self, name: str) -> bool: ... + + +class _Runtime(Protocol): + backend_registry: _BackendRegistry + + def load_program(self, data: bytes) -> Any: ... + + +class DelegateCompatibilityError(ImportError): + """The runtime wheel is incompatible with the active native runtime.""" + + +def activate() -> ModuleType: + """Make the delegate-enabled portable runtime back ``executorch.runtime``.""" + existing = sys.modules.get(_NATIVE_NAME) + if existing is not None and existing.__name__ == __name__ + "._portable_lib": + return existing + if existing is not None or _WRAPPER_NAME in sys.modules: + raise DelegateCompatibilityError( + "ExecuTorch's stock runtime was imported first. Call " + 'torch_tensorrt.load(..., format="executorch") before importing ' + "executorch.runtime." + ) + previous_data_loader = sys.modules.get(_DATA_LOADER_NAME) + try: + data_loader = importlib.import_module(__name__ + ".data_loader") + # _portable_lib imports this canonical name while its module initializer + # runs. Install our binding first so Python does not load ExecuTorch's + # stock data_loader and register PyDataLoader a second time. + sys.modules[_DATA_LOADER_NAME] = data_loader + native = importlib.import_module(__name__ + "._portable_lib") + except ImportError as error: + if previous_data_loader is None: + sys.modules.pop(_DATA_LOADER_NAME, None) + else: + sys.modules[_DATA_LOADER_NAME] = previous_data_loader + raise DelegateCompatibilityError( + "Could not load the prebuilt Torch-TensorRT ExecuTorch runtime. " + "Install torch, executorch, torch-tensorrt, and the runtime package from " + "the same release matrix." + ) from error + sys.modules[_NATIVE_NAME] = native + sys.modules.pop(_WRAPPER_NAME, None) + return native + + +def get_runtime() -> _Runtime: + """Return the activated ExecuTorch Runtime singleton.""" + activate() + from executorch.runtime import Runtime + + value = cast(_Runtime, Runtime.get()) + if not value.backend_registry.is_available(BACKEND_NAME): + raise DelegateCompatibilityError(f"{BACKEND_NAME} is not registered") + return value + + +__all__ = ["BACKEND_NAME", "DelegateCompatibilityError", "activate", "get_runtime"] diff --git a/py/torch-tensorrt-executorch-runtime/torch_tensorrt_executorch_runtime/runtime.py b/py/torch-tensorrt-executorch-runtime/torch_tensorrt_executorch_runtime/runtime.py new file mode 100644 index 0000000000..deb9600e64 --- /dev/null +++ b/py/torch-tensorrt-executorch-runtime/torch_tensorrt_executorch_runtime/runtime.py @@ -0,0 +1,81 @@ +"""Python inference API for Torch-TensorRT ExecuTorch programs.""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any, Collection, Sequence, Union, cast + +if TYPE_CHECKING: + from torch_tensorrt_executorch_runtime import _Runtime + + +def _get_runtime() -> _Runtime: + try: + from torch_tensorrt_executorch_runtime import get_runtime + except ImportError as error: + raise ImportError( + "ExecuTorch Python inference requires the prebuilt delegate. " + 'Install it with: pip install "torch-tensorrt[executorch]"' + ) from error + return get_runtime() + + +class Program: + """A loaded ExecuTorch program backed by TensorRTBackend. + + The ExecuTorch Python portable runtime executes across a CPU tensor + boundary: CUDA inputs are copied to CPU before dispatch and outputs are + returned on CPU. TensorRT still executes the delegated graph on GPU, but + the device-resident input/output fast path is available only through the + ExecuTorch C++ runner. + """ + + def __init__(self, program: Any, data: bytes) -> None: + # ExecuTorch's BufferDataLoader references this memory without copying it. + self._data = data + self._program = program + + @property + def method_names(self) -> Collection[str]: + return cast(Collection[str], self._program.method_names) + + def run(self, inputs: Sequence[Any], method: str = "forward") -> Sequence[Any]: + """Run a method using CPU inputs and return CPU outputs. + + CUDA tensor inputs are copied to CPU before entering the portable + Python runtime. Use the C++ runner when inputs and outputs must remain + device-resident. + """ + import torch + + inputs = tuple( + value.cpu() if isinstance(value, torch.Tensor) and value.is_cuda else value + for value in inputs + ) + if method not in self.method_names: + raise ValueError( + f"Unknown method {method!r}; available methods: {sorted(self.method_names)}" + ) + loaded = self._program.load_method(method) + if loaded is None: + raise RuntimeError(f"ExecuTorch failed to load method {method!r}") + return cast(Sequence[Any], loaded.execute(inputs)) + + def forward(self, *inputs: Any) -> Sequence[Any]: + return self.run(inputs, "forward") + + +def load(path: Union[str, Path]) -> Program: + """Load a `.pte` with the delegate-enabled ExecuTorch Python runtime. + + External `.pdf` weight files are not supported; weights must be embedded + in the `.pte` file. + """ + model_path = Path(path) + if not model_path.is_file(): + raise FileNotFoundError(f"ExecuTorch model not found: {model_path}") + data = model_path.read_bytes() + return Program(_get_runtime().load_program(data), data) + + +__all__ = ["Program", "load"] diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index a4d1feaabf..f34384d900 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -18,6 +18,7 @@ Set, Tuple, Union, + cast, ) import torch @@ -89,6 +90,13 @@ def _has_executorch_exir() -> bool: return False +def _has_executorch_runtime() -> bool: + try: + return importlib.util.find_spec("torch_tensorrt_executorch_runtime") is not None + except ModuleNotFoundError: + return False + + def _non_fx_input_interface( inputs: Sequence[Input | torch.Tensor], ) -> TypeGuard[List[Input | torch.Tensor]]: @@ -316,7 +324,7 @@ def _fx_input_interface( "'arg_inputs' and 'inputs' should not be used at the same time." ) if inputs is not None: - arg_inputs = inputs # type: ignore[assignment] + arg_inputs = inputs if kwarg_inputs is None: kwarg_inputs = {} @@ -415,7 +423,7 @@ def cross_compile_for_windows( "'arg_inputs' and 'inputs' should not be used at the same time." ) - arg_inputs = inputs or arg_inputs # type: ignore[assignment] + arg_inputs = inputs or arg_inputs if kwarg_inputs is None: kwarg_inputs = {} @@ -515,7 +523,7 @@ def convert_method_to_trt_engine( raise AssertionError( "'arg_inputs' and 'inputs' should not be used at the same time." ) - arg_inputs = arg_inputs or inputs # type: ignore[assignment] + arg_inputs = arg_inputs or inputs module_type = _parse_module_type(module) target_ir = _get_target_fe(module_type, ir) @@ -559,11 +567,14 @@ def convert_method_to_trt_engine( module, torchtrt_arg_inputs, kwarg_inputs=torchtrt_kwarg_inputs, **kwargs ) - return dynamo_convert_exported_program_to_serialized_trt_engine( - exp_program, - arg_inputs=tuple(normalized_arg_inputs), - kwarg_inputs=torchtrt_kwarg_inputs, - **kwargs, + return cast( + bytes, + dynamo_convert_exported_program_to_serialized_trt_engine( + exp_program, + arg_inputs=tuple(normalized_arg_inputs), + kwarg_inputs=torchtrt_kwarg_inputs, + **kwargs, + ), ) elif target_ir == _IRType.torch_compile: raise RuntimeError( @@ -587,16 +598,23 @@ def load_cross_compiled_exported_program(file_path: str = "") -> Any: def load( - file_path: str = "", extra_files: Optional[dict[str, Any]] = None, **kwargs: Any + file_path: str = "", + extra_files: Optional[dict[str, Any]] = None, + *, + format: Optional[str] = None, + **kwargs: Any, ) -> Any: """ - Load either a Torchscript model or ExportedProgram. + Load a TorchScript, ExportedProgram, or ExecuTorch program. - Loads a TorchScript or ExportedProgram file from disk. File type will be detect the type using try, except. + By default, detects TorchScript and ExportedProgram files. Set + ``format="executorch"`` explicitly for an ExecuTorch ``.pte`` file. Arguments: file_path (str): Path to file on the disk extra_files (dict[str, Any]): Extra files to load with the model + format (Optional[str]): Set to ``"executorch"`` to load a ``.pte`` file + using the separately installed ExecuTorch runtime package. Example: # Load with extra files. @@ -605,8 +623,24 @@ def load( print(extra_files["foo.txt"]) Raises: - ValueError: If there is no file or the file is not either a TorchScript file or ExportedProgram file + ImportError: If ExecuTorch format is requested without the runtime package + ValueError: If the format is unsupported or the file is not a TorchScript or ExportedProgram file """ + if format == "executorch": + if not _has_executorch_runtime(): + raise ImportError( + "Loading an ExecuTorch program requires the prebuilt " + "Torch-TensorRT ExecuTorch delegate. Install it with: " + 'pip install "torch-tensorrt[executorch]"' + ) + from torch_tensorrt_executorch_runtime.runtime import load as load_executorch + + return load_executorch(file_path) + if format is not None: + raise ValueError( + f"Unsupported format {format!r}; expected None or 'executorch'" + ) + # Ensure Python TRT engine ops are registered so torch.export.load can # resolve tensorrt::execute_engine when the C++ runtime is absent. if not ENABLED_FEATURES.torch_tensorrt_runtime: @@ -756,7 +790,7 @@ def save( CUDA ``.pte``. See :ref:`the ExecuTorch save guide ` for the ``CudaPartitioner`` recipe, its export-time requirements (CUDA backend + nvcc), and the external - ``.ptd`` weight caveats. + ``.pdf`` weight caveats. """ if isinstance(module, CudaGraphsTorchTensorRTModule): module = module.compiled_module @@ -1336,14 +1370,14 @@ def _replace_execute_engine_for_executorch(exp_program: Any) -> Any: def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None: - """Write an ExecuTorch program's external named tensor data (``.ptd``) next to the ``.pte``. + """Write an ExecuTorch program's external named tensor data (``.pdf``) next to the ``.pte``. The CUDA (AOTInductor) backend emits its weights as external named data (``save_data_externally``), which ExecuTorch serializes only via ``write_tensor_data_to_file`` -- ``ExecutorchProgram.write_to_file`` does not persist it. So a partition that carries external weights (e.g. a ``CudaPartitioner`` delegate) would lose its blob and the ``.pte`` could not - load. This writes the ``.ptd`` data file(s) into the ``.pte``'s directory when + load. This writes the ``.pdf`` data file(s) into the ``.pte``'s directory when the program has any; the runtime must then be given those data file(s) (e.g. via the ExecuTorch ``Module`` data-files argument) to load the weights. """ @@ -1352,7 +1386,7 @@ def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None out_dir = os.path.dirname(os.path.abspath(file_path)) executorch_program.write_tensor_data_to_file(out_dir) logger.info( - "Wrote external delegate weights (.ptd) to %s; point the runtime's " + "Wrote external delegate weights (.pdf) to %s; point the runtime's " "data-files at this directory to load them.", out_dir, ) diff --git a/py/torch_tensorrt/executorch/__init__.py b/py/torch_tensorrt/executorch/__init__.py index 123eee846d..28a84a6878 100644 --- a/py/torch_tensorrt/executorch/__init__.py +++ b/py/torch_tensorrt/executorch/__init__.py @@ -1,3 +1,10 @@ +"""ExecuTorch compilation and export integration. + +Runtime loading is provided by the optional +``torch-tensorrt-executorch-runtime`` distribution and dispatched through +``torch_tensorrt.load(..., format="executorch")``. +""" + import importlib from typing import TYPE_CHECKING, NoReturn diff --git a/setup.py b/setup.py index 91651c7e15..6007ef15cc 100644 --- a/setup.py +++ b/setup.py @@ -99,12 +99,6 @@ def load_dep_info(): CI_BUILD = False USE_TRT_RTX = False -EXECUTORCH_REQUIREMENT = "executorch>=1.3.1" -EXTRAS_REQUIRE = { - "executorch": [EXECUTORCH_REQUIREMENT], - "all": [EXECUTORCH_REQUIREMENT], -} - if "--use-rtx" in sys.argv: USE_TRT_RTX = True @@ -189,6 +183,15 @@ def load_dep_info(): else: __version__ = f"{get_base_version()}.dev0+{get_git_revision_short_hash()}" +EXECUTORCH_REQUIREMENT = "executorch>=1.3.1" +EXECUTORCH_RUNTIME_REQUIREMENT = ( + f"torch-tensorrt-executorch-runtime=={__version__}; " "platform_system == 'Linux'" +) +EXTRAS_REQUIRE = { + "executorch": [EXECUTORCH_REQUIREMENT, EXECUTORCH_RUNTIME_REQUIREMENT], + "all": [EXECUTORCH_REQUIREMENT, EXECUTORCH_RUNTIME_REQUIREMENT], +} + if "--ci" in sys.argv: sys.argv.remove("--ci") if RELEASE: @@ -612,7 +615,6 @@ def run(self): ) package_data = {} -executorch_header_package_data = ["include/torch_tensorrt/executorch/*.h"] if not (PY_ONLY or NO_TS): tensorrt_x86_64_external_dir = ( @@ -779,7 +781,6 @@ def run(self): { "torch_tensorrt": [ "include/torch_tensorrt/*.h", - *executorch_header_package_data, "include/torch_tensorrt/core/*.h", "include/torch_tensorrt/core/conversion/*.h", "include/torch_tensorrt/core/conversion/conversionctx/*.h", @@ -809,7 +810,6 @@ def run(self): { "torch_tensorrt": [ "include/torch_tensorrt/*.h", - *executorch_header_package_data, "include/torch_tensorrt/core/*.h", "include/torch_tensorrt/core/runtime/*.h", "lib/*", diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 5d88822716..4ac9cdcf04 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -53,22 +53,79 @@ def fake_find_spec(name, package=None): ) +@pytest.mark.unit +def test_load_executorch_error_when_delegate_missing(monkeypatch): + from torch_tensorrt import _compile + + monkeypatch.setattr(_compile, "_has_executorch_runtime", lambda: False) + + with pytest.raises(ImportError, match=r"torch-tensorrt\[executorch\]"): + _compile.load("model.pte", format="executorch") + + +@pytest.mark.unit +def test_load_executorch_dispatches_to_delegate(monkeypatch): + from torch_tensorrt import _compile + + delegate = types.ModuleType("torch_tensorrt_executorch_runtime") + delegate.__path__ = [] + runtime = types.ModuleType("torch_tensorrt_executorch_runtime.runtime") + sentinel = object() + runtime.load = lambda path: (sentinel, path) + monkeypatch.setitem(sys.modules, delegate.__name__, delegate) + monkeypatch.setitem(sys.modules, runtime.__name__, runtime) + monkeypatch.setattr(_compile, "_has_executorch_runtime", lambda: True) + + assert _compile.load("model.pte", format="executorch") == ( + sentinel, + "model.pte", + ) + + @pytest.mark.unit def test_public_api_symbols_present(): module = importlib.import_module("torch_tensorrt.executorch") assert "get_edge_compile_config" in module.__all__ assert "TensorRTPartitioner" in module.__all__ assert "TensorRTBackend" in module.__all__ + assert "Program" not in module.__all__ + assert "load" not in module.__all__ _REPO_ROOT = Path(__file__).resolve().parents[4] _SETUP_PY = _REPO_ROOT / "setup.py" +_RUNTIME_SETUP_PY = _REPO_ROOT / "py/torch-tensorrt-executorch-runtime/setup.py" + + +@pytest.mark.unit +def test_runtime_implementation_is_owned_by_runtime_package(): + assert not (_REPO_ROOT / "py/torch_tensorrt/executorch/runtime.py").exists() + assert ( + _REPO_ROOT + / "py/torch-tensorrt-executorch-runtime" + / "torch_tensorrt_executorch_runtime/runtime.py" + ).is_file() + + +@pytest.mark.unit +def test_runtime_extension_has_dependency_wheel_rpaths(): + cmake = ( + _REPO_ROOT / "py/torch-tensorrt-executorch-runtime/native/CMakeLists.txt" + ).read_text(encoding="utf-8") + assert "BUILD_WITH_INSTALL_RPATH ON" in cmake + assert "$ORIGIN/../tensorrt_libs" in cmake + assert "$ORIGIN/../nvidia/cuda_runtime/lib" in cmake + assert "$ORIGIN/../nvidia/cu13/lib" in cmake def _setup_tree(): return ast.parse(_SETUP_PY.read_text(encoding="utf-8")) +def _runtime_setup_tree(): + return ast.parse(_RUNTIME_SETUP_PY.read_text(encoding="utf-8")) + + def _assignment_value(tree, name): for node in tree.body: if isinstance(node, ast.Assign) and any( @@ -86,6 +143,31 @@ def _function_def(tree, name): raise AssertionError(f"Could not find function {name}") +@pytest.mark.unit +def test_runtime_wheel_uses_public_torch_version(): + function = _function_def(_runtime_setup_tree(), "public_version") + namespace = {} + exec( + compile(ast.Module(body=[function], type_ignores=[]), "", "exec"), + namespace, + ) + + assert namespace["public_version"]("2.14.0.dev20260726+cu132") == ( + "2.14.0.dev20260726" + ) + + +@pytest.mark.unit +def test_runtime_wheel_pins_cuda_13_native_dependencies(): + setup_source = _RUNTIME_SETUP_PY.read_text(encoding="utf-8") + assert 'TENSORRT_DISTRIBUTION = "tensorrt-cu13"' in setup_source + assert 'CUDA_RUNTIME_DISTRIBUTION = "nvidia-cuda-runtime"' in setup_source + assert "torch=={public_version(torch.__version__)}" in setup_source + assert "{TENSORRT_DISTRIBUTION}=={tensorrt_version}" in setup_source + assert "{CUDA_RUNTIME_DISTRIBUTION}=={cuda_runtime_version}" in setup_source + assert "nvidia-cuda-runtime-cu12" not in setup_source + + @pytest.mark.unit def test_packaging_declares_executorch_extra(): tree = _setup_tree() @@ -101,11 +183,14 @@ def test_packaging_declares_executorch_extra(): assert extra_name in extras_by_name requirements = extras_by_name[extra_name] assert isinstance(requirements, ast.List) - assert any( - isinstance(requirement, ast.Name) - and requirement.id == "EXECUTORCH_REQUIREMENT" - for requirement in requirements.elts - ) + for requirement_name in ( + "EXECUTORCH_REQUIREMENT", + "EXECUTORCH_RUNTIME_REQUIREMENT", + ): + assert any( + isinstance(requirement, ast.Name) and requirement.id == requirement_name + for requirement in requirements.elts + ) setup_call = next( node @@ -134,19 +219,20 @@ def test_executorch_is_not_base_install_requirement(): ): function = _function_def(tree, function_name) assert not any( - isinstance(node, ast.Name) and node.id == "EXECUTORCH_REQUIREMENT" + isinstance(node, ast.Name) + and node.id + in { + "EXECUTORCH_REQUIREMENT", + "EXECUTORCH_RUNTIME_REQUIREMENT", + } for node in ast.walk(function) ) @pytest.mark.unit -def test_executorch_headers_are_not_dlfw_gated(): - tree = _setup_tree() - header_package_data = _assignment_value(tree, "executorch_header_package_data") - assert isinstance(header_package_data, ast.List) - assert not any( - isinstance(node, ast.Name) and node.id == "IS_DLFW_CI" - for node in ast.walk(header_package_data) +def test_main_wheel_does_not_package_executorch_delegate_headers(): + assert "include/torch_tensorrt/executorch/*.h" not in _SETUP_PY.read_text( + encoding="utf-8" ) diff --git a/tests/py/dynamo/executorch/test_python_runtime.py b/tests/py/dynamo/executorch/test_python_runtime.py new file mode 100644 index 0000000000..6520dd6159 --- /dev/null +++ b/tests/py/dynamo/executorch/test_python_runtime.py @@ -0,0 +1,153 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +RUNTIME_PATH = ( + Path(__file__).parents[4] + / "py/torch-tensorrt-executorch-runtime/torch_tensorrt_executorch_runtime/runtime.py" +) +DELEGATE_PATH = ( + Path(__file__).parents[4] + / "py/torch-tensorrt-executorch-runtime/torch_tensorrt_executorch_runtime/__init__.py" +) + + +def load_runtime_module(): + spec = importlib.util.spec_from_file_location( + "torchtrt_et_runtime_test", RUNTIME_PATH + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def load_delegate_module(): + spec = importlib.util.spec_from_file_location( + "torchtrt_et_delegate_test", DELEGATE_PATH + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class FakeMethod: + def execute(self, inputs): + return [inputs[0] + 1] + + +class FakeProgram: + method_names = {"forward"} + + def load_method(self, name): + return FakeMethod() if name == "forward" else None + + +class FakeRuntime: + def load_program(self, data): + self.data = data + return FakeProgram() + + +def test_load_and_forward(monkeypatch, tmp_path): + delegate = types.ModuleType("torch_tensorrt_executorch_runtime") + fake_runtime = FakeRuntime() + delegate.get_runtime = lambda: fake_runtime + monkeypatch.setitem(sys.modules, delegate.__name__, delegate) + model = tmp_path / "model.pte" + model.write_bytes(b"pte") + + program = load_runtime_module().load(model) + + assert program.forward(2) == [3] + assert program._data is fake_runtime.data + + +def test_unknown_method(monkeypatch, tmp_path): + delegate = types.ModuleType("torch_tensorrt_executorch_runtime") + delegate.get_runtime = FakeRuntime + monkeypatch.setitem(sys.modules, delegate.__name__, delegate) + model = tmp_path / "model.pte" + model.write_bytes(b"pte") + + with pytest.raises(ValueError, match="Unknown method"): + load_runtime_module().load(model).run([], "missing") + + +def test_missing_model(): + with pytest.raises(FileNotFoundError): + load_runtime_module().load("does-not-exist.pte") + + +def test_activate_twice_is_safe(monkeypatch): + delegate = load_delegate_module() + data_loader = types.ModuleType(delegate.__name__ + ".data_loader") + native = types.ModuleType(delegate.__name__ + "._portable_lib") + imported = [] + + def fake_import(name): + imported.append(name) + return { + data_loader.__name__: data_loader, + native.__name__: native, + }[name] + + monkeypatch.setattr( + delegate, "importlib", types.SimpleNamespace(import_module=fake_import) + ) + assert delegate.activate() is native + wrapper = types.ModuleType(delegate._WRAPPER_NAME) + monkeypatch.setitem(sys.modules, delegate._WRAPPER_NAME, wrapper) + assert delegate.activate() is native + assert imported == [data_loader.__name__, native.__name__] + assert sys.modules[delegate._NATIVE_NAME] is native + assert sys.modules[delegate._DATA_LOADER_NAME] is data_loader + assert sys.modules[delegate._WRAPPER_NAME] is wrapper + + +def test_activate_rejects_preloaded_stock_runtime(monkeypatch): + delegate = load_delegate_module() + stock_runtime = types.ModuleType(delegate._NATIVE_NAME) + monkeypatch.setitem(sys.modules, delegate._NATIVE_NAME, stock_runtime) + + with pytest.raises(delegate.DelegateCompatibilityError, match="stock runtime"): + delegate.activate() + + +def test_activate_rejects_preloaded_stock_wrapper(monkeypatch): + delegate = load_delegate_module() + stock_wrapper = types.ModuleType(delegate._WRAPPER_NAME) + monkeypatch.delitem(sys.modules, delegate._NATIVE_NAME, raising=False) + monkeypatch.setitem(sys.modules, delegate._WRAPPER_NAME, stock_wrapper) + + with pytest.raises( + delegate.DelegateCompatibilityError, + match=r"torch_tensorrt\.load", + ): + delegate.activate() + + +def test_activate_cleans_up_data_loader_when_native_import_fails(monkeypatch): + delegate = load_delegate_module() + data_loader = types.ModuleType(delegate.__name__ + ".data_loader") + + def fake_import(name): + if name == data_loader.__name__: + return data_loader + assert sys.modules[delegate._DATA_LOADER_NAME] is data_loader + raise ImportError("native module failed to load") + + monkeypatch.setattr( + delegate, "importlib", types.SimpleNamespace(import_module=fake_import) + ) + monkeypatch.delitem(sys.modules, delegate._NATIVE_NAME, raising=False) + monkeypatch.delitem(sys.modules, delegate._DATA_LOADER_NAME, raising=False) + + with pytest.raises(delegate.DelegateCompatibilityError): + delegate.activate() + + assert delegate._DATA_LOADER_NAME not in sys.modules diff --git a/tests/py/utils/ci_helpers.sh b/tests/py/utils/ci_helpers.sh index 466f864f7f..7f3b4e5861 100755 --- a/tests/py/utils/ci_helpers.sh +++ b/tests/py/utils/ci_helpers.sh @@ -155,12 +155,12 @@ trt_tier_l2_dynamo_compile() { trt_tier_l2_dynamo_core() { ( cd "${TRT_REPO_ROOT}/tests/py/dynamo" - _trt_py -m pytest -ra $(_trt_nproc auto) --junitxml="$(_trt_xml l2_dynamo_core_tests_results)" -k "not test_000_ and not test_001_" runtime/* "$@" - if [ "${USE_TRT_RTX:-false}" != "true" ]; then - # ExecuTorch integration is standard-TRT only. - _trt_py -m pip install pyyaml "executorch>=1.3.1" - _trt_py -m pytest -ra $(_trt_nproc auto) --junitxml="$(_trt_xml l2_dynamo_executorch_tests_results)" executorch/ "$@" - fi ) + _trt_py -m pytest -ra $(_trt_nproc auto) --junitxml="$(_trt_xml l2_dynamo_core_tests_results)" -k "not test_000_ and not test_001_" runtime/* "$@" ) +} + +trt_tier_executorch() { + ( cd "${TRT_REPO_ROOT}/tests/py/dynamo" + _trt_py -m pytest -ra $(_trt_nproc auto) --junitxml="$(_trt_xml executorch_tests_results)" executorch/ "$@" ) } trt_tier_l2_plugin() { diff --git a/toolchains/ci_workspaces/MODULE.bazel.tmpl b/toolchains/ci_workspaces/MODULE.bazel.tmpl index 46479dc114..1aa929c506 100644 --- a/toolchains/ci_workspaces/MODULE.bazel.tmpl +++ b/toolchains/ci_workspaces/MODULE.bazel.tmpl @@ -172,13 +172,12 @@ new_local_repository( build_file = "third_party/libtorch/BUILD" ) -# Pinned to the ExecuTorch release/1.3 branch head. +# Pinned to ExecuTorch main for PyTorch 2.14 native ABI compatibility. new_git_repository( name = "executorch", build_file = "@//third_party/executorch:BUILD", - # latest commit in release/1.3 branch - commit = "6118688a095fd9697224f5cad72ce42db641c9cd", - recursive_init_submodules = True, + # executorch==1.3.1 + commit = "e2f18eb23c45bd22ca332b0b8b49a81de304b472", patch_cmds = [ "find . -mindepth 2 \\( -name BUILD -o -name BUILD.bazel \\) -delete", "mkdir executorch && find . -mindepth 1 -maxdepth 1 ! -name executorch ! -name BUILD ! -name BUILD.bazel ! -name REPO.bazel -exec cp -a {} executorch/ \\;", @@ -187,6 +186,7 @@ new_git_repository( "Get-ChildItem -Recurse -Include BUILD,BUILD.bazel | Where-Object { $_.DirectoryName -ne (Get-Location).Path } | Remove-Item -Force", "New-Item -ItemType Directory executorch | Out-Null; Get-ChildItem -Force | Where-Object { $_.Name -notin @('executorch','BUILD','BUILD.bazel','REPO.bazel') } | Copy-Item -Destination executorch -Recurse -Force", ], + recursive_init_submodules = True, remote = "https://github.com/pytorch/executorch.git", )