[STF] Add graph_ctx constructors for explicit graph with user stream (#8094)

The impl(cudaGraph_t) constructor was missing base class initialization
with the async_resources_handle. Fix this and add a new constructor
taking (cudaGraph_t, cudaStream_t, async_resources_handle) that launches
the graph on a user-provided stream with non-blocking finalize.

Also add explicit_graph_async test exercising the new constructor.

Author: caugonnet

cudax/include/cuda/experimental/__stf/graph/graph_ctx.cuh

@@ -187,13 +187,27 @@ class graph_ctx : public backend_ctx<graph_ctx>
     }
 
     // Note that graph contexts with an explicit graph passed by the user cannot use stages
-    impl(cudaGraph_t g)
-        : _graph(wrap_cuda_graph(g))
+    impl(cudaGraph_t g, async_resources_handle _async_resources = async_resources_handle(nullptr))
+        : backend_ctx<graph_ctx>::impl(mv(_async_resources))
+        , _graph(wrap_cuda_graph(g))
         , explicit_graph(true)
     {
       reserved::backend_ctx_setup_allocators<impl, uncached_graph_allocator>(*this);
     }
 
+    // Constructor with explicit graph, user stream, and async resources
+    impl(cudaGraph_t g,
+         cudaStream_t user_stream,
+         async_resources_handle _async_resources = async_resources_handle(nullptr))
+        : backend_ctx<graph_ctx>::impl(mv(_async_resources))
+        , submitted_stream(user_stream)
+        , _graph(wrap_cuda_graph(g))
+        , explicit_graph(true)
+        , blocking_finalize(false)
+    {
+      reserved::backend_ctx_setup_allocators<impl, uncached_graph_allocator>(*this);
+    }
+
     ~impl() override {}
 
     ::std::string to_string() const override
@@ -285,8 +299,13 @@ public:
   }
 
   /// @brief Constructor taking a user-provided graph. User code is not supposed to destroy the graph later.
-  graph_ctx(cudaGraph_t g)
-      : backend_ctx<graph_ctx>(::std::make_shared<impl>(g))
+  graph_ctx(cudaGraph_t g, async_resources_handle handle = async_resources_handle(nullptr))
+      : backend_ctx<graph_ctx>(::std::make_shared<impl>(g, mv(handle)))
+  {}
+
+  /// @brief Constructor with explicit graph, support stream, and async resources
+  graph_ctx(cudaGraph_t g, cudaStream_t user_stream, async_resources_handle handle = async_resources_handle(nullptr))
+      : backend_ctx<graph_ctx>(::std::make_shared<impl>(g, user_stream, mv(handle)))
   {}
   ///@}
 

cudax/test/stf/CMakeLists.txt

@@ -33,6 +33,7 @@ set(
   gnu/include_only.cpp
   graph/concurrency_test.cu
   graph/explicit_graph.cu
+  graph/explicit_graph_async.cu
   graph/explicit_graph_while.cu
   graph/explicit_graph_while-kernels.cu
   graph/get_cache_stats.cu

cudax/test/stf/graph/explicit_graph_async.cu

@@ -0,0 +1,57 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of CUDASTF in CUDA C++ Core Libraries,
+// under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES.
+//
+//===----------------------------------------------------------------------===//
+
+//! @file
+//! @brief Add tasks to a user-provided graph and launch on a user-provided stream.
+//!        Exercises graph_ctx(cudaGraph_t, cudaStream_t): finalize() submits the
+//!        graph on the given stream and does not block; the caller synchronizes.
+
+#include <cuda/experimental/__stf/graph/graph_ctx.cuh>
+
+using namespace cuda::experimental::stf;
+
+__global__ void dummy() {}
+
+int main()
+{
+  cudaGraph_t graph;
+  cudaStream_t stream;
+
+  cuda_safe_call(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+  cuda_safe_call(cudaGraphCreate(&graph, 0));
+
+  graph_ctx ctx(graph, stream);
+
+  auto lX = ctx.token();
+  auto lY = ctx.token();
+  auto lZ = ctx.token();
+
+  ctx.task(lX.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lX.read(), lY.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lX.read(), lZ.write())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  ctx.task(lY.rw(), lZ.rw())->*[](cudaStream_t s) {
+    dummy<<<1, 1, 0, s>>>();
+  };
+
+  // Non-blocking: submits the graph on the user-provided stream
+  ctx.finalize();
+
+  cuda_safe_call(cudaStreamSynchronize(stream));
+  cuda_safe_call(cudaStreamDestroy(stream));
+}