diff --git a/cudax/include/cuda/experimental/__stf/graph/graph_ctx.cuh b/cudax/include/cuda/experimental/__stf/graph/graph_ctx.cuh index 3b687ee9da2f..5735749b9442 100644 --- a/cudax/include/cuda/experimental/__stf/graph/graph_ctx.cuh +++ b/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 diff --git a/cudax/include/cuda/experimental/__stf/internal/async_resources_handle.cuh b/cudax/include/cuda/experimental/__stf/internal/async_resources_handle.cuh index 3d641625dec5..0827d7956111 100644 --- a/cudax/include/cuda/experimental/__stf/internal/async_resources_handle.cuh +++ b/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, 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, bool> cached_graphs_query(cudaGraph_t g) + ::cuda::std::pair<::std::shared_ptr, bool> cached_graphs_query(cudaGraph_t g, cudaStream_t stream) { const size_t nnodes = cuda_try(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) diff --git a/cudax/include/cuda/experimental/__stf/internal/executable_graph_cache.cuh b/cudax/include/cuda/experimental/__stf/internal/executable_graph_cache.cuh index 71fb4adc5d65..d0c09fbd062c 100644 --- a/cudax/include/cuda/experimental/__stf/internal/executable_graph_cache.cuh +++ b/cudax/include/cuda/experimental/__stf/internal/executable_graph_cache.cuh @@ -30,7 +30,7 @@ #include #include -#include // for ::std::priority_queue +#include #include namespace cuda::experimental::stf @@ -119,12 +119,18 @@ public: // One entry of the cache struct entry { - entry(executable_graph_cache* cache, ::std::shared_ptr exec_g_, size_t footprint) + entry(executable_graph_cache* cache, + ::std::shared_ptr 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 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, bool> query(size_t nnodes, size_t nedges, cudaGraph_t g) + ::cuda::std::pair<::std::shared_ptr, 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(); _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,7 +222,8 @@ 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; } @@ -199,47 +231,84 @@ public: } 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; - - 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, 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 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 diff --git a/cudax/include/cuda/experimental/__stf/stackable/stackable_ctx_impl.cuh b/cudax/include/cuda/experimental/__stf/stackable/stackable_ctx_impl.cuh index f4473fc86a62..f1160562bd62 100644 --- a/cudax/include/cuda/experimental/__stf/stackable/stackable_ctx_impl.cuh +++ b/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) diff --git a/cudax/test/stf/CMakeLists.txt b/cudax/test/stf/CMakeLists.txt index 4125830d5ba0..aeb6a35d7c33 100644 --- a/cudax/test/stf/CMakeLists.txt +++ b/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 diff --git a/cudax/test/stf/graph/get_cache_stats.cu b/cudax/test/stf/graph/get_cache_stats.cu index 37b99ec3ba58..a080b21717d5 100644 --- a/cudax/test/stf/graph/get_cache_stats.cu +++ b/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(); { - graph_ctx ctx(handle); - auto lA = ctx.logical_data(shape_of>(64)); - ctx.launch(lA.write())->*[] _CCCL_DEVICE(auto t, slice 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>(64)); + ctx.launch(lA.write())->*[] _CCCL_DEVICE(auto t, slice 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)); } diff --git a/cudax/test/stf/graph/graph_cache_policy.cu b/cudax/test/stf/graph/graph_cache_policy.cu index 1cc663422b09..644646208590 100644 --- a/cudax/test/stf/graph/graph_cache_policy.cu +++ b/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(); { - 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>(64)); - ctx.launch(lA.write())->*[] _CCCL_DEVICE(auto t, slice 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>(64)); + ctx.launch(lA.write())->*[] _CCCL_DEVICE(auto t, slice 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)); } diff --git a/cudax/test/stf/stackable/executable_graph_cache_streams.cu b/cudax/test/stf/stackable/executable_graph_cache_streams.cu new file mode 100644 index 000000000000..ba1ab890d3a6 --- /dev/null +++ b/cudax/test/stf/stackable/executable_graph_cache_streams.cu @@ -0,0 +1,150 @@ +//===----------------------------------------------------------------------===// +// +// 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) 2026 NVIDIA CORPORATION & AFFILIATES. +// +//===----------------------------------------------------------------------===// + +#include + +#include +#include + +using namespace cuda::experimental::stf; + +__global__ void keep_stream_busy(unsigned long long cycles) +{ + const unsigned long long start = clock64(); + while (clock64() - start < cycles) + { + } +} + +void test_cache_stream_affinity() +{ + cudaStream_t stream_a = cuda_try(); + cudaStream_t stream_b = cuda_try(); + + { + async_resources_handle handle; + cudaGraph_t graph_a = cuda_try(0); + cudaGraph_t graph_b = cuda_try(0); + + cuda_try(graph_a, nullptr, 0); + cuda_try(graph_b, nullptr, 0); + + auto [exec_a, hit_a] = handle.cached_graphs_query(1, 0, graph_a, stream_a); + EXPECT(!hit_a); + + auto [exec_a_again, hit_a_again] = handle.cached_graphs_query(1, 0, graph_b, stream_a); + EXPECT(hit_a_again); + EXPECT(exec_a_again == exec_a); + + auto [exec_b, hit_b] = handle.cached_graphs_query(1, 0, graph_b, stream_b); + EXPECT(!hit_b); + EXPECT(exec_b != exec_a); + + keep_stream_busy<<<1, 1, 0, stream_a>>>(100000000); + cuda_try(cudaPeekAtLastError()); + + auto [exec_busy, hit_busy] = handle.cached_graphs_query(1, 0, graph_b, stream_a); + EXPECT(!hit_busy); + EXPECT(exec_busy != exec_a); + + cuda_try(cudaStreamSynchronize(stream_a)); + cuda_try(cudaGraphDestroy(graph_b)); + cuda_try(cudaGraphDestroy(graph_a)); + } + + cuda_try(cudaStreamDestroy(stream_b)); + cuda_try(cudaStreamDestroy(stream_a)); +} + +template +void build_scope( + stackable_ctx& ctx, + Data data, + int value, + int* active, + int* max_active, + ::std::atomic& ready, + ::std::atomic& launch) +{ + ctx.set_head_offset(ctx.get_root_offset()); + auto scope = ctx.graph_scope(); + ctx.parallel_for(data.shape(), data.write())->*[active, max_active, value] __device__(size_t i, auto data) { + if (i == 0) + { + const int concurrent = atomicAdd(active, 1) + 1; + atomicMax(max_active, concurrent); + const unsigned long long start = clock64(); + while (clock64() - start < 200000000) + { + } + atomicSub(active, 1); + } + data(i) = value; + }; + + ready.fetch_add(1, ::std::memory_order_release); + while (!launch.load(::std::memory_order_acquire)) + { + } +} + +void test_sibling_graph_scopes_overlap() +{ + constexpr size_t count = 128; + int output_a[count] = {}; + int output_b[count] = {}; + + int* active = nullptr; + int* max_active = nullptr; + cuda_try(cudaMallocManaged(&active, sizeof(int))); + cuda_try(cudaMallocManaged(&max_active, sizeof(int))); + *active = 0; + *max_active = 0; + + stackable_ctx ctx; + auto data_a = ctx.logical_data(output_a); + auto data_b = ctx.logical_data(output_b); + + ::std::atomic ready{0}; + ::std::atomic launch{false}; + + ::std::thread thread_a([&] { + build_scope(ctx, data_a, 17, active, max_active, ready, launch); + }); + ::std::thread thread_b([&] { + build_scope(ctx, data_b, 29, active, max_active, ready, launch); + }); + + while (ready.load(::std::memory_order_acquire) != 2) + { + } + launch.store(true, ::std::memory_order_release); + + thread_a.join(); + thread_b.join(); + + ctx.finalize(); + + EXPECT(*max_active == 2); + for (size_t i = 0; i < count; ++i) + { + EXPECT(output_a[i] == 17); + EXPECT(output_b[i] == 29); + } + + cuda_try(cudaFree(max_active)); + cuda_try(cudaFree(active)); +} + +int main() +{ + test_cache_stream_affinity(); + test_sibling_graph_scopes_overlap(); +}