Merge pull request NVIDIA#1740 from ericniebler/fix-multi-gpu-stream-scheduler-get-completion-scheduler-query

fix code that was C-style casting a multi_`gpu_stream_scheduler` to a `stream_scheduler`

Author: ericniebler

include/nvexec/multi_gpu_context.cuh

@@ -27,7 +27,7 @@ STDEXEC_PRAGMA_IGNORE_EDG(cuda_compile)
 
 namespace nvexec {
   namespace _strm {
-    struct multi_gpu_stream_scheduler : private stream_scheduler_env {
+    struct multi_gpu_stream_scheduler : private stream_scheduler_env<multi_gpu_stream_scheduler> {
       using __t = multi_gpu_stream_scheduler;
       using __id = multi_gpu_stream_scheduler;
 
@@ -138,6 +138,13 @@ namespace nvexec {
       int num_devices_{};
       context_state_t context_state_;
     };
+
+    template <>
+    STDEXEC_ATTRIBUTE(nodiscard)
+    inline auto stream_scheduler_env<multi_gpu_stream_scheduler>::query(
+      get_completion_scheduler_t<set_value_t>) const noexcept -> multi_gpu_stream_scheduler {
+      return stdexec::__c_downcast<multi_gpu_stream_scheduler>(*this);
+    }
   } // namespace _strm
 
   using _strm::multi_gpu_stream_scheduler;

include/nvexec/stream_context.cuh

@@ -45,22 +45,23 @@ namespace nvexec {
   namespace _strm {
     struct stream_scheduler;
 
-    struct stream_scheduler_env {
+    template <class StreamScheduler>
+    struct stream_scheduler_env { // NOLINT(bugprone-crtp-constructor-accessibility)
       STDEXEC_ATTRIBUTE(nodiscard)
       static auto query(get_forward_progress_guarantee_t) noexcept -> forward_progress_guarantee {
         return forward_progress_guarantee::weakly_parallel;
       }
 
       STDEXEC_ATTRIBUTE(nodiscard)
-      auto query(get_completion_scheduler_t<set_value_t>) const noexcept -> stream_scheduler;
+      auto query(get_completion_scheduler_t<set_value_t>) const noexcept -> StreamScheduler;
 
       STDEXEC_ATTRIBUTE(nodiscard)
       constexpr auto query(get_completion_domain_t<set_value_t>) const noexcept -> stream_domain {
         return {};
       }
     };
 
-    struct stream_scheduler : private stream_scheduler_env {
+    struct stream_scheduler : private stream_scheduler_env<stream_scheduler> {
       using __t = stream_scheduler;
       using __id = stream_scheduler;
 
@@ -153,10 +154,11 @@ namespace nvexec {
       context_state_t context_state_;
     };
 
+    template <>
     STDEXEC_ATTRIBUTE(nodiscard)
-    inline auto stream_scheduler_env::query(get_completion_scheduler_t<set_value_t>) const noexcept
-      -> stream_scheduler {
-      return (const stream_scheduler&) *this;
+    inline auto stream_scheduler_env<stream_scheduler>::query(
+      get_completion_scheduler_t<set_value_t>) const noexcept -> stream_scheduler {
+      return stdexec::__c_downcast<stream_scheduler>(*this);
     }
   } // namespace _strm