Run the CUDA delegate on a caller-selected CUDA stream (#20082)

Summary:

Problem: The CUDA/AOTI backend always runs a model on a CUDA stream it creates for itself, so an application cannot make it run on a stream the application picked -- for example a CUDA green-context stream that confines the work to part of the GPU.

Fix: Add a small thread-local handshake to the CUDA backend's stream layer in `backends/aoti/slim/cuda/guard`: `CallerStreamGuard`, an RAII scope that records (for the calling thread) the CUDA stream the backend should run on and restores the previous choice when it goes out of scope, and `getCallerStream()`, which returns that stream or nothing if no guard is active.

The CUDA/AOTI backend now consults it: when a caller stream is set, `execute()` runs the kernels and the input and output copies on that stream; when none is set, it uses its own stream exactly as before, so existing behavior is unchanged. CUDA graph capture and replay is refused while a caller stream is set, because a captured graph is bound to its own stream.

The handshake lives next to the existing stream registry and device guards in the same `guard` unit, so the delegate uses it without taking on a new dependency.

Differential Revision: D107698747

Author: shoumikhin

backends/aoti/slim/core/storage.h

@@ -177,7 +177,18 @@ struct DeviceTraits<c10::DeviceType::CUDA> {
           static_cast<int>(dst_device.index()));
     }
 
-    ET_CUDA_CHECK(cudaMemcpy(dst, src, nbytes, direction));
+    // Plain cudaMemcpy is host-synchronous on the default stream, which a
+    // green context would not confine. When a caller stream is active, copy
+    // on it asynchronously and synchronize it to preserve blocking
+    // semantics; otherwise fall back to the plain synchronous copy.
+    const auto caller_stream = executorch::backends::cuda::getCallerStream();
+    if (caller_stream) {
+      ET_CUDA_CHECK(
+          cudaMemcpyAsync(dst, src, nbytes, direction, *caller_stream));
+      ET_CUDA_CHECK(cudaStreamSynchronize(*caller_stream));
+    } else {
+      ET_CUDA_CHECK(cudaMemcpy(dst, src, nbytes, direction));
+    }
   }
 };
 #else

backends/aoti/slim/cuda/guard.cpp

@@ -9,13 +9,15 @@
 #include <executorch/backends/aoti/slim/cuda/guard.h>
 #include <executorch/runtime/platform/log.h>
 #include <limits>
+#include <optional>
 #include <unordered_map>
 
 namespace executorch::backends::cuda {
 
 namespace {
 // Thread-local stream storage (private to this file)
 thread_local std::unordered_map<DeviceIndex, cudaStream_t> current_streams_;
+thread_local std::optional<cudaStream_t> caller_stream_;
 } // namespace
 
 Error setCurrentCUDAStream(cudaStream_t stream, DeviceIndex device_index) {
@@ -52,6 +54,46 @@ Result<cudaStream_t> getCurrentCUDAStream(DeviceIndex device_index) {
   return stream;
 }
 
+std::optional<cudaStream_t> peekCurrentCUDAStream(DeviceIndex device_index) {
+  if (device_index == -1) {
+    int tmp_device = -1;
+    if (cudaGetDevice(&tmp_device) != cudaSuccess) {
+      return std::nullopt;
+    }
+    device_index = static_cast<DeviceIndex>(tmp_device);
+  }
+
+  auto it = current_streams_.find(device_index);
+  if (it == current_streams_.end()) {
+    return std::nullopt;
+  }
+  return it->second;
+}
+
+void clearCurrentCUDAStream(DeviceIndex device_index) {
+  if (device_index == -1) {
+    int tmp_device = -1;
+    if (cudaGetDevice(&tmp_device) != cudaSuccess) {
+      return;
+    }
+    device_index = static_cast<DeviceIndex>(tmp_device);
+  }
+  current_streams_.erase(device_index);
+}
+
+std::optional<cudaStream_t> getCallerStream() {
+  return caller_stream_;
+}
+
+CallerStreamGuard::CallerStreamGuard(cudaStream_t stream)
+    : previous_(caller_stream_) {
+  caller_stream_ = stream;
+}
+
+CallerStreamGuard::~CallerStreamGuard() {
+  caller_stream_ = previous_;
+}
+
 CUDAGuard::CUDAGuard(CUDAGuard&& other) noexcept
     : original_device_index_(other.original_device_index_),
       current_device_index_(other.current_device_index_) {

backends/aoti/slim/cuda/guard.h

@@ -9,6 +9,8 @@
 #pragma once
 
 #include <cuda_runtime.h>
+#include <optional>
+
 #include <executorch/backends/aoti/slim/c10/core/Device.h>
 #include <executorch/backends/aoti/slim/c10/cuda/Exception.h>
 #include <executorch/runtime/core/error.h>
@@ -43,6 +45,54 @@ Error setCurrentCUDAStream(cudaStream_t stream, DeviceIndex device_index = -1);
  */
 Result<cudaStream_t> getCurrentCUDAStream(DeviceIndex device_index = -1);
 
+/**
+ * The CUDA stream registered for the specified device, or std::nullopt if none
+ * is set. Unlike getCurrentCUDAStream, it never creates one, so it can snapshot
+ * the current selection without side effects. Also returns std::nullopt if the
+ * current device cannot be queried (device_index -1), so nullopt does not
+ * distinguish "no stream set" from "device query failed".
+ *
+ * @param device_index The device index (-1 to use current device)
+ */
+std::optional<cudaStream_t> peekCurrentCUDAStream(
+    DeviceIndex device_index = -1);
+
+/**
+ * Clears any CUDA stream registered for the specified device, restoring the
+ * "no stream selected" state. Best-effort: if device_index is -1 and the
+ * current device cannot be queried, it silently does nothing.
+ *
+ * @param device_index The device index (-1 to use current device)
+ */
+void clearCurrentCUDAStream(DeviceIndex device_index = -1);
+
+/**
+ * The CUDA stream the caller selected for this thread (via CallerStreamGuard),
+ * or std::nullopt if none. The CUDA backend runs on it when set, otherwise it
+ * uses its own stream. Kept separate from getCurrentCUDAStream so an explicit
+ * caller choice is distinguishable from a lazily-created stream.
+ */
+std::optional<cudaStream_t> getCallerStream();
+
+/**
+ * Scopes the CUDA stream the backend should run on for the calling thread, and
+ * restores the previous selection on destruction. One value per thread; a
+ * cuGreenCtxStreamCreate stream confines work to that green context's SM
+ * partition.
+ */
+class CallerStreamGuard {
+ public:
+  explicit CallerStreamGuard(cudaStream_t stream);
+  ~CallerStreamGuard();
+  CallerStreamGuard(const CallerStreamGuard&) = delete;
+  CallerStreamGuard& operator=(const CallerStreamGuard&) = delete;
+  CallerStreamGuard(CallerStreamGuard&&) = delete;
+  CallerStreamGuard& operator=(CallerStreamGuard&&) = delete;
+
+ private:
+  std::optional<cudaStream_t> previous_;
+};
+
 /**
  * RAII guard that sets the current CUDA device and restores it on destruction.
  * This ensures that the device is properly restored even if an exception

backends/aoti/slim/cuda/test/test_cuda_stream_guard.cpp

@@ -11,6 +11,8 @@
 #include <executorch/runtime/platform/platform.h>
 #include <gtest/gtest.h>
 
+#include <type_traits>
+
 using namespace executorch::backends::cuda;
 using namespace executorch::runtime;
 
@@ -265,3 +267,60 @@ TEST_F(CUDAStreamGuardTest, NullStreamPointer) {
   auto current_stream_result = getCurrentCUDAStream(0);
   ASSERT_TRUE(current_stream_result.ok());
 }
+
+// CallerStreamGuard / getCallerStream select the backend's stream through pure
+// thread-local state and never touch a device. They still need the CUDA headers
+// to build, but no CUDA device at runtime, so they run outside the device-gated
+// fixture above using opaque (fake) stream values.
+namespace {
+// Opaque, distinct, never-dereferenced stream handles; using object addresses
+// avoids an int-to-pointer cast.
+cudaStream_t fake_stream(int index) {
+  static char storage[3];
+  return reinterpret_cast<cudaStream_t>(&storage[index]);
+}
+} // namespace
+
+TEST(CallerStreamGuardTest, NoGuardReportsNullopt) {
+  EXPECT_FALSE(getCallerStream().has_value());
+}
+
+TEST(CallerStreamGuardTest, GuardSelectsThenRestores) {
+  const cudaStream_t selected = fake_stream(0);
+  {
+    CallerStreamGuard guard(selected);
+    EXPECT_EQ(getCallerStream(), selected);
+  }
+  EXPECT_FALSE(getCallerStream().has_value());
+}
+
+TEST(CallerStreamGuardTest, NestedGuardsRestoreOuter) {
+  const cudaStream_t outer = fake_stream(1);
+  const cudaStream_t inner = fake_stream(2);
+  CallerStreamGuard outer_guard(outer);
+  {
+    CallerStreamGuard inner_guard(inner);
+    EXPECT_EQ(getCallerStream(), inner);
+  }
+  EXPECT_EQ(getCallerStream(), outer);
+}
+
+TEST(CallerStreamGuardCompileTimeTest, NotCopyable) {
+  static_assert(
+      !std::is_copy_constructible_v<CallerStreamGuard>,
+      "CallerStreamGuard should not be copy constructible");
+  static_assert(
+      !std::is_copy_assignable_v<CallerStreamGuard>,
+      "CallerStreamGuard should not be copy assignable");
+}
+
+TEST(CUDAStreamRegistryTest, PeekDoesNotCreateAndClearResets) {
+  // An explicit index skips the cudaGetDevice path, so this needs no device;
+  // use an index no other test touches.
+  constexpr DeviceIndex kIdx = 5;
+  EXPECT_FALSE(peekCurrentCUDAStream(kIdx).has_value());
+  ASSERT_EQ(setCurrentCUDAStream(fake_stream(0), kIdx), Error::Ok);
+  EXPECT_EQ(peekCurrentCUDAStream(kIdx), fake_stream(0));
+  clearCurrentCUDAStream(kIdx);
+  EXPECT_FALSE(peekCurrentCUDAStream(kIdx).has_value());
+}

backends/cuda/runtime/cuda_backend.cpp

@@ -22,6 +22,7 @@
 #include <filesystem>
 #include <fstream>
 #include <mutex>
+#include <optional>
 #include <string>
 #include <string_view>
 #include <unordered_map>
@@ -32,6 +33,7 @@
 #include <executorch/backends/aoti/slim/c10/cuda/Exception.h>
 #include <executorch/backends/aoti/slim/core/slim_tensor.h>
 #include <executorch/backends/aoti/slim/core/storage.h>
+#include <executorch/backends/aoti/slim/cuda/guard.h>
 #include <executorch/backends/aoti/slim/factory/empty.h>
 #include <executorch/backends/aoti/slim/factory/from_blob.h>
 #include <executorch/backends/aoti/slim/factory/from_etensor.h>
@@ -482,7 +484,39 @@ class ET_EXPERIMENTAL CudaBackend final
     size_t n_outputs;
     handle->get_num_outputs(handle->container_handle, &n_outputs);
 
-    setCurrentCUDAStream(handle->get_cuda_stream(), 0);
+    // Run on the caller-selected stream when one is active on this thread (e.g.
+    // a CUDA green-context stream), otherwise the handle's own stream. Every
+    // kernel and boundary copy reads getCurrentCUDAStream, so installing the
+    // choice here routes the whole execution; restore the prior selection on
+    // return so a caller stream does not linger for later work on this thread.
+    const std::optional<cudaStream_t> caller_stream =
+        executorch::backends::cuda::getCallerStream();
+
+    // A captured CUDA graph is bound to its capture stream and cannot be safely
+    // replayed on a different, caller-provided stream.
+    ET_CHECK_OR_RETURN_ERROR(
+        !(caller_stream &&
+          handle->cuda_graph_state.phase != CudaGraphPhase::Disabled),
+        NotSupported,
+        "CUDA graph is not supported together with a caller-provided CUDA stream.");
+
+    // Snapshot the prior selection without creating one (peek, not get), so the
+    // restore is exact and we don't leak a stream just to snapshot.
+    std::optional<cudaStream_t> prev_stream;
+    if (caller_stream) {
+      prev_stream = peekCurrentCUDAStream(0);
+    }
+    setCurrentCUDAStream(caller_stream.value_or(handle->get_cuda_stream()), 0);
+    executorch::backends::aoti::ScopeGuard restore_stream([&]() noexcept {
+      if (!caller_stream) {
+        return;
+      }
+      if (prev_stream) {
+        setCurrentCUDAStream(*prev_stream, 0);
+      } else {
+        clearCurrentCUDAStream(0);
+      }
+    });
 
     size_t n_io_sum = 0;
     ET_CHECK_OR_RETURN_ERROR(