[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,9 +119,11 @@ 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_, size_t footprint)
         : cache(cache)
         , exec_g(mv(exec_g_))
+        , stream(stream_)
         , footprint(footprint)
     {
       last_use = cache->index;
@@ -135,6 +137,7 @@ public:
 
     executable_graph_cache* cache;
     ::std::shared_ptr<cudaGraphExec_t> exec_g;
+    cudaStream_t stream;
     size_t last_use;
     size_t footprint;
   };
@@ -156,15 +159,27 @@ 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");
 
     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.
+      if (e.stream != stream || !stream_is_idle(stream))
+      {
+        continue;
+      }
+
       if (reserved::try_updating_executable_graph(*e.exec_g, g))
       {
         // update the last use index for the LRU algorithm
@@ -191,55 +206,58 @@ 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, 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)
+  static bool stream_is_idle(cudaStream_t stream)
   {
-    size_t reclaimed = 0;
+    const cudaError_t status = cudaStreamQuery(stream);
+    if (status == cudaSuccess)
+    {
+      return true;
+    }
+    if (status == cudaErrorNotReady)
+    {
+      return false;
+    }
 
-    // Use a priority queue (min-heap) to track least recently used entries
-    using key_type = ::std::pair<size_t, size_t>;
+    cuda_try(status);
+    return false;
+  }
 
+  void reclaim(int dev_id, size_t to_reclaim)
+  {
+    size_t reclaimed   = 0;
     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;
-    };
-
-    // Priority queue storing keys, ordered by least recently used
-    ::std::priority_queue<key_type, ::std::vector<key_type>, decltype(cmp)> lru_queue(cmp);
-
-    // Populate queue with keys from the cache
-    for (const auto& kv : device_cache)
-    {
-      lru_queue.push(kv.first);
-    }
-
-    // Reclaim least recently used entries
-    while (!lru_queue.empty() && reclaimed < to_reclaim)
+    // 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.
+    while (reclaimed < to_reclaim)
     {
-      key_type key = lru_queue.top();
-      lru_queue.pop();
-
-      // Find the entry before erasing
-      auto it = device_cache.find(key);
-      if (it != device_cache.end())
+      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 (stream_is_idle(it->second.stream)
+            && (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 +271,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));
 }

cudax/test/stf/graph/graph_cache_policy.cu

@@ -19,39 +19,44 @@ 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);
-
-    // If i is a multiple of 3 we enable the cache, the first iteration will fill the cache
-    ctx.set_graph_cache_policy([i]() {
-      return (i % 3) == 0;
-    });
-
-    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);
+
+      // If i is a multiple of 3 we enable the cache, the first iteration will fill the cache
+      ctx.set_graph_cache_policy([i]() {
+        return (i % 3) == 0;
+      });
+
+      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();
+
+      // For the first iteration, or non multiple of 3 we have to instantiate, otherwise we should have a cache hit
+      if (i == 0 || (i % 3) != 0)
       {
-        A(i) = 2 * i;
+        EXPECT(st->instantiate_cnt == 1);
+        EXPECT(st->update_cnt == 0);
+      }
+      else
+      {
+        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();
-
-    // For the first iteration, or non multiple of 3 we have to instantiate, otherwise we should have a cache hit
-    if (i == 0 || (i % 3) != 0)
-    {
-      EXPECT(st->instantiate_cnt == 1);
-      EXPECT(st->update_cnt == 0);
-    }
-    else
-    {
-      EXPECT(st->instantiate_cnt == 0);
-      EXPECT(st->update_cnt == 1);
     }
   }
+  cuda_try(cudaStreamDestroy(stream));
 }