cudax/stf: migrate stackable/ from cuda_safe_call to cuda_try (#9165)

* cudax/stf: migrate stackable/ from cuda_safe_call to cuda_try

In stackable_ctx_impl.cuh, replace cuda_safe_call with cuda_try in the
graph_ctx_node constructor and finalize() so CUDA errors are reported as
exceptions rather than aborting the process.

The constructor builds a CUDA graph in stages, so add transactional
cleanup:

  - In the nested non-conditional branch, the freshly created
    dummy_graph is destroyed intentionally mid-block.  Guard it with a
    SCOPE(fail) that frees it only while dummy_graph_owned is true, and
    disarm the flag right after the intentional destroy.

  - The outer `graph` is owned by us only in the non-nested case (in the
    nested cases it is either parent_graph or a child of parent_graph,
    both owned upstream).  A SCOPE(fail) destroys it on early throw and
    is disarmed the instant graph_ctx adopts it via
    `auto gctx = graph_ctx(sub_graph, ...);`, matching graph_ctx's
    documented ownership contract ("User code is not supposed to destroy
    the graph later").

  - The conditional handle (cudaGraphConditionalHandleCreate) and any
    nodes added to `graph` (cudaGraphAddNode, cudaGraphAddKernelNode)
    are implicitly cleaned up by the outer SCOPE(fail) destroying
    `graph`.

Two residual hazards are intentionally documented inline rather than
fixed in this commit:

  - cudaGraphAddChildGraphNode leaves an orphaned child node inside
    parent_graph on later throw; cleanly removing it would need
    cudaGraphDestroyNode and dependency rewiring.

  - cudaGraphConditionalHandleCreate writes a handle into a caller-owned
    pointer; CUDA has no destroy API for conditional handles, so on
    throw the handle is left invalid (its backing graph is destroyed).

Both are no worse than the prior behavior (which aborted).

The four cuda_safe_call sites in finalize() (cudaGraphAddDependencies
on both CTK branches, cudaGraphDebugDotPrint, cudaGraphLaunch) become
plain cuda_try; no resource rollback applies.

The two cuda_safe_call sites inside the new SCOPE(fail) bodies are
intentional: SCOPE bodies are noexcept, so cuda_safe_call is the
correct tool there.

In stackable_ctx.cuh, the two cuda_safe_call sites inside
UNITTEST host-task lambdas are kept and annotated.  Those lambdas
are dispatched by the STF host-task path, whose exception-safety has
not been audited, so an abort remains safer than an unannotated throw.

* cudax/stf: use templated cuda_try<F> form for graph out-params in stackable/

Where a graph_ctx_node site captures a CUDA out-parameter, switch from
the runtime-status form `cuda_try(cudaFn(&out, ...))` to the templated
form `out = cuda_try<cudaFn>(...)`.  This allows the captured handle to
be a single const-initialized local instead of declare-then-fill.

Converted (each capturing one output handle):

  - cudaGraphCreate (x2)            -> dummy_graph / graph
  - cudaGraphAddChildGraphNode      -> const n
  - cudaGraphChildGraphNodeGetGraph -> graph (last-output-param form)
  - cudaGraphAddNode (both CTK      -> const conditionalNode
    branches)
  - cudaGraphAddKernelNode          -> const reset_node

dummy_graph stays non-const because it is reset to nullptr to disarm its
SCOPE(fail) guard after the intentional destroy.

Left in runtime-status form on purpose:

  - cudaGraphConditionalHandleCreate: its output is written into the
    caller-owned config.conditional_handle, not a synthesizable local;
    the templated form would create a throwaway local and lose the write.
  - cudaGraphDestroy, cudaGraphAddDependencies, cudaGraphDebugDotPrint,
    cudaGraphLaunch: no captured output, so the templated form adds
    nothing.

cudaGraphAddNode is convertible despite its trailing non-const
cudaGraphNodeParams* because the last-output interpretation fails to
typecheck (cudaGraph_t is not convertible to cudaGraphNode_t* once the
synthesized pointer is appended), so cuda_try selects the first-output
form unambiguously.

Author: andralex

cudax/include/cuda/experimental/__stf/stackable/stackable_ctx.cuh

@@ -1638,6 +1638,10 @@ UNITTEST("stackable task on exec_place::host()")
   stackable_ctx ctx;
   auto lA = ctx.logical_data(shape_of<slice<int>>(1024));
   ctx.task(exec_place::host(), lA.write())->*[](cudaStream_t stream, auto) {
+    // cuda_safe_call (not cuda_try) on purpose: this lambda body is invoked from
+    // the STF runtime under host-task dispatch, where exception safety has not
+    // been audited. An abort here is preferable to an unannotated throw escaping
+    // into the runtime.
     cuda_safe_call(cudaStreamSynchronize(stream));
   };
   ctx.finalize();
@@ -1648,6 +1652,8 @@ UNITTEST("stackable task with set_symbol and set_exec_place")
   stackable_ctx ctx;
   auto lA = ctx.logical_data(shape_of<slice<int>>(1024));
   ctx.task(lA.write()).set_symbol("task").set_exec_place(exec_place::host())->*[](cudaStream_t stream, auto) {
+    // Same rationale as the previous test: keep cuda_safe_call inside this
+    // host-task lambda until the dispatch path is audited for exception safety.
     cuda_safe_call(cudaStreamSynchronize(stream));
   };
   ctx.finalize();

cudax/include/cuda/experimental/__stf/stackable/stackable_ctx_impl.cuh

@@ -37,6 +37,7 @@
 #include "cuda/experimental/__stf/stackable/stackable_node_hierarchy.cuh"
 #include "cuda/experimental/__stf/stackable/stackable_task_dep.cuh"
 #include "cuda/experimental/__stf/utility/hash.cuh"
+#include "cuda/experimental/__stf/utility/scope_guard.cuh"
 #include "cuda/experimental/__stf/utility/source_location.cuh"
 
 namespace cuda::experimental::stf
@@ -453,38 +454,76 @@ public:
           else
 #endif // _CCCL_CTK_AT_LEAST(12, 4) && !defined(CUDASTF_DISABLE_CODE_GENERATION) && defined(__CUDACC__)
           {
-            cudaGraph_t dummy_graph;
-            cuda_safe_call(cudaGraphCreate(&dummy_graph, 0));
+            cudaGraph_t dummy_graph = cuda_try<cudaGraphCreate>(0);
+
+            // dummy_graph is intentionally destroyed below. Until that destroy
+            // succeeds, we own it; this guard releases it if any of the following
+            // calls throws. The handle itself carries the ownership bit -- we null
+            // it after the intentional destroy to disarm.
+            SCOPE(fail)
+            {
+              if (dummy_graph != nullptr)
+              {
+                cuda_safe_call(cudaGraphDestroy(dummy_graph));
+              }
+            };
 
-            // The dependencies to this child graph will be added later
-            cudaGraphNode_t n;
-            cuda_safe_call(cudaGraphAddChildGraphNode(&n, parent_graph, nullptr, 0, dummy_graph));
+            // The dependencies to this child graph will be added later.
+            // NOTE: cudaGraphAddChildGraphNode adds `n` into parent_graph. If the
+            // constructor throws after this point, `n` remains as an orphaned child
+            // node inside parent_graph. Removing it cleanly would require
+            // cudaGraphDestroyNode plus careful dependency rewiring; we accept the
+            // orphan because parent_graph's lifetime extends well beyond this scope
+            // and the orphan is harmless until parent_graph is destroyed.
+            const cudaGraphNode_t n = cuda_try<cudaGraphAddChildGraphNode>(parent_graph, nullptr, 0, dummy_graph);
 
             // cudaGraphAddChildGraphNode clones dummy_graph into the parent
-            cuda_safe_call(cudaGraphDestroy(dummy_graph));
+            cuda_try(cudaGraphDestroy(dummy_graph));
+            dummy_graph = nullptr;
 
             // Get the graph described by the child, not the graph that was
             // cloned into the child graph node so that changes are reflected
             // in it.
-            cuda_safe_call(cudaGraphChildGraphNodeGetGraph(n, &graph));
+            graph       = cuda_try<cudaGraphChildGraphNodeGetGraph>(n);
             input_node  = n;
             output_node = n;
           }
         }
         else
         {
-          cuda_safe_call(cudaGraphCreate(&graph, 0));
+          graph = cuda_try<cudaGraphCreate>(0);
         }
 
+        // We own `graph` (the raw cudaGraph_t member) only when we freshly allocated
+        // it above (the non-nested case). In the nested case, `graph` is either
+        // parent_graph itself or a child graph inside parent_graph, both owned by
+        // parent_graph -- destroying it would be incorrect. The graph_ctx built
+        // below (`gctx`) takes ownership when constructed successfully (see
+        // graph_ctx's ctor doc: "User code is not supposed to destroy the graph
+        // later."), so we disarm this guard right after that point.
+        bool graph_owned_by_us = !nested_graph;
+        SCOPE(fail)
+        {
+          if (graph_owned_by_us)
+          {
+            cuda_safe_call(cudaGraphDestroy(graph));
+          }
+        };
+
         // This is the graph which will be used by our STF context. If we need
         // to use a conditional node later, will will change that value.
         cudaGraph_t sub_graph = graph;
 
 #if _CCCL_CTK_AT_LEAST(12, 4) && !defined(CUDASTF_DISABLE_CODE_GENERATION) && defined(__CUDACC__)
         if (config.conditional_handle != nullptr)
         {
-          // Create the conditional handle and store it in the provided pointer
-          cuda_safe_call(cudaGraphConditionalHandleCreate(
+          // Create the conditional handle and store it in the provided pointer.
+          // NOTE: on a thrown cuda_try below, *config.conditional_handle is left
+          // pointing at a handle whose backing graph will be destroyed by the
+          // SCOPE(fail) above. CUDA has no destroy API for conditional handles
+          // (they are tied to their graph), so the handle simply becomes invalid
+          // and the caller must not use it after catching the exception.
+          cuda_try(cudaGraphConditionalHandleCreate(
             config.conditional_handle, graph, config.default_launch_value, config.flags));
 
           // Create conditional node parameters
@@ -494,12 +533,13 @@ public:
           cParams.conditional.type    = config.conditional_type;
           cParams.conditional.size    = 1;
 
-          // Add conditional node to parent graph
-          cudaGraphNode_t conditionalNode;
+          // Add conditional node to parent graph. The node lives inside `graph`;
+          // if construction fails later, the SCOPE(fail) above destroys `graph`
+          // and the node is cleaned up implicitly.
 #  if _CCCL_CTK_AT_LEAST(13, 0)
-          cuda_safe_call(cudaGraphAddNode(&conditionalNode, graph, nullptr, nullptr, 0, &cParams));
+          const cudaGraphNode_t conditionalNode = cuda_try<cudaGraphAddNode>(graph, nullptr, nullptr, 0, &cParams);
 #  else
-          cuda_safe_call(cudaGraphAddNode(&conditionalNode, graph, nullptr, 0, &cParams));
+          const cudaGraphNode_t conditionalNode = cuda_try<cudaGraphAddNode>(graph, nullptr, 0, &cParams);
 #  endif
 
           // Get the body graph from the conditional node
@@ -518,8 +558,7 @@ public:
           kconfig.kernelParams   = kconfig_args;
           kconfig.sharedMemBytes = 0;
 
-          cudaGraphNode_t reset_node;
-          cuda_safe_call(cudaGraphAddKernelNode(&reset_node, graph, &conditionalNode, 1, &kconfig));
+          const cudaGraphNode_t reset_node = cuda_try<cudaGraphAddKernelNode>(graph, &conditionalNode, 1, &kconfig);
 
           input_node  = conditionalNode;
           output_node = reset_node;
@@ -540,6 +579,11 @@ public:
         }
 
         auto gctx = graph_ctx(sub_graph, support_stream, handle);
+        // graph_ctx's ctor took ownership of sub_graph (and, in the non-nested
+        // non-conditional case, of `graph` since sub_graph == graph). From here on,
+        // gctx's destructor handles cleanup, including if any of the following lines
+        // throws before `ctx = gctx`.
+        graph_owned_by_us = false;
 
         // Set up context properties
         gctx.set_parent_ctx(parent_ctx);
@@ -605,7 +649,7 @@ public:
         {
           static ::std::atomic<int> debug_graph_cnt{0};
           ::std::string filename = "instantiated_graph" + ::std::to_string(debug_graph_cnt++) + ".dot";
-          cuda_safe_call(cudaGraphDebugDotPrint(graph, filename.c_str(), cudaGraphDebugDotFlags(0)));
+          cuda_try(cudaGraphDebugDotPrint(graph, filename.c_str(), cudaGraphDebugDotFlags(0)));
           ::std::cout << "Debug: Stackable graph DOT output written to " << filename << '\n';
         }
       }
@@ -628,11 +672,11 @@ public:
 
         size_t nnodes;
         size_t nedges;
-        cuda_safe_call(cudaGraphGetNodes(graph, nullptr, &nnodes));
+        cuda_try(cudaGraphGetNodes(graph, nullptr, &nnodes));
 #if _CCCL_CTK_AT_LEAST(13, 0)
-        cuda_safe_call(cudaGraphGetEdges(graph, nullptr, nullptr, nullptr, &nedges));
+        cuda_try(cudaGraphGetEdges(graph, nullptr, nullptr, nullptr, &nedges));
 #else
-        cuda_safe_call(cudaGraphGetEdges(graph, nullptr, nullptr, &nedges));
+        cuda_try(cudaGraphGetEdges(graph, nullptr, nullptr, &nedges));
 #endif
 
         auto [cached_exec, cache_hit] = ctx.async_resources().cached_graphs_query(nnodes, nedges, graph);
@@ -680,7 +724,7 @@ public:
         _CCCL_ASSERT(exec_graph_, "launch_once called before ensure_instantiated");
         _CCCL_ASSERT(stream == support_stream, "launch_once only supports the node's support stream");
         ensure_prereqs_synced();
-        cuda_safe_call(cudaGraphLaunch(*exec_graph_, stream));
+        cuda_try(cudaGraphLaunch(*exec_graph_, stream));
       }
 
       // Release resources and build the finalize_prereqs event list that the
@@ -740,11 +784,11 @@ public:
         auto* cache_stat = ctx.graph_get_cache_stat();
         if (cache_stat)
         {
-          cuda_safe_call(cudaGraphGetNodes(body, nullptr, &cache_stat->nnodes));
+          cuda_try(cudaGraphGetNodes(body, nullptr, &cache_stat->nnodes));
 #if _CCCL_CTK_AT_LEAST(13, 0)
-          cuda_safe_call(cudaGraphGetEdges(body, nullptr, nullptr, nullptr, &cache_stat->nedges));
+          cuda_try(cudaGraphGetEdges(body, nullptr, nullptr, nullptr, &cache_stat->nedges));
 #else
-          cuda_safe_call(cudaGraphGetEdges(body, nullptr, nullptr, &cache_stat->nedges));
+          cuda_try(cudaGraphGetEdges(body, nullptr, nullptr, &cache_stat->nedges));
 #endif
         }
 
@@ -754,7 +798,7 @@ public:
         {
           static ::std::atomic<int> nested_cnt{0};
           ::std::string filename = "nested_graph" + ::std::to_string(nested_cnt++) + ".dot";
-          cuda_safe_call(cudaGraphDebugDotPrint(body, filename.c_str(), cudaGraphDebugDotFlags(0)));
+          cuda_try(cudaGraphDebugDotPrint(body, filename.c_str(), cudaGraphDebugDotFlags(0)));
         }
 
         cudaGraph_t support_graph = parent_ctx.graph();
@@ -773,10 +817,10 @@ public:
           // Create a vector of input_node repeated for each dependency
           ::std::vector<cudaGraphNode_t> to_nodes(ctx_ready_nodes.size(), input_node);
 #if _CCCL_CTK_AT_LEAST(13, 0)
-          cuda_safe_call(cudaGraphAddDependencies(
+          cuda_try(cudaGraphAddDependencies(
             support_graph, ctx_ready_nodes.data(), to_nodes.data(), nullptr, ctx_ready_nodes.size()));
 #else // _CCCL_CTK_AT_LEAST(13, 0)
-          cuda_safe_call(
+          cuda_try(
             cudaGraphAddDependencies(support_graph, ctx_ready_nodes.data(), to_nodes.data(), ctx_ready_nodes.size()));
 #endif // _CCCL_CTK_AT_LEAST(13, 0)
         }