Fix nvexec upon_stopped completion delivery (#2183)

* Fix nvexec upon_stopped completions

* working around nvhpc ICE and warnings

---------

Co-authored-by: Eric Niebler <eniebler@nvidia.com>

Author: fallintoplace

examples/nvexec/maxwell/common.cuh

@@ -39,6 +39,9 @@
 #  include "nvexec/detail/throw_on_cuda_error.cuh"
 #endif
 
+STDEXEC_PRAGMA_PUSH()
+STDEXEC_PRAGMA_IGNORE_EDG(is_constant_evaluated_in_nonconstexpr_context)
+
 struct deleter_t
 {
   bool on_gpu{};
@@ -538,3 +541,5 @@ auto value(std::map<std::string_view, std::size_t> const &params,
   }
   return default_value;
 }
+
+STDEXEC_PRAGMA_POP()

examples/nvexec/maxwell/stdpar.cuh

@@ -31,6 +31,9 @@
 
 #  include <algorithm>
 
+STDEXEC_PRAGMA_PUSH()
+STDEXEC_PRAGMA_IGNORE_EDG(is_constant_evaluated_in_nonconstexpr_context)
+
 template <class Policy>
 auto is_gpu_policy([[maybe_unused]] Policy&& policy) -> bool
 {
@@ -81,4 +84,6 @@ void run_stdpar(float            dt,
                      });
 }
 
+STDEXEC_PRAGMA_POP()
+
 #endif  // !STDEXEC_NO_STDCPP_PARALLEL_ALGORITHMS()

include/exec/static_thread_pool.hpp

@@ -62,6 +62,9 @@ import stdexec;
 #include "sequence/iterate.hpp"
 #include "sequence_senders.hpp"
 
+STDEXEC_PRAGMA_PUSH()
+STDEXEC_PRAGMA_IGNORE_EDG(is_constant_evaluated_in_nonconstexpr_context)
+
 namespace experimental::execution
 {
   struct bwos_params
@@ -1847,3 +1850,5 @@ namespace experimental::execution
 
 STDEXEC_MODULE_EXPORT
 namespace exec = experimental::execution;
+
+STDEXEC_PRAGMA_POP()

include/nvexec/stream/common.cuh

@@ -1013,6 +1013,13 @@ namespace nv::execution
     template <class Fun, class... Args>
     using __f = STDEXEC::__msize_t<_sizeof_v<STDEXEC::__call_result_t<Fun, Args...>>>;
   };
+
+  using _cuda_error_completion_t =
+    STDEXEC::completion_signatures<STDEXEC::set_error_t(cudaError_t)>;
+
+  template <class _NoExcept>
+  using _cuda_error_completion_unless_t =
+    STDEXEC::__if<_NoExcept, STDEXEC::completion_signatures<>, _cuda_error_completion_t>;
 }  // namespace nv::execution
 
 namespace nvexec = nv::execution;

include/nvexec/stream/upon_stopped.cuh

@@ -19,14 +19,16 @@
 #pragma once
 
 #include "../../stdexec/execution.hpp"
-#include <cstddef>
-#include <exception>
-#include <type_traits>
 
-#include <cuda/std/utility>
+#include "../../exec/completion_signatures.hpp"
 
 #include "common.cuh"
 
+#include <cuda/std/utility>
+
+#include <cstddef>
+#include <type_traits>
+
 STDEXEC_PRAGMA_PUSH()
 STDEXEC_PRAGMA_IGNORE_EDG(cuda_compile)
 
@@ -106,7 +108,7 @@ namespace nv::execution::_strm
               status == cudaSuccess)
           {
             opstate_.defer_temp_storage_destruction(d_result);
-            opstate_.propagate_completion_signal(STDEXEC::set_value, *d_result);
+            opstate_.propagate_completion_signal(STDEXEC::set_value, std::move(*d_result));
           }
           else
           {
@@ -125,28 +127,24 @@ namespace nv::execution::_strm
       Fun                            fun_;
       _strm::opstate_base<Receiver>& opstate_;
     };
+
+    template <class Fun>
+    consteval auto _get_completions_fun() noexcept
+    {
+      return []() noexcept
+      {
+        return __set_value_from_t<Fun>();
+      };
+    }
   }  // namespace _upon_stopped
 
   template <class Sender, class Fun>
   struct upon_stopped_sender : stream_sender_base
   {
     using sender_concept = STDEXEC::sender_tag;
-    using _set_error_t   = completion_signatures<set_error_t(std::exception_ptr)>;
 
     template <class Receiver>
-    using receiver_t = _upon_stopped::receiver<Receiver, Fun>;
-
-    template <class Self, class... Env>
-    using completion_signatures = __transform_completion_signatures_t<
-      __completion_signatures_of_t<__copy_cvref_t<Self, Sender>, Env...>,
-      __with_error_invoke_t<__mbind_front_q<__callable_error_t, upon_stopped_t>,
-                            set_stopped_t,
-                            Fun,
-                            __copy_cvref_t<Self, Sender>,
-                            Env...>,
-      __cmplsigs::__default_set_value,
-      __cmplsigs::__default_set_error,
-      __set_value_from_t<Fun>>;
+    using _receiver_t = _upon_stopped::receiver<Receiver, Fun>;
 
     explicit upon_stopped_sender(Sender sndr, Fun fun)
       noexcept(__nothrow_move_constructible<Sender, Fun>)
@@ -155,22 +153,26 @@ namespace nv::execution::_strm
     {}
 
     template <__decays_to<upon_stopped_sender> Self, STDEXEC::receiver Receiver>
-      requires receiver_of<Receiver, completion_signatures<Self, env_of_t<Receiver>>>
     STDEXEC_EXPLICIT_THIS_BEGIN(auto connect)(this Self&& self, Receiver rcvr)
-      -> stream_opstate_t<__copy_cvref_t<Self, Sender>, receiver_t<Receiver>, Receiver>
+      -> stream_opstate_t<__copy_cvref_t<Self, Sender>, _receiver_t<Receiver>, Receiver>
     {
       return stream_opstate<__copy_cvref_t<Self, Sender>>(
         static_cast<Self&&>(self).sndr_,
         static_cast<Receiver&&>(rcvr),
-        [&](_strm::opstate_base<Receiver>& stream_provider) -> receiver_t<Receiver>
-        { return receiver_t<Receiver>(self.fun_, stream_provider); });
+        [&](_strm::opstate_base<Receiver>& stream_provider) -> _receiver_t<Receiver>
+        { return _receiver_t<Receiver>(self.fun_, stream_provider); });
     }
     STDEXEC_EXPLICIT_THIS_END(connect)
 
     template <__decays_to<upon_stopped_sender> Self, class... Env>
-    static consteval auto get_completion_signatures() -> completion_signatures<Self, Env...>
+    static consteval auto get_completion_signatures()
     {
-      return {};
+      return exec::transform_completion_signatures(
+        STDEXEC::get_completion_signatures<__copy_cvref_t<Self, Sender>, Env...>(),
+        {},
+        {},
+        _upon_stopped::_get_completions_fun<Fun>(),
+        _cuda_error_completion_t());
     }
 
     auto get_env() const noexcept -> stream_sender_attrs<Sender>

test/nvexec/upon_stopped.cpp

@@ -1,5 +1,7 @@
 #include <stdexec/execution.hpp>
 #include <test_common/catch2.hpp>
+#include <test_common/senders.hpp>
+#include <test_common/type_helpers.hpp>
 
 #include "common.cuh"
 #include "nvexec/stream_context.cuh"
@@ -10,6 +12,45 @@ using nvexec::is_on_gpu;
 
 namespace
 {
+  struct move_only_result
+  {
+    STDEXEC_ATTRIBUTE(host, device)
+    explicit move_only_result(int value) noexcept
+      : value_(value)
+    {}
+
+    STDEXEC_ATTRIBUTE(host, device)
+    move_only_result(move_only_result&& other) noexcept
+      : value_(other.value_)
+    {
+      other.value_ = 0;
+    }
+
+    move_only_result(move_only_result const &) = delete;
+
+    STDEXEC_ATTRIBUTE(host, device)
+    ~move_only_result() = default;
+
+    STDEXEC_ATTRIBUTE(host, device)
+    auto value() const noexcept -> int
+    {
+      return value_;
+    }
+
+   private:
+    int value_;
+  };
+
+  TEST_CASE("nvexec upon_stopped advertises CUDA launch errors",
+            "[cuda][stream][adaptors][upon_stopped]")
+  {
+    auto fun = []() noexcept {};
+    using sender_t =
+      nvexec::_strm::upon_stopped_sender<a_sender_of<ex::set_stopped_t()>, decltype(fun)>;
+    sender_t snd{a_sender_of<ex::set_stopped_t()>{}, std::move(fun)};
+
+    check_err_types<ex::__mset<cudaError_t>>(snd);
+  }
 
   TEST_CASE("nvexec upon_stopped returns a sender", "[cuda][stream][adaptors][upon_stopped]")
   {
@@ -41,4 +82,16 @@ namespace
 
     REQUIRE(flags_storage.all_set_once());
   }
+
+  TEST_CASE("nvexec upon_stopped moves its result", "[cuda][stream][adaptors][upon_stopped]")
+  {
+    nvexec::stream_context stream_ctx{};
+
+    auto snd = ex::just_stopped() | ex::continues_on(stream_ctx.get_scheduler())
+             | ex::upon_stopped([] { return move_only_result{42}; });
+
+    auto [result] = STDEXEC::sync_wait(std::move(snd)).value();
+
+    REQUIRE(result.value() == 42);
+  }
 }  // namespace