implement type-erased allocator for parallel_scheduler_backend

Author: ericniebler

include/exec/any_sender_of.hpp

@@ -342,15 +342,15 @@ namespace exec {
         static_assert(sizeof(_Tp) <= __buffer_size && alignof(_Tp) <= __alignment);
         _Tp* __pointer = reinterpret_cast<_Tp*>(&__buffer_[0]);
         using _Alloc = std::allocator_traits<_Allocator>::template rebind_alloc<_Tp>;
-        _Alloc __alloc{__allocator_};
+        _Alloc __alloc{__alloc_};
         std::allocator_traits<_Alloc>::construct(__alloc, __pointer, static_cast<_As&&>(__args)...);
         __object_pointer_ = __pointer;
       }
 
       template <class _Tp, class... _As>
       void __construct_large(_As&&... __args) {
         using _Alloc = std::allocator_traits<_Allocator>::template rebind_alloc<_Tp>;
-        _Alloc __alloc{__allocator_};
+        _Alloc __alloc{__alloc_};
         _Tp* __pointer = std::allocator_traits<_Alloc>::allocate(__alloc, 1);
         STDEXEC_TRY {
           std::allocator_traits<_Alloc>::construct(
@@ -369,7 +369,7 @@ namespace exec {
           return;
         }
         using _Alloc = std::allocator_traits<_Allocator>::template rebind_alloc<_Tp>;
-        _Alloc __alloc{__allocator_};
+        _Alloc __alloc{__alloc_};
         _Tp* __pointer = static_cast<_Tp*>(std::exchange(__object_pointer_, nullptr));
         std::allocator_traits<_Alloc>::destroy(__alloc, __pointer);
         if constexpr (!__is_small<_Tp>) {
@@ -382,7 +382,7 @@ namespace exec {
       void* __object_pointer_{nullptr};
       alignas(__alignment) std::byte __buffer_[__buffer_size]{};
       STDEXEC_IMMOVABLE_NO_UNIQUE_ADDRESS
-      _Allocator __allocator_{};
+      _Allocator __alloc_{};
     };
 
     template <
@@ -510,15 +510,15 @@ namespace exec {
         static_assert(sizeof(_Tp) <= __buffer_size && alignof(_Tp) <= __alignment);
         _Tp* __pointer = reinterpret_cast<_Tp*>(&__buffer_[0]);
         using _Alloc = std::allocator_traits<_Allocator>::template rebind_alloc<_Tp>;
-        _Alloc __alloc{__allocator_};
+        _Alloc __alloc{__alloc_};
         std::allocator_traits<_Alloc>::construct(__alloc, __pointer, static_cast<_As&&>(__args)...);
         __object_pointer_ = __pointer;
       }
 
       template <class _Tp, class... _As>
       void __construct_large(_As&&... __args) {
         using _Alloc = std::allocator_traits<_Allocator>::template rebind_alloc<_Tp>;
-        _Alloc __alloc{__allocator_};
+        _Alloc __alloc{__alloc_};
         _Tp* __pointer = std::allocator_traits<_Alloc>::allocate(__alloc, 1);
         STDEXEC_TRY {
           std::allocator_traits<_Alloc>::construct(
@@ -537,7 +537,7 @@ namespace exec {
           return;
         }
         using _Alloc = std::allocator_traits<_Allocator>::template rebind_alloc<_Tp>;
-        _Alloc __alloc{__allocator_};
+        _Alloc __alloc{__alloc_};
         _Tp* __pointer = static_cast<_Tp*>(std::exchange(__object_pointer_, nullptr));
         std::allocator_traits<_Alloc>::destroy(__alloc, __pointer);
         if constexpr (!__is_small<_Tp>) {
@@ -555,7 +555,7 @@ namespace exec {
           _Tp& __other_object = *__pointer;
           this->template __construct_small<_Tp>(static_cast<_Tp&&>(__other_object));
           using _Alloc = std::allocator_traits<_Allocator>::template rebind_alloc<_Tp>;
-          _Alloc __alloc{__allocator_};
+          _Alloc __alloc{__alloc_};
           std::allocator_traits<_Alloc>::destroy(__alloc, __pointer);
         } else {
           __object_pointer_ = __pointer;
@@ -582,7 +582,7 @@ namespace exec {
       const __vtable_t* __vtable_{__default_storage_vtable(static_cast<__vtable_t*>(nullptr))};
       void* __object_pointer_{nullptr};
       alignas(__alignment) std::byte __buffer_[__buffer_size]{};
-      STDEXEC_ATTRIBUTE(no_unique_address) _Allocator __allocator_ { };
+      STDEXEC_ATTRIBUTE(no_unique_address) _Allocator __alloc_ { };
     };
 
     struct __empty_vtable {

include/stdexec/__detail/__any.hpp

@@ -1210,7 +1210,10 @@ namespace STDEXEC::__any {
   // __bad_any_cast
   struct __bad_any_cast : std::exception {
     [[nodiscard]]
-    constexpr char const *what() const noexcept override {
+#if __cpp_lib_constexpr_exceptions >= 2025'02L // constexpr support for std::exception
+    constexpr
+#endif
+      char const *what() const noexcept override {
       return "__bad_any_cast";
     }
   };

include/stdexec/__detail/__any_allocator.hpp

@@ -0,0 +1,107 @@
+/*
+ * Copyright (c) 2026 NVIDIA Corporation
+ *
+ * Licensed under the Apache License Version 2.0 with LLVM Exceptions
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *   https://llvm.org/LICENSE.txt
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "__execution_fwd.hpp"
+
+#include "__any.hpp"
+#include "__concepts.hpp"
+#include "__memory.hpp"
+#include "__typeinfo.hpp"
+
+namespace STDEXEC {
+  namespace __detail {
+    template <class _Base>
+    struct __byte_allocator;
+
+    template <class _Base>
+    using __byte_allocator_interface_t = __any::interface<
+      __byte_allocator,
+      _Base,
+      __any::__extends<__any::__icopyable, __any::__iequality_comparable>
+    >;
+
+    // NOLINTBEGIN(modernize-use-override)
+    template <class _Base>
+    struct __byte_allocator : __byte_allocator_interface_t<_Base> {
+      using __byte_allocator_interface_t<_Base>::interface::interface;
+
+      [[nodiscard]]
+      constexpr virtual auto allocate(size_t __n) -> std::byte* {
+        return __any::__value(*this).allocate(__n);
+      }
+
+      constexpr virtual void deallocate(std::byte* __byte, size_t __n) noexcept {
+        __any::__value(*this).deallocate(__byte, __n);
+      }
+    };
+    // NOLINTEND(modernize-use-override)
+  } // namespace __detail
+
+  template <class _Ty>
+  struct __any_allocator {
+    using value_type = _Ty;
+
+    __any_allocator() = default;
+
+    template <__not_same_as<__any_allocator> _Alloc>
+      requires __is_not_instance_of<_Alloc, __any_allocator> && __simple_allocator<_Alloc>
+    __any_allocator(_Alloc __alloc) noexcept {
+      using __value_t = std::allocator_traits<_Alloc>::value_type;
+      static_assert(
+        __same_as<_Ty, __value_t>,
+        "__any_allocator<T> must be constructed with an allocator of the same value type");
+      __alloc_.emplace(STDEXEC::__rebind_allocator<std::byte>(__alloc));
+    }
+
+    template <__not_same_as<_Ty> _Uy>
+    __any_allocator(__any_allocator<_Uy> __other) noexcept
+      : __alloc_(std::move(__other.__alloc_)) {
+    }
+
+    [[nodiscard]]
+    constexpr bool has_value() const noexcept {
+      return !__any::__empty(__alloc_);
+    }
+
+    [[nodiscard]]
+    constexpr auto type() const noexcept -> __type_index const & {
+      return __any::__type(__alloc_);
+    }
+
+    [[nodiscard]]
+    constexpr auto allocate(size_t __n) -> _Ty* {
+      void* __void_ptr = __alloc_.allocate(__n * sizeof(_Ty));
+      return static_cast<_Ty*>(__void_ptr);
+    }
+
+    constexpr virtual void deallocate(_Ty* __ptr, size_t __n) noexcept {
+      void* __void_ptr = static_cast<void*>(__ptr);
+      __alloc_.deallocate(static_cast<std::byte*>(__void_ptr), __n * sizeof(_Ty));
+    }
+
+   private:
+    __any::__any<__detail::__byte_allocator> __alloc_{};
+  };
+
+  template <class _Alloc>
+  STDEXEC_HOST_DEVICE_DEDUCTION_GUIDE
+    __any_allocator(_Alloc) -> __any_allocator<typename _Alloc::value_type>;
+
+  STDEXEC_HOST_DEVICE_DEDUCTION_GUIDE
+  __any_allocator(std::allocator<void>) -> __any_allocator<std::byte>;
+} // namespace STDEXEC

include/stdexec/__detail/__concepts.hpp

@@ -294,22 +294,13 @@ namespace STDEXEC {
   template <class _Ty, class _Up>
   concept __decays_to_derived_from = __std::derived_from<__decay_t<_Ty>, _Up>;
 
-  namespace __detail {
-    template <class _Alloc>
-    constexpr auto __test_alloc_pointer(int) -> _Alloc::pointer;
-    template <class _Alloc>
-    constexpr auto __test_alloc_pointer(long) -> _Alloc::value_type*;
-
-    template <class _Alloc>
-    using __alloc_pointer_t = decltype(__detail::__test_alloc_pointer<__decay_t<_Alloc>>(0));
-  } // namespace __detail
-
+  // See [allocator.requirements.general]/p99 (https://eel.is/c++draft/allocator.requirements.general#99)
   template <class _Alloc>
-  concept __allocator_ = //
-    requires(__decay_t<_Alloc>& __alloc, std::size_t __bytes) {
-      { __alloc.allocate(__bytes) } -> __std::same_as<__detail::__alloc_pointer_t<_Alloc>>;
-      __alloc.deallocate(__alloc.allocate(__bytes), __bytes);
+  concept __simple_allocator = //
+    requires(_Alloc __alloc, std::size_t __count) {
+      { *__alloc.allocate(__count) } -> __std::same_as<typename _Alloc::value_type&>;
+      __alloc.deallocate(__alloc.allocate(__count), __count);
     } //
-    && __std::copy_constructible<__decay_t<_Alloc>> //
-    && __std::equality_comparable<__decay_t<_Alloc>>;
+    && __std::copy_constructible<_Alloc> //
+    && __std::equality_comparable<_Alloc>;
 } // namespace STDEXEC

include/stdexec/__detail/__parallel_scheduler_backend.hpp

@@ -22,7 +22,8 @@
 // include these after __execution_fwd.hpp
 // #include "any_allocator.cuh"
 #include "../functional.hpp" // IWYU pragma: keep for __with_default
-#include "../stop_token.hpp"
+#include "../stop_token.hpp" // IWYU pragma: keep for get_stop_token_t
+#include "__any_allocator.hpp"
 #include "__queries.hpp"
 #include "__typeinfo.hpp"
 
@@ -34,34 +35,14 @@ STDEXEC_PRAGMA_PUSH()
 STDEXEC_PRAGMA_IGNORE_MSVC(4702) // warning C4702: unreachable code
 
 namespace STDEXEC {
-  template <class _Ty>
-  class any_allocator : public std::allocator<_Ty> {
-   public:
-    template <class _OtherTy>
-    struct rebind {
-      using other = any_allocator<_OtherTy>;
-    };
-
-    template <__not_same_as<any_allocator> _Alloc>
-    any_allocator(const _Alloc&) noexcept {
-    }
-  };
-
-  template <class _Alloc>
-  STDEXEC_HOST_DEVICE_DEDUCTION_GUIDE
-    any_allocator(_Alloc) -> any_allocator<typename _Alloc::value_type>;
-
-  STDEXEC_HOST_DEVICE_DEDUCTION_GUIDE
-  any_allocator(std::allocator<void>) -> any_allocator<std::byte>;
-
   class task_scheduler;
 
   // namespace __detail {
   //   struct __env_proxy : __immovable {
   //     [[nodiscard]]
   //     virtual auto query(const get_stop_token_t&) const noexcept -> inplace_stop_token = 0;
   //     [[nodiscard]]
-  //     virtual auto query(const get_allocator_t&) const noexcept -> any_allocator<std::byte> = 0;
+  //     virtual auto query(const get_allocator_t&) const noexcept -> __any_allocator<std::byte> = 0;
   //     [[nodiscard]]
   //     virtual auto query(const get_scheduler_t&) const noexcept -> task_scheduler = 0;
   //   };
@@ -168,10 +149,12 @@ namespace STDEXEC {
       }
 
       void __query(get_allocator_t, __type_index __value_type, void* __dest) const noexcept {
-        if (__value_type == __mtypeid<any_allocator<std::byte>>) {
-          using __dest_t = std::optional<any_allocator<std::byte>>;
-          *static_cast<__dest_t*>(__dest) = any_allocator{
-            __with_default(get_allocator, std::allocator<std::byte>())(STDEXEC::get_env(__rcvr_))};
+        if (__value_type == __mtypeid<__any_allocator<std::byte>>) {
+          using __dest_t = std::optional<__any_allocator<std::byte>>;
+          constexpr auto __get_alloc = __with_default(get_allocator, std::allocator<std::byte>());
+          auto __alloc = STDEXEC::__rebind_allocator<std::byte>(
+            __get_alloc(STDEXEC::get_env(__rcvr_)));
+          *static_cast<__dest_t*>(__dest) = __any_allocator{std::move(__alloc)};
         }
       }
 
@@ -212,9 +195,10 @@ namespace STDEXEC {
       }
     };
 
-    // A receiver type that forwards its completion operations to a _RcvrProxy member held by
-    // reference (where _RcvrProxy is one of receiver_proxy or bulk_item_receiver_proxy). It
-    // is also responsible to destroying and, if necessary, deallocating the operation state.
+    // A receiver type that forwards its completion operations to a _RcvrProxy member held
+    // by reference (where _RcvrProxy is one of receiver_proxy or
+    // bulk_item_receiver_proxy). It is also responsible for destroying and, if necessary,
+    // deallocating the operation state.
     template <class _RcvrProxy>
     struct __proxy_receiver {
       using receiver_concept = receiver_t;

include/stdexec/__detail/__queries.hpp

@@ -129,7 +129,8 @@ struct get_allocator_t : STDEXEC::__query<get_allocator_t> {
   STDEXEC_ATTRIBUTE(always_inline, host, device)
   static constexpr void __validate() noexcept {
     static_assert(STDEXEC::__nothrow_callable<get_allocator_t, const _Env&>);
-    static_assert(STDEXEC::__allocator_<STDEXEC::__call_result_t<get_allocator_t, const _Env&>>);
+    using __alloc_t = STDEXEC::__call_result_t<get_allocator_t, const _Env&>;
+    static_assert(STDEXEC::__simple_allocator<STDEXEC::__decay_t<__alloc_t>>);
   }
 
   STDEXEC_ATTRIBUTE(nodiscard, always_inline, host, device)