[V4] Join (#60)

* Tidy up comments

* deduplicate transforms

* convert to macros

* comment

* deduplicate fork/call awaitables

* less initialization

* api todo

* propagate cancel

* fix constexpr

* return parent on exception

* revert switch api

* reorder immovable

* comments

* rethrow from await_suspend

* spell

* add atomic ops to frame

* correct frame init

* regular constraint on checkpoint

* full final suspend

* comment

* spell

* joins ops

* join primer

* sketch of a join

* join basics

* add constexpr

* optimise fast path

* why does this slow me down?

* small refactor

* tweaks for perf

* golf

* remove no inline

* set no inline

* spell

* additional asserts

* drop todo

* constexpr gaurd

* init cancel in root

* Update src/core/promise.cxx

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update src/core/promise.cxx

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* spell

* Update src/core/promise.cxx

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update src/core/concepts.cxx

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* rename reset -> reset_counters

* rm false asserts

* add todo

* spell

* fix bugged indexing

Author: ConorWilliams

benchmark/src/libfork_benchmark/fib/fib.hpp

@@ -73,6 +73,8 @@ struct vector_ctx {
 
   auto alloc() noexcept -> Alloc & { return allocator; }
 
+  // TODO: try LF_NO_INLINE for final allocator
+  LF_NO_INLINE
   void push(handle_type handle) { work.push_back(handle); }
 
   auto pop() noexcept -> handle_type {

benchmark/src/libfork_benchmark/fib/lf_parts.cpp

@@ -14,12 +14,15 @@ namespace {
 
 struct global_allocator {
 
-  struct empty {};
+  struct empty {
+    auto operator==(empty const &) const -> bool = default;
+  };
 
   constexpr static auto push(std::size_t sz) -> void * { return ::operator new(sz); }
   constexpr static auto pop(void *p, std::size_t sz) noexcept -> void { ::operator delete(p, sz); }
   constexpr static auto checkpoint() noexcept -> empty { return {}; }
-  constexpr static auto switch_to(empty) noexcept -> void {}
+  constexpr static auto release() noexcept -> void {}
+  constexpr static auto acquire(empty) noexcept -> void {}
 };
 
 static_assert(lf::stack_allocator<global_allocator>);
@@ -37,8 +40,8 @@ struct linear_allocator {
   constexpr auto pop(void *p, std::size_t) noexcept -> void { ptr = static_cast<std::byte *>(p); }
 
   constexpr auto checkpoint() noexcept -> std::byte * { return data.get(); }
-
-  constexpr static auto switch_to(std::byte *) noexcept -> void {}
+  constexpr auto release() noexcept -> void {}
+  constexpr auto acquire(std::byte *) noexcept -> void {}
 };
 
 static_assert(lf::stack_allocator<linear_allocator>);
@@ -57,10 +60,14 @@ constexpr auto no_await = [](this auto fib, std::int64_t *ret, std::int64_t n) -
 
   auto t1 = fib(&lhs, n - 1);
   t1.promise->frame.kind = lf::category::root;
+  t1.promise->frame.stack_ckpt = lf::thread_context<T>->alloc().checkpoint();
+  t1.promise->frame.cancel = nullptr;
   t1.promise->handle().resume();
 
   auto t2 = fib(&rhs, n - 2);
   t2.promise->frame.kind = lf::category::root;
+  t2.promise->frame.stack_ckpt = lf::thread_context<T>->alloc().checkpoint();
+  t2.promise->frame.cancel = nullptr;
   t2.promise->handle().resume();
 
   *ret = lhs + rhs;
@@ -97,7 +104,7 @@ constexpr auto ret = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t,
   co_return lhs + rhs;
 };
 
-template <typename T>
+template <typename T, bool Join = false>
 constexpr auto fork_call = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t, T> {
   if (n < 2) {
     co_return n;
@@ -109,6 +116,10 @@ constexpr auto fork_call = [](this auto fib, std::int64_t n) -> lf::task<std::in
   co_await lf::fork(&rhs, fib, n - 2);
   co_await lf::call(&lhs, fib, n - 1);
 
+  if constexpr (Join) {
+    co_await lf::join();
+  }
+
   co_return lhs + rhs;
 };
 
@@ -138,11 +149,15 @@ void fib(benchmark::State &state) {
     if constexpr (requires { Fn(&result, n); }) {
       auto task = Fn(&result, n);
       task.promise->frame.kind = lf::category::root;
+      task.promise->frame.cancel = nullptr;
+      task.promise->frame.stack_ckpt = lf::thread_context<U>->alloc().checkpoint();
       task.promise->handle().resume();
     } else {
       auto task = Fn(n);
       task.promise->frame.kind = lf::category::root;
+      task.promise->frame.cancel = nullptr;
       task.promise->return_address = &result;
+      task.promise->frame.stack_ckpt = lf::thread_context<U>->alloc().checkpoint();
       task.promise->handle().resume();
     }
 
@@ -184,9 +199,21 @@ BENCHMARK(fib<ret<linear_alloc>, linear_alloc>)->Name("base/libfork/fib/bump/ret
 BENCHMARK(fib<fork_call<linear_alloc>, linear_alloc>)->Name("test/libfork/fib/vector_ctx")->Arg(fib_test);
 BENCHMARK(fib<fork_call<linear_alloc>, linear_alloc>)->Name("base/libfork/fib/vector_ctx")->Arg(fib_base);
 
+// Same as above but with join.
+BENCHMARK(fib<fork_call<linear_alloc, true>, linear_alloc>)
+    ->Name("test/libfork/fib/vector_ctx/join")
+    ->Arg(fib_test);
+BENCHMARK(fib<fork_call<linear_alloc, true>, linear_alloc>)
+    ->Name("base/libfork/fib/vector_ctx/join")
+    ->Arg(fib_base);
+
 using A = poly_vector_ctx<linear_allocator>;
 using B = lf::polymorphic_context<linear_allocator>;
 
 // Same as above but with polymorphic contexts.
 BENCHMARK(fib<fork_call<B>, A, B>)->Name("test/libfork/fib/poly_vector_ctx")->Arg(fib_test);
 BENCHMARK(fib<fork_call<B>, A, B>)->Name("base/libfork/fib/poly_vector_ctx")->Arg(fib_base);
+
+// Same as above but with join.
+BENCHMARK(fib<fork_call<B, true>, A, B>)->Name("test/libfork/fib/poly_vector_ctx/join")->Arg(fib_test);
+BENCHMARK(fib<fork_call<B, true>, A, B>)->Name("base/libfork/fib/poly_vector_ctx/join")->Arg(fib_base);

include/libfork/__impl/utils.hpp

@@ -25,3 +25,8 @@
  * @brief Use like `std::forward` to perfectly forward an expression.
  */
 #define LF_FWD(...) ::std::forward<decltype(__VA_ARGS__)>(__VA_ARGS__)
+
+/**
+ * @brief Use to define a `T` that is aligned to the required alignment of `std::atomic_ref<T>`.
+ */
+#define ATOMIC_ALIGN(T) alignas(std::atomic_ref<T>::required_alignment) T

src/core/concepts.cxx

@@ -40,25 +40,28 @@ consteval auto constify(T &&x) noexcept -> std::add_const_t<T> &;
 /**
  * @brief Defines the API for a libfork compatible stack allocator.
  *
- * - After construction push is valid.
+ * - After construction `this` is in the empty state and push is valid.
  * - Pop is valid provided the FILO order is respected.
  * - Destruction is expected to only occur when the stack is empty.
- * - Result of `.checkpoint()` is expected to be "cheap to copy".
- * - Switch releases the current stack and resumes from the checkpoint:
- *     - This is a noop if the checkpoint is from this stack.
- *     - If the checkpoint is default-constructed it is expected to switch to a new stack.
+ * - Result of `.checkpoint()` is expected to:
+ *     - Be "cheap to copy".
+ *     - Compare equal if they belong to the same stack.
+ * - Release detaches the current stack and leaves `this` in the empty state.
+ * - Acquire attaches to the stack that the checkpoint came from:
+ *     - This is a noop if the checkpoint is from the current stack.
+ *     - Otherwise `this` is empty.
  *
  * Fast-path operations: empty, push, pop, checkpoint
- * Slow-path operations: switch
+ * Slow-path operations: release, acquire
  */
 export template <typename T>
 concept stack_allocator = std::is_object_v<T> && requires (T alloc, std::size_t n, void *ptr) {
   // { alloc.empty() } noexcept -> std::same_as<bool>;
   { alloc.push(n) } -> std::same_as<void *>;
   { alloc.pop(ptr, n) } noexcept -> std::same_as<void>;
-  { alloc.checkpoint() } noexcept -> std::semiregular;
-  { alloc.switch_to({}) } noexcept -> std::same_as<void>;
-  { alloc.switch_to(constify(alloc.checkpoint())) } noexcept -> std::same_as<void>;
+  { alloc.checkpoint() } noexcept -> std::regular;
+  { alloc.release() } noexcept -> std::same_as<void>;
+  { alloc.acquire(constify(alloc.checkpoint())) } noexcept -> std::same_as<void>;
 };
 
 /**

src/core/constants.cxx

@@ -8,6 +8,8 @@ namespace lf {
 constexpr std::size_t k_kilobyte = 1024;
 constexpr std::size_t k_megabyte = 1024 * k_kilobyte;
 
+constexpr std::uint32_t k_u16_max = std::numeric_limits<std::uint16_t>::max();
+
 export constexpr std::size_t k_new_align = __STDCPP_DEFAULT_NEW_ALIGNMENT__;
 export constexpr std::size_t k_cache_line = std::hardware_destructive_interference_size;
 

src/core/frame.cxx

@@ -5,6 +5,7 @@ export module libfork.core:frame;
 import std;
 
 import :concepts;
+import :constants;
 
 namespace lf {
 
@@ -26,18 +27,39 @@ struct frame_type {
   using allocator_type = allocator_t<Context>;
   using checkpoint_type = checkpoint_t<allocator_type>;
 
-  frame_type *parent = nullptr;
-  cancellation *cancel = nullptr;
+  frame_type *parent;
+  cancellation *cancel;
+
   [[no_unique_address]]
   checkpoint_type stack_ckpt;
 
-  std::uint32_t merges = 0;       // Atomic is 32 bits for speed
-  std::uint16_t steals = 0;       // In debug do overflow checking
-  category kind = category::call; // Fork/Call/Just/Root
-  std::uint8_t exception_bit = 0; // Atomically set
+  ATOMIC_ALIGN(std::uint32_t) joins = 0;        // Atomic is 32 bits for speed
+  std::uint16_t steals = 0;                     // In debug do overflow checking
+  category kind = static_cast<category>(0);     // Fork/Call/Just/Root
+  ATOMIC_ALIGN(std::uint8_t) exception_bit = 0; // Atomically set
+
+  // Explicitly post construction, this allows the compiler to emit a single
+  // instruction for the zero init then an instruction for the joins init,
+  // instead of three instructions.
+  constexpr frame_type() noexcept { joins = k_u16_max; }
 
   [[nodiscard]]
   constexpr auto handle() LF_HOF(std::coroutine_handle<frame_type>::from_promise(*this))
+
+  [[nodiscard]]
+  constexpr auto atomic_joins() noexcept -> std::atomic_ref<std::uint32_t> {
+    return std::atomic_ref{joins};
+  }
+
+  [[nodiscard]]
+  constexpr auto atomic_except() noexcept -> std::atomic_ref<std::uint8_t> {
+    return std::atomic_ref{exception_bit};
+  }
+
+  constexpr void reset_counters() noexcept {
+    joins = k_u16_max;
+    steals = 0;
+  }
 };
 
 // =================== Handle =================== //

src/core/ops.cxx

@@ -13,47 +13,56 @@ namespace lf {
 // clang-format off
 
 template <typename R, typename Fn, typename... Args>
-struct pkg {
+struct pkg : immovable {
   R *return_address;
   [[no_unique_address]] Fn fn;
   [[no_unique_address]] tuple<Args...> args;
 };
 
 template <typename Fn, typename... Args>
-struct pkg<void, Fn, Args...> {
+struct pkg<void, Fn, Args...> : immovable {
   [[no_unique_address]] Fn fn;
   [[no_unique_address]] tuple<Args...> args;
 };
 
 // clang-format on
 
+// TODO: consider a prelude namespace for these ops + task
+// TODO: consider neibloids to block ADL
+
 // ======== Fork ======== //
 
 template <typename R, typename Fn, typename... Args>
-struct [[nodiscard("You should immediately co_await this!")]] fork_pkg : pkg<R, Fn, Args...>, immovable {};
+struct [[nodiscard("You should immediately co_await this!")]] fork_pkg : pkg<R, Fn, Args...> {};
 
 export template <typename... Args, async_invocable_to<void, Args...> Fn>
 constexpr auto fork(Fn &&fn, Args &&...args) noexcept -> fork_pkg<void, Fn, Args &&...> {
-  return {LF_FWD(fn), {LF_FWD(args)...}};
+  return {{.fn = LF_FWD(fn), .args = {LF_FWD(args)...}}};
 }
 
 export template <typename R, typename... Args, async_invocable_to<R, Args...> Fn>
 constexpr auto fork(R *ret, Fn &&fn, Args &&...args) noexcept -> fork_pkg<R, Fn, Args &&...> {
-  return {ret, LF_FWD(fn), {LF_FWD(args)...}};
+  return {{.return_address = ret, .fn = LF_FWD(fn), .args = {LF_FWD(args)...}}};
 }
 // ======== Call ======== //
 
 template <typename R, typename Fn, typename... Args>
-struct [[nodiscard("You should immediately co_await this!")]] call_pkg : pkg<R, Fn, Args...>, immovable {};
+struct [[nodiscard("You should immediately co_await this!")]] call_pkg : pkg<R, Fn, Args...> {};
 
 export template <typename... Args, async_invocable_to<void, Args...> Fn>
 constexpr auto call(Fn &&fn, Args &&...args) noexcept -> call_pkg<void, Fn, Args &&...> {
-  return {LF_FWD(fn), {LF_FWD(args)...}};
+  return {{.fn = LF_FWD(fn), .args = {LF_FWD(args)...}}};
 }
 
 export template <typename R, typename... Args, async_invocable_to<R, Args...> Fn>
 constexpr auto call(R *ret, Fn &&fn, Args &&...args) noexcept -> call_pkg<R, Fn, Args &&...> {
-  return {ret, LF_FWD(fn), {LF_FWD(args)...}};
+  return {{.return_address = ret, .fn = LF_FWD(fn), .args = {LF_FWD(args)...}}};
 }
 
+// =============== Join =============== //
+
+struct [[nodiscard("You should immediately co_await this!")]] join_type {};
+
+export constexpr auto join() noexcept -> join_type { return {}; }
+
 } // namespace lf