From b5fce104c40bd85dc09f97f65cc2cfbb27e5e17b Mon Sep 17 00:00:00 2001 From: alwaysprince05 Date: Thu, 5 Nov 2026 05:28:16 +0530 Subject: [PATCH] nvexec: forward reference results from then, upon_error and upon_stopped When the callable returns a reference, the kernel now stores a pointer to the referenced object in temporary storage, and the reference is passed on from that pointer. Before, the result was decay-copied, so the receiver got a prvalue while the sender advertised set_value_t(T&), and code like then(f -> int&) | then(g(int&)) did not compile. The address is only known after the kernel has run, so a reference result waits for the stream (cudaStreamSynchronize) before it is passed on. Value results are unchanged. then's set_value constraint now checks the callable against the arguments as they are passed, like upon_error already does. Fixes #958 --- include/nvexec/stream/common.cuh | 54 +++++++++++++++++++++++++- include/nvexec/stream/then.cuh | 27 +++++++------ include/nvexec/stream/upon_error.cuh | 23 ++++++----- include/nvexec/stream/upon_stopped.cuh | 25 +++++++----- test/nvexec/then.cpp | 34 ++++++++++++++++ test/nvexec/upon_error.cpp | 21 ++++++++++ test/nvexec/upon_stopped.cpp | 18 +++++++++ 7 files changed, 171 insertions(+), 31 deletions(-) diff --git a/include/nvexec/stream/common.cuh b/include/nvexec/stream/common.cuh index d715da9b6..6d4332814 100644 --- a/include/nvexec/stream/common.cuh +++ b/include/nvexec/stream/common.cuh @@ -1023,10 +1023,62 @@ namespace nv::execution using __f = STDEXEC::__msize_t; }; + // How a callable's result is kept in temporary storage: a reference result is + // kept as a pointer to the referenced object, anything else as a decayed value. + template + using _stored_result_t = + STDEXEC::__if_c, std::add_pointer_t, STDEXEC::__decay_t>; + + template + STDEXEC_ATTRIBUTE(host, device) + void _store_result(_stored_result_t* storage, Fun&& fun, Args&&... args) + { + if constexpr (std::is_reference_v) + { + new (storage) _stored_result_t( + __builtin_addressof(static_cast(fun)(static_cast(args)...))); + } + else + { + new (storage) + _stored_result_t(static_cast(fun)(static_cast(args)...)); + } + } + + // A referenced object's address is only known once the kernel has run, so a + // reference result needs the stream to finish before it can be passed on. + template + auto _wait_for_result(cudaStream_t stream) noexcept -> cudaError_t + { + if constexpr (std::is_reference_v) + { + return STDEXEC_LOG_CUDA_API(cudaStreamSynchronize(stream)); + } + else + { + return cudaSuccess; + } + } + + // The stored result as it is passed to the next receiver. + template + auto _load_result(_stored_result_t* storage) noexcept -> decltype(auto) + { + if constexpr (std::is_reference_v) + { + return static_cast(**storage); + } + else + { + return std::move(*storage); + } + } + struct result_size_for { template - using __f = STDEXEC::__msize_t<_sizeof_v>>; + using __f = + STDEXEC::__msize_t<_sizeof_v<_stored_result_t>>>; }; using _cuda_error_completion_t = diff --git a/include/nvexec/stream/then.cuh b/include/nvexec/stream/then.cuh index 4f7b0dcae..084fbdf10 100644 --- a/include/nvexec/stream/then.cuh +++ b/include/nvexec/stream/then.cuh @@ -47,12 +47,13 @@ namespace nv::execution::_strm ::cuda::std::move(fn)(static_cast(args)...); } - template + template STDEXEC_ATTRIBUTE(launch_bounds(1)) - __global__ void _then_kernel_with_result(Fun fn, ResultT* result, Args... args) + __global__ + void _then_kernel_with_result(Fun fn, _stored_result_t* result, Args... args) { static_assert(trivially_copyable); - new (result) ResultT(::cuda::std::move(fn)(static_cast(args)...)); + _store_result(result, ::cuda::std::move(fn), static_cast(args)...); } template @@ -70,10 +71,10 @@ namespace nv::execution::_strm {} template - requires std::invocable...> + requires std::invocable void set_value(Args&&... args) noexcept { - using result_t = std::invoke_result_t...>; + using result_t = std::invoke_result_t; constexpr bool does_not_return_a_value = std::is_same_v; cudaStream_t stream = opstate_.get_stream(); @@ -93,16 +94,20 @@ namespace nv::execution::_strm } else { - using decayed_result_t = __decay_t; - auto* d_result = static_cast(opstate_.temp_storage_); - _then_kernel_with_result + auto* d_result = static_cast<_stored_result_t*>(opstate_.temp_storage_); + _then_kernel_with_result <<<1, 1, 0, stream>>>(std::move(f_), d_result, static_cast(args)...); opstate_.defer_temp_storage_destruction(d_result); - if (cudaError_t status = STDEXEC_LOG_CUDA_API(cudaPeekAtLastError()); - status == cudaSuccess) + cudaError_t status = STDEXEC_LOG_CUDA_API(cudaPeekAtLastError()); + if (status == cudaSuccess) + { + status = _wait_for_result(stream); + } + if (status == cudaSuccess) { - opstate_.propagate_completion_signal(STDEXEC::set_value, std::move(*d_result)); + opstate_.propagate_completion_signal(STDEXEC::set_value, + _load_result(d_result)); } else { diff --git a/include/nvexec/stream/upon_error.cuh b/include/nvexec/stream/upon_error.cuh index b994f08f0..2086164b1 100644 --- a/include/nvexec/stream/upon_error.cuh +++ b/include/nvexec/stream/upon_error.cuh @@ -46,12 +46,13 @@ namespace nv::execution::_strm ::cuda::std::move(fn)(static_cast(args)...); } - template + template STDEXEC_ATTRIBUTE(launch_bounds(1)) - __global__ void _upon_error_kernel_with_result(Fun fn, ResultT* result, Args... args) + __global__ + void _upon_error_kernel_with_result(Fun fn, _stored_result_t* result, Args... args) { static_assert(trivially_copyable); - new (result) ResultT(::cuda::std::move(fn)(static_cast(args)...)); + _store_result(result, ::cuda::std::move(fn), static_cast(args)...); } template @@ -100,15 +101,19 @@ namespace nv::execution::_strm } else { - using decayed_result_t = __decay_t; - auto* d_result = static_cast(opstate_.temp_storage_); - _upon_error_kernel_with_result + auto* d_result = static_cast<_stored_result_t*>(opstate_.temp_storage_); + _upon_error_kernel_with_result <<<1, 1, 0, stream>>>(std::move(f_), d_result, static_cast(error)); - if (cudaError_t status = STDEXEC_LOG_CUDA_API(cudaPeekAtLastError()); - status == cudaSuccess) + cudaError_t status = STDEXEC_LOG_CUDA_API(cudaPeekAtLastError()); + if (status == cudaSuccess) + { + status = _wait_for_result(stream); + } + if (status == cudaSuccess) { opstate_.defer_temp_storage_destruction(d_result); - opstate_.propagate_completion_signal(STDEXEC::set_value, std::move(*d_result)); + opstate_.propagate_completion_signal(STDEXEC::set_value, + _load_result(d_result)); } else { diff --git a/include/nvexec/stream/upon_stopped.cuh b/include/nvexec/stream/upon_stopped.cuh index ffc4118d4..20ea281ff 100644 --- a/include/nvexec/stream/upon_stopped.cuh +++ b/include/nvexec/stream/upon_stopped.cuh @@ -44,12 +44,12 @@ namespace nv::execution::_strm ::cuda::std::move(fn)(); } - template + template STDEXEC_ATTRIBUTE(launch_bounds(1)) - __global__ void _upon_stopped_kernel_with_result(Fun fn, ResultT* result) + __global__ void _upon_stopped_kernel_with_result(Fun fn, _stored_result_t* result) { static_assert(trivially_copyable); - new (result) ResultT(::cuda::std::move(fn)()); + _store_result(result, ::cuda::std::move(fn)); } template @@ -61,7 +61,7 @@ namespace nv::execution::_strm static constexpr std::size_t memory_allocation_size() noexcept { - return _sizeof_v<_result_t>; + return _sizeof_v<_stored_result_t<_result_t>>; } explicit receiver(Fun fun, _strm::opstate_base& opstate) @@ -101,14 +101,19 @@ namespace nv::execution::_strm } else { - using decayed_result_t = __decay_t<_result_t>; - auto* d_result = static_cast(opstate_.temp_storage_); - _upon_stopped_kernel_with_result<<<1, 1, 0, stream>>>(std::move(fun_), d_result); - if (cudaError_t status = STDEXEC_LOG_CUDA_API(cudaPeekAtLastError()); - status == cudaSuccess) + auto* d_result = static_cast<_stored_result_t<_result_t>*>(opstate_.temp_storage_); + _upon_stopped_kernel_with_result<_result_t> + <<<1, 1, 0, stream>>>(std::move(fun_), d_result); + cudaError_t status = STDEXEC_LOG_CUDA_API(cudaPeekAtLastError()); + if (status == cudaSuccess) + { + status = _wait_for_result<_result_t>(stream); + } + if (status == cudaSuccess) { opstate_.defer_temp_storage_destruction(d_result); - opstate_.propagate_completion_signal(STDEXEC::set_value, std::move(*d_result)); + opstate_.propagate_completion_signal(STDEXEC::set_value, + _load_result<_result_t>(d_result)); } else { diff --git a/test/nvexec/then.cpp b/test/nvexec/then.cpp index f19e25a6f..ef76dab95 100644 --- a/test/nvexec/then.cpp +++ b/test/nvexec/then.cpp @@ -371,4 +371,38 @@ namespace REQUIRE(handle.alive() == 0); } + + TEST_CASE("nvexec then can return a reference", "[cuda][stream][adaptors][then]") + { + nvexec::stream_context stream_ctx{}; + + int *data{}; + STDEXEC_TRY_CUDA_API(cudaMallocManaged(&data, sizeof(int))); + *data = 0; + + auto snd = ex::schedule(stream_ctx.get_scheduler()) + | ex::then([data]() -> int & { return *data; }) | ex::then([](int &ref) { ref = 7; }); + STDEXEC::sync_wait(std::move(snd)); + + int const result = *data; + STDEXEC_ASSERT_CUDA_API(cudaFree(data)); + REQUIRE(result == 7); + } + + TEST_CASE("nvexec then can return a const reference", "[cuda][stream][adaptors][then]") + { + nvexec::stream_context stream_ctx{}; + + int *data{}; + STDEXEC_TRY_CUDA_API(cudaMallocManaged(&data, sizeof(int))); + *data = 5; + + auto snd = ex::schedule(stream_ctx.get_scheduler()) + | ex::then([data]() -> int const & { return *data; }) + | ex::then([](int const &ref) { return ref * 2; }); + auto [result] = STDEXEC::sync_wait(std::move(snd)).value(); + + STDEXEC_ASSERT_CUDA_API(cudaFree(data)); + REQUIRE(result == 10); + } } // namespace diff --git a/test/nvexec/upon_error.cpp b/test/nvexec/upon_error.cpp index 8d049571b..709c750f7 100644 --- a/test/nvexec/upon_error.cpp +++ b/test/nvexec/upon_error.cpp @@ -125,4 +125,25 @@ namespace REQUIRE(flags_storage.all_set_once()); } + + TEST_CASE("nvexec upon_error can return a reference", "[cuda][stream][adaptors][upon_error]") + { + nvexec::stream_context stream_ctx{}; + + int *data{}; + STDEXEC_TRY_CUDA_API(cudaMallocManaged(&data, 2 * sizeof(int))); + data[0] = 0; + data[1] = 0; + + auto snd = ex::just_error(1) | ex::continues_on(stream_ctx.get_scheduler()) + | ex::upon_error([data](int i) -> int & { return data[i]; }) + | ex::then([](int &ref) { ref = 9; }); + STDEXEC::sync_wait(std::move(snd)); + + int const first = data[0]; + int const second = data[1]; + STDEXEC_ASSERT_CUDA_API(cudaFree(data)); + REQUIRE(first == 0); + REQUIRE(second == 9); + } } // namespace diff --git a/test/nvexec/upon_stopped.cpp b/test/nvexec/upon_stopped.cpp index 1cac5a096..5972b4bca 100644 --- a/test/nvexec/upon_stopped.cpp +++ b/test/nvexec/upon_stopped.cpp @@ -125,4 +125,22 @@ namespace REQUIRE(result.value() == 42); } + + TEST_CASE("nvexec upon_stopped can return a reference", "[cuda][stream][adaptors][upon_stopped]") + { + nvexec::stream_context stream_ctx{}; + + int *data{}; + STDEXEC_TRY_CUDA_API(cudaMallocManaged(&data, sizeof(int))); + *data = 0; + + auto snd = ex::just_stopped() | ex::continues_on(stream_ctx.get_scheduler()) + | ex::upon_stopped([data]() -> int & { return *data; }) + | ex::then([](int &ref) { ref = 11; }); + STDEXEC::sync_wait(std::move(snd)); + + int const result = *data; + STDEXEC_ASSERT_CUDA_API(cudaFree(data)); + REQUIRE(result == 11); + } } // namespace