Merge pull request #1458 from NVIDIA/no-alloc-co_await-for-exec-task

increase the size of the small-object buffer for type-erased operation states
NVIDIA · Dec 30, 2024 · 2b91c3f · 2b91c3f
2 parents 5f01147 + b5bdbbf
commit 2b91c3f
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 10 deletions.
diff --git a/include/exec/any_sender_of.hpp b/include/exec/any_sender_of.hpp
@@ -567,8 +567,13 @@ namespace exec {
       }
     };
 
-    template <class _VTable = __empty_vtable, class _Allocator = std::allocator<std::byte>>
-    using __immovable_storage_t = __t<__immovable_storage<_VTable, _Allocator>>;
+    template <
+      class _VTable = __empty_vtable,
+      class _Allocator = std::allocator<std::byte>,
+      std::size_t _InlineSize = 3 * sizeof(void*),
+      std::size_t _Alignment = alignof(std::max_align_t)>
+    using __immovable_storage_t =
+      __t<__immovable_storage<_VTable, _Allocator, _InlineSize, _Alignment>>;
 
     template <class _VTable, class _Allocator = std::allocator<std::byte>>
     using __unique_storage_t = __t<__storage<_VTable, _Allocator>>;
@@ -791,7 +796,8 @@ namespace exec {
       }
     };
 
-    using __immovable_operation_storage = __immovable_storage_t<__operation_vtable>;
+    using __immovable_operation_storage =
+      __immovable_storage_t<__operation_vtable, std::allocator<std::byte>, 6 * sizeof(void*)>;
 
     template <class _Sigs, class _Queries>
     using __receiver_ref = __mapply<__mbind_front<__q<__rec::__ref>, _Sigs>, _Queries>;
@@ -1205,6 +1211,7 @@ namespace exec {
 
       template <auto... _SchedulerQueries>
       class any_scheduler {
+        // Add the required set_value_t() completions to the schedule-sender.
         using __schedule_completions = stdexec::__concat_completion_signatures<
           _Completions,
           stdexec::completion_signatures<stdexec::set_value_t()>>;
@@ -1216,10 +1223,11 @@ namespace exec {
         template <class _Tag>
         struct __ret_equals_to {
           template <class _Sig>
-          using __f = std::is_same<_Tag, decltype(__ret_fn(static_cast<_Sig>(nullptr)))>;
+          using __f =
+            stdexec::__mbool<STDEXEC_IS_SAME(_Tag, decltype(__ret_fn(static_cast<_Sig>(nullptr))))>;
         };
 
-        using schedule_sender_queries = stdexec::__minvoke<
+        using __schedule_sender_queries = stdexec::__minvoke<
           stdexec::__mremove_if<
             __ret_equals_to<stdexec::get_completion_scheduler_t<stdexec::set_value_t>>>,
           decltype(_SenderQueries)...>;
@@ -1238,7 +1246,7 @@ namespace exec {
           stdexec::get_completion_scheduler<stdexec::set_value_t>.template signature<any_scheduler() noexcept>>;
 #endif
         using __schedule_sender =
-          stdexec::__mapply<stdexec::__q<__schedule_sender_fn>, schedule_sender_queries>;
+          stdexec::__mapply<stdexec::__q<__schedule_sender_fn>, __schedule_sender_queries>;
 
         using __scheduler_base =
           __any::__scheduler<__schedule_sender, queries<_SchedulerQueries...>>;

diff --git a/include/exec/task.hpp b/include/exec/task.hpp
@@ -38,10 +38,14 @@ namespace exec {
   namespace __task {
     using namespace stdexec;
 
-    using __any_scheduler =                                                     //
-      any_receiver_ref<                                                         //
-        completion_signatures<set_error_t(std::exception_ptr), set_stopped_t()> //
-        >::any_sender<>::any_scheduler<>;
+    // The required set_value_t() scheduler-sender completion signature is added in
+    // any_receiver_ref::any_sender::any_scheduler.
+    using __any_scheduler_completions =
+      completion_signatures<set_error_t(std::exception_ptr), set_stopped_t()>;
+
+    using __any_scheduler =
+      any_receiver_ref<__any_scheduler_completions>::any_sender<>::any_scheduler<>;
+
     static_assert(scheduler<__any_scheduler>);
 
     template <class _Ty>