refactor(executorch): use shared caller stream

[shoumikhin]

Stacked on #4434

This change needs the ExecuTorch 1.4 caller-stream APIs, which #4434 provides by
bumping the pinned ExecuTorch commit. That bump is a separate concern and is reviewed
separately.

Both pull requests come from a fork, and GitHub only accepts a base branch that lives
in the upstream repository, so this cannot be expressed as a GitHub-level stack. The
commits are stacked correctly in git instead:

42bf7d4d2  refactor(executorch): use shared caller stream   <- this pull request
8002a42a3  Pin ExecuTorch to the release/1.4 branch head    <- #4434
c60720959  main

Because the base here has to be main, this diff also shows the four pin files owned
by #4434:

MODULE.bazel
docker/MODULE.bazel.docker
docker/MODULE.bazel.ngc
toolchains/ci_workspaces/MODULE.bazel.tmpl

Please review only the caller-stream changes here, and #4434 for the pin.

The two commits touch no files in common, so there is nothing to reconcile between
them. Once #4434 lands I will rebase, and those four files will drop out of this diff
with no change to the caller-stream code. Reviewing this as 17 files rather than 21 is
equivalent to reviewing the post-rebase state.

The problem

An ExecuTorch program can mix delegates. One subgraph may run on TensorRT while
another runs on the CUDA/AOTI backend. If the application wants both to run on a
CUDA stream it owns, both delegates have to agree on which stream that is.

Today they cannot agree. Torch-TensorRT keeps its own private "caller stream"
value, and ExecuTorch's CUDA backend keeps a different one. Selecting a stream for
one backend leaves the other running somewhere else, which breaks ordering.

The fix

Delete Torch-TensorRT's private caller-stream state and read ExecuTorch's shared
one instead.

Before:

torch_tensorrt::executorch_backend::CudaStreamGuard guard(stream);
module.forward(inputs);

After:

#include <executorch/extension/cuda/caller_stream.h>

executorch::extension::cuda::CallerStreamGuard guard(stream);
module.forward(inputs);

TensorRT and other CUDA-capable ExecuTorch delegates now read the same selection.

Why one shared library matters

The selection lives in a thread_local variable inside a single shared library,
libextension_cuda.so. A shared library is a .so file that a program loads at
run time, and every part of the program that loads the same one sees the same
variable.

If a second copy of that variable ends up in the process, for example because
something linked a static archive instead, then the runner writes to one copy and
the delegate reads the other. Nothing crashes. The delegate silently falls back to
a different stream:

runner selected stream 0x1056930
delegate observed      0x2          <- cudaStreamPerThread, the fallback

So the build must guarantee exactly one shared library. This PR enforces that in
three ways:

• CMake accepts a prebuilt override only after reading the file's ELF header and
  confirming it is a shared object. Checking the file name is not enough, because
  the linker happily links a static archive that has been renamed to .so, which
  produces exactly the duplicate above.
• When building from an ExecuTorch source checkout, CMake calls
  add_subdirectory on ExecuTorch's own extension/cuda instead of re-declaring
  the target. This keeps the two builds identical, and if something else also
  declares the target, CMake fails loudly rather than producing two libraries.
• CI inspects the built and the packaged runner and requires a real
  DT_NEEDED entry for libextension_cuda.so plus dynamic imports of both
  getCallerStream and CallerStreamGuard. A private copy satisfies those
  references at link time and leaves no import, so its absence is the signal.

What else changes

• Use the selected stream for the TensorRT enqueue and for host staging copies.
• Fall back to cudaStreamPerThread when no guard is active.
• Read the selection once, then derive both the stream and whether the backend may
  return with work still in flight. Two separate reads could drift apart.
• Package exactly one shared libextension_cuda.so.
• Run the ExecuTorch backend C++ tests in CI, which were previously only built.
• Keep the native reference runner free of libtorch, and run real inference inside
  a CallerStreamGuard.

Compatibility

Removing torch_tensorrt::executorch_backend::CudaStreamGuard is an intentional
source-level C++ API change. Native callers switch to CallerStreamGuard as shown
above. The replacement is behavior-preserving: an explicitly selected null stream
still counts as a caller selection, matching the old two-field encoding.

This PR validates ordinary CUDA streams. On the discrete-GPU CI configuration the
reference runner exercises the synchronized host-staging path. The device-resident
asynchronous path is not covered end to end by that runner. CUDA green-context
streams need context-aware completion-event handling and are not claimed as
supported here.

Python export-only usage is unchanged.

Testing

Verified on Linux x86_64 with an NVIDIA H100, CUDA 12.8, and TensorRT 11.0,
against the pinned ExecuTorch source checkout:

• Every CMake selection branch, 10 cases: source build, prebuilt shared library,
  empty override, missing override, and rejection of a static archive, a linker
  script, and an executable renamed to .so. A shared object with no .so
  extension is correctly accepted, since the header decides, not the name.
• The linkage assertions, checked against a deliberately broken runner that
  embeds a private copy. The gate rejects it, including after stripping, and
  accepts a correctly linked one.
• The caller-stream unit tests, 6 cases, including an explicit null stream, an
  explicit cudaStreamPerThread, nesting, and per-thread isolation.
• One shared thread-local across two independently linked shared libraries plus
  the executable, and the reverse case confirming a duplicate copy is detectable.
• ET_CHECK_MSG survives -O2 -DNDEBUG, so release-build checks are not
  compiled out.
• The backend builds for both runtime flavors. C++ formatting and shell syntax
  are clean.

[shoumikhin]

CI caught a real mistake in my previous push, now fixed.

I had removed the LD_LIBRARY_PATH setup around the Bazel test step, on the assumption
that only the caller-stream test needed TensorRT at run time. That was wrong.
test_executorch_binding_names also depends on @tensorrt//:nvinfer, so it failed to
start:

//tests/cpp/executorch:test_caller_stream           PASSED
//tests/cpp/executorch:test_executorch_blob_header  PASSED
//tests/cpp/executorch:test_executorch_binding_names FAILED
  error while loading shared libraries: libnvinfer.so.11:
  cannot open shared object file: No such file or directory

The underlying reason is that Bazel's cc_import provides lib/libnvinfer.so, but a
binary linked against it records a DT_NEEDED entry for the versioned soname
libnvinfer.so.11, which is not in the test runfiles. I reproduced this locally
against the same TensorRT SDK and got the identical message, then confirmed that
pointing LD_LIBRARY_PATH at the SDK's lib directory resolves it.

So the test step now locates the versioned library Bazel materialized and passes it
through:

tensorrt_lib_dir="$(
  find -L "$(bazel info output_base)/external" \
    -path '*tensorrt*/lib/libnvinfer.so.*' -printf '%h\n' 2>/dev/null |
    sort -u | head -n1
)"
bazel test //tests/cpp/executorch:executorch_backend_tests \
  --compilation_mode opt --config=linux --test_output=errors \
  --test_env=LD_LIBRARY_PATH="${tensorrt_lib_dir}"

This is still simpler than before. It keeps the sandbox enabled, since
--test_strategy=standalone is no longer needed, and it drops the
CUDA_VISIBLE_DEVICES passthrough, since none of these tests run CUDA work. The glob
only matches versioned sonames, which is what the loader actually looks for.

Worth noting from the same run: test_caller_stream passed, and the full
bazel build //:libtorchtrt completed successfully with the reworked CMake, so the
ELF check and the switch to ExecuTorch's own extension/cuda build are exercised and
working.

[shoumikhin]

Second attempt at the test-runtime fix. My previous one was wrong in a way the logs
made obvious.

I set LD_LIBRARY_PATH to only the TensorRT directory, which replaced the value
instead of adding to it. That fixed libnvinfer.so.11 and immediately broke
libcudart.so.13, so the run went from one failing test to two:

test_executorch_blob_header    PASSED
test_caller_stream             FAILED  <- regressed, previously passed
test_executorch_binding_names  FAILED
  error while loading shared libraries: libcudart.so.13

Both libraries have the same root cause. Bazel's cc_import ships the unversioned
libnvinfer.so and libcudart.so, but a binary linked against them records
DT_NEEDED entries for the versioned sonames, which are not in the test runfiles.

The step now locates the directory holding each versioned soname and appends to
LD_LIBRARY_PATH, keeping the toolchain entries that were already there:

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

I verified this locally by reproducing both failure modes rather than reasoning about
them. Copying the two libraries somewhere off the default search path and running the
loader with --inhibit-cache, so the system library cache cannot mask the problem:

no LD_LIBRARY_PATH        -> libnvinfer.so.11: cannot open shared object file
TensorRT directory only   -> libcudart.so.12: cannot open shared object file
both directories appended -> both libraries loaded OK

The middle line is the failure this run hit. The guard also names the specific library
if either directory cannot be found, instead of failing later with a loader error.

This is still simpler than the original version of the step: the sandbox stays enabled,
since --test_strategy=standalone is not needed, and the CUDA_VISIBLE_DEVICES
passthrough stays removed, since none of these tests run CUDA work.

[shoumikhin]

CI is green on the ExecuTorch gate. All three C++ tests pass:

//tests/cpp/executorch:test_caller_stream                PASSED in 0.1s
//tests/cpp/executorch:test_executorch_binding_names     PASSED in 0.1s
//tests/cpp/executorch:test_executorch_blob_header       PASSED in 0.1s
Executed 3 out of 3 tests: 3 tests pass.

The same job also exercised the rest of this change end to end:

• Built libextension_cuda.so through ExecuTorch's own extension/cuda CMake, which
  is the definition this change now reuses instead of re-declaring the target.
• Verified both the CMake-built and the packaged runner carry a real DT_NEEDED entry
  for libextension_cuda.so and import the caller-stream symbols rather than defining
  private copies.
• Ran real inference in both runners inside a CallerStreamGuard, each producing the
  expected values:

output[0] shape=[2,3,4,4] numel=96 dtype=6
  first 8 values: 2.0000 2.0000 2.0000 2.0000 2.0000 2.0000 2.0000 2.0000

Remaining failures on this commit are pre-existing on main and unrelated to this
change:

Check	Here	On main
Python-only dynamo runtime, two variants	fail	fail
RTX Python-only dynamo runtime, two variants	fail	fail / cancelled
L0 core python	fail	flaky

The dynamo jobs stop before running any test with HTTP 403 FORBIDDEN for channel pkgs/main, a package-index problem in the job image. The L0 failure is
test_libtorchtrt_linkage.py, which fails on a ctypes open of
libnvinfer_plugin.so.11 that is missing from the runner image; this change touches no
files under tests/py.

Note that a large number of pytorch/executorch checks are also attached to this
commit SHA because it appears in that repository's CI as well. Those are not from this
pull request. Filtering by originating repository, the counts are 65 checks from
pytorch/TensorRT and 100 from pytorch/executorch.

[cehongwang]

This means that eventually enqueueV3 is happening on the caller stream, right? Is it going to be a problem @narendasan

[shoumikhin]

Good question, and the short answer is yes, but that is not new here.

enqueueV3 already ran on the caller-selected stream before this change. On the base
commit the same function had:

cudaStream_t stream = g_user_stream_set ? g_user_stream : cudaStreamPerThread;  // line 382
...
if (!ctx->enqueueV3(stream)) {                                                 // line 571

This change only swaps where that stream value is read from. It was a thread-local
private to this delegate, and it is now ExecuTorch's shared one. The
enqueueV3 call, the staging copies, the synchronize, and the completion event are all
untouched.

What happens with no guard active

Nothing changes. getCallerStream() returns empty, so the stream is
cudaStreamPerThread and execute() still synchronizes before returning, exactly as
before. Existing callers see identical behavior.

When execute() can return with work still in flight

This is the part worth being precise about. The decision is:

const bool must_sync = output_staged_to_host || input_staged_from_host || !caller_stream_set;

execute() skips the end-of-call synchronize only when all of these hold:

1. the caller scoped a CallerStreamGuard, and
2. every input is already GPU-accessible, and
3. every output is already GPU-accessible.

If any tensor is host-backed, the backend staged it through a temporary device buffer,
so it synchronizes before returning to keep the "results are ready on return" contract.
And with no guard at all it always synchronizes. So the asynchronous path only happens
when the caller explicitly asked for stream semantics and no host memory is involved.

Ordering and lifetime safety

Within execute(), stream is assigned once and never reassigned, and every CUDA
operation uses that same variable: the host-to-device copy, enqueueV3, the
device-to-host copy, the synchronize, and the event record. There is no mixing of the
caller's stream with the default stream.

For the asynchronous path, the backend records a completion event on that stream and
waits on it in two places before it could disturb TensorRT state:

• at the start of the next execute(), before any setInputShape or
  setTensorAddress, since those are host-side calls on the execution context;
• in the handle destructor, before freeing staging buffers and releasing the execution
  context.

It waits on the event rather than the stream, so teardown stays correct even if the
caller has already destroyed their stream.

This is the same pattern ExecuTorch's CUDA backend uses

The caller-stream primitive exists precisely so several CUDA delegates can share one
caller-provided stream. ExecuTorch's own CUDA/AOTI backend reads the same value and
installs it for its whole execution:

// backends/cuda/runtime/cuda_backend.cpp
const std::optional<cudaStream_t> caller_stream =
    executorch::extension::cuda::getCallerStream();
...
setCurrentCUDAStream(caller_stream.value_or(handle->get_cuda_stream()), 0);

If the two delegates read different stream values, a program that mixes them cannot
order their work, which is the problem this change fixes.

Known limits, stated plainly

• The selected stream must be on the TensorRT engine's device. That is not validated in
  code today; a mismatch surfaces as an enqueueV3 failure. Pre-existing, and this
  change makes the error message name the likely cause.
• CUDA green-context streams need context-aware completion-event handling and are not
  claimed as supported. This change narrows the previous claim rather than widening it.
• The device-resident asynchronous path is not covered end to end by CI, because the
  reference runner's tensors are host-backed and therefore take the synchronized path.

Happy to gate the asynchronous return behind an explicit backend option instead, if you
would prefer that a backend-neutral stream selection never changes when execute()
returns. That would be a small follow-up and would leave the default fully synchronous.