fix: make the packaged ExecuTorch reference runner run TensorRT delegated models (#4573)

Author: shoumikhin

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

@@ -441,6 +441,24 @@ if [[ -n "${extra_defs}" ]]; then
   exit 1
 fi
 
+# ET_CHECK_MSG and ET_LOG hand their text to the core archive's vlogf, so an
+# archive compiled with logging off turns every runtime failure in the packaged
+# runner into a bare exit code with no output at all. Prove the runner can still
+# say why it failed, on a path that needs neither a GPU nor a valid model.
+diagnostic_log="${verify_root}/packaged_runner_diagnostic.log"
+if "${packaged_runner}" \
+  --model_path="${verify_root}/no-such-model.pte" >"${diagnostic_log}" 2>&1; then
+  echo "Packaged runner exited 0 on a missing model file" >&2
+  exit 1
+fi
+if ! grep -q "FileDataLoader::from" "${diagnostic_log}"; then
+  echo "Packaged runner gave no diagnostic for a missing model file." >&2
+  echo "ET_LOG_ENABLED likely disagrees between libexecutorch_core.a and the" >&2
+  echo "Bazel targets that call into it. Output was:" >&2
+  cat "${diagnostic_log}" >&2
+  exit 1
+fi
+
 "${runner_path}" \
   --model_path="${model_path}" \
   --num_runs=1 2>&1 | tee "${runner_log}"
@@ -454,9 +472,19 @@ packaged_runner_log="${verify_root}/packaged_runner.log"
 # 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. ET_LOG output is not part of the packaged runner contract
-# and may be compiled out, so nothing here depends on it.
+# 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.
+  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
     echo "Unexpected output shape in ${_log}:" >&2
     grep 'output\[0\] shape=' "${_log}" >&2 || echo "  no shape line at all" >&2

cpp/BUILD

@@ -141,6 +141,37 @@ cc_library(
     ],
 )
 
+# Registration only, kept out of :tensorrt_executorch_backend so that linking
+# the backend into both a shared library and an executable cannot register the
+# allocator twice. Executables that run a TensorRT delegated program depend on
+# this directly.
+#
+# Not part of :executorch_backend_source_files. The CMake build of the reference
+# runner enables the CUDA/AOTI delegate, which already registers this allocator.
+cc_library(
+    name = "tensorrt_executorch_cuda_device_allocator",
+    srcs = [
+        "src/torch_tensorrt/executorch/RegisterCudaDeviceAllocator.cpp",
+    ],
+    alwayslink = True,
+    target_compatible_with = select({
+        ":linux_x86_64": [],
+        ":sbsa": [],
+        "//conditions:default": ["@platforms//:incompatible"],
+    }),
+    deps = select({
+        ":linux_x86_64": [
+            "@executorch//:executorch_cuda_allocator",
+            "@executorch//:executorch_headers",
+        ],
+        ":sbsa": [
+            "@executorch//:executorch_cuda_allocator",
+            "@executorch//:executorch_headers",
+        ],
+        "//conditions:default": [],
+    }),
+)
+
 cc_library(
     name = "tensorrt_executorch_backend",
     srcs = [

cpp/src/torch_tensorrt/executorch/RegisterCudaDeviceAllocator.cpp

@@ -0,0 +1,20 @@
+#include <executorch/backends/cuda/runtime/cuda_allocator.h>
+#include <executorch/runtime/core/device_allocator.h>
+
+namespace {
+
+// A program delegated to TensorRT still asks the ExecuTorch runtime for its
+// device planned buffers, and the runtime can only serve those once something
+// has registered a CUDA DeviceAllocator. The pinned ExecuTorch release does
+// that from inside the CUDA/AOTI delegate, which an application that delegates
+// to TensorRT alone has no reason to link. Register it here instead.
+//
+// Deliberately unguarded. If the CUDA/AOTI delegate is ever linked into the
+// same binary it registers the same singleton a second time, and the registry
+// is meant to abort on that rather than silently pick one.
+[[maybe_unused]] const bool cuda_device_allocator_registered = [] {
+  executorch::runtime::register_device_allocator(&executorch::backends::cuda::CudaAllocator::instance());
+  return true;
+}();
+
+} // namespace

examples/executorch_reference_runner/BUILD

@@ -25,6 +25,7 @@ cc_binary(
     ],
     deps = [
         "//cpp:tensorrt_executorch_backend",
+        "//cpp:tensorrt_executorch_cuda_device_allocator",
         "@executorch//:executorch_core",
         "@executorch//:executorch_file_data_loader",
         "@executorch//:extension_cuda",
@@ -36,6 +37,7 @@ cc_binary(
     srcs = ["kv_cache_decode_check.cpp"],
     deps = [
         "//cpp:tensorrt_executorch_backend",
+        "//cpp:tensorrt_executorch_cuda_device_allocator",
         "@cuda//:cudart",
         "@executorch//:executorch_core",
         "@executorch//:executorch_file_data_loader",

third_party/executorch/BUILD

@@ -19,6 +19,14 @@ filegroup(
     ),
 )
 
+# EXECUTORCH_ENABLE_LOGGING is pinned rather than left to follow
+# CMAKE_BUILD_TYPE. Its default is off in Release, which compiles the archive
+# with ET_LOG_ENABLED=0, but that is a CMake compile definition and
+# rules_foreign_cc harvests only the .a, so the Bazel targets that call into it
+# keep runtime/platform/log.h's default of 1. logf is an inline function whose
+# body is guarded by that macro, so the two halves disagree on its definition
+# and every ET_LOG and ET_CHECK_MSG message reaches an empty vlogf. A runner
+# built that way dies without saying why.
 cmake(
     name = "executorch_core",
     cache_entries = {
@@ -31,6 +39,7 @@ cmake(
         "EXECUTORCH_BUILD_EXTENSION_MODULE": "ON",
         "EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP": "ON",
         "EXECUTORCH_BUILD_PYBIND": "OFF",
+        "EXECUTORCH_ENABLE_LOGGING": "ON",
     },
     install = False,
     lib_source = ":executorch_sources",
@@ -116,6 +125,28 @@ cc_library(
     ],
 )
 
+# The CUDA DeviceAllocator that serves device planned buffers. In the pinned
+# ExecuTorch release this source is compiled only into the CUDA/AOTI delegate,
+# so a binary that delegates to TensorRT alone cannot get one without linking a
+# delegate it never calls. Compile it on its own instead. It has no delegate
+# dependencies: cuda_runtime.h, the caller-stream header already in
+# libextension_cuda.so, and the ExecuTorch runtime headers.
+#
+# This target names a path that goes away when the allocator stops living under
+# backends/cuda upstream, so the pin cannot move forward while it is still here.
+cc_library(
+    name = "executorch_cuda_allocator",
+    srcs = ["executorch/backends/cuda/runtime/cuda_allocator.cpp"],
+    hdrs = ["executorch/backends/cuda/runtime/cuda_allocator.h"],
+    include_prefix = "executorch",
+    strip_include_prefix = "executorch",
+    deps = [
+        ":executorch_headers",
+        ":extension_cuda",
+        "@cuda//:cudart",
+    ],
+)
+
 cc_library(
     name = "executorch_headers",
     hdrs = glob(