[STF] Make executable graph cache stream-affine

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

@@ -503,7 +503,7 @@ public:
     {
       /* This will lookup in the cache (if any) and update an existing entry, or
        * instantiate a graph if none is found. */
-      auto query_result = async_resources().cached_graphs_query(nnodes, nedges, *g);
+      auto query_result = async_resources().cached_graphs_query(nnodes, nedges, *g, state.submitted_stream);
       state.exec_graph  = query_result.first;
 
       hit = query_result.second; // indicate if this was a hit or miss in the cache

cudax/include/cuda/experimental/__stf/internal/async_resources_handle.cuh

@@ -226,13 +226,13 @@ public:
   // The graph is only used during the call (to update or instantiate); it is never stored, so the
   // caller only needs to keep it valid for the duration of the call.
   ::cuda::std::pair<::std::shared_ptr<cudaGraphExec_t>, bool>
-  cached_graphs_query(size_t nnodes, size_t nedges, cudaGraph_t g)
+  cached_graphs_query(size_t nnodes, size_t nedges, cudaGraph_t g, cudaStream_t stream)
   {
     _CCCL_ASSERT(pimpl, "async_resources_handle is not initialized");
-    return pimpl->cached_graphs.query(nnodes, nedges, g);
+    return pimpl->cached_graphs.query(nnodes, nedges, g, stream);
   }
 
-  ::cuda::std::pair<::std::shared_ptr<cudaGraphExec_t>, bool> cached_graphs_query(cudaGraph_t g)
+  ::cuda::std::pair<::std::shared_ptr<cudaGraphExec_t>, bool> cached_graphs_query(cudaGraph_t g, cudaStream_t stream)
   {
     const size_t nnodes = cuda_try<cudaGraphGetNodes>(g, nullptr);
 #if _CCCL_CTK_AT_LEAST(13, 0)
@@ -242,7 +242,7 @@ public:
 #endif // _CCCL_CTK_AT_LEAST(13, 0)
 
     _CCCL_ASSERT(pimpl, "async_resources_handle is not initialized");
-    return cached_graphs_query(nnodes, nedges, g);
+    return cached_graphs_query(nnodes, nedges, g, stream);
   }
 
 #if _CCCL_CTK_AT_LEAST(12, 4)

cudax/include/cuda/experimental/__stf/internal/executable_graph_cache.cuh

@@ -30,7 +30,7 @@
 #include <cuda/experimental/__stf/utility/pretty_print.cuh>
 #include <cuda/experimental/__stf/utility/source_location.cuh>
 
-#include <queue> // for ::std::priority_queue
+#include <mutex>
 #include <unordered_map>
 
 namespace cuda::experimental::stf
@@ -119,12 +119,18 @@ public:
   // One entry of the cache
   struct entry
   {
-    entry(executable_graph_cache* cache, ::std::shared_ptr<cudaGraphExec_t> exec_g_, size_t footprint)
+    entry(executable_graph_cache* cache,
+          ::std::shared_ptr<cudaGraphExec_t> exec_g_,
+          cudaStream_t stream_,
+          unsigned long long stream_id_,
+          size_t footprint)
         : cache(cache)
         , exec_g(mv(exec_g_))
+        , stream(stream_)
+        , stream_id(stream_id_)
         , footprint(footprint)
     {
-      last_use = cache->index;
+      last_use = cache->index++;
     }
 
     // Update the last_use field to mark that this entry was used recently
@@ -135,6 +141,14 @@ public:
 
     executable_graph_cache* cache;
     ::std::shared_ptr<cudaGraphExec_t> exec_g;
+    // The binding identity is the driver-assigned stream id, which is unique
+    // for the lifetime of the process: a cudaStream_t handle value can be
+    // recycled after cudaStreamDestroy, so comparing handles could falsely
+    // match an entry bound to a dead stream against an unrelated new one.
+    // The raw handle is kept only to probe idleness, which is meaningful
+    // only while the bound stream is alive (see query_stream_state).
+    cudaStream_t stream;
+    unsigned long long stream_id;
     size_t last_use;
     size_t footprint;
   };
@@ -156,15 +170,32 @@ public:
   // Check if there is a matching entry (and update it if necessary)
   // the returned bool indicate is this is a cache hit (true = cache hit, false = cache miss)
   // The graph g is only used during this call (for update or instantiate); it is never stored.
-  ::cuda::std::pair<::std::shared_ptr<cudaGraphExec_t>, bool> query(size_t nnodes, size_t nedges, cudaGraph_t g)
+  ::cuda::std::pair<::std::shared_ptr<cudaGraphExec_t>, bool>
+  query(size_t nnodes, size_t nedges, cudaGraph_t g, cudaStream_t stream)
   {
+    ::std::lock_guard<::std::mutex> guard(mutex);
+
     int dev_id = cuda_try<cudaGetDevice>();
     _CCCL_ASSERT(dev_id < int(cached_graphs.size()), "invalid device id value");
 
+    const unsigned long long stream_id = stream_unique_id(stream);
+
     auto range = cached_graphs[dev_id].equal_range({nnodes, nedges});
     for (auto it = range.first; it != range.second; ++it)
     {
       auto& e = it->second;
+      // Executable graphs are only reused on the stream to which the cache
+      // entry is bound. In addition to preventing CUDA from serializing
+      // concurrent launches of one executable on different streams, this
+      // gives us an explicit completion check before the host-side update.
+      // The caller's stream is alive by definition, so probing it is safe;
+      // a caller stream in capture reads as busy (a query would invalidate
+      // the capture), falling through to a fresh instantiation.
+      if (e.stream_id != stream_id || query_stream_state(stream) != stream_state::idle)
+      {
+        continue;
+      }
+
       if (reserved::try_updating_executable_graph(*e.exec_g, g))
       {
         // update the last use index for the LRU algorithm
@@ -191,55 +222,93 @@ public:
     // If we maintain a cache, store the executable graph
     if (cache_size_limit != 0)
     {
-      cached_graphs[dev_id].insert({::std::make_pair(nnodes, nedges), entry(this, exec_g, footprint)});
+      cached_graphs[dev_id].insert(
+        {::std::make_pair(nnodes, nedges), entry(this, exec_g, stream, stream_id, footprint)});
       total_cache_footprint[dev_id] += footprint;
     }
 
     return ::cuda::std::make_pair(exec_g, false);
   }
 
 private:
-  void reclaim(int dev_id, size_t to_reclaim)
+  // The driver-assigned stream id: unique for the process lifetime, unlike
+  // the handle value (see entry::stream_id).
+  static unsigned long long stream_unique_id(cudaStream_t stream)
   {
-    size_t reclaimed = 0;
-
-    // Use a priority queue (min-heap) to track least recently used entries
-    using key_type = ::std::pair<size_t, size_t>;
-
-    auto& device_cache = cached_graphs[dev_id];
-
-    auto cmp = [&device_cache](const key_type& key_a, const key_type& key_b) {
-      auto iter_a = device_cache.find(key_a);
-      auto iter_b = device_cache.find(key_b);
-
-      // Directly compare last_use timestamps
-      return iter_a->second.last_use > iter_b->second.last_use;
-    };
+    unsigned long long id = 0;
+    cuda_safe_call(cudaStreamGetId(stream, &id));
+    return id;
+  }
 
-    // Priority queue storing keys, ordered by least recently used
-    ::std::priority_queue<key_type, ::std::vector<key_type>, decltype(cmp)> lru_queue(cmp);
+  enum class stream_state
+  {
+    idle,
+    busy,
+    unavailable
+  };
 
-    // Populate queue with keys from the cache
-    for (const auto& kv : device_cache)
+  // Probe a stream without ever throwing and without touching a capture:
+  // cudaStreamQuery on a capturing stream would invalidate that capture (a
+  // cross-thread hazard when reclaim probes another context's stream), so
+  // capture status is checked first with the capture-legal API. Errors from
+  // either call (e.g. a destroyed handle for an entry whose bound stream the
+  // cache does not own) read as `unavailable`: such an entry is neither
+  // reusable nor provably safe to destroy.
+  static stream_state query_stream_state(cudaStream_t stream)
+  {
+    cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone;
+    if (cudaStreamIsCapturing(stream, &capture) != cudaSuccess)
     {
-      lru_queue.push(kv.first);
+      cudaGetLastError();
+      return stream_state::unavailable;
     }
-
-    // Reclaim least recently used entries
-    while (!lru_queue.empty() && reclaimed < to_reclaim)
+    if (capture != cudaStreamCaptureStatusNone)
+    {
+      return stream_state::busy;
+    }
+    const cudaError_t status = cudaStreamQuery(stream);
+    if (status == cudaSuccess)
     {
-      key_type key = lru_queue.top();
-      lru_queue.pop();
+      return stream_state::idle;
+    }
+    cudaGetLastError();
+    return (status == cudaErrorNotReady) ? stream_state::busy : stream_state::unavailable;
+  }
 
-      // Find the entry before erasing
-      auto it = device_cache.find(key);
-      if (it != device_cache.end())
+  void reclaim(int dev_id, size_t to_reclaim)
+  {
+    size_t reclaimed   = 0;
+    auto& device_cache = cached_graphs[dev_id];
+
+    // Reclaim the least-recently-used idle entries. cudaGraphExecDestroy must
+    // not race an in-flight launch, so a busy entry remains cached even if
+    // that temporarily leaves the cache above its configured size. An
+    // `unavailable` entry (bound stream destroyed) is skipped too: its final
+    // launch may still be draining, so destroying the executable is not
+    // provably safe, and the entry stays as an unreclaimable zombie. This is
+    // benign when cache-bound streams outlive the cache (the pool streams
+    // handed out by async_resources_handle do); binding entries to streams
+    // with independent lifetimes is what makes zombies possible at all.
+    while (reclaimed < to_reclaim)
+    {
+      auto victim = device_cache.end();
+      for (auto it = device_cache.begin(); it != device_cache.end(); ++it)
       {
-        reclaimed += it->second.footprint;
-        total_cache_footprint[dev_id] -= it->second.footprint;
+        if (query_stream_state(it->second.stream) == stream_state::idle
+            && (victim == device_cache.end() || it->second.last_use < victim->second.last_use))
+        {
+          victim = it;
+        }
+      }
 
-        device_cache.erase(it);
+      if (victim == device_cache.end())
+      {
+        break;
       }
+
+      reclaimed += victim->second.footprint;
+      total_cache_footprint[dev_id] -= victim->second.footprint;
+      device_cache.erase(victim);
     }
   }
 
@@ -253,5 +322,10 @@ private:
   ::std::vector<size_t> total_cache_footprint;
 
   size_t cache_size_limit;
+
+  // A handle may be shared by multiple host threads. Serialize cache lookup,
+  // update, insertion, and reclaim so one executable cannot be updated by two
+  // queries concurrently.
+  ::std::mutex mutex;
 };
 } // namespace cuda::experimental::stf

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

@@ -685,8 +685,9 @@ public:
         cuda_try(cudaGraphGetEdges(graph, nullptr, nullptr, &nedges));
 #endif
 
-        auto [cached_exec, cache_hit] = ctx.async_resources().cached_graphs_query(nnodes, nedges, graph);
-        exec_graph_                   = mv(cached_exec);
+        auto [cached_exec,
+              cache_hit] = ctx.async_resources().cached_graphs_query(nnodes, nedges, graph, support_stream);
+        exec_graph_      = mv(cached_exec);
 
         auto* cache_stat = ctx.graph_get_cache_stat();
         if (cache_stat)

cudax/test/stf/CMakeLists.txt

@@ -114,6 +114,7 @@ set(
   reductions/sum_multiple_places_no_refvalue.cu
   slice/pinning.cu
   stackable/composite_conditional.cu
+  stackable/executable_graph_cache_streams.cu
   stackable/graph_scope_test.cu
   stencil/stencil-1D.cu
   stress/empty_tasks.cu

cudax/test/stf/graph/get_cache_stats.cu

@@ -19,33 +19,38 @@ using namespace cuda::experimental::stf;
 
 int main()
 {
-  async_resources_handle handle;
-  for (size_t i = 0; i < 10; i++)
+  cudaStream_t stream = cuda_try<cudaStreamCreate>();
   {
-    graph_ctx ctx(handle);
-    auto lA = ctx.logical_data(shape_of<slice<size_t>>(64));
-    ctx.launch(lA.write())->*[] _CCCL_DEVICE(auto t, slice<size_t> A) {
-      for (auto i : t.apply_partition(shape(A)))
+    async_resources_handle handle;
+    for (size_t i = 0; i < 10; i++)
+    {
+      graph_ctx ctx(stream, handle);
+      auto lA = ctx.logical_data(shape_of<slice<size_t>>(64));
+      ctx.launch(lA.write())->*[] _CCCL_DEVICE(auto t, slice<size_t> A) {
+        for (auto i : t.apply_partition(shape(A)))
+        {
+          A(i) = 2 * i;
+        }
+      };
+      ctx.finalize();
+      cuda_try(cudaStreamSynchronize(stream));
+
+      // Query statistics about the graph context : the first iteration needs to
+      // instantiate the graph, then we will reuse graphs saved in the handle.
+      auto* st = ctx.graph_get_cache_stat();
+      if (i == 0)
+      {
+        EXPECT(st->instantiate_cnt == 1);
+        EXPECT(st->update_cnt == 0);
+      }
+      else
       {
-        A(i) = 2 * i;
+        EXPECT(st->instantiate_cnt == 0);
+        EXPECT(st->update_cnt == 1);
       }
-    };
-    ctx.finalize();
 
-    // Query statistics about the graph context : the first iteration needs to
-    // instantiate the graph, then we will reuse graphs saved in the handle.
-    auto* st = ctx.graph_get_cache_stat();
-    if (i == 0)
-    {
-      EXPECT(st->instantiate_cnt == 1);
-      EXPECT(st->update_cnt == 0);
-    }
-    else
-    {
-      EXPECT(st->instantiate_cnt == 0);
-      EXPECT(st->update_cnt == 1);
+      // fprintf(stderr, "nnodes %ld nedges %ld\n", st->nnodes, st->nedges);
     }
-
-    // fprintf(stderr, "nnodes %ld nedges %ld\n", st->nnodes, st->nedges);
   }
+  cuda_try(cudaStreamDestroy(stream));
 }