add torch-tensorrt-executorch-delegate wheel build

.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

.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.

.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

.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

.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/

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/ \\;",

examples/executorch_reference_runner/BUILD

@@ -7,6 +7,7 @@ filegroup(
     srcs = [
         "CMakeLists.txt",
         "README.md",
+        "load_model.py",
         "main.cpp",
     ],
 )

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

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()

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)))
+```