[libc++][PSTL] Simplify the partitioning algorithm until we have better data to know how to chunk better

[philnik]

Summary

The current chunking strategy is very bad for sorting, and we don't really know how to chunk in general. This fixes the performance problem for sorting.

[ldionne]

libcxx/src/pstl/libdispatch.cpp

12

You can likely minimize some includes.

libcxx/test/libcxx/algorithms/pstl.libdispatch.chunk_partitions.pass.cpp

11

This shouldn't be needed even in your WIP patches. This should be enabled by default on the apple-system job.

[philnik]

Address comments

[philnik]

Try to fix CI

[ldionne]

Requesting changes until the tests are fixed.

[philnik]

• Rebased
• Try to fix CI

[ldionne]

LGTM w/ test and passing CI, assuming fixing the CI is just a matter of rebasing and poking the CI.

libcxx/src/pstl/libdispatch.cpp

31–33

Can you add tests that check the exact values for element_count == 0 and 1?

[philnik]

Address comments

libcxx/src/pstl/libdispatch.cpp

31–33

We already have that.

libcxx/include/__algorithm/pstl_backends/cpu_backends/libdispatch.h

 
 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 pmr::memory_resource* __get_memory_resource();
-[[__gnu__::__const__]] _LIBCPP_EXPORTED_FROM_ABI __chunk_partitions __partition_chunks(ptrdiff_t __size);
+[[__gnu__::__const__]] _LIBCPP_EXPORTED_FROM_ABI __chunk_partitions __partition_chunks(ptrdiff_t __size) noexcept;
 
 template <class _RandomAccessIterator, class _Functor>
 _LIBCPP_HIDE_FROM_ABI void
 __parallel_for(_RandomAccessIterator __first, _RandomAccessIterator __last, _Functor __func) {
   auto __partitions = __libdispatch::__partition_chunks(__last - __first);
 
   // Perform the chunked execution.
   __libdispatch::__dispatch_apply(__partitions.__chunk_count_, [&](size_t __chunk) {

libcxx/src/pstl/libdispatch.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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/min.h>
 #include <__algorithm/pstl_backends/cpu_backends/libdispatch.h>
 #include <__config>
 #include <dispatch/dispatch.h>
 #include <memory_resource>
 #include <thread>
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
 namespace __par_backend::inline __libdispatch {
 
 pmr::memory_resource* __get_memory_resource() {
   static std::pmr::synchronized_pool_resource pool{pmr::new_delete_resource()};
   return &pool;
 }
 
 void __dispatch_apply(size_t chunk_count, void* context, void (*func)(void* context, size_t chunk)) noexcept {
   ::dispatch_apply_f(chunk_count, DISPATCH_APPLY_AUTO, context, func);
 }
 
-__chunk_partitions __partition_chunks(ptrdiff_t element_count) {
+__chunk_partitions __partition_chunks(ptrdiff_t element_count) noexcept {
   __chunk_partitions partitions;
-  partitions.__chunk_count_ = [&] {
-    ptrdiff_t cores = std::max(1u, thread::hardware_concurrency());
-
-    auto medium = [&](ptrdiff_t n) { return cores + ((n - cores) / cores); };
-
-    // This is an approximation of `log(1.01, sqrt(n))` which seemes to be reasonable for `n` larger than 500 and tops
-    // at 800 tasks for n ~ 8 million
-    auto large = [](ptrdiff_t n) { return static_cast<ptrdiff_t>(100.499 * std::log(std::sqrt(n))); };
-
-    if (element_count < cores)
-      return element_count;
-    else if (element_count < 500)
-      return medium(element_count);
-    else
-      return std::min(medium(element_count), large(element_count)); // provide a "smooth" transition
-  }();
-  partitions.__chunk_size_       = element_count / partitions.__chunk_count_;
-  partitions.__first_chunk_size_ = partitions.__chunk_size_;
-
-  const ptrdiff_t leftover_item_count = element_count - (partitions.__chunk_count_ * partitions.__chunk_size_);
-
-  if (leftover_item_count == 0)
-    return partitions;
-
-  if (leftover_item_count == partitions.__chunk_size_) {
-    partitions.__chunk_count_ += 1;
-    return partitions;
-  }
-
-  const ptrdiff_t n_extra_items_per_chunk = leftover_item_count / partitions.__chunk_count_;
-  const ptrdiff_t n_final_leftover_items  = leftover_item_count - (n_extra_items_per_chunk * partitions.__chunk_count_);
-
-  partitions.__chunk_size_ += n_extra_items_per_chunk;
-  partitions.__first_chunk_size_ = partitions.__chunk_size_ + n_final_leftover_items;
+  partitions.__chunk_count_      = element_count / 256;
+  partitions.__chunk_size_       = partitions.__chunk_count_ * 256;
+  partitions.__first_chunk_size_ = element_count - (partitions.__chunk_count_ - 1) * partitions.__chunk_size_;
   return partitions;
 }
 
 // NOLINTNEXTLINE(llvm-namespace-comment) // This is https://llvm.org/PR56804
 } // namespace __par_backend::inline __libdispatch
 
 _LIBCPP_END_NAMESPACE_STD

libcxx/test/libcxx/algorithms/pstl.libdispatch.chunk_partitions.pass.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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-cpu-backend-libdispatch
+// REQUIRES-: libcpp-pstl-cpu-backend-libdispatch
 
 // ADDITIONAL_COMPILE_FLAGS: -Wno-private-header
 
 // __chunk_partitions __partition_chunks(ptrdiff_t);
 
 #include <__algorithm/pstl_backends/cpu_backends/libdispatch.h>
 #include <cassert>
 #include <cstddef>