ci(executorch): run a coalesced TensorRT + CUDA program in the reference runner gate

.github/scripts/verify-executorch-reference-runner.sh

@@ -18,6 +18,11 @@ set +x
 #   examples/torchtrt_executorch_example/export_kv_cache_decode.py). When given,
 #   kv_cache_decode_check is built and run against it as well.
 #
+# Optional third argument: path to a coalesced TensorRT + CUDA .pte (see
+#   examples/torchtrt_executorch_example/export_coalesced.py). When given, both
+#   runners are run against it and their output is compared to the eager
+#   reference that export script wrote next to the model.
+#
 # Optional:
 #   TensorRT_ROOT=/path/to/extracted/TensorRT
 #     If unset, the script reuses Bazel's fetched TensorRT SDK when available
@@ -33,8 +38,8 @@ set +x
 repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
 cd "${repo_root}"
 
-if [[ $# -lt 1 || $# -gt 2 ]]; then
-  echo "Usage: $0 PATH_TO_MODEL.pte [PATH_TO_KV_CACHE_DECODE.pte]" >&2
+if [[ $# -lt 1 || $# -gt 3 ]]; then
+  echo "Usage: $0 PATH_TO_MODEL.pte [PATH_TO_KV_CACHE_DECODE.pte [PATH_TO_COALESCED.pte]]" >&2
   exit 1
 fi
 model_path="$1"
@@ -47,6 +52,11 @@ if [[ -n "${kv_model_path}" && ! -f "${kv_model_path}" ]]; then
   echo "KV-cache decode model not found: ${kv_model_path}" >&2
   exit 1
 fi
+coalesced_model_path="${3:-}"
+if [[ -n "${coalesced_model_path}" && ! -f "${coalesced_model_path}" ]]; then
+  echo "Coalesced model not found: ${coalesced_model_path}" >&2
+  exit 1
+fi
 
 python_executable="${PYTHON_EXECUTABLE:-}"
 if [[ -z "${python_executable}" ]]; then
@@ -467,25 +477,35 @@ packaged_runner_log="${verify_root}/packaged_runner.log"
   --model_path="${model_path}" \
   --num_runs=1 2>&1 | tee "${packaged_runner_log}"
 
-# The sample model is x + 1 on a (2,3,4,4) input and both runners fill inputs with
-# 1.0f, so the shape is exactly [2,3,4,4] and every printed value is exactly 2.0000.
-# Assert both precisely. Matching only "shape=" accepts any shape, and matching one
-# 2.0000 anywhere on the values line accepts a line of wrong numbers that happens to
-# contain one right one, so neither catches a stream-ordering regression returning
-# stale or partial output. Both lines come from fprintf in the runner, so these
-# assertions hold whatever ET_LOG_ENABLED is set to.
-for _log in "${runner_log}" "${packaged_runner_log}"; do
-  # A right answer produced entirely on the host would not prove much here: the
-  # program is delegated to TensorRT, so at least one planned buffer has to be
-  # served by a registered CUDA DeviceAllocator. Pin that, otherwise a model or
-  # a planning change could quietly turn this into a CPU-only test.
+# Assert the printed shape, and that EVERY value on the "first N values:" line is
+# the expected one. Matching only "shape=" accepts any shape, and matching one
+# right value anywhere on the values line accepts a line of wrong numbers that
+# happens to contain one, so neither on its own catches a stream-ordering
+# regression returning stale or partial output. Both lines come from fprintf in
+# the runner, so these assertions hold whatever ET_LOG_ENABLED is set to. The
+# models used here are elementwise on an all-ones input, so one number describes
+# the whole expected output.
+assert_runner_output() {
+  local _log="$1"
+  local _shape="$2"
+  local _expected="$3"
+  local _tolerance="$4"
+  local _values
+  local _value
+
+  # A right answer produced entirely on the host would not prove much: the
+  # programs here are delegated, so at least one planned buffer has to be served
+  # by a registered CUDA DeviceAllocator. Pin that, otherwise a model or a
+  # planning change could quietly turn this into a CPU-only test.
   if ! grep -q 'planned buffer\[[0-9]*\] = [0-9]* bytes on device_type 1' "${_log}"; then
     echo "No CUDA planned buffer was allocated in ${_log}:" >&2
     grep 'planned buffer' "${_log}" >&2 || echo "  no planned buffer line at all" >&2
     exit 1
   fi
 
-  if ! grep -q 'output\[0\] shape=\[2,3,4,4\]' "${_log}"; then
+  # -F: the shape is bracketed, and an unescaped [2,3,4,4] is a regex character
+  # class that would match any single one of those characters.
+  if ! grep -qF "output[0] shape=${_shape}" "${_log}"; then
     echo "Unexpected output shape in ${_log}:" >&2
     grep 'output\[0\] shape=' "${_log}" >&2 || echo "  no shape line at all" >&2
     exit 1
@@ -498,11 +518,19 @@ for _log in "${runner_log}" "${packaged_runner_log}"; do
     exit 1
   fi
   for _value in ${_values}; do
-    if [[ "${_value}" != "2.0000" ]]; then
-      echo "Unexpected output value '${_value}' in ${_log}: ${_values}" >&2
+    if ! awk -v got="${_value}" -v want="${_expected}" -v tol="${_tolerance}" \
+        'BEGIN { d = got - want; if (d < 0) d = -d; exit !(d <= tol) }'; then
+      echo "Unexpected output value '${_value}' in ${_log}" \
+        "(expected ${_expected} within ${_tolerance}): ${_values}" >&2
       exit 1
     fi
   done
+}
+
+# The sample model is x + 1 on a (2,3,4,4) input, so every output value is exactly
+# 2. That is exact in float32, hence a zero tolerance.
+for _log in "${runner_log}" "${packaged_runner_log}"; do
+  assert_runner_output "${_log}" "[2,3,4,4]" "2.0000" 0
 done
 
 if [[ -n "${kv_model_path}" ]]; then
@@ -521,3 +549,37 @@ if [[ -n "${kv_model_path}" ]]; then
   "${kv_check_path}" --model_path="${kv_model_path}" 2>&1 | tee "${kv_check_log}"
   grep -q "PASS: decode at pos=1 observed the KV written at pos=0" "${kv_check_log}"
 fi
+
+if [[ -n "${coalesced_model_path}" ]]; then
+  # A coalesced program splits one graph across the TensorRT delegate and
+  # ExecuTorch's CUDA delegate, so a value produced by one delegate is consumed by
+  # the other on the device, inside one method. The checks above use a program with
+  # a single delegate, so they exercise neither the second backend nor the handover
+  # between the two.
+  coalesced_expected_path="${coalesced_model_path%.pte}.expected"
+  if [[ ! -f "${coalesced_expected_path}" ]]; then
+    echo "Coalesced reference output not found: ${coalesced_expected_path}" >&2
+    echo "It is written by examples/torchtrt_executorch_example/export_coalesced.py" >&2
+    exit 1
+  fi
+  coalesced_shape="$(sed -n '1p' "${coalesced_expected_path}")"
+  coalesced_value="$(sed -n '2p' "${coalesced_expected_path}")"
+  if [[ -z "${coalesced_shape}" || -z "${coalesced_value}" ]]; then
+    echo "Malformed coalesced reference output in ${coalesced_expected_path}" >&2
+    exit 1
+  fi
+
+  # Only the from-source runner runs the coalesced program. It links every
+  # ExecuTorch delegate, including the CUDA/AOTI backend that the CUDA partition
+  # of a coalesced .pte is handed to. The packaged runner in the wheel ships the
+  # TensorRT delegate alone, so it has no CudaBackend to run that partition and
+  # cannot execute this model.
+  coalesced_runner_log="${verify_root}/coalesced_my_runner.log"
+  "${runner_path}" \
+    --model_path="${coalesced_model_path}" \
+    --num_runs=1 2>&1 | tee "${coalesced_runner_log}"
+
+  # TensorRT, AOTInductor and eager PyTorch compute the same math with different
+  # kernels, so compare within a tolerance instead of on the printed digits.
+  assert_runner_output "${coalesced_runner_log}" "${coalesced_shape}" "${coalesced_value}" 0.001
+fi

.github/workflows/executorch-test-linux.yml

@@ -95,8 +95,11 @@ jobs:
           --model_path="${RUNNER_TEMP}/torchtrt-python.pte"
         python examples/torchtrt_executorch_example/export_kv_cache_decode.py \
           --model_path="${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte"
+        python examples/torchtrt_executorch_example/export_coalesced.py \
+          --model_path="${RUNNER_TEMP}/torchtrt-coalesced.pte"
         .github/scripts/verify-executorch-reference-runner.sh \
           "${RUNNER_TEMP}/torchtrt-python.pte" \
-          "${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte"
+          "${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" \
+          "${RUNNER_TEMP}/torchtrt-coalesced.pte"
         python examples/executorch_reference_runner/load_model.py \
           --model_path="${RUNNER_TEMP}/torchtrt-python.pte" --num_runs=1

examples/torchtrt_executorch_example/export_coalesced.py

@@ -0,0 +1,120 @@
+"""
+.. _executorch_export_coalesced:
+
+Exporting a Coalesced TensorRT + CUDA Model to ExecuTorch (.pte)
+================================================================
+
+This example exports one graph split across two ExecuTorch backends. TensorRT
+takes the operators it can convert, and everything left over goes to
+ExecuTorch's own CUDA backend, which compiles it with AOTInductor. Both
+delegates end up inside a single ``.pte`` and run in the same method.
+
+The model is ``cos(erfinv(tanh(x)))``. TensorRT has no converter for
+``erfinv``, so that operator is the one the CUDA backend has to claim. That
+makes the split real rather than incidental.
+
+A value produced by one delegate is consumed by the other on the device, inside
+a single method, so this is what proves the two backends can complete each
+other rather than only work on their own.
+
+The CUDA backend also writes an ``aoti_cuda_blob.ptd`` next to the ``.pte`` for
+its external weights. This model has no weights, so that file is empty of
+tensors and the reference runner does not need it.
+
+Besides the ``.pte`` this writes ``<model_path>.expected``, holding the output
+shape and the eager reference value for an all-ones input. The reference runner
+gate reads that file instead of hard-coding a number, so the expected value
+cannot drift away from the model.
+
+Prerequisites
+-------------
+Install Torch-TensorRT with the ExecuTorch extra before running this example::
+
+    pip install -e ".[executorch]"
+
+ExecuTorch's CUDA backend also needs a CUDA toolkit (``nvcc``) at export time,
+for the AOTInductor compile.
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+import torch_tensorrt
+from executorch.backends.cuda.cuda_backend import CudaBackend
+from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+from executorch.exir._serialize._program import deserialize_pte_binary
+
+SHAPE = (64, 64)
+
+
+class CoalescedModel(torch.nn.Module):
+    def forward(self, x):
+        return torch.cos(torch.erfinv(torch.tanh(x)))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_path", default="coalesced.pte", help="Path to save the .pte"
+    )
+    args = parser.parse_args()
+    model_path = Path(args.model_path)
+
+    with torch.no_grad():
+        model = CoalescedModel().eval().cuda()
+        example_input = (torch.randn(SHAPE).cuda(),)
+
+        exported_program = torch.export.export(model, example_input)
+        trt_gm = torch_tensorrt.dynamo.compile(
+            exported_program,
+            arg_inputs=example_input,
+            min_block_size=1,
+            truncate_double=True,
+        )
+        torch_tensorrt.save(
+            trt_gm,
+            str(model_path),
+            output_format="executorch",
+            arg_inputs=example_input,
+            retrace=False,
+            # Catch-all, so every operator TensorRT rejected goes to the CUDA
+            # backend instead of falling back to a portable CPU kernel.
+            partitioners=[
+                CudaPartitioner(
+                    [CudaBackend.generate_method_name_compile_spec("forward")]
+                )
+            ],
+        )
+
+        # Both delegates must really be in the file. A partitioner change that
+        # quietly routed the whole graph to TensorRT would otherwise leave a
+        # green job that no longer tests the coalesced path at all.
+        program = deserialize_pte_binary(model_path.read_bytes()).program
+        delegates = [d.id for plan in program.execution_plan for d in plan.delegates]
+        missing = [
+            name for name in ("TensorRTBackend", "CudaBackend") if name not in delegates
+        ]
+        if missing:
+            sys.exit(
+                f"{model_path} is not coalesced: missing {missing}, found {delegates}"
+            )
+
+        # The reference runners fill every input element with 1.0, and this model
+        # is elementwise, so a single number describes the whole expected output.
+        reference = model(torch.ones(SHAPE).cuda())
+        expected_path = model_path.with_suffix(".expected")
+        expected_path.write_text(
+            "[{}]\n{:.4f}\n".format(
+                ",".join(str(dim) for dim in reference.shape),
+                reference.flatten()[0].item(),
+            )
+        )
+
+    print(f"Saved {model_path} with delegates {delegates}.")
+    print(f"Saved {expected_path} with the eager reference output.")
+
+
+if __name__ == "__main__":
+    main()

tests/py/dynamo/executorch/test_cuda_partitioner_composition.py

@@ -89,11 +89,10 @@ def _delegate_ids(pte_path):
 
 # NOTE: these tests are COMPOSITION-ONLY. They assert the serialized .pte carries
 # both a TensorRTBackend and a CudaBackend delegate (and, below, that external
-# CUDA weights are persisted as a .ptd). They do NOT load or run the program: a
-# coalesced ATen-mode .pte cannot be loaded yet because memory-planned CUDA
-# buffers get a CPU data pointer, so Method::init fails the CUDA backend's device
-# check (tensor_parser_aten hardcodes CPU). A load-run-allclose test should be
-# added once that runtime fix lands (follow-up).
+# CUDA weights are persisted as a .ptd). They do NOT load or run the program.
+# Execution of a coalesced program is covered by the reference runner gate, see
+# examples/torchtrt_executorch_example/export_coalesced.py and the third argument
+# of .github/scripts/verify-executorch-reference-runner.sh.
 
 
 def test_erfinv_routes_to_cuda_backend(tmp_path):