Skip to content

make task_scheduler allocator-aware #1961

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ericniebler merged 1 commit into from
Mar 26, 2026
+186 −45
Merged

make task_scheduler allocator-aware #1961

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 109 additions & 36 deletions include/stdexec/__detail/__any.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

#include "__concepts.hpp"
#include "__config.hpp"
#include "__memory.hpp"
#include "__type_traits.hpp"
#include "__typeinfo.hpp"
#include "__utility.hpp"
Expand All @@ -28,6 +29,7 @@
#include <bit>
#include <exception>
#include <memory>
#include <new>
#include <span>
#include <type_traits>
#include <utility>
Expand Down Expand Up @@ -211,7 +213,7 @@ namespace STDEXEC::__any
{
template <class _Ty>
STDEXEC_ATTRIBUTE(nodiscard, always_inline)
constexpr auto &operator()(_Ty &&__arg) const noexcept
constexpr auto &&operator()(_Ty &&__arg) const noexcept
{
return __arg.__value_(static_cast<_Ty &&>(__arg));
}
Expand Down Expand Up @@ -395,6 +397,11 @@ namespace STDEXEC::__any
// needed by MSVC for EBO to work for some reason:
constexpr virtual ~__iroot() = default;

constexpr virtual void __delete_() noexcept
{
this->~__iroot();
}

private:
template <template <class> class, class, class, size_t, size_t>
friend struct __interface_base;
Expand Down Expand Up @@ -506,8 +513,9 @@ namespace STDEXEC::__any

//////////////////////////////////////////////////////////////////////////////////////////
// __emplace_into
template <class _Model, class... _Args>
constexpr _Model &__emplace_into([[maybe_unused]] __iroot *&__root_ptr,
template <class _Model, class _Allocator, class... _Args>
constexpr _Model &__emplace_into([[maybe_unused]] _Allocator const &__alloc,
[[maybe_unused]] __iroot *&__root_ptr,
[[maybe_unused]] std::span<std::byte> __buff,
_Args &&...__args)
{
Expand All @@ -526,20 +534,27 @@ namespace STDEXEC::__any
}
else
{
auto *const __model = ::new _Model(static_cast<_Args &&>(__args)...);
auto __alloc2 = STDEXEC::__rebind_allocator<_Model>(__alloc);
using __traits_t = std::allocator_traits<decltype(__alloc2)>;
auto *const __model = __traits_t::allocate(__alloc2, 1);
__scope_guard __guard{[&]() noexcept { __traits_t::deallocate(__alloc2, __model, 1); }};
__traits_t::construct(__alloc2, __model, static_cast<_Args &&>(__args)...);
__guard.__dismiss();
*__std::start_lifetime_as<__tagged_ptr>(__buff.data()) = __tagged_ptr(__model);
return *__model;
}
}
}

template <int = 0, class _CvRefValue, class _Value = std::decay_t<_CvRefValue>>
template <int = 0, class _CvRefValue, class _Allocator, class _Value = std::decay_t<_CvRefValue>>
STDEXEC_ATTRIBUTE(always_inline)
constexpr _Value &__emplace_into(__iroot *&__root_ptr,
constexpr _Value &__emplace_into(_Allocator const &__alloc,
__iroot *&__root_ptr,
std::span<std::byte> __buff,
_CvRefValue &&__value)
{
return STDEXEC::__any::__emplace_into<_Value>(__root_ptr,
return STDEXEC::__any::__emplace_into<_Value>(__alloc,
__root_ptr,
__buff,
static_cast<_CvRefValue &&>(__value));
}
Expand Down Expand Up @@ -574,23 +589,39 @@ namespace STDEXEC::__any
using __reference_proxy_model = __reference<_Interface>;

// __value
template <template <class> class _Interface, class _Value>
struct __value_root;
template <template <class> class _Interface, class _Value, class _Allocator>
struct __value_root_with_allocator;

template <template <class> class _Interface, class _Value>
template <template <class> class _Interface, class _Value, class _Allocator>
struct __value_model final
: _Interface<__mcall1<__bases_of<_Interface>, __value_root<_Interface, _Value>>>
: _Interface<__mcall1<__bases_of<_Interface>,
__value_root_with_allocator<_Interface, _Value, _Allocator>>>
{
using __base_t = _Interface<__mcall1<__bases_of<_Interface>, __value_root<_Interface, _Value>>>;
using __base_t =
_Interface<__mcall1<__bases_of<_Interface>,
__value_root_with_allocator<_Interface, _Value, _Allocator>>>;
using __base_t::__base_t;

constexpr void __delete_() noexcept final
{
auto __alloc = STDEXEC::__rebind_allocator<__value_model>(this->__get_allocator());
using __traits_t = std::allocator_traits<decltype(__alloc)>;
__traits_t::destroy(__alloc, this);
__traits_t::deallocate(__alloc, this, 1);
}

// This is a virtual override if _Interface extends __imovable
//! @pre __is_small<__value_model>(__buff.size())
constexpr void __move_to(__iroot *&__ptr, std::span<std::byte> __buff) noexcept
{
static_assert(__extension_of<__iabstract<_Interface>, __imovable>);
STDEXEC_ASSERT(STDEXEC::__any::__is_small<__value_model>(__buff.size()));
STDEXEC::__any::__emplace_into(__ptr, __buff, std::move(*this));
STDEXEC::__any::__emplace_into<__value_model>(this->__get_allocator(),
__ptr,
__buff,
std::allocator_arg,
this->__get_allocator(),
__value(std::move(*this)));
__reset(*this);
}

Expand All @@ -599,7 +630,12 @@ namespace STDEXEC::__any
{
static_assert(__extension_of<__iabstract<_Interface>, __icopyable>);
STDEXEC_ASSERT(!__empty(*this));
STDEXEC::__any::__emplace_into(__ptr, __buff, *this);
STDEXEC::__any::__emplace_into<__value_model>(this->__get_allocator(),
__ptr,
__buff,
std::allocator_arg,
this->__get_allocator(),
__value(*this));
}
};

Expand Down Expand Up @@ -739,6 +775,34 @@ namespace STDEXEC::__any
}
};

//////////////////////////////////////////////////////////////////////////////////////////
// __value_root_with_allocator
template <template <class> class _Interface, class _Value, class _Allocator>
struct STDEXEC_ATTRIBUTE(empty_bases) __value_root_with_allocator
: __value_root<_Interface, _Value>
, private _Allocator
{
constexpr explicit __value_root_with_allocator(_Value __value,
_Allocator const &__alloc = _Allocator())
: __value_root<_Interface, _Value>(std::move(__value))
, _Allocator(__alloc)
{}

template <class... _Args>
constexpr explicit __value_root_with_allocator(std::allocator_arg_t,
_Allocator const &__alloc,
_Args &&...__args)
: __value_root<_Interface, _Value>(static_cast<_Args &&>(__args)...)
, _Allocator(__alloc)
{}

[[nodiscard]]
constexpr auto __get_allocator() const noexcept -> _Allocator const &
{
return *this;
}
};

//////////////////////////////////////////////////////////////////////////////////////////
// __value_proxy_root
template <template <class> class _Interface>
Expand Down Expand Up @@ -869,22 +933,27 @@ namespace STDEXEC::__any
friend struct __any;
friend struct __access;

template <class _Value, class... _Args>
constexpr _Value &__emplace_(_Args &&...__args)
template <class _Value, class _Allocator, class... _Args>
constexpr _Value &__emplace_(std::allocator_arg_t, _Allocator const &__alloc, _Args &&...__args)
{
static_assert(__decays_to<_Value, _Value>, "Value must be an object type.");
using __model_type = __value_model<_Interface, _Value>;
auto &__model = STDEXEC::__any::__emplace_into<__model_type>(__root_ptr_,
__buff_,
static_cast<_Args &&>(
__args)...);
static_assert(__simple_allocator<_Allocator>);
using __model_type = __value_model<_Interface, _Value, _Allocator>;
auto &__model = //
STDEXEC::__any::__emplace_into<__model_type>(__alloc,
__root_ptr_,
__buff_,
std::allocator_arg,
__alloc,
static_cast<_Args &&>(__args)...);
return __value(__model);
}

template <int = 0, class _CvRefValue, class _Value = std::decay_t<_CvRefValue>>
constexpr _Value &__emplace_(_CvRefValue &&__value)
template <class _Value, class... _Args>
constexpr _Value &__emplace_(_Args &&...__args)
{
return __emplace_<_Value>(static_cast<_CvRefValue &&>(__value));
std::allocator<_Value> const __alloc{};
return __emplace_<_Value>(std::allocator_arg, __alloc, static_cast<_Args &&>(__args)...);
}

template <class _Self>
Expand Down Expand Up @@ -938,7 +1007,7 @@ namespace STDEXEC::__any
else if (!__ptr.__is_tagged())
std::destroy_at(std::addressof(__value(*this)));
else
delete std::addressof(__value(*this));
__value(*this).__delete_();

__ptr = __tagged_ptr();
}
Expand Down Expand Up @@ -1238,15 +1307,17 @@ namespace STDEXEC::__any
using __model_type = __reference_model<_Interface, __value_type, __extension_type>;
if constexpr (_CvModel::__root_kind == __root_kind::__reference)
{
STDEXEC::__any::__emplace_into<__model_type>(__root_ptr_,
STDEXEC::__any::__emplace_into<__model_type>(std::allocator<std::byte>{}, // not used
__root_ptr_,
__buff_,
__model.__get_value_ptr_(),
__model.__get_root_ptr_());
}
else
{
__iroot *__root_ptr = std::addressof(STDEXEC::__unconst(__model));
STDEXEC::__any::__emplace_into<__model_type>(__root_ptr_,
STDEXEC::__any::__emplace_into<__model_type>(std::allocator<std::byte>{}, // not used
__root_ptr_,
__buff_,
static_cast<__value_type *>(nullptr),
STDEXEC_DECAY_COPY(__root_ptr));
Expand All @@ -1258,7 +1329,8 @@ namespace STDEXEC::__any
{
static_assert(!__extension_of<_CvValue, _Interface>);
using __model_type = __reference_model<_Interface, _CvValue>;
STDEXEC::__any::__emplace_into<__model_type>(__root_ptr_,
STDEXEC::__any::__emplace_into<__model_type>(std::allocator<std::byte>{}, // not used
__root_ptr_,
__buff_,
std::addressof(__value),
static_cast<__iroot *>(nullptr));
Expand Down Expand Up @@ -1597,18 +1669,19 @@ namespace STDEXEC::__any

// Construct from an object that implements the interface (and is not an any<>
// itself)
template <__model_of<_Interface> _Value>
constexpr __any(_Value __value)
template <__model_of<_Interface> _Value, class _Allocator = std::allocator<_Value>>
constexpr __any(_Value __value, _Allocator const &__alloc = _Allocator())
: __any()
{
(*this).__emplace_(std::move(__value));
static_assert(__simple_allocator<_Allocator>);
(*this).template __emplace_<_Value>(std::allocator_arg, __alloc, std::move(__value));
}

template <class _Type, class... _Args>
constexpr explicit __any(std::in_place_type_t<_Type>, _Args &&...__args)
template <class _Value, class... _Args>
constexpr explicit __any(std::in_place_type_t<_Value>, _Args &&...__args)
: __any()
{
(*this).template __emplace_<_Type>(static_cast<_Args &&>(__args)...);
(*this).template __emplace_<_Value>(static_cast<_Args &&>(__args)...);
}

// Implicit derived-to-base conversion constructor
Expand All @@ -1634,7 +1707,7 @@ namespace STDEXEC::__any
constexpr __any &operator=(_Value __value)
{
__reset(*this);
(*this).__emplace_(std::move(__value));
(*this).template __emplace_<_Value>(std::move(__value));
return *this;
}

Expand All @@ -1654,7 +1727,7 @@ namespace STDEXEC::__any
&& (_Other::__root_kind == __root_kind::__reference)
constexpr __any &operator=(_Interface<_Other> const &__other)
{
// Guard against __self-assignment when __other is a reference to *this
// Guard against self-assignment when __other is a reference to *this
if (__data(__other) == __data(*this))
return *this;

Expand Down
20 changes: 11 additions & 9 deletions include/stdexec/__detail/__task_scheduler.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,10 +227,11 @@ namespace STDEXEC
public:
using scheduler_concept = scheduler_t;

template <__not_same_as<task_scheduler> _Sch, class _Alloc = std::allocator<std::byte>>
template <__not_same_as<task_scheduler> _Sch, class _Alloc = std::allocator<_Sch>>
requires __infallible_scheduler<_Sch, __task::__env_t<true>>
constexpr explicit task_scheduler(_Sch __sch, _Alloc __alloc = {})
: __backend_(__backend_for{std::move(__sch), __alloc})
constexpr explicit task_scheduler(_Sch __sch, _Alloc const & __alloc = {})
: __backend_(__backend_for{std::move(__sch), STDEXEC::__rebind_allocator<_Sch>(__alloc)},
STDEXEC::__rebind_allocator<__backend_for<_Sch, _Alloc>>(__alloc))
{
using __opstate_t = __task::__opstate<_Alloc, schedule_result_t<_Sch>, __task::__env_t<true>>;
static_assert(sizeof(__opstate_t) <= STDEXEC_TASK_SCHEDULE_OPSTATE_SIZE,
Expand Down Expand Up @@ -668,10 +669,11 @@ namespace STDEXEC
// __emplace_into
template <class _Ty, class _Alloc, class... _Args>
constexpr auto
__emplace_into(std::span<std::byte> __storage, _Alloc& __alloc, _Args&&... __args) -> _Ty&
__emplace_into(std::span<std::byte> __storage, _Alloc const & __alloc, _Args&&... __args)
-> _Ty&
{
using __traits_t = std::allocator_traits<_Alloc>::template rebind_traits<_Ty>;
auto __alloc_copy = STDEXEC::__rebind_allocator<_Ty>(__alloc);
using __traits_t = std::allocator_traits<decltype(__alloc_copy)>;

bool const __in_situ = __storage.size() >= sizeof(_Ty);
auto* __ptr = __in_situ ? reinterpret_cast<_Ty*>(__storage.data())
Expand Down Expand Up @@ -756,10 +758,10 @@ namespace STDEXEC
{
STDEXEC_TRY
{
using __opstate_t = __task::__opstate<_Alloc, _Sndr, _Env>;
bool const __in_situ = __storage.size() >= sizeof(__opstate_t);
_Alloc& __alloc = *this;
auto& __opstate = __task::__emplace_into<__opstate_t>(__storage,
using __opstate_t = __task::__opstate<_Alloc, _Sndr, _Env>;
bool const __in_situ = __storage.size() >= sizeof(__opstate_t);
_Alloc const & __alloc = *this;
auto& __opstate = __task::__emplace_into<__opstate_t>(__storage,
__alloc,
__alloc,
static_cast<_Sndr&&>(__sndr),
Expand Down
Loading
Loading