refactor(executorch): use shared caller stream

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

@@ -153,7 +153,18 @@ download_tensorrt_root() {
   if [[ ! -f "${tensorrt_archive}" ]]; then
     curl -fL "${tensorrt_url}" -o "${tensorrt_archive}" || return 1
   fi
-  tar -xzf "${tensorrt_archive}" -C "${tensorrt_extract_dir}" || return 1
+  case "${tensorrt_archive}" in
+    *.tar.zst)
+      tar --zstd -xf "${tensorrt_archive}" -C "${tensorrt_extract_dir}" || return 1
+      ;;
+    *.tar.gz | *.tgz)
+      tar -xzf "${tensorrt_archive}" -C "${tensorrt_extract_dir}" || return 1
+      ;;
+    *)
+      echo "Unsupported TensorRT archive format: ${tensorrt_archive}" >&2
+      return 1
+      ;;
+  esac
 
   if [[ -n "${tensorrt_strip_prefix}" ]]; then
     tensorrt_root="${tensorrt_extract_dir}/${tensorrt_strip_prefix}"
@@ -290,6 +301,8 @@ require_tar_entry() {
 
 require_tar_entry "torch_tensorrt/src/torch_tensorrt/executorch/CMakeLists.txt"
 require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/CMakeLists.txt"
+require_tar_entry "torch_tensorrt/bin/example_executorch_runner"
+require_tar_entry "torch_tensorrt/lib/libextension_cuda.so"
 require_tar_entry "torch_tensorrt/BUILD"
 
 export TORCH_TENSORRT_ROOT="${verify_root}/torch_tensorrt"
@@ -317,19 +330,135 @@ cmake --build "${verify_root}/build-executorch-reference-runner" \
 
 runner_log="${verify_root}/my_runner.log"
 runner_path="${verify_root}/build-executorch-reference-runner/example_executorch_runner"
-if command -v ldd >/dev/null 2>&1 &&
-  ldd "${runner_path}" |
+
+# Symbol/linkage inspection tools are mandatory on this Linux gate: silently
+# skipping them would let a broken single-TLS layout pass unnoticed.
+for _tool in ldd readelf nm; do
+  if ! command -v "${_tool}" >/dev/null 2>&1; then
+    echo "Required tool '${_tool}' not found; cannot verify caller-stream linkage" >&2
+    exit 1
+  fi
+done
+
+# The runner must not pull in libtorch: this native path is libtorch-free.
+if ldd "${runner_path}" |
     grep -E "libtorch|libtorch_cpu|libtorch_cuda|libc10" >&2; then
   echo "example_executorch_runner links PyTorch/libtorch shared libraries" >&2
   exit 1
 fi
 
+# The runner must declare a real DT_NEEDED dependency on libextension_cuda.so
+# (an ldd filename match alone would also accept a "=> not found" line).
+if ! readelf -d "${runner_path}" |
+    grep -E "\(NEEDED\).*libextension_cuda\.so" >&2; then
+  echo "example_executorch_runner has no DT_NEEDED entry for libextension_cuda.so" >&2
+  exit 1
+fi
+
+# ...and that dependency must actually resolve at load time.
+if ldd "${runner_path}" | grep -E "libextension_cuda\.so.*=>.*not found" >&2; then
+  echo "example_executorch_runner cannot resolve libextension_cuda.so at runtime" >&2
+  exit 1
+fi
+
+# The runner must import the caller-stream API from the shared library rather
+# than define it privately. A private definition means a second copy of the
+# thread-local, which silently breaks the cross-backend handshake. Assert the
+# import (in .dynsym) rather than the absence of a definition: absence-of-symbol
+# checks read .symtab, which is stripped from release binaries and would make
+# the assertion pass vacuously. A private definition would satisfy the reference
+# at link time and leave no import here.
+for _symbol in getCallerStream CallerStreamGuard; do
+  if ! nm -D --undefined-only "${runner_path}" 2>/dev/null | grep -q "${_symbol}"; then
+    echo "example_executorch_runner does not import ${_symbol} from libextension_cuda.so" >&2
+    exit 1
+  fi
+done
+
+# Validate the .so the runner ACTUALLY loads (resolved via ldd), not just a
+# packaged copy. The runner is CMake-built and may link the CMake-built
+# extension_cuda; whichever .so the loader binds to must export the accessor and
+# must be the sole definer the runner sees.
+loaded_extension_cuda="$(
+  ldd "${runner_path}" 2>/dev/null |
+    sed -n 's/.*libextension_cuda\.so[^ ]* => \([^ ]*\).*/\1/p' |
+    head -n1
+)"
+if [[ -z "${loaded_extension_cuda}" || ! -f "${loaded_extension_cuda}" ]]; then
+  echo "Could not resolve the libextension_cuda.so the runner loads" >&2
+  exit 1
+fi
+if ! nm --defined-only --dynamic "${loaded_extension_cuda}" 2>/dev/null |
+    grep -q "getCallerStream"; then
+  echo "Loaded ${loaded_extension_cuda} does not export getCallerStream" >&2
+  exit 1
+fi
+
+# Packaging integrity (independent of the CMake runner): the Bazel-packaged .so
+# must exist and export the accessor, and no other packaged ELF may define the
+# caller-stream symbols -- a second definition would reintroduce a duplicate
+# thread-local in the shipped artifact.
+packaged_runner="${TORCH_TENSORRT_ROOT}/bin/example_executorch_runner"
+packaged_extension_cuda="${TORCH_TENSORRT_ROOT}/lib/libextension_cuda.so"
+if [[ ! -x "${packaged_runner}" ]]; then
+  echo "Packaged example_executorch_runner missing or not executable: ${packaged_runner}" >&2
+  exit 1
+fi
+if [[ ! -f "${packaged_extension_cuda}" ]]; then
+  echo "Packaged libextension_cuda.so missing at ${packaged_extension_cuda}" >&2
+  exit 1
+fi
+if ! nm --defined-only --dynamic "${packaged_extension_cuda}" 2>/dev/null |
+    grep -q "getCallerStream"; then
+  echo "Packaged libextension_cuda.so does not export getCallerStream" >&2
+  exit 1
+fi
+if ! readelf -d "${packaged_runner}" |
+    grep -E "\(NEEDED\).*libextension_cuda\.so" >&2; then
+  echo "Packaged runner has no DT_NEEDED entry for libextension_cuda.so" >&2
+  exit 1
+fi
+if ldd "${packaged_runner}" | grep -E "libextension_cuda\.so.*=>.*not found" >&2; then
+  echo "Packaged runner cannot resolve libextension_cuda.so" >&2
+  exit 1
+fi
+for _symbol in getCallerStream CallerStreamGuard; do
+  if ! nm -D --undefined-only "${packaged_runner}" 2>/dev/null | grep -q "${_symbol}"; then
+    echo "Packaged runner does not import ${_symbol} from libextension_cuda.so" >&2
+    exit 1
+  fi
+done
+
+# No other packaged ELF may define the caller-stream symbols: a second
+# definition would reintroduce a duplicate thread-local.
+extra_defs="$(
+  find "${TORCH_TENSORRT_ROOT}/lib" -maxdepth 1 -type f -name '*.so*' \
+    ! -name 'libextension_cuda.so' -print0 2>/dev/null |
+    while IFS= read -r -d '' _so; do
+      if nm --defined-only --dynamic "${_so}" 2>/dev/null |
+          grep -qE "getCallerStream|CallerStreamGuard"; then
+        echo "${_so}"
+      fi
+    done
+)"
+if [[ -n "${extra_defs}" ]]; then
+  echo "Unexpected caller-stream definitions outside libextension_cuda.so:" >&2
+  echo "${extra_defs}" >&2
+  exit 1
+fi
+
 "${runner_path}" \
   --model_path="${model_path}" \
   --num_runs=1 2>&1 | tee "${runner_log}"
-
-# The sample model is x + 1, and the reference runner fills inputs with 1.0f,
-# so the output sample should contain 2.0000.
-grep -q "Inference completed" "${runner_log}"
-grep -q "output\\[0\\] shape=" "${runner_log}"
-grep -Eq "first [0-9]+ values:.* 2\\.0000" "${runner_log}"
+packaged_runner_log="${verify_root}/packaged_runner.log"
+"${packaged_runner}" \
+  --model_path="${model_path}" \
+  --num_runs=1 2>&1 | tee "${packaged_runner_log}"
+
+# The sample model is x + 1, and both runners fill inputs with 1.0f, so each
+# output sample must report the expected shape and values. ET_LOG output is not
+# part of the packaged runner contract and may be compiled out.
+for _log in "${runner_log}" "${packaged_runner_log}"; do
+  grep -q "output\\[0\\] shape=" "${_log}"
+  grep -Eq "first [0-9]+ values:.* 2\\.0000" "${_log}"
+done

.github/workflows/_linux-x86_64-core.yml

@@ -639,6 +639,7 @@ jobs:
     needs:
       [
         build,
+        executorch-static-build,
         L0-dynamo-converter-tests,
         L0-dynamo-core-tests,
         L0-py-core-tests,

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

@@ -92,6 +92,27 @@ jobs:
 
         # this is to build the libtorchtrt.tar.gz
         bazel build //:libtorchtrt --compilation_mode opt --config=linux
+        # Run the ExecuTorch backend C++ unit tests, which are otherwise only
+        # built, never executed. These tests link TensorRT and CUDA, but Bazel's
+        # cc_import ships the unversioned libnvinfer.so and libcudart.so while the
+        # loader asks for the versioned sonames, so add the directories holding
+        # those to LD_LIBRARY_PATH. Append rather than replace: the toolchain
+        # entries already present are still needed.
+        bazel_external="$(bazel info output_base)/external"
+        for _soname in libnvinfer.so libcudart.so; do
+          _dir="$(
+            find -L "${bazel_external}" -path "*/lib*/${_soname}.*" -printf '%h\n' 2>/dev/null |
+              sort -u | head -n1
+          )"
+          if [[ -z "${_dir}" ]]; then
+            echo "Could not locate a versioned ${_soname} under ${bazel_external}" >&2
+            exit 1
+          fi
+          export LD_LIBRARY_PATH="${_dir}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
+        done
+        bazel test //tests/cpp/executorch:executorch_backend_tests \
+          --compilation_mode opt --config=linux --test_output=errors \
+          --test_env=LD_LIBRARY_PATH
         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}"

BUILD.bazel

@@ -163,20 +163,23 @@ alias(
     actual = ":executorch_trt_backend_archive",
 )
 
+# Ship only the shared caller-stream library. The packaged reference runner
+# builds the ExecuTorch core and TensorRT backend from source (its CMake uses
+# add_subdirectory), so prebuilt .a archives are not consumed by that flow -- but
+# libextension_cuda.so must ship so the runner has a single shared caller-stream
+# TLS at runtime. Package the .so target directly so the tarball never contains a
+# static caller-stream archive.
 pkg_files(
-    name = "executorch_lib_pkg_files",
-    srcs = [
-        ":executorch_core_archive",
-        ":executorch_trt_backend_archive",
-    ],
+    name = "executorch_extension_cuda_pkg_files",
+    srcs = ["@executorch//:libextension_cuda.so"],
     prefix = "lib/",
     visibility = ["//visibility:public"],
 )
 
 pkg_tar(
-    name = "executorch_lib",
-    srcs = [":executorch_lib_pkg_files"],
-    mode = "0644",
+    name = "executorch_extension_cuda_lib",
+    srcs = [":executorch_extension_cuda_pkg_files"],
+    mode = "0755",
     package_dir = "",
 )
 
@@ -275,6 +278,7 @@ pkg_tar(
         ":rtx_sbsa": [],
         ":windows": [],
         "//conditions:default": [
+            ":executorch_extension_cuda_lib",
             ":executorch_source_package",
             ":include_executorch",
         ],

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 the ExecuTorch release/1.4 branch head.
 new_git_repository(
     name = "executorch",
     build_file = "@//third_party/executorch:BUILD",
-    # latest commit in release/1.3 branch
-    commit = "6118688a095fd9697224f5cad72ce42db641c9cd",
+    # latest commit in release/1.4 branch
+    commit = "cd380e7aefd18c171271cc228d3a155455095219",
     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/ \\;",

cpp/BUILD

@@ -141,11 +141,13 @@ cc_library(
     ] + select({
         ":linux_x86_64": [
             "@executorch//:executorch_headers",
+            "@executorch//:extension_cuda",
             "@cuda//:cudart",
             "@tensorrt//:nvinfer",
         ],
         ":sbsa": [
             "@executorch//:executorch_headers",
+            "@executorch//:extension_cuda",
             "@cuda//:cudart",
             "@tensorrt_sbsa//:nvinfer",
         ],

cpp/include/torch_tensorrt/executorch/TensorRTBackend.h

@@ -90,34 +90,5 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface {
   void destroy(::executorch::runtime::DelegateHandle* handle) const override;
 };
 
-// Selects, for the calling thread, the CUDA stream the delegate runs TensorRT on;
-// scope it around execution.
-//
-// Confines inference to a CUDA green context's SM partition when the caller
-// passes a cuGreenCtxStreamCreate stream: confinement rides the stream (the green
-// context need not be current), and cudaStreamPerThread — the no-guard default —
-// is rejected while a green context is current. While active, device-resident
-// outputs are left enqueued on the stream (no end sync) to compose with later GPU
-// work.
-//
-// Contract: the stream is on the engine's device and outlives the guard; a handle
-// is executed by one thread at a time. On the no-end-sync path (guard active, all
-// I/O device-resident) execute() returns with the TensorRT enqueue still in flight
-// on the stream; the delegate itself orders the next execute() on, and the
-// destruction of, that handle after the work completes (via an internal completion
-// event), so the caller need only synchronize the stream before reading
-// device-resident outputs.
-class CudaStreamGuard {
- public:
-  explicit CudaStreamGuard(cudaStream_t stream);
-  ~CudaStreamGuard();
-  CudaStreamGuard(const CudaStreamGuard&) = delete;
-  CudaStreamGuard& operator=(const CudaStreamGuard&) = delete;
-
- private:
-  cudaStream_t prev_stream_;
-  bool prev_set_;
-};
-
 } // namespace executorch_backend
 } // namespace torch_tensorrt