implement proposed resolution of P3718 (#1542)

Author: ericniebler

include/nvexec/nvtx.cuh

@@ -121,8 +121,8 @@ namespace nvexec {
           return {};
         }
 
-        auto get_env() const noexcept -> env_of_t<const Sender&> {
-          return stdexec::get_env(sndr_);
+        auto get_env() const noexcept -> stream_sender_attrs<Sender> {
+          return {&sndr_};
         }
       };
     };

include/nvexec/stream/bulk.cuh

@@ -147,8 +147,8 @@ namespace nvexec::_strm {
         return {};
       }
 
-      auto get_env() const noexcept -> env_of_t<const Sender&> {
-        return stdexec::get_env(sndr_);
+      auto get_env() const noexcept -> stream_sender_attrs<Sender> {
+        return {&sndr_};
       }
     };
   };
@@ -383,8 +383,8 @@ namespace nvexec::_strm {
         return {};
       }
 
-      auto get_env() const noexcept -> env_of_t<const Sender&> {
-        return stdexec::get_env(sndr_);
+      auto get_env() const noexcept -> stream_sender_attrs<Sender> {
+        return {&sndr_};
       }
     };
   };

include/nvexec/stream/common.cuh

@@ -68,12 +68,39 @@ namespace nvexec {
   inline STDEXEC_ATTRIBUTE((host, device)) auto is_on_gpu() noexcept -> bool {
     return get_device_type() == device_type::device;
   }
+
+  namespace _strm {
+    // Used by stream_domain to late-customize senders for execution
+    // on the stream_scheduler.
+    template <class Tag, class... Env>
+    struct transform_sender_for;
+
+    template <class Tag>
+    struct apply_sender_for;
+  } // namespace _strm
 } // namespace nvexec
 
 namespace nvexec {
   struct stream_context;
 
-  struct stream_domain;
+  // The stream_domain is how the stream scheduler customizes the sender algorithms. All of the
+  // algorithms use the current scheduler's domain to transform senders before starting them.
+  struct stream_domain : stdexec::default_domain {
+    template <stdexec::sender_expr Sender, class Tag = stdexec::tag_of_t<Sender>, class... Env>
+      requires stdexec::
+        __callable<stdexec::__sexpr_apply_t, Sender, _strm::transform_sender_for<Tag, Env...>>
+      static auto transform_sender(Sender&& sndr, const Env&... env) {
+      return stdexec::__sexpr_apply(
+        static_cast<Sender&&>(sndr), _strm::transform_sender_for<Tag, Env...>{env...});
+    }
+
+    template <class Tag, stdexec::sender Sender, class... Args>
+      requires stdexec::__callable<_strm::apply_sender_for<Tag>, Sender, Args...>
+    static auto apply_sender(Tag, Sender&& sndr, Args&&... args) {
+      return _strm::apply_sender_for<Tag>{}(
+        static_cast<Sender&&>(sndr), static_cast<Args&&>(args)...);
+    }
+  };
 
   namespace _strm {
 
@@ -86,14 +113,6 @@ namespace nvexec {
       ((STDEXEC_IS_TRIVIALLY_COPYABLE(Ts) || std::is_reference_v<Ts>) && ...);
 #endif
 
-    // Used by stream_domain to late-customize senders for execution
-    // on the stream_scheduler.
-    template <class Tag, class... Env>
-    struct transform_sender_for;
-
-    template <class Tag>
-    struct apply_sender_for;
-
     inline auto get_stream_priority(stream_priority priority) -> std::pair<int, cudaError_t> {
       int least{};
       int greatest{};
@@ -336,6 +355,26 @@ namespace nvexec {
       }
     };
 
+    template <class Sender>
+    struct stream_sender_attrs {
+      using __t = stream_sender_attrs;
+      using __id = stream_sender_attrs;
+
+      STDEXEC_ATTRIBUTE((nodiscard)) constexpr auto query(get_domain_late_t) const noexcept -> stream_domain {
+        return {};
+      }
+
+      template <__forwarding_query Query>
+        requires __env::__queryable<env_of_t<Sender>, Query>
+      STDEXEC_ATTRIBUTE((nodiscard)) constexpr auto query(Query) const //
+        noexcept(__env::__nothrow_queryable<env_of_t<Sender>, Query>)
+          -> __env::__query_result_t<env_of_t<Sender>, Query> {
+        return stdexec::get_env(*child_).query(Query{});
+      }
+
+      const Sender* child_{};
+    };
+
     template <class BaseEnv>
     auto make_stream_env(BaseEnv&& base_env, stream_provider_t* stream_provider) noexcept {
       return __env::__join(

include/nvexec/stream/continues_on.cuh

@@ -116,9 +116,10 @@ namespace nvexec::_strm {
     };
   } // namespace _continues_on
 
-  template <class SenderId>
+  template <class Scheduler, class SenderId>
   struct continues_on_sender_t {
     using Sender = stdexec::__t<SenderId>;
+    using LateDomain = __detail::__early_domain_of_t<Sender, __none_such>;
 
     struct __t : stream_sender_base {
       using __id = continues_on_sender_t;
@@ -128,6 +129,7 @@ namespace nvexec::_strm {
         stdexec::__t<
           _continues_on::operation_state_t<__cvref_id<Self, Sender>, stdexec::__id<Receiver>>>;
 
+      Scheduler sched_;
       context_state_t context_state_;
       Sender sndr_;
 
@@ -155,17 +157,18 @@ namespace nvexec::_strm {
       }
 
       template <__decays_to<__t> Self, class... Env>
-      static auto
-        get_completion_signatures(Self&&, Env&&...) -> _completion_signatures_t<Self, Env...> {
+      static auto get_completion_signatures(Self&&, Env&&...) //
+        -> _completion_signatures_t<Self, Env...> {
         return {};
       }
 
-      auto get_env() const noexcept -> env_of_t<const Sender&> {
-        return stdexec::get_env(sndr_);
+      auto get_env() const noexcept -> __sched_attrs<Scheduler, LateDomain> {
+        return {sched_, {}};
       }
 
-      __t(context_state_t context_state, Sender sndr)
-        : context_state_(context_state)
+      __t(Scheduler sched, context_state_t context_state, Sender sndr)
+        : sched_(sched)
+        , context_state_(context_state)
         , sndr_{static_cast<Sender&&>(sndr)} {
       }
     };
@@ -180,18 +183,18 @@ namespace nvexec::_strm {
     template <class Sched, class Sender>
       requires gpu_stream_scheduler<_current_scheduler_t<Sender>>
     auto operator()(__ignore, Sched sched, Sender&& sndr) const {
-      using _sender_t = __t<continues_on_sender_t<__id<__decay_t<Sender>>>>;
+      using _sender_t = __t<continues_on_sender_t<Sched, __id<__decay_t<Sender>>>>;
       auto stream_sched = get_completion_scheduler<set_value_t>(get_env(sndr));
       return schedule_from(
         static_cast<Sched&&>(sched),
-        _sender_t{stream_sched.context_state_, static_cast<Sender&&>(sndr)});
+        _sender_t{sched, stream_sched.context_state_, static_cast<Sender&&>(sndr)});
     }
   };
 
 } // namespace nvexec::_strm
 
 namespace stdexec::__detail {
-  template <class SenderId>
-  inline constexpr __mconst<nvexec::_strm::continues_on_sender_t<__name_of<__t<SenderId>>>>
-    __name_of_v<nvexec::_strm::continues_on_sender_t<SenderId>>{};
+  template <class Scheduler, class SenderId>
+  inline constexpr __mconst<nvexec::_strm::continues_on_sender_t<Scheduler, __name_of<__t<SenderId>>>>
+    __name_of_v<nvexec::_strm::continues_on_sender_t<Scheduler, SenderId>>{};
 } // namespace stdexec::__detail

include/nvexec/stream/ensure_started.cuh

@@ -45,9 +45,7 @@ namespace nvexec::_strm {
       const inplace_stop_source& stop_source,
       stream_provider_t* stream_provider) noexcept {
       return make_stream_env(
-        __env::__from{[&](get_stop_token_t) noexcept {
-          return stop_source.get_token();
-        }},
+        __env::__from{[&](get_stop_token_t) noexcept { return stop_source.get_token(); }},
         stream_provider);
     }
 
@@ -355,8 +353,8 @@ namespace nvexec::_strm {
         return operation_t<Receiver>{static_cast<Receiver&&>(rcvr), std::move(shared_state_)};
       }
 
-      auto get_env() const noexcept -> env_of_t<const Sender&> {
-        return stdexec::get_env(sndr_);
+      auto get_env() const noexcept -> stream_sender_attrs<Sender> {
+        return {&sndr_};
       }
 
       template <class... Tys>
@@ -398,7 +396,7 @@ namespace nvexec::_strm {
     using _sender_t = __t<ensure_started_sender_t<__id<__decay_t<Sender>>>>;
 
     template <class Env, stream_completing_sender Sender>
-    auto operator()(__ignore, Env&& /*env*/, Sender&& sndr) const -> _sender_t<Sender> {
+    auto operator()(__ignore, Env&&, Sender&& sndr) const -> _sender_t<Sender> {
       auto sched = get_completion_scheduler<set_value_t>(get_env(sndr));
       return _sender_t<Sender>{sched.context_state_, static_cast<Sender&&>(sndr)};
     }

include/nvexec/stream/launch.cuh

@@ -151,8 +151,8 @@ namespace nvexec {
           return {};
         }
 
-        auto get_env() const noexcept -> env_of_t<const Sender&> {
-          return stdexec::get_env(sndr_);
+        auto get_env() const noexcept -> stream_sender_attrs<Sender> {
+          return {&sndr_};
         }
       };
     };