refactor(executorch): use shared caller stream - #4421
Conversation
764b1ca to
5a2ddd9
Compare
|
CI note: the failing checks here are unrelated to this change and appear to be pre-existing trunk flakiness.
The branch is already based on the latest |
801cdcf to
c7e3a09
Compare
c7e3a09 to
a2dc104
Compare
|
CI update after the rebase/force-push:
The caller-stream-specific ExecuTorch static workflow and the remaining build matrix are still running. |
a2dc104 to
27b28a0
Compare
b18ebc2 to
9e0294f
Compare
Replace the TensorRT ExecuTorch backend private caller-stream TLS with ExecuTorch CallerStreamGuard/getCallerStream so CUDA-capable delegates share one process-wide selection. Link and package one shared extension_cuda instance, add ordinary caller-stream inference coverage in the reference runner, and verify both CMake-built and packaged runners consume the shared TLS without libtorch. The ExecuTorch release/1.4 pin is intentionally owned by the preceding version-bump commit.
9e0294f to
6e51eef
Compare
|
Pushed a revision that tightens the single-shared-library guarantee and removes some The shared-object check was not actually checking. The prebuilt-override path Two libraries could be created silently. If a project added this package before An empty override bricked a working configuration. Removed the one-line public wrapper. Corrected the CI symbol assertions. The check used Smaller build and CI surface. The unit test no longer links TensorRT, so the Documentation. Softened a claim that the runner "validates stream selection". It Net effect is 27 fewer lines than the previous revision. Full verification, on Linux Two things I did not change, and want to flag rather than hide:
|
|
Closing note: this pull request was auto-closed by GitHub, not intentionally. I force-pushed an update from a shallow clone ( The branch has been repaired: the commit now sits on the correct parent and the file Continued in #4454. |
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:
After:
TensorRT and other CUDA-capable ExecuTorch delegates now read the same selection.
Why one shared library matters
The selection lives in a
thread_localvariable inside a single shared library,libextension_cuda.so. A shared library is a.sofile that a program loads atrun 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:
So the build must guarantee exactly one shared library. This PR enforces that in
three ways:
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, whichproduces exactly the duplicate above.
add_subdirectoryon ExecuTorch's ownextension/cudainstead of re-declaringthe target. This keeps the two builds identical, and if something else also
declares the target, CMake fails loudly rather than producing two libraries.
DT_NEEDEDentry forlibextension_cuda.soplus dynamic imports of bothgetCallerStreamandCallerStreamGuard. A private copy satisfies thosereferences at link time and leaves no import, so its absence is the signal.
What else changes
cudaStreamPerThreadwhen no guard is active.return with work still in flight. Two separate reads could drift apart.
libextension_cuda.so.a
CallerStreamGuard.Compatibility
Removing
torch_tensorrt::executorch_backend::CudaStreamGuardis an intentionalsource-level C++ API change. Native callers switch to
CallerStreamGuardas shownabove. 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:
empty override, missing override, and rejection of a static archive, a linker
script, and an executable renamed to
.so. A shared object with no.soextension is correctly accepted, since the header decides, not the name.
embeds a private copy. The gate rejects it, including after stripping, and
accepts a correctly linked one.
explicit
cudaStreamPerThread, nesting, and per-thread isolation.the executable, and the reverse case confirming a duplicate copy is detectable.
ET_CHECK_MSGsurvives-O2 -DNDEBUG, so release-build checks are notcompiled out.
are clean.
Dependency
Native caller-stream support needs ExecuTorch 1.4 APIs. The version bump is owned
by a separate change and is temporarily included in this branch's diff against
main; those pin files disappear after rebasing once it lands.