[V4] Fork infrastructure (#56)

* move ops to new file

* add frame members

* add fork pkg

* fix catagory

* context (not compiling)

* interfaces

* temp fill in type

* spell

* rename

* todo

* trim deps

* take 3

* use kind

* add noexcept

* comment

* todo

* force colour

* todo (use llm agent)

* use call by default

* all initialized

* fixup tuple impl

* add tests for tuple

* add template kw

* very rough dispatch on final suspend

* defer in final suspend

* assume -> ensure

* non_null and safe_cast

* some fork handling code (unused by bench)

* todo

* more comments

* re-order parts

* missing contexts in task/promise

* export/tidy context

* vector context in bench header

* add forking version

* todo

* renames

* benchmark tls on stack alloc

* more checks

* rm dead type

* rm defer comment

* rename

* rename

* fix missing include

* specify void return type

Author: ConorWilliams

CMakeLists.txt

@@ -62,6 +62,8 @@ target_sources(libfork_libfork
       src/core/frame.cxx
       src/core/constants.cxx
       src/core/tuple.cxx
+      src/core/ops.cxx
+      src/core/context.cxx
     PRIVATE
       src/exception.cpp
 )

CMakeUserPresets.json

@@ -5,13 +5,19 @@
       "name": "dev",
       "inherits": "ci-hardened",
       "displayName": "Hardened development build",
-      "toolchainFile": "${sourceDir}/cmake/llvm-brew-toolchain.cmake"
+      "toolchainFile": "${sourceDir}/cmake/llvm-brew-toolchain.cmake",
+      "cacheVariables": {
+        "CMAKE_COLOR_DIAGNOSTICS": "ON"
+      }
     },
     {
       "name": "bench",
       "inherits": "ci-release",
       "displayName": "Release build for benchmarks",
-      "toolchainFile": "${sourceDir}/cmake/llvm-brew-toolchain.cmake"
+      "toolchainFile": "${sourceDir}/cmake/llvm-brew-toolchain.cmake",
+      "cacheVariables": {
+        "CMAKE_COLOR_DIAGNOSTICS": "ON"
+      }
     }
   ],
   "buildPresets": [

benchmark/src/libfork_benchmark/fib/baremetal.cpp

@@ -12,7 +12,7 @@
 namespace {
 
 struct task {
-  struct promise_type : fib_bump_allocator {
+  struct promise_type : tls_bump {
 
     auto get_return_object() -> task { return {std::coroutine_handle<promise_type>::from_promise(*this)}; }
 
@@ -81,7 +81,7 @@ void fib_coro_no_queue(benchmark::State &state) {
 
   // 8MB stack
   std::unique_ptr buffer = std::make_unique<std::byte[]>(1024 * 1024 * 8);
-  fib_bump_ptr = buffer.get();
+  tls_bump_ptr = buffer.get();
 
   for (auto _ : state) {
     benchmark::DoNotOptimize(n);
@@ -91,7 +91,7 @@ void fib_coro_no_queue(benchmark::State &state) {
     benchmark::DoNotOptimize(result);
   }
 
-  if (fib_bump_ptr != buffer.get()) {
+  if (tls_bump_ptr != buffer.get()) {
     std::terminate(); // Stack leak
   }
 }

benchmark/src/libfork_benchmark/fib/fib.hpp

@@ -3,9 +3,14 @@
 #include <bit>
 #include <cstddef>
 #include <cstdint>
+#include <vector>
+
+#include "libfork/__impl/compiler.hpp"
 
 #include "libfork_benchmark/common.hpp"
 
+import libfork.core;
+
 inline constexpr int fib_test = 3;
 inline constexpr int fib_base = 37;
 
@@ -39,17 +44,49 @@ inline auto fib_align_size(std::size_t n) -> std::size_t {
   return (n + k_fib_align - 1) & ~(k_fib_align - 1);
 }
 
-inline thread_local std::byte *fib_bump_ptr = nullptr;
+constinit inline thread_local std::byte *tls_bump_ptr = nullptr;
+
+struct tls_bump {
+
+  static auto operator new(std::size_t sz) -> void * {
+    auto *prev = tls_bump_ptr;
+    tls_bump_ptr += fib_align_size(sz);
+    return prev;
+  }
+
+  static auto operator delete(void *p, [[maybe_unused]] std::size_t sz) noexcept -> void {
+    tls_bump_ptr = std::bit_cast<std::byte *>(p);
+  }
+};
+
+constinit inline std::byte *bump_ptr = nullptr;
 
-struct fib_bump_allocator {
+struct global_bump {
 
   static auto operator new(std::size_t sz) -> void * {
-    auto *prev = fib_bump_ptr;
-    fib_bump_ptr += fib_align_size(sz);
+    auto *prev = bump_ptr;
+    bump_ptr += fib_align_size(sz);
     return prev;
   }
 
   static auto operator delete(void *p, [[maybe_unused]] std::size_t sz) noexcept -> void {
-    fib_bump_ptr = std::bit_cast<std::byte *>(p);
+    bump_ptr = std::bit_cast<std::byte *>(p);
+  }
+};
+
+// === Shared Context Logic ===
+
+struct vector_ctx final : lf::polymorphic_context {
+
+  std::vector<lf::work_handle> work;
+
+  vector_ctx() { work.reserve(1024); }
+
+  void push(lf::work_handle handle) override { work.push_back(handle); }
+
+  auto pop() noexcept -> lf::work_handle override {
+    auto handle = work.back();
+    work.pop_back();
+    return handle;
   }
 };

benchmark/src/libfork_benchmark/fib/lf_parts.cpp

@@ -19,9 +19,8 @@ struct stack_on_heap {
   }
 };
 
-template <lf::alloc_mixin StackPolicy>
-constexpr auto no_await =
-    [](this auto fib, std::int64_t *ret, std::int64_t n) -> lf::task<void, StackPolicy> {
+template <lf::alloc_mixin Stack>
+constexpr auto no_await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> lf::task<void, Stack> {
   if (n < 2) {
     *ret = n;
     co_return;
@@ -30,14 +29,19 @@ constexpr auto no_await =
   std::int64_t lhs = 0;
   std::int64_t rhs = 0;
 
-  fib(&lhs, n - 1).promise->handle().resume();
-  fib(&rhs, n - 2).promise->handle().resume();
+  auto t1 = fib(&lhs, n - 1);
+  t1.promise->frame.kind = lf::category::root;
+  t1.promise->handle().resume();
+
+  auto t2 = fib(&rhs, n - 2);
+  t2.promise->frame.kind = lf::category::root;
+  t2.promise->handle().resume();
 
   *ret = lhs + rhs;
 };
 
-template <lf::alloc_mixin StackPolicy>
-constexpr auto await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> lf::task<void, StackPolicy> {
+template <lf::alloc_mixin Stack>
+constexpr auto await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> lf::task<void, Stack> {
   if (n < 2) {
     *ret = n;
     co_return;
@@ -52,6 +56,35 @@ constexpr auto await = [](this auto fib, std::int64_t *ret, std::int64_t n) -> l
   *ret = lhs + rhs;
 };
 
+constexpr auto ret = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t, tls_bump> {
+  if (n < 2) {
+    co_return n;
+  }
+
+  std::int64_t lhs = 0;
+  std::int64_t rhs = 0;
+
+  co_await lf::call(&lhs, fib, n - 1);
+  co_await lf::call(&rhs, fib, n - 2);
+
+  co_return lhs + rhs;
+};
+
+template <typename Ctx, typename A = tls_bump>
+constexpr auto fork_call = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t, A, Ctx> {
+  if (n < 2) {
+    co_return n;
+  }
+
+  std::int64_t lhs = 0;
+  std::int64_t rhs = 0;
+
+  co_await lf::fork(&rhs, fib, n - 2);
+  co_await lf::call(&lhs, fib, n - 1);
+
+  co_return lhs + rhs;
+};
+
 template <auto Fn>
 void fib(benchmark::State &state) {
 
@@ -60,18 +93,27 @@ void fib(benchmark::State &state) {
 
   state.counters["n"] = static_cast<double>(n);
 
-  std::unique_ptr buffer = std::make_unique<std::byte[]>(1024 * 1024);
+  // Set bump allocator buffer
+  std::unique_ptr buf = std::make_unique<std::byte[]>(1024 * 1024);
+  tls_bump_ptr = buf.get();
+  bump_ptr = buf.get();
 
-  fib_bump_ptr = buffer.get();
+  // Set both context and poly context
+  std::unique_ptr ctx = std::make_unique<vector_ctx>();
+  lf::thread_context<vector_ctx> = ctx.get();
+  lf::thread_context<lf::polymorphic_context> = ctx.get();
 
   for (auto _ : state) {
     benchmark::DoNotOptimize(n);
     std::int64_t result = 0;
 
     if constexpr (requires { Fn(&result, n); }) {
-      Fn(&result, n).promise->handle().resume();
+      auto task = Fn(&result, n);
+      task.promise->frame.kind = lf::category::root;
+      task.promise->handle().resume();
     } else {
       auto task = Fn(n);
+      task.promise->frame.kind = lf::category::root;
       task.promise->return_address = &result;
       task.promise->handle().resume();
     }
@@ -80,39 +122,54 @@ void fib(benchmark::State &state) {
     benchmark::DoNotOptimize(result);
   }
 
-  if (fib_bump_ptr != buffer.get()) {
+  if (tls_bump_ptr != buf.get() || bump_ptr != buf.get()) {
     LF_TERMINATE("Stack leak detected");
   }
-}
-
-template <lf::alloc_mixin StackPolicy>
-constexpr auto ret = [](this auto fib, std::int64_t n) -> lf::task<std::int64_t, StackPolicy> {
-  if (n < 2) {
-    co_return n;
-  }
-
-  std::int64_t lhs = 0;
-  std::int64_t rhs = 0;
 
-  co_await lf::call(&lhs, fib, n - 1);
-  co_await lf::call(&rhs, fib, n - 2);
-
-  co_return lhs + rhs;
-};
+  tls_bump_ptr = nullptr;
+  bump_ptr = nullptr;
+  lf::thread_context<vector_ctx> = nullptr;
+  lf::thread_context<lf::polymorphic_context> = nullptr;
+}
 
 } // namespace
 
+// Return by ref-arg, test direct root, no co-await, direct resumes, uses new/delete for alloc
 BENCHMARK(fib<no_await<stack_on_heap>>)->Name("test/libfork/fib/heap/no_await")->Arg(fib_test);
 BENCHMARK(fib<no_await<stack_on_heap>>)->Name("base/libfork/fib/heap/no_await")->Arg(fib_base);
 
+// Same as above but uses tls bump allocator
+BENCHMARK(fib<no_await<tls_bump>>)->Name("test/libfork/fib/tls_bump/no_await")->Arg(fib_test);
+BENCHMARK(fib<no_await<tls_bump>>)->Name("base/libfork/fib/tls_bump/no_await")->Arg(fib_base);
+
+// Same as above but with global bump allocator
+BENCHMARK(fib<no_await<global_bump>>)->Name("test/libfork/fib/global_bump/no_await")->Arg(fib_test);
+BENCHMARK(fib<no_await<global_bump>>)->Name("base/libfork/fib/global_bump/no_await")->Arg(fib_base);
+
+// TODO: no_await with segmented stack allocator?
+
+// Return by ref-arg, libfork call/call with co-await, uses new/delete for alloc
 BENCHMARK(fib<await<stack_on_heap>>)->Name("test/libfork/fib/heap/await")->Arg(fib_test);
 BENCHMARK(fib<await<stack_on_heap>>)->Name("base/libfork/fib/heap/await")->Arg(fib_base);
 
-BENCHMARK(fib<no_await<fib_bump_allocator>>)->Name("test/libfork/fib/bump_alloc/no_await")->Arg(fib_test);
-BENCHMARK(fib<no_await<fib_bump_allocator>>)->Name("base/libfork/fib/bump_alloc/no_await")->Arg(fib_base);
+// Same as above but uses tls bump allocator
+BENCHMARK(fib<await<tls_bump>>)->Name("test/libfork/fib/tls_bump/await")->Arg(fib_test);
+BENCHMARK(fib<await<tls_bump>>)->Name("base/libfork/fib/tls_bump/await")->Arg(fib_base);
+
+// Same as above but with global bump allocator
+BENCHMARK(fib<await<global_bump>>)->Name("test/libfork/fib/global_bump/await")->Arg(fib_test);
+BENCHMARK(fib<await<global_bump>>)->Name("base/libfork/fib/global_bump/await")->Arg(fib_base);
+
+// Return by value
+// libfork call/call with co-await
+BENCHMARK(fib<ret>)->Name("test/libfork/fib/tls_bump/return")->Arg(fib_test);
+BENCHMARK(fib<ret>)->Name("base/libfork/fib/tls_bump/return")->Arg(fib_base);
 
-BENCHMARK(fib<await<fib_bump_allocator>>)->Name("test/libfork/fib/bump_alloc/await")->Arg(fib_test);
-BENCHMARK(fib<await<fib_bump_allocator>>)->Name("base/libfork/fib/bump_alloc/await")->Arg(fib_base);
+// Return by value
+// libfork call/fork (no join)
+// Non-polymorphic vector-backed context
+BENCHMARK(fib<fork_call<vector_ctx>>)->Name("test/libfork/fib/vector_ctx")->Arg(fib_test);
+BENCHMARK(fib<fork_call<vector_ctx>>)->Name("base/libfork/fib/vector_ctx")->Arg(fib_base);
 
-BENCHMARK(fib<ret<fib_bump_allocator>>)->Name("test/libfork/fib/bump_alloc/return")->Arg(fib_test);
-BENCHMARK(fib<ret<fib_bump_allocator>>)->Name("base/libfork/fib/bump_alloc/return")->Arg(fib_base);
+BENCHMARK(fib<fork_call<lf::polymorphic_context>>)->Name("test/libfork/fib/poly_vector_ctx")->Arg(fib_test);
+BENCHMARK(fib<fork_call<lf::polymorphic_context>>)->Name("base/libfork/fib/poly_vector_ctx")->Arg(fib_base);

include/libfork/__impl/assume.hpp

@@ -15,7 +15,7 @@
  *
  * This macro is always active, regardless of optimization settings or `NDEBUG`.
  */
-#define LF_ASSERT(expr)                                                                                      \
+#define LF_ENSURE(expr)                                                                                      \
   do {                                                                                                       \
     if (!(expr)) {                                                                                           \
       LF_TERMINATE("Assumption '" #expr "' failed!");                                                        \
@@ -43,5 +43,5 @@
       }                                                                                                      \
     } while (false)
 #else
-  #define LF_ASSUME(expr) LF_ASSERT(expr)
+  #define LF_ASSUME(expr) LF_ENSURE(expr)
 #endif

src/core/concepts.cxx

@@ -35,7 +35,7 @@ template <typename T, template <typename...> typename Template>
 concept specialization_of = is_specialization_of<std::remove_cvref_t<T>, Template>::value;
 
 // Forward-decl
-export template <returnable T, alloc_mixin Stack>
+export template <returnable T, alloc_mixin Stack, typename Context>
 struct task;
 
 /**

src/core/context.cxx

@@ -0,0 +1,34 @@
+module;
+export module libfork.core:context;
+
+import std;
+
+import :frame;
+
+namespace lf {
+
+// TODO: private bits / split
+export struct work_handle {
+  frame_type *frame;
+};
+
+static_assert(std::atomic<work_handle>::is_always_lock_free);
+
+template <typename T>
+concept context = requires (T &ctx, work_handle h) {
+  { ctx.push(h) } -> std::same_as<void>;
+  { ctx.pop() } noexcept -> std::same_as<work_handle>;
+};
+
+export struct polymorphic_context {
+  virtual void push(work_handle h) = 0;
+  virtual auto pop() noexcept -> work_handle = 0;
+  virtual ~polymorphic_context() = default;
+};
+
+static_assert(context<polymorphic_context>);
+
+export template <typename Context>
+constinit inline thread_local Context *thread_context = nullptr;
+
+} // namespace lf

src/core/core.cxx

@@ -1,9 +1,15 @@
 export module libfork.core;
 
-// All partitions
-export import :promise;
+// T1 partitions
 export import :concepts;
-export import :utility;
-export import :frame;
 export import :constants;
+export import :frame;
+export import :utility;
 export import :tuple;
+
+// T2 partitions
+export import :ops;     // concepts, tuple, utility
+export import :context; // frame
+
+// T3 partitions
+export import :promise; // context, ops