Share the CUDA device guard through the extension so other backends can use it by shoumikhin · Pull Request #22970 · pytorch/executorch

shoumikhin · 2026-09-21T16:23:50Z

The problem

Selecting a CUDA device and putting the caller's device back is written out by hand in two places in the CUDA
backend. One is a struct inside a source file, private to that file. The other is the same save and restore
spelled inline in a destructor. A comment on a third copy says it follows the first. The TensorRT delegate,
which lives in another repository and loads into the same process, had written a fourth.

The change

This tree already had the right guard, sitting in a backend directory. It restores on destruction, it is
built through a factory that returns a result so a failure cannot be ignored by accident, it logs a failed
restore, it leaves a moved-from copy looking already restored so only one of the two puts the device back,
and it has tests. What it lacked was a home another backend could reach.

So it moves to the CUDA extension, beside the caller-stream guard, for the same reason that one is there. One
program can run several backends whose engines sit on different devices, so each has to leave the device as
it found it, or the next one inherits a selection it never made. Kernel launches follow the stream they are
given, but allocation follows the current device, so a backend that allocates has to select one even when the
caller chose the stream.

Only the guard moves. The per-device stream helpers stay where their callers are, because the extension
already answers the stream question a different way and two mechanisms side by side would leave a reader
guessing which to use. Their header keeps the name in scope, which it needs regardless, since its own stream
guard holds one.

Two details were carried across rather than rewritten. The destructor decides whether to restore from what
the caller asked for, not from whether the switch reported success, because a failed selection can still move
the calling thread once an error is pending on the device. And a negative index is refused rather than
ignored, so the one call site that can reach one asks first, exactly as it did before.

The device index is a plain int here. The old signature used an alias from a backend compatibility layer, and
that alias was the only thing tying this guard to that backend.

Testing

Measured on four machines: two with several cards each, and two with a single card, on both processor
architectures. A selection that needs a switch restores the caller's device, one that does not changes
nothing, a negative index and an absent device are both refused and leave the device alone, and moving a
guard restores once rather than twice. The single-card machines take the path where the multi-device cases
are skipped, which is the same path the new test takes there.

A new test pins the move. It takes the borrow over, lets the taker go out of scope, moves the device
somewhere a second restore would be visible, and requires the moved-from guard to leave it alone. Deleting
the one line that marks a moved-from guard as already restored turns it red, which matters because a
defaulted move would have restored in both scopes.

The moved header and the file it came from both compile against the real runtime headers on Arm.

Every place that reaches the header is wired: a target of its own in the extension, a dependency from the
guard's former target, and one from each backend target whose sources include it. No public job builds that
graph, so this part was checked by reading rather than by a green tick.

The guard's tests, and the stream guard's tests beside them, had a target in one build system and none in
the other, so the build a pull request runs never compiled either file. Both have a target now, and the job
that runs the other CUDA runtime tests builds and runs them, so twenty nine cases that compiled nowhere
now compile and run.

pytorch-bot · 2026-09-21T16:23:54Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22970

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

linux.dgx.b200 runner queuing exceeds 4 hours and dgxb200-05 "CUDA not available"

⏳ No Failures, 5 Pending

As of commit 087bf0d with merge base bb2683b ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Copilot