[libcxx-commits] [libcxx] [libc++][pstl] std_thread scheduler (PR #212121)

Michael G. Kazakov via libcxx-commits libcxx-commits at lists.llvm.org
Sun Jul 26 07:06:13 PDT 2026


https://github.com/mikekazakov created https://github.com/llvm/llvm-project/pull/212121

None

>From 976602afdf179acc02b62fa67bc2a8acd176555a Mon Sep 17 00:00:00 2001
From: Michael Kazakov <mike.kazakov at gmail.com>
Date: Sun, 26 Jul 2026 15:01:05 +0100
Subject: [PATCH] Draft of std_thread scheduler

---
 libcxx/include/__pstl/backends/std_thread.h   | 325 ++++++++++++-
 libcxx/src/CMakeLists.txt                     |   6 +
 libcxx/src/pstl/std_thread.cpp                | 433 ++++++++++++++++++
 .../algorithms/pstl.std_thread.apply.pass.cpp |  52 +++
 .../libcxx/test/features/libcxx_macros.py     |   1 +
 5 files changed, 792 insertions(+), 25 deletions(-)
 create mode 100644 libcxx/src/pstl/std_thread.cpp
 create mode 100644 libcxx/test/libcxx/algorithms/pstl.std_thread.apply.pass.cpp

diff --git a/libcxx/include/__pstl/backends/std_thread.h b/libcxx/include/__pstl/backends/std_thread.h
index dd2c3f15403e3..9b4d5831a8787 100644
--- a/libcxx/include/__pstl/backends/std_thread.h
+++ b/libcxx/include/__pstl/backends/std_thread.h
@@ -9,7 +9,23 @@
 #ifndef _LIBCPP___PSTL_BACKENDS_STD_THREAD_H
 #define _LIBCPP___PSTL_BACKENDS_STD_THREAD_H
 
+#include <__algorithm/inplace_merge.h>
+#include <__algorithm/lower_bound.h>
+#include <__algorithm/max.h>
+#include <__algorithm/merge.h>
+#include <__algorithm/upper_bound.h>
+#include <__atomic/atomic.h>
 #include <__config>
+#include <__cstddef/ptrdiff_t.h>
+#include <__exception/terminate.h>
+#include <__iterator/iterator_traits.h>
+#include <__iterator/move_iterator.h>
+#include <__memory/allocator.h>
+#include <__memory/construct_at.h>
+#include <__memory/destroy.h>
+#include <__memory/unique_ptr.h>
+#include <__new/exceptions.h>
+#include <__numeric/reduce.h>
 #include <__pstl/backend_fwd.h>
 #include <__pstl/cpu_algos/any_of.h>
 #include <__pstl/cpu_algos/cpu_traits.h>
@@ -21,7 +37,9 @@
 #include <__pstl/cpu_algos/transform.h>
 #include <__pstl/cpu_algos/transform_reduce.h>
 #include <__utility/empty.h>
+#include <__utility/exception_guard.h>
 #include <__utility/move.h>
+#include <__utility/pair.h>
 #include <optional>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
@@ -34,8 +52,47 @@ _LIBCPP_PUSH_MACROS
 #if _LIBCPP_STD_VER >= 17
 
 _LIBCPP_BEGIN_NAMESPACE_STD
+_LIBCPP_BEGIN_EXPLICIT_ABI_ANNOTATIONS
 namespace __pstl {
 
+namespace __std_thread {
+
+_LIBCPP_EXPORTED_FROM_ABI void
+__apply(size_t __iterations, void* __context, void (*__func)(void* __context, size_t __iteration)) noexcept;
+
+struct __chunk_partitions {
+  ptrdiff_t __chunk_count_; // includes the first chunk
+  ptrdiff_t __chunk_size_;
+  ptrdiff_t __first_chunk_size_;
+};
+
+[[__gnu__::__const__]] _LIBCPP_EXPORTED_FROM_ABI __chunk_partitions __partition_chunks(ptrdiff_t __size) noexcept;
+
+template <class _Func>
+_LIBCPP_HIDE_FROM_ABI void __apply(size_t __chunk_count, _Func __func) noexcept {
+  __std_thread::__apply(__chunk_count, &__func, [](void* __context, size_t __chunk) {
+    (*static_cast<_Func*>(__context))(__chunk);
+  });
+}
+
+template <class _RandomAccessIterator, class _Functor>
+_LIBCPP_HIDE_FROM_ABI optional<__empty>
+__parallel_for(__chunk_partitions __partitions, _RandomAccessIterator __first, _Functor __func) {
+  // Perform the chunked execution.
+  __std_thread::__apply(__partitions.__chunk_count_, [&](size_t __chunk) {
+    auto __this_chunk_size = __chunk == 0 ? __partitions.__first_chunk_size_ : __partitions.__chunk_size_;
+    auto __index =
+        __chunk == 0
+            ? 0
+            : (__chunk * __partitions.__chunk_size_) + (__partitions.__first_chunk_size_ - __partitions.__chunk_size_);
+    __func(__first + __index, __first + __index + __this_chunk_size);
+  });
+
+  return __empty{};
+}
+
+} // namespace __std_thread
+
 //
 // This partial backend implementation is for testing purposes only and not meant for production use. This will be
 // replaced by a proper implementation once the PSTL implementation is somewhat stable.
@@ -45,45 +102,262 @@ namespace __pstl {
 
 template <>
 struct __cpu_traits<__std_thread_backend_tag> {
-  template <class _RandomAccessIterator, class _Fp>
+  template <class _RandomAccessIterator, class _Functor>
   _LIBCPP_HIDE_FROM_ABI static optional<__empty>
-  __for_each(_RandomAccessIterator __first, _RandomAccessIterator __last, _Fp __f) {
-    __f(__first, __last);
-    return __empty{};
+  __for_each(_RandomAccessIterator __first, _RandomAccessIterator __last, _Functor __func) {
+    return __std_thread::__parallel_for(
+        __std_thread::__partition_chunks(__last - __first), std::move(__first), std::move(__func));
   }
 
-  template <class _Index, class _UnaryOp, class _Tp, class _BinaryOp, class _Reduce>
-  _LIBCPP_HIDE_FROM_ABI static optional<_Tp>
-  __transform_reduce(_Index __first, _Index __last, _UnaryOp, _Tp __init, _BinaryOp, _Reduce __reduce) {
-    return __reduce(std::move(__first), std::move(__last), std::move(__init));
-  }
+  template <class _RandomAccessIterator1, class _RandomAccessIterator2, class _RandomAccessIteratorOut>
+  struct __merge_range {
+    __merge_range(_RandomAccessIterator1 __mid1, _RandomAccessIterator2 __mid2, _RandomAccessIteratorOut __result)
+        : __mid1_(__mid1), __mid2_(__mid2), __result_(__result) {}
 
-  template <class _RandomAccessIterator, class _Compare, class _LeafSort>
-  _LIBCPP_HIDE_FROM_ABI static optional<__empty>
-  __stable_sort(_RandomAccessIterator __first, _RandomAccessIterator __last, _Compare __comp, _LeafSort __leaf_sort) {
-    __leaf_sort(__first, __last, __comp);
-    return __empty{};
-  }
+    _RandomAccessIterator1 __mid1_;
+    _RandomAccessIterator2 __mid2_;
+    _RandomAccessIteratorOut __result_;
+  };
 
-  _LIBCPP_HIDE_FROM_ABI static void __cancel_execution() {}
-
-  template <class _RandomAccessIterator1,
-            class _RandomAccessIterator2,
-            class _RandomAccessIterator3,
-            class _Compare,
-            class _LeafMerge>
+  template <typename _RandomAccessIterator1,
+            typename _RandomAccessIterator2,
+            typename _RandomAccessIterator3,
+            typename _Compare,
+            typename _LeafMerge>
   _LIBCPP_HIDE_FROM_ABI static optional<__empty>
   __merge(_RandomAccessIterator1 __first1,
           _RandomAccessIterator1 __last1,
           _RandomAccessIterator2 __first2,
           _RandomAccessIterator2 __last2,
-          _RandomAccessIterator3 __outit,
+          _RandomAccessIterator3 __result,
           _Compare __comp,
-          _LeafMerge __leaf_merge) {
-    __leaf_merge(__first1, __last1, __first2, __last2, __outit, __comp);
+          _LeafMerge __leaf_merge) noexcept {
+    __std_thread::__chunk_partitions __partitions =
+        __std_thread::__partition_chunks(std::max<ptrdiff_t>(__last1 - __first1, __last2 - __first2));
+
+    if (__partitions.__chunk_count_ == 0)
+      return __empty{};
+
+    if (__partitions.__chunk_count_ == 1) {
+      __leaf_merge(__first1, __last1, __first2, __last2, __result, __comp);
+      return __empty{};
+    }
+
+    using __merge_range_t = __merge_range<_RandomAccessIterator1, _RandomAccessIterator2, _RandomAccessIterator3>;
+    auto const __n_ranges = __partitions.__chunk_count_ + 1;
+
+    // TODO: use __uninitialized_buffer
+    auto __destroy = [=](__merge_range_t* __ptr) {
+      std::destroy_n(__ptr, __n_ranges);
+      std::allocator<__merge_range_t>().deallocate(__ptr, __n_ranges);
+    };
+
+    unique_ptr<__merge_range_t[], decltype(__destroy)> __ranges(
+        [&]() -> __merge_range_t* {
+#  if _LIBCPP_HAS_EXCEPTIONS
+          try {
+#  endif
+            return std::allocator<__merge_range_t>().allocate(__n_ranges);
+#  if _LIBCPP_HAS_EXCEPTIONS
+          } catch (const std::bad_alloc&) {
+            return nullptr;
+          }
+#  endif
+        }(),
+        __destroy);
+
+    if (!__ranges)
+      return nullopt;
+
+    // TODO: Improve the case where the smaller range is merged into just a few (or even one) chunks of the larger case
+    __merge_range_t* __r = __ranges.get();
+    std::__construct_at(__r++, __first1, __first2, __result);
+
+    bool __iterate_first_range = __last1 - __first1 > __last2 - __first2;
+
+    auto __compute_chunk = [&](size_t __chunk_size) -> __merge_range_t {
+      auto [__mid1, __mid2] = [&] {
+        if (__iterate_first_range) {
+          auto __m1 = __first1 + __chunk_size;
+          auto __m2 = std::lower_bound(__first2, __last2, __m1[-1], __comp);
+          return std::make_pair(__m1, __m2);
+        } else {
+          auto __m2 = __first2 + __chunk_size;
+          auto __m1 = std::lower_bound(__first1, __last1, __m2[-1], __comp);
+          return std::make_pair(__m1, __m2);
+        }
+      }();
+
+      __result += (__mid1 - __first1) + (__mid2 - __first2);
+      __first1 = __mid1;
+      __first2 = __mid2;
+      return {std::move(__mid1), std::move(__mid2), __result};
+    };
+
+    // handle first chunk
+    std::__construct_at(__r++, __compute_chunk(__partitions.__first_chunk_size_));
+
+    // handle 2 -> N - 1 chunks
+    for (ptrdiff_t __i = 0; __i != __partitions.__chunk_count_ - 2; ++__i)
+      std::__construct_at(__r++, __compute_chunk(__partitions.__chunk_size_));
+
+    // handle last chunk
+    std::__construct_at(__r, __last1, __last2, __result);
+
+    __std_thread::__apply(__partitions.__chunk_count_, [&](size_t __index) {
+      auto __first_iters = __ranges[__index];
+      auto __last_iters  = __ranges[__index + 1];
+      __leaf_merge(
+          __first_iters.__mid1_,
+          __last_iters.__mid1_,
+          __first_iters.__mid2_,
+          __last_iters.__mid2_,
+          __first_iters.__result_,
+          __comp);
+    });
+
+    return __empty{};
+  }
+
+  template <class _RandomAccessIterator, class _Transform, class _Value, class _Combiner, class _Reduction>
+  _LIBCPP_HIDE_FROM_ABI static optional<_Value> __transform_reduce(
+      _RandomAccessIterator __first,
+      _RandomAccessIterator __last,
+      _Transform __transform,
+      _Value __init,
+      _Combiner __combiner,
+      _Reduction __reduction) {
+    if (__first == __last)
+      return __init;
+
+    auto __partitions = __std_thread::__partition_chunks(__last - __first);
+
+    auto __destroy = [__count = __partitions.__chunk_count_](_Value* __ptr) {
+      std::destroy_n(__ptr, __count);
+      std::allocator<_Value>().deallocate(__ptr, __count);
+    };
+
+    // TODO: use __uninitialized_buffer
+    // TODO: allocate one element per worker instead of one element per chunk
+    unique_ptr<_Value[], decltype(__destroy)> __values(
+        std::allocator<_Value>().allocate(__partitions.__chunk_count_), __destroy);
+
+    // __apply is noexcept
+    __std_thread::__apply(__partitions.__chunk_count_, [&](size_t __chunk) {
+      auto __this_chunk_size = __chunk == 0 ? __partitions.__first_chunk_size_ : __partitions.__chunk_size_;
+      auto __index           = __chunk == 0 ? 0
+                                            : (__chunk * __partitions.__chunk_size_) +
+                                                  (__partitions.__first_chunk_size_ - __partitions.__chunk_size_);
+      if (__this_chunk_size != 1) {
+        std::__construct_at(
+            __values.get() + __chunk,
+            __reduction(__first + __index + 2,
+                        __first + __index + __this_chunk_size,
+                        __combiner(__transform(__first + __index), __transform(__first + __index + 1))));
+      } else {
+        std::__construct_at(__values.get() + __chunk, __transform(__first + __index));
+      }
+    });
+
+    return std::reduce(
+        std::make_move_iterator(__values.get()),
+        std::make_move_iterator(__values.get() + __partitions.__chunk_count_),
+        std::move(__init),
+        __combiner);
+  }
+
+  template <class _RandomAccessIterator, class _Comp, class _LeafSort>
+  _LIBCPP_HIDE_FROM_ABI static optional<__empty>
+  __stable_sort(_RandomAccessIterator __first, _RandomAccessIterator __last, _Comp __comp, _LeafSort __leaf_sort) {
+    const auto __size = __last - __first;
+    auto __partitions = __std_thread::__partition_chunks(__size);
+
+    if (__partitions.__chunk_count_ == 0)
+      return __empty{};
+
+    if (__partitions.__chunk_count_ == 1) {
+      __leaf_sort(__first, __last, __comp);
+      return __empty{};
+    }
+
+    using _Value = __iterator_value_type<_RandomAccessIterator>;
+
+    auto __destroy = [__size](_Value* __ptr) {
+      std::destroy_n(__ptr, __size);
+      std::allocator<_Value>().deallocate(__ptr, __size);
+    };
+
+    // TODO: use __uninitialized_buffer
+    unique_ptr<_Value[], decltype(__destroy)> __values(std::allocator<_Value>().allocate(__size), __destroy);
+
+    // Initialize all elements to a moved-from state
+    // TODO: Don't do this - this can be done in the first merge - see https://llvm.org/PR63928
+    std::__construct_at(__values.get(), std::move(*__first));
+    for (__iterator_difference_type<_RandomAccessIterator> __i = 1; __i != __size; ++__i) {
+      std::__construct_at(__values.get() + __i, std::move(__values.get()[__i - 1]));
+    }
+    *__first = std::move(__values.get()[__size - 1]);
+
+    __std_thread::__parallel_for(
+        __partitions,
+        __first,
+        [&__leaf_sort, &__comp](_RandomAccessIterator __chunk_first, _RandomAccessIterator __chunk_last) {
+          __leaf_sort(std::move(__chunk_first), std::move(__chunk_last), __comp);
+        });
+
+    bool __objects_are_in_buffer = false;
+    do {
+      const auto __old_chunk_size = __partitions.__chunk_size_;
+      if (__partitions.__chunk_count_ % 2 == 1) {
+        auto __inplace_merge_chunks = [&__comp, &__partitions](auto __first_chunk_begin) {
+          std::inplace_merge(
+              __first_chunk_begin,
+              __first_chunk_begin + __partitions.__first_chunk_size_,
+              __first_chunk_begin + __partitions.__first_chunk_size_ + __partitions.__chunk_size_,
+              __comp);
+        };
+        if (__objects_are_in_buffer)
+          __inplace_merge_chunks(__values.get());
+        else
+          __inplace_merge_chunks(__first);
+        __partitions.__first_chunk_size_ += 2 * __partitions.__chunk_size_;
+      } else {
+        __partitions.__first_chunk_size_ += __partitions.__chunk_size_;
+      }
+
+      __partitions.__chunk_size_ *= 2;
+      __partitions.__chunk_count_ /= 2;
+
+      auto __merge_chunks = [__partitions, __old_chunk_size, &__comp](auto __from_first, auto __to_first) {
+        __std_thread::__parallel_for(
+            __partitions,
+            __from_first,
+            [__old_chunk_size, &__from_first, &__to_first, &__comp](auto __chunk_first, auto __chunk_last) {
+              std::merge(std::make_move_iterator(__chunk_first),
+                         std::make_move_iterator(__chunk_last - __old_chunk_size),
+                         std::make_move_iterator(__chunk_last - __old_chunk_size),
+                         std::make_move_iterator(__chunk_last),
+                         __to_first + (__chunk_first - __from_first),
+                         __comp);
+            });
+      };
+
+      if (__objects_are_in_buffer)
+        __merge_chunks(__values.get(), __first);
+      else
+        __merge_chunks(__first, __values.get());
+      __objects_are_in_buffer = !__objects_are_in_buffer;
+    } while (__partitions.__chunk_count_ > 1);
+
+    if (__objects_are_in_buffer) {
+      std::move(__values.get(), __values.get() + __size, __first);
+    }
+
     return __empty{};
   }
 
+  _LIBCPP_HIDE_FROM_ABI static void __cancel_execution() {}
+
   static constexpr size_t __lane_size = 64;
 };
 
@@ -130,6 +404,7 @@ struct __fill<__std_thread_backend_tag, _ExecutionPolicy>
     : __cpu_parallel_fill<__std_thread_backend_tag, _ExecutionPolicy> {};
 
 } // namespace __pstl
+_LIBCPP_END_EXPLICIT_ABI_ANNOTATIONS
 _LIBCPP_END_NAMESPACE_STD
 
 #endif // _LIBCPP_STD_VER >= 17
diff --git a/libcxx/src/CMakeLists.txt b/libcxx/src/CMakeLists.txt
index e029f9f29d5ea..472552afb6057 100644
--- a/libcxx/src/CMakeLists.txt
+++ b/libcxx/src/CMakeLists.txt
@@ -310,6 +310,12 @@ if (LIBCXX_PSTL_BACKEND STREQUAL "libdispatch")
     )
 endif()
 
+if (LIBCXX_PSTL_BACKEND STREQUAL "std_thread")
+  list(APPEND LIBCXX_EXPERIMENTAL_SOURCES
+    pstl/std_thread.cpp
+    )
+endif()
+
 if (LIBCXX_ENABLE_LOCALIZATION AND LIBCXX_ENABLE_FILESYSTEM AND LIBCXX_ENABLE_TIME_ZONE_DATABASE)
   list(APPEND LIBCXX_EXPERIMENTAL_SOURCES
     experimental/include/tzdb/time_zone_private.h
diff --git a/libcxx/src/pstl/std_thread.cpp b/libcxx/src/pstl/std_thread.cpp
new file mode 100644
index 0000000000000..626ee6754000f
--- /dev/null
+++ b/libcxx/src/pstl/std_thread.cpp
@@ -0,0 +1,433 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <__algorithm/max.h>
+#include <__atomic/aliases.h>
+#include <__atomic/atomic.h>
+#include <__config>
+#include <__mutex/lock_guard.h>
+#include <__mutex/mutex.h>
+#include <__mutex/once_flag.h>
+#include <__mutex/unique_lock.h>
+#include <__new/interference_size.h>
+#include <__pstl/backends/std_thread.h>
+#include <__thread/this_thread.h>
+#include <__thread/thread.h>
+#include <__type_traits/is_trivial.h>
+
+// #define WITH_LOGGING 1
+#ifdef WITH_LOGGING
+#  include <stdio.h>
+#endif
+
+_LIBCPP_BEGIN_NAMESPACE_STD
+_LIBCPP_BEGIN_EXPLICIT_ABI_ANNOTATIONS
+namespace __pstl::__std_thread {
+
+#ifdef WITH_LOGGING
+[[clang::no_destroy]] static std::mutex LOG_MUTEX;
+#  define LOG(...)                                                                                                     \
+    do {                                                                                                               \
+      std::lock_guard lock{LOG_MUTEX};                                                                                 \
+      fprintf(stderr, __VA_ARGS__);                                                                                    \
+    } while (0)
+#else
+#  define LOG(...)                                                                                                     \
+    do {                                                                                                               \
+    } while (0)
+#endif
+
+struct Task;
+
+struct alignas(std::hardware_destructive_interference_size) SPMCQueue {
+  static constexpr unsigned log_initial_size = 6; // 64 elements by default
+
+  SPMCQueue() noexcept : bottom(0), top(0), buffer(alloc_buffer(log_initial_size)) {}
+  SPMCQueue(const SPMCQueue&)            = delete;
+  SPMCQueue& operator=(const SPMCQueue&) = delete;
+  ~SPMCQueue() { ::operator delete(buffer); }
+
+  bool push(Task* task) noexcept {
+    std::uint64_t b    = bottom.load();
+    std::uint64_t t    = top.load();
+    std::uint64_t size = b - t;
+    if (size >= buffer_size(buffer) - 1) {
+      Buffer* grown = grow_buffer(buffer, b, t);
+      if (grown == nullptr) {
+        return false;
+      }
+      Buffer* current;
+      {
+        std::lock_guard lock{mut};
+        current = std::exchange(buffer, grown);
+      }
+      if (current->ref_count.fetch_sub(1) == 1) {
+        ::operator delete(current);
+      }
+    }
+    put(buffer, b, task);
+    bottom.store(b + 1);
+    return true;
+  }
+
+  Task* pop() noexcept {
+    std::uint64_t b = bottom.load();
+    if (b == 0)
+      return nullptr;
+    --b;
+    bottom.store(b);
+    std::uint64_t t = top.load();
+    if (t <= b) {
+      Task* task = get(buffer, b);
+      if (b > t)
+        return task;
+      if (top.compare_exchange_strong(t, t + 1)) {
+        bottom.store(t + 1);
+        return task;
+      } else {
+        bottom.store(t);
+        return nullptr;
+      }
+    } else {
+      bottom.store(t);
+      return nullptr;
+    }
+  }
+
+  Task* steal() noexcept {
+    std::uint64_t t = top.load();
+    std::uint64_t b = bottom.load();
+    if (t >= b)
+      return nullptr;
+
+    Buffer* current;
+    {
+      std::lock_guard lock{mut};
+      current = buffer;
+      current->ref_count.fetch_add(1);
+    }
+    Task* task = get(current, t);
+    if (current->ref_count.fetch_sub(1) == 1) {
+      ::operator delete(current);
+    }
+
+    if (top.compare_exchange_strong(t, t + 1))
+      return task;
+    else
+      return nullptr;
+
+    return nullptr;
+  }
+
+private:
+  struct Buffer {
+    std::uint32_t log_size;
+    std::atomic_int32_t ref_count;
+  };
+
+  static Buffer* alloc_buffer(unsigned log_size) noexcept {
+    std::size_t size  = 1Z << log_size;
+    std::size_t bytes = sizeof(SPMCQueue) + sizeof(Task*) * size;
+    Buffer* buffer    = static_cast<Buffer*>(::operator new(bytes, std::nothrow));
+    if (buffer == nullptr)
+      return nullptr;
+    buffer->log_size  = log_size;
+    buffer->ref_count = 1;
+    return buffer;
+  }
+
+  static Buffer* grow_buffer(Buffer* existing, std::uint64_t bottom, std::uint64_t top) noexcept {
+    Buffer* grown = alloc_buffer(existing->log_size + 1);
+    if (grown == nullptr)
+      return nullptr;
+    for (std::uint64_t i = top; i != bottom; ++i) {
+      put(grown, i, get(existing, i));
+    }
+    return grown;
+  }
+
+  static std::size_t buffer_size(Buffer* buffer) noexcept { return 1Z << buffer->log_size; }
+
+  static Task* get(Buffer* buffer, std::size_t index) noexcept {
+    Task** elements  = reinterpret_cast<Task**>(reinterpret_cast<std::byte*>(buffer) + sizeof(Buffer));
+    std::size_t mask = (1Z << buffer->log_size) - 1;
+    return elements[index & mask];
+  }
+
+  static void put(Buffer* buffer, std::size_t index, Task* task) noexcept {
+    Task** elements        = reinterpret_cast<Task**>(reinterpret_cast<std::byte*>(buffer) + sizeof(Buffer));
+    std::size_t mask       = (1Z << buffer->log_size) - 1;
+    elements[index & mask] = task;
+  }
+
+  std::atomic_uint64_t bottom;
+  std::atomic_uint64_t top;
+  Buffer* buffer;
+  std::mutex mut;
+};
+
+struct Task {
+  void* ctxt;
+  void (*func)(void*, std::size_t);
+  size_t iterations;
+  alignas(std::hardware_destructive_interference_size) std::atomic_size_t index;
+  alignas(std::hardware_destructive_interference_size) std::atomic_size_t requests;
+  alignas(std::hardware_destructive_interference_size) std::mutex m;
+  bool completed = false;
+  std::condition_variable cv;
+};
+
+struct DequeID {
+  static_assert(std::atomic<std::thread::id>::is_always_lock_free);
+  std::atomic<std::thread::id> id = {};
+  std::atomic_size_t aquire_count = 0;
+};
+
+struct Sched {
+private:
+  std::unique_ptr<std::thread[]> m_workers;
+  std::unique_ptr<SPMCQueue[]> m_deques;
+  std::unique_ptr<DequeID[]> m_deque_ids;
+
+  size_t m_workers_count;
+  size_t m_deques_count;
+
+  std::mutex m_workers_sleep_mutex;
+  std::condition_variable m_workers_sleep_cv;
+
+  void worker_thread(size_t worker_num) noexcept;
+
+public:
+  Sched() {
+    size_t cpu_threads = std::thread::hardware_concurrency();
+    if (cpu_threads == 0)
+      cpu_threads = 1; // safe fallback
+
+    m_deques_count = cpu_threads * 2;
+    m_deques       = std::make_unique<SPMCQueue[]>(m_deques_count);
+    m_deque_ids    = std::make_unique<DequeID[]>(m_deques_count);
+
+    m_workers_count = cpu_threads - 1;
+    m_workers       = std::make_unique<std::thread[]>(m_workers_count);
+
+    for (size_t i = 0; i < m_workers_count; ++i) {
+      std::thread worker(&Sched::worker_thread, this, i);
+      m_deque_ids[i].id.store(worker.get_id());
+      m_deque_ids[i].aquire_count.store(1);
+      m_workers[i] = std::move(worker);
+    }
+  }
+
+  void apply(std::size_t iterations, void* ctxt, void (*func)(void*, std::size_t)) noexcept;
+
+  size_t num_workers() const noexcept { return m_workers_count; }
+
+  void wake_workers() noexcept { m_workers_sleep_cv.notify_all(); }
+
+  SPMCQueue* acquire_deque(std::thread::id thread_id) {
+    for (size_t i = 0; i < m_deques_count; ++i) {
+      std::thread::id id = m_deque_ids[i].id.load();
+      if (id == thread_id) {
+        bool is_guest = i >= m_workers_count;
+        if (is_guest) {
+          m_deque_ids[i].aquire_count.fetch_add(1);
+        }
+        return &m_deques[i];
+      }
+    }
+    for (size_t i = 0; i < m_deques_count; ++i) {
+      std::thread::id id = m_deque_ids[i].id.load();
+      if (id == std::thread::id()) {
+        std::thread::id expected;
+        if (m_deque_ids[i].id.compare_exchange_strong(expected, thread_id)) {
+          // Acquired by a new guest thread => set aquire_count to 1
+          m_deque_ids[i].aquire_count.store(1);
+          return &m_deques[i];
+        }
+      }
+    }
+    return nullptr;
+  }
+
+  void release_deque(SPMCQueue* deque) {
+    size_t idx    = static_cast<size_t>(deque - m_deques.get());
+    bool is_guest = idx >= m_workers_count;
+    if (is_guest) {
+      if (m_deque_ids[idx].aquire_count.fetch_sub(1) == 1) {
+        m_deque_ids[idx].id.store(std::thread::id());
+      }
+    }
+  }
+};
+
+static Sched* g_sched = nullptr;
+static std::once_flag g_sched_once_flag;
+
+static Sched* get_sched() {
+  std::call_once(g_sched_once_flag, []() { g_sched = new Sched(); });
+  return g_sched;
+}
+
+static void apply_serial(std::size_t iterations, void* ctxt, void (*func)(void*, std::size_t)) noexcept {
+  for (std::size_t i = 0; i < iterations; ++i) {
+    func(ctxt, i);
+  }
+}
+
+static void process_request(Task* task) {
+  size_t const iterations                = task->iterations;
+  void* const ctxt                       = task->ctxt;
+  void (*const func)(void*, std::size_t) = task->func;
+
+  size_t index;
+  while ((index = task->index.fetch_add(1, std::memory_order_relaxed)) < iterations) {
+    func(ctxt, index);
+  }
+
+  if (task->requests.fetch_sub(1, std::memory_order_release) == 1) {
+    std::lock_guard lock{task->m};
+    task->completed = true;
+    task->cv.notify_one();
+  }
+}
+
+void Sched::worker_thread(size_t worker_num) noexcept {
+  const size_t deques_count    = m_deques_count;
+  constexpr int max_spin_count = 16;
+  int spin_count               = 0;
+
+  while (true) {
+    Task* task = nullptr;
+    //    bool has_task = false;
+
+    // 1st - try to pop from own deque
+    if ((task = m_deques[worker_num].pop())) {
+      LOG("Worker %zu popped task %p from own deque\n", worker_num, task);
+    } else {
+      // 2nd - try to steal from other deques
+      for (size_t i = 0; i != m_deques_count; ++i) {
+        size_t idx = (worker_num + i) % deques_count;
+        if ((task = m_deques[idx].steal())) {
+          LOG("Worker %zu stole task %p from deque %zu\n", worker_num, task, idx);
+          break;
+        }
+      }
+    }
+
+    if (task) {
+      // Do the work
+      process_request(task);
+    } else {
+      // Wait and sleep
+      if (spin_count < max_spin_count) {
+        ++spin_count;
+        std::this_thread::yield();
+      } else {
+        LOG("Worker %zu is going to sleep\n", worker_num);
+        std::unique_lock lock{m_workers_sleep_mutex};
+        m_workers_sleep_cv.wait(lock);
+        spin_count = 0;
+        LOG("Worker %zu woke up\n", worker_num);
+      }
+    }
+  }
+}
+
+void Sched::apply(std::size_t iterations, void* ctxt, void (*func)(void*, std::size_t)) noexcept {
+  if (m_workers_count == 0) {
+    // No worker threads available, fallback to serial execution
+    apply_serial(iterations, ctxt, func);
+    return;
+  }
+
+  std::thread::id thread_id = std::this_thread::get_id();
+  SPMCQueue* deque          = acquire_deque(thread_id);
+  if (deque == nullptr) {
+    apply_serial(iterations, ctxt, func);
+    return;
+  }
+
+  size_t num_requests = std::min(m_workers_count + 1, iterations);
+  Task root_task;
+  root_task.ctxt       = ctxt;
+  root_task.func       = func;
+  root_task.iterations = iterations;
+  root_task.index      = 0;
+  root_task.requests   = num_requests;
+  LOG("New root task %p of %zu iterations with %zu requests\n", &root_task, iterations, num_requests);
+
+  // Push all but one request to the deque, and process the last one directly in this thread.
+  while (num_requests > 1) {
+    if (!deque->push(&root_task)) {
+      break;
+    }
+    --num_requests;
+  }
+
+  wake_workers();
+
+  // Process the last request directly in this thread
+  while (num_requests > 0) {
+    LOG("Guest thread is directly processing task %p, num_requests left: %zu\n", &root_task, num_requests);
+    process_request(&root_task);
+    --num_requests;
+  }
+
+  while (root_task.requests.load() > 0) {
+    Task* task = nullptr;
+
+    // 1st - try to pop from own deque
+    if ((task = deque->pop())) {
+      LOG("Guest thread popped task %p from own deque\n", task);
+    } else {
+      // 2nd - try to steal from other deques
+      for (size_t i = 0; i != m_deques_count; ++i) {
+        if (&m_deques[i] != deque && (task = m_deques[i].steal())) {
+          LOG("Guest stole task %p from deque %zu\n", task, i);
+          break;
+        }
+      }
+    }
+
+    if (task) {
+      // Do the work
+      process_request(task);
+    } else {
+      // Wait a bit and try again
+      std::this_thread::yield();
+    }
+  }
+
+  std::unique_lock lock{root_task.m};
+  root_task.cv.wait(lock, [&root_task] { return root_task.completed; });
+
+  release_deque(deque);
+}
+
+void __apply(size_t __iterations, void* __context, void (*__func)(void* __context, size_t __iteration)) noexcept {
+  if (__iterations < 2) {
+    apply_serial(__iterations, __context, __func);
+    return;
+  }
+  Sched* sched = get_sched();
+  sched->apply(__iterations, __context, __func);
+}
+
+__chunk_partitions __partition_chunks(ptrdiff_t element_count) noexcept {
+  __chunk_partitions partitions;
+  partitions.__chunk_count_      = std::max<ptrdiff_t>(1, element_count / 256);
+  partitions.__chunk_size_       = element_count / partitions.__chunk_count_;
+  partitions.__first_chunk_size_ = element_count - (partitions.__chunk_count_ - 1) * partitions.__chunk_size_;
+  if (partitions.__chunk_count_ == 0 && element_count > 0)
+    partitions.__chunk_count_ = 1;
+  return partitions;
+}
+
+} // namespace __pstl::__std_thread
+_LIBCPP_END_EXPLICIT_ABI_ANNOTATIONS
+_LIBCPP_END_NAMESPACE_STD
diff --git a/libcxx/test/libcxx/algorithms/pstl.std_thread.apply.pass.cpp b/libcxx/test/libcxx/algorithms/pstl.std_thread.apply.pass.cpp
new file mode 100644
index 0000000000000..021504a6eb9cd
--- /dev/null
+++ b/libcxx/test/libcxx/algorithms/pstl.std_thread.apply.pass.cpp
@@ -0,0 +1,52 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// <algorithm>
+
+// REQUIRES: libcpp-pstl-backend-std-thread
+
+// void __apply(size_t __iterations, void* __context, void (*__func)(void* __context, size_t __iteration)) noexcept;
+
+#include <__pstl/backends/std_thread.h>
+#include <algorithm>
+#include <cassert>
+#include <cstddef>
+#include <vector>
+#include <cstdio>
+
+// Fork-bomb Fibonacci implementation.
+int fibb(int n) {
+  if (n <= 1)
+    return n;
+
+  struct Ctx {
+    int n;
+    int n12[2];
+  } ctx{.n = n};
+  std::__pstl::__std_thread::__apply(2, &ctx, [](void* ctxt, std::size_t i) {
+    auto& c  = *static_cast<Ctx*>(ctxt);
+    c.n12[i] = fibb(c.n - 1 - i);
+  });
+
+  return ctx.n12[0] + ctx.n12[1];
+}
+
+// Flat fork-join style application.
+bool flat_fork_join() {
+  std::vector<int> v(1'000'000, 0);
+  std::__pstl::__std_thread::__apply(v.size(), v.data(), [](void* data, std::size_t i) {
+    static_cast<int*>(data)[i] = 42;
+  });
+  return std::all_of(v.begin(), v.end(), [](int x) { return x == 42; });
+}
+
+int main(int, char**) {
+  assert(fibb(25) == 75025);
+  assert(flat_fork_join());
+  return 0;
+}
diff --git a/libcxx/utils/libcxx/test/features/libcxx_macros.py b/libcxx/utils/libcxx/test/features/libcxx_macros.py
index bcc081e2dc2de..c864df4481e05 100644
--- a/libcxx/utils/libcxx/test/features/libcxx_macros.py
+++ b/libcxx/utils/libcxx/test/features/libcxx_macros.py
@@ -34,6 +34,7 @@
     "_LIBCPP_DEPRECATED_ABI_DISABLE_PAIR_TRIVIAL_COPY_CTOR": "libcpp-deprecated-abi-disable-pair-trivial-copy-ctor",
     "_LIBCPP_ABI_NO_COMPRESSED_PAIR_PADDING": "libcpp-abi-no-compressed-pair-padding",
     "_LIBCPP_PSTL_BACKEND_LIBDISPATCH": "libcpp-pstl-backend-libdispatch",
+    "_LIBCPP_PSTL_BACKEND_STD_THREAD": "libcpp-pstl-backend-std-thread",
     "_LIBCPP_ABI_USE_SMALL_DEQUE_BLOCK_SIZE": "libcpp-abi-use-small-deque-block-size",
     "_LIBCPP_ABI_VECTORIZED_MERSENNE_TWISTER_ENGINE": "libcpp-abi-vectorized-mersenne-twister-engine",
 }



More information about the libcxx-commits mailing list