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..c32fc55bdc 100644 --- a/.github/workflows/_linux-x86_64-core.yml +++ b/.github/workflows/_linux-x86_64-core.yml @@ -96,6 +96,7 @@ jobs: trigger-event: ${{ github.event_name }} architecture: "x86_64" use-rtx: ${{ inputs.use-rtx }} + build-executorch-delegate: ${{ !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. diff --git a/.github/workflows/build_linux.yml b/.github/workflows/build_linux.yml index 90333aedf8..109c439ab5 100644 --- a/.github/workflows/build_linux.yml +++ b/.github/workflows/build_linux.yml @@ -90,6 +90,11 @@ on: required: false default: "python -m build --wheel" type: string + build-executorch-delegate: + description: Build and publish the matching ExecuTorch delegate wheel + required: false + default: false + type: boolean pip-install-torch-extra-args: # NOTE: Why does this exist? # Well setuptools / python packaging doesn't actually allow you to specify dependencies @@ -152,6 +157,7 @@ jobs: ARCH: ${{ inputs.architecture }} BUILD_TARGET: ${{ inputs.build-target }} USE_TRT_RTX: ${{ inputs.use-rtx }} + BUILD_EXECUTORCH_DELEGATE: ${{ inputs.build-executorch-delegate }} name: build-wheel-${{ matrix.python_version }}-${{ matrix.desired_cuda }}-${{ matrix.gpu_arch_type }}-${{ inputs.architecture }}-${{ inputs.use-rtx }}-${{ inputs.is-jetpack }} 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' || ''}} @@ -313,6 +319,19 @@ jobs: ${CONDA_RUN} python setup.py bdist_wheel ${param} fi fi + - name: Build the ExecuTorch delegate wheel + if: ${{ inputs.build-executorch-delegate && inputs.use-rtx == false && inputs.is-jetpack == false && inputs.is-release-tarball == false }} + working-directory: ${{ inputs.repository }} + shell: bash -l {0} + run: | + set -euxo pipefail + source "${BUILD_ENV_FILE}" + ${CONDA_RUN} python -m pip install pyyaml executorch + 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}" + ${CONDA_RUN} python -m pip wheel --no-build-isolation --no-deps \ + --wheel-dir dist py/torch-tensorrt-executorch-delegate - name: Repair Manylinux_2_28 Wheel shell: bash -l {0} env: @@ -346,16 +365,31 @@ jobs: run: | set -euxo pipefail source "${BUILD_ENV_FILE}" - WHEEL_NAME=$(ls "${{ inputs.repository }}/dist/") - echo "$WHEEL_NAME" + mapfile -t WHEEL_PATHS < <(find "${{ inputs.repository }}/dist" -maxdepth 1 -name "*.whl" -print | sort) + if [[ ${#WHEEL_PATHS[@]} -eq 0 ]]; then + echo "No wheels found in ${{ inputs.repository }}/dist" >&2 + exit 1 + fi + printf "Wheel: %s\n" "${WHEEL_PATHS[@]}" + MAIN_WHEEL="" + for wheel_path in "${WHEEL_PATHS[@]}"; do + if [[ $(basename "${wheel_path}") == torch_tensorrt-[0-9]*.whl ]]; then + MAIN_WHEEL="${wheel_path}" + break + fi + done + if [[ -z "${MAIN_WHEEL}" ]]; then + echo "Could not identify the torch_tensorrt wheel" >&2 + exit 1 + fi if [[ ${{ inputs.architecture }} == "aarch64" ]]; then echo "Skipping smoke test for aarch64, since it is not an actual gpu runner" else - ${CONDA_RUN} pip install "${{ inputs.repository }}/dist/$WHEEL_NAME" --use-deprecated=legacy-resolver --extra-index-url https://download.pytorch.org/${CHANNEL}/${CU_VERSION} + ${CONDA_RUN} pip install "${WHEEL_PATHS[@]}" --use-deprecated=legacy-resolver --extra-index-url https://download.pytorch.org/${CHANNEL}/${CU_VERSION} # 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" + unzip -o "${MAIN_WHEEL}" # Ensure that pytorch version is pinned, should output file where it was found grep "Requires-Dist: torch (==.*)" -r . ) @@ -374,6 +408,9 @@ jobs: echo "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT} found" ${CONDA_RUN} python "${{ inputs.repository }}/${SMOKE_TEST_SCRIPT}" fi + if [[ "${BUILD_EXECUTORCH_DELEGATE}" == "true" ]]; then + ${CONDA_RUN} python -c 'from torch_tensorrt_executorch_delegate import runtime; assert runtime().backend_registry.is_available("TensorRTBackend")' + fi fi # NB: Only upload to GitHub after passing smoke tests diff --git a/.github/workflows/executorch-static-linux.yml b/.github/workflows/executorch-static-linux.yml index fe53fdb12c..b98c1a4dac 100644 --- a/.github/workflows/executorch-static-linux.yml +++ b/.github/workflows/executorch-static-linux.yml @@ -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-delegate script: | set -euo pipefail BAZELISK_VERSION="1.26.0" @@ -96,5 +97,15 @@ jobs: 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 + + # export the model to a .pte file + python examples/torchtrt_executorch_example/export_static_shape.py --model_path="${RUNNER_TEMP}/torchtrt-python.pte" + # verify the c++ reference runner can run the model + .github/scripts/verify-executorch-reference-runner.sh "${RUNNER_TEMP}/torchtrt-python.pte" + + # Build the no-compile-for-users Python runtime/delegate wheel. + python -m pip wheel --no-build-isolation --no-deps --wheel-dir dist py/torch-tensorrt-executorch-delegate + python -m pip install --no-deps --force-reinstall dist/torch_tensorrt_executorch_delegate-*.whl + python -c 'from torch_tensorrt_executorch_delegate import runtime; assert runtime().backend_registry.is_available("TensorRTBackend")' + python examples/executorch_reference_runner/load_model.py --model_path="${RUNNER_TEMP}/torchtrt-python.pte" --num_runs=1 diff --git a/.gitignore b/.gitignore index a9fb3b5ebf..662a6fc267 100644 --- a/.gitignore +++ b/.gitignore @@ -86,3 +86,5 @@ CLAUDE.md /compile_commands.json /external/ build-executorch/ +py/torch-tensorrt-executorch-delegate/build/ +py/torch-tensorrt-executorch-delegate/dist/ \ No newline at end of file diff --git a/MODULE.bazel b/MODULE.bazel index 2f384e8066..99c22600bc 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -42,12 +42,12 @@ 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. +# 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", + # PyTorch 2.14-compatible ExecuTorch main commit + commit = "a6d812a082df57898b8608f56c867140cc9da32c", 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/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..171fd7c246 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-delegate` 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..7497744743 --- /dev/null +++ b/examples/executorch_reference_runner/load_model.py @@ -0,0 +1,61 @@ +"""Load and optionally run a Torch-TensorRT ExecuTorch ``.pte`` model. + +The default input matches ``export_static_shape.py``: one CUDA float tensor +with shape ``(2, 3, 4, 4)`` filled with ones. +""" + +import argparse +from pathlib import Path + +import torch +from torch_tensorrt.executorch.runtime import load + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model_path", type=Path, default=Path("model.pte")) + parser.add_argument("--method", default="forward") + parser.add_argument("--num_runs", type=int, default=1) + parser.add_argument( + "--load_only", + action="store_true", + help="Parse the program and print its methods without loading/executing one", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if not args.model_path.is_file(): + raise FileNotFoundError(f"ExecuTorch model not found: {args.model_path}") + if args.num_runs < 1: + raise ValueError("--num_runs must be at least 1") + + program = load(args.model_path) + print(f"Loaded {args.model_path}") + print(f"Program methods: {sorted(program.method_names)}") + + if args.load_only: + return + if args.method not in program.method_names: + raise ValueError( + f"Method {args.method!r} is not in the program; " + f"available methods: {sorted(program.method_names)}" + ) + inputs = (torch.ones((2, 3, 4, 4), dtype=torch.float32),) + for run in range(args.num_runs): + outputs = program.run(inputs, args.method) + print(f"Run {run + 1} outputs:") + for index, output in enumerate(outputs): + if isinstance(output, torch.Tensor): + print( + f" output[{index}]: shape={tuple(output.shape)}, " + f"dtype={output.dtype}, device={output.device}, " + f"sample={output.flatten()[:8]}" + ) + else: + print(f" output[{index}]: {output!r}") + + +if __name__ == "__main__": + main() diff --git a/py/torch-tensorrt-executorch-delegate/README.md b/py/torch-tensorrt-executorch-delegate/README.md new file mode 100644 index 0000000000..c2f93dcb19 --- /dev/null +++ b/py/torch-tensorrt-executorch-delegate/README.md @@ -0,0 +1,44 @@ +# Torch-TensorRT ExecuTorch Delegate Wheel + +This directory builds `torch-tensorrt-executorch-delegate`. 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. + +## 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 +executorch_cmake_location="$(bazel query \ + @executorch//:executorch/CMakeLists.txt --output=location)" +export EXECUTORCH_SOURCE_DIR="$(dirname "${executorch_cmake_location%%:*}")" +export TensorRT_ROOT=/path/to/TensorRT + +python -m pip install executorch==1.3.1 +python -m pip wheel --no-build-isolation --no-deps \ + --wheel-dir dist py/torch-tensorrt-executorch-delegate +``` + +The static ExecuTorch and delegate archives are intermediate build inputs; +users receive the final native Python module and do not compile anything. + +## Use + +```bash +pip install "torch-tensorrt[executorch]" +``` + +```python +import torch +from torch_tensorrt.executorch.runtime import load + +program = load("model.pte") +outputs = program.forward(torch.ones((2, 3, 4, 4))) +``` diff --git a/py/torch-tensorrt-executorch-delegate/native/CMakeLists.txt b/py/torch-tensorrt-executorch-delegate/native/CMakeLists.txt new file mode 100644 index 0000000000..1b61d3f767 --- /dev/null +++ b/py/torch-tensorrt-executorch-delegate/native/CMakeLists.txt @@ -0,0 +1,68 @@ +cmake_minimum_required(VERSION 3.24) +project(torch_tensorrt_executorch_delegate LANGUAGES CXX) + +if(NOT EXECUTORCH_SOURCE_DIR OR NOT TORCH_TENSORRT_SOURCE_DIR) + message(FATAL_ERROR "EXECUTORCH_SOURCE_DIR and TORCH_TENSORRT_SOURCE_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() + +# ExecuTorch portable c10 delegates selected headers to torch/headeronly. The +# pybind bridge must resolve those delegated headers from the exact PyTorch +# wheel it links, not ExecuTorch's vendored snapshot. +file(MAKE_DIRECTORY "${CMAKE_BINARY_DIR}/torch_header_shim/torch") +file(CREATE_LINK + "${TORCH_PRIMARY_INCLUDE_DIR}/torch/headeronly" + "${CMAKE_BINARY_DIR}/torch_header_shim/torch/headeronly" + SYMBOLIC COPY_ON_ERROR) +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) +set_target_properties(portable_lib PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${DELEGATE_LIBRARY_OUTPUT_DIRECTORY}" + OUTPUT_NAME "_portable_lib") +set_target_properties(data_loader PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${DELEGATE_LIBRARY_OUTPUT_DIRECTORY}") + +add_custom_target(torch_tensorrt_executorch_portable_lib DEPENDS portable_lib data_loader) diff --git a/py/torch-tensorrt-executorch-delegate/pyproject.toml b/py/torch-tensorrt-executorch-delegate/pyproject.toml new file mode 100644 index 0000000000..90cff97b26 --- /dev/null +++ b/py/torch-tensorrt-executorch-delegate/pyproject.toml @@ -0,0 +1,9 @@ +[build-system] +requires = [ + "setuptools>=68", + "wheel>=0.40", + # Intentionally supplied by the active Torch-TensorRT build environment. + # Builds must use --no-build-isolation; see README.md. + "torch", +] +build-backend = "setuptools.build_meta" diff --git a/py/torch-tensorrt-executorch-delegate/setup.py b/py/torch-tensorrt-executorch-delegate/setup.py new file mode 100644 index 0000000000..83a645a721 --- /dev/null +++ b/py/torch-tensorrt-executorch-delegate/setup.py @@ -0,0 +1,109 @@ +"""Build the precompiled Torch-TensorRT backend for ExecuTorch Python.""" + +from __future__ import annotations + +import importlib.metadata +import os +import pathlib +import re +import subprocess +import sys + +import torch +from torch.utils.cpp_extension import include_paths +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] + +def torchtrt_version() -> str: + if value := os.getenv("TORCH_TENSORRT_EXECUTORCH_DELEGATE_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() + + +class CMakeExtension(Extension): + def __init__(self, name: str) -> None: + super().__init__(name, sources=[]) + + +class CMakeBuild(build_ext): + def build_extension(self, ext: Extension) -> None: + if sys.platform != "linux": + raise RuntimeError("The ExecuTorch TensorRT delegate supports Linux only") + executorch_source = os.getenv("EXECUTORCH_SOURCE_DIR") + if not executorch_source: + raise RuntimeError( + "Set EXECUTORCH_SOURCE_DIR to the pinned ExecuTorch source tree" + ) + + output = pathlib.Path(self.get_ext_fullpath(ext.name)).resolve() + if output.exists(): + return + build_dir = pathlib.Path(self.build_temp) / "executorch_delegate_native" + output.parent.mkdir(parents=True, exist_ok=True) + build_dir.mkdir(parents=True, exist_ok=True) + configure = [ + "cmake", + "-S", + str(HERE / "native"), + "-B", + str(build_dir), + f"-DEXECUTORCH_SOURCE_DIR={pathlib.Path(executorch_source).resolve()}", + f"-DTORCH_TENSORRT_SOURCE_DIR={REPO_ROOT}", + f"-DPYTHON_EXECUTABLE={sys.executable}", + f"-DTORCH_PRIMARY_INCLUDE_DIR={include_paths()[0]}", + f"-DDELEGATE_LIBRARY_OUTPUT_DIRECTORY={output.parent}", + f"-DCMAKE_BUILD_TYPE={'Debug' if self.debug else 'Release'}", + ] + if value := os.getenv("TensorRT_ROOT"): + configure.append(f"-DTensorRT_ROOT={value}") + configure.extend(os.getenv("CMAKE_ARGS", "").split()) + subprocess.run(configure, check=True) + subprocess.run( + [ + "cmake", + "--build", + str(build_dir), + "--target", + "torch_tensorrt_executorch_portable_lib", + "--parallel", + str(self.parallel or os.cpu_count() or 1), + ], + check=True, + ) + library_stem = ( + "_portable_lib" if ext.name.endswith("._portable_lib") else "data_loader" + ) + built = next(output.parent.glob(f"{library_stem}*.so"), None) + if built is None: + raise RuntimeError( + f"CMake did not produce {library_stem} in {output.parent}" + ) + if built != output: + built.replace(output) + + +executorch_version = importlib.metadata.version("executorch") +setup( + name="torch-tensorrt-executorch-delegate", + version=torchtrt_version(), + description="Torch-TensorRT delegate for the ExecuTorch Python runtime", + packages=find_packages(), + ext_modules=[ + CMakeExtension("torch_tensorrt_executorch_delegate._portable_lib"), + CMakeExtension("torch_tensorrt_executorch_delegate.data_loader"), + ], + cmdclass={"build_ext": CMakeBuild}, + python_requires=">=3.10", + install_requires=[ + f"torch=={torch.__version__}", + f"executorch=={executorch_version}", + ], + zip_safe=False, +) diff --git a/py/torch-tensorrt-executorch-delegate/torch_tensorrt_executorch_delegate/__init__.py b/py/torch-tensorrt-executorch-delegate/torch_tensorrt_executorch_delegate/__init__.py new file mode 100644 index 0000000000..343873d789 --- /dev/null +++ b/py/torch-tensorrt-executorch-delegate/torch_tensorrt_executorch_delegate/__init__.py @@ -0,0 +1,62 @@ +"""Activate the TensorRT delegate-enabled ExecuTorch Python runtime.""" + +from __future__ import annotations + +import importlib +import sys +from types import ModuleType + +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 DelegateCompatibilityError(ImportError): + """The delegate 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.executorch.load(...) before importing " + "executorch.runtime." + ) + missing = object() + previous_data_loader = sys.modules.get(_DATA_LOADER_NAME, missing) + try: + data_loader = importlib.import_module(__name__ + ".data_loader") + native = importlib.import_module(__name__ + "._portable_lib") + except ImportError as error: + if previous_data_loader is missing: + 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 delegate from " + "the same release matrix." + ) from error + sys.modules[_DATA_LOADER_NAME] = data_loader + sys.modules[_NATIVE_NAME] = native + sys.modules.pop(_WRAPPER_NAME, None) + return native + + +def runtime(): + """Return the activated ExecuTorch Runtime singleton.""" + activate() + from executorch.runtime import Runtime + + value = 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", "runtime"] diff --git a/py/torch_tensorrt/executorch/__init__.py b/py/torch_tensorrt/executorch/__init__.py index 123eee846d..a068329e8b 100644 --- a/py/torch_tensorrt/executorch/__init__.py +++ b/py/torch_tensorrt/executorch/__init__.py @@ -29,6 +29,7 @@ def __getattr__(name: str) -> NoReturn: else: from torch_tensorrt.executorch.backend import TensorRTBackend from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + from torch_tensorrt.executorch.runtime import Program, load def get_edge_compile_config() -> "EdgeCompileConfig": """Return the EdgeCompileConfig used for Torch-TensorRT ExecuTorch export.""" @@ -40,4 +41,6 @@ def get_edge_compile_config() -> "EdgeCompileConfig": "get_edge_compile_config", "TensorRTPartitioner", "TensorRTBackend", + "Program", + "load", ] diff --git a/py/torch_tensorrt/executorch/runtime.py b/py/torch_tensorrt/executorch/runtime.py new file mode 100644 index 0000000000..fb78b54d77 --- /dev/null +++ b/py/torch_tensorrt/executorch/runtime.py @@ -0,0 +1,61 @@ +"""Python inference API for Torch-TensorRT ExecuTorch programs.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Sequence, Union + + +def _runtime(): + try: + from torch_tensorrt_executorch_delegate import 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 runtime() + + +class Program: + """A loaded ExecuTorch program backed by TensorRTBackend.""" + + 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): + return self._program.method_names + + def run(self, inputs: Sequence[Any], method: str = "forward") -> Sequence[Any]: + 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 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.""" + 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(_runtime().load_program(data), data) + + +__all__ = ["Program", "load"] diff --git a/setup.py b/setup.py index 91651c7e15..1563d41c33 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,16 @@ def load_dep_info(): else: __version__ = f"{get_base_version()}.dev0+{get_git_revision_short_hash()}" +EXECUTORCH_REQUIREMENT = "executorch>=1.3.1" +EXECUTORCH_DELEGATE_REQUIREMENT = ( + f"torch-tensorrt-executorch-delegate=={__version__}; " + "platform_system == 'Linux' and platform_machine == 'x86_64'" +) +EXTRAS_REQUIRE = { + "executorch": [EXECUTORCH_REQUIREMENT, EXECUTORCH_DELEGATE_REQUIREMENT], + "all": [EXECUTORCH_REQUIREMENT, EXECUTORCH_DELEGATE_REQUIREMENT], +} + if "--ci" in sys.argv: sys.argv.remove("--ci") if RELEASE: 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..648607fa54 --- /dev/null +++ b/tests/py/dynamo/executorch/test_python_runtime.py @@ -0,0 +1,150 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +RUNTIME_PATH = Path(__file__).parents[4] / "py/torch_tensorrt/executorch/runtime.py" +DELEGATE_PATH = ( + Path(__file__).parents[4] + / "py/torch-tensorrt-executorch-delegate/torch_tensorrt_executorch_delegate/__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_delegate") + fake_runtime = FakeRuntime() + delegate.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_delegate") + delegate.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"Call torch_tensorrt\.executorch\.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 delegate._DATA_LOADER_NAME not in sys.modules + 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/toolchains/ci_workspaces/MODULE.bazel.tmpl b/toolchains/ci_workspaces/MODULE.bazel.tmpl index 46479dc114..9c22f41291 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, + # PyTorch 2.14-compatible ExecuTorch main commit + commit = "a6d812a082df57898b8608f56c867140cc9da32c", 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", )