Skip to content

nvexec: forward reference results from then, upon_error and upon_stopped #2410

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ericniebler merged 1 commit into from
Nov 5, 2026
Merged

nvexec: forward reference results from then, upon_error and upon_stopped #2410

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 53 additions & 1 deletion include/nvexec/stream/common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1023,10 +1023,62 @@ namespace nv::execution
using __f = STDEXEC::__msize_t<STDEXEC::__umax({std::size_t(0), Sizes::value...})>;
};

// 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 <class Ty>
using _stored_result_t =
STDEXEC::__if_c<std::is_reference_v<Ty>, std::add_pointer_t<Ty>, STDEXEC::__decay_t<Ty>>;

template <class ResultT, class Fun, class... Args>
STDEXEC_ATTRIBUTE(host, device)
void _store_result(_stored_result_t<ResultT>* storage, Fun&& fun, Args&&... args)
{
if constexpr (std::is_reference_v<ResultT>)
{
new (storage) _stored_result_t<ResultT>(
__builtin_addressof(static_cast<Fun&&>(fun)(static_cast<Args&&>(args)...)));
}
else
{
new (storage)
_stored_result_t<ResultT>(static_cast<Fun&&>(fun)(static_cast<Args&&>(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 <class ResultT>
auto _wait_for_result(cudaStream_t stream) noexcept -> cudaError_t
{
if constexpr (std::is_reference_v<ResultT>)
{
return STDEXEC_LOG_CUDA_API(cudaStreamSynchronize(stream));
}
else
{
return cudaSuccess;
}
}

// The stored result as it is passed to the next receiver.
template <class ResultT>
auto _load_result(_stored_result_t<ResultT>* storage) noexcept -> decltype(auto)
{
if constexpr (std::is_reference_v<ResultT>)
{
return static_cast<ResultT>(**storage);
}
else
{
return std::move(*storage);
}
}

struct result_size_for
{
template <class Fun, class... Args>
using __f = STDEXEC::__msize_t<_sizeof_v<STDEXEC::__call_result_t<Fun, Args...>>>;
using __f =
STDEXEC::__msize_t<_sizeof_v<_stored_result_t<STDEXEC::__call_result_t<Fun, Args...>>>>;
};

using _cuda_error_completion_t =
Expand Down
27 changes: 16 additions & 11 deletions include/nvexec/stream/then.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,13 @@ namespace nv::execution::_strm
::cuda::std::move(fn)(static_cast<Args&&>(args)...);
}

template <class... Args, class Fun, class ResultT>
template <class ResultT, class... Args, class Fun>
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<ResultT>* result, Args... args)
{
static_assert(trivially_copyable<Fun, Args...>);
new (result) ResultT(::cuda::std::move(fn)(static_cast<Args&&>(args)...));
_store_result<ResultT>(result, ::cuda::std::move(fn), static_cast<Args&&>(args)...);
}

template <std::size_t MemoryAllocationSize, class Receiver, class Fun>
Expand All @@ -70,10 +71,10 @@ namespace nv::execution::_strm
{}

template <class... Args>
requires std::invocable<Fun, __decay_t<Args>...>
requires std::invocable<Fun, Args...>
void set_value(Args&&... args) noexcept
{
using result_t = std::invoke_result_t<Fun, __decay_t<Args>...>;
using result_t = std::invoke_result_t<Fun, Args...>;
constexpr bool does_not_return_a_value = std::is_same_v<void, result_t>;
cudaStream_t stream = opstate_.get_stream();

Expand All @@ -93,16 +94,20 @@ namespace nv::execution::_strm
}
else
{
using decayed_result_t = __decay_t<result_t>;
auto* d_result = static_cast<decayed_result_t*>(opstate_.temp_storage_);
_then_kernel_with_result<Args&&...>
auto* d_result = static_cast<_stored_result_t<result_t>*>(opstate_.temp_storage_);
_then_kernel_with_result<result_t, Args&&...>
<<<1, 1, 0, stream>>>(std::move(f_), d_result, static_cast<Args&&>(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<result_t>(stream);
}
if (status == cudaSuccess)
{
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
{
Expand Down
23 changes: 14 additions & 9 deletions include/nvexec/stream/upon_error.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,12 +46,13 @@ namespace nv::execution::_strm
::cuda::std::move(fn)(static_cast<Args&&>(args)...);
}

template <class... Args, class Fun, class ResultT>
template <class ResultT, class... Args, class Fun>
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<ResultT>* result, Args... args)
{
static_assert(trivially_copyable<Fun, Args...>);
new (result) ResultT(::cuda::std::move(fn)(static_cast<Args&&>(args)...));
_store_result<ResultT>(result, ::cuda::std::move(fn), static_cast<Args&&>(args)...);
}

template <std::size_t MemoryAllocationSize, class Receiver, class Fun>
Expand Down Expand Up @@ -100,15 +101,19 @@ namespace nv::execution::_strm
}
else
{
using decayed_result_t = __decay_t<result_t>;
auto* d_result = static_cast<decayed_result_t*>(opstate_.temp_storage_);
_upon_error_kernel_with_result<Error&&>
auto* d_result = static_cast<_stored_result_t<result_t>*>(opstate_.temp_storage_);
_upon_error_kernel_with_result<result_t, Error&&>
<<<1, 1, 0, stream>>>(std::move(f_), d_result, static_cast<Error&&>(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<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
{
Expand Down
25 changes: 15 additions & 10 deletions include/nvexec/stream/upon_stopped.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,12 +44,12 @@ namespace nv::execution::_strm
::cuda::std::move(fn)();
}

template <class Fun, class ResultT>
template <class ResultT, class Fun>
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<ResultT>* result)
{
static_assert(trivially_copyable<Fun>);
new (result) ResultT(::cuda::std::move(fn)());
_store_result<ResultT>(result, ::cuda::std::move(fn));
}

template <class Receiver, class Fun>
Expand All @@ -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<Receiver>& opstate)
Expand Down Expand Up @@ -101,14 +101,19 @@ namespace nv::execution::_strm
}
else
{
using decayed_result_t = __decay_t<_result_t>;
auto* d_result = static_cast<decayed_result_t*>(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
{
Expand Down
34 changes: 34 additions & 0 deletions test/nvexec/then.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
21 changes: 21 additions & 0 deletions test/nvexec/upon_error.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
18 changes: 18 additions & 0 deletions test/nvexec/upon_stopped.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading