[libc++] nth_element change to partial sort when range is really small

[Backl1ght]

Summary

For really small range, to make the element pointed at by __nth correct, we don't need to selection sort the whole range, we only need to do partial sort on the range so that [__first, __nth] contains the sorted __nth - __first + 1 smallest elements in the range [__first, __last).

And this can be done by only modifying the terminate condition of outer loop in selection sort since element before __first will not be changed anymore in the future round.

The difference is that the latter will end up earlier and make less comparison unless __nth equal to __last - 1 in which case there are same number of comparisons.

[Mordante]

Thanks for your contribution!

Can you show the benchmark results of the improvement, if needed by adding a new benchmark.

libcxx/include/__algorithm/nth_element.h

29

Please make this a _LIBCPP_ASSERT instead.

30

I really love to have some documentation for this function.

33

Why not use __nth? The new variable name does not help me what this does.

34

If __nth == __last this will fail.

37

Is _Compare not deduced?

[Backl1ght]

refine code

[Backl1ght]

For previous diff I just copy code of __selection_sort from sort.h and modify the terminate condition.

[Backl1ght]

It seems like the improvement is too small to make sense in benchmark. But from the output of code bellow we can tell that this reduce some calculations.

#include <algorithm>
#include <iostream>
#include <random>
#include <vector>

int main(int argc, char** argv) {
  const int REPETITIONS = 100'000;
  const int N = 7;

  int seed = 998244353;
  std::mt19937 rng(seed);
  int num_comp = 0;
  std::vector<int> a(N);
  for (int i = 0; i < REPETITIONS; ++i) {
    std::iota(a.begin(), a.end(), 0);
    std::shuffle(a.begin(), a.end(), rng);

    int k = rng() % N;
    std::nth_element(a.begin(), a.begin() + k, a.end(),
                     [&](int x, int y) -> bool {
                       ++num_comp;
                       return x < y;
                     });
  }
  std::cout << "num comp: " << num_comp << std::endl;
  return 0;
}

output before: num comp: 2100000
output after: num comp: 1597554

[Backl1ght]

fix build error by replace U+2212 with U+002d

[Backl1ght]

ping

[Backl1ght]

ping?

[philnik]

In D141614#4053789, @Backl1ght wrote:

It seems like the improvement is too small to make sense in benchmark. But from the output of code bellow we can tell that this reduce some calculations.

#include <algorithm>
#include <iostream>
#include <random>
#include <vector>

int main(int argc, char** argv) {
  const int REPETITIONS = 100'000;
  const int N = 7;

  int seed = 998244353;
  std::mt19937 rng(seed);
  int num_comp = 0;
  std::vector<int> a(N);
  for (int i = 0; i < REPETITIONS; ++i) {
    std::iota(a.begin(), a.end(), 0);
    std::shuffle(a.begin(), a.end(), rng);

    int k = rng() % N;
    std::nth_element(a.begin(), a.begin() + k, a.end(),
                     [&](int x, int y) -> bool {
                       ++num_comp;
                       return x < y;
                     });
  }
  std::cout << "num comp: " << num_comp << std::endl;
  return 0;
}

output before: num comp: 2100000
output after: num comp: 1597554

I the improvement is too small to show up in a benchmark, I don't think it's actually an improvement. Maybe you could use a comparator that isn't trivial (e.g. use std::string)? This part at least shows that it isn't a pessimization for trivial comparators.

libcxx/include/__algorithm/nth_element.h

42

Why not min_element? Using the public API makes it always easier to refactor internal stuff, since that part is stable.

[Backl1ght]

libcxx/include/__algorithm/nth_element.h

42

I use __min_element since selection_sort in sort.h uses __min_element.

I agree with that using the public API will be better. So I tried using min_element but got fail in a unittest defined in ranges_robust_against_proxy_iterators.pass.cpp.

I find out that selection_sort in sort.h used min_element before and changed to __min_element in D129823 to make range algorithms support proxy iterators

[Backl1ght]

benchmark_6454 KBDownload

This is the benchmark result.

[Backl1ght]

benchmark_81928 KBDownload

This is the benchmark result for BM_NthElement_string* after I changed length of each string to 8192.

libcxx/benchmarks/CMakeLists.txt

 #==============================================================================
 set(BENCHMARK_TESTS
     algorithms.partition_point.bench.cpp
     algorithms/equal.bench.cpp
     algorithms/lower_bound.bench.cpp
     algorithms/make_heap.bench.cpp
     algorithms/make_heap_then_sort_heap.bench.cpp
     algorithms/min_max_element.bench.cpp
+    algorithms/nth_element.bench.cpp
     algorithms/pop_heap.bench.cpp
     algorithms/push_heap.bench.cpp
     algorithms/ranges_make_heap.bench.cpp
     algorithms/ranges_make_heap_then_sort_heap.bench.cpp
     algorithms/ranges_pop_heap.bench.cpp
     algorithms/ranges_push_heap.bench.cpp
     algorithms/ranges_sort.bench.cpp
     algorithms/ranges_sort_heap.bench.cpp

libcxx/benchmarks/algorithms/nth_element.bench.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>
+
+#include "common.h"
+
+namespace {
+template <class ValueType, class Order>
+struct NthElement {
+  size_t Quantity;
+
+  void run(benchmark::State& state) const {
+    runOpOnCopies<ValueType>(state, Quantity, Order(), BatchSize::CountElements, [](auto& Copy) {
+      size_t k = Copy.size() / 2;
+      std::nth_element(Copy.begin(), Copy.begin() + k, Copy.end());
+    });
+  }
+
+  std::string name() const {
+    return "BM_NthElement" + ValueType::name() + Order::name() + "_" + std::to_string(Quantity);
+  };
+};
+} // namespace
+
+int main(int argc, char** argv) {
+  benchmark::Initialize(&argc, argv);
+  if (benchmark::ReportUnrecognizedArguments(argc, argv))
+    return 1;
+  makeCartesianProductBenchmark<NthElement, AllValueTypes, AllOrders>(Quantities);
+  benchmark::RunSpecifiedBenchmarks();
+}
 No newline at end of file

libcxx/include/__algorithm/nth_element.h

 //===----------------------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef _LIBCPP___ALGORITHM_NTH_ELEMENT_H
 #define _LIBCPP___ALGORITHM_NTH_ELEMENT_H
 
 #include <__algorithm/comp.h>
 #include <__algorithm/comp_ref_type.h>
 #include <__algorithm/iterator_operations.h>
+#include <__algorithm/min_element.h>
 #include <__algorithm/sort.h>
 #include <__config>
 #include <__debug>
 #include <__debug_utils/randomize_range.h>
 #include <__iterator/iterator_traits.h>
 #include <__utility/move.h>
 
 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
 #  pragma GCC system_header
 #endif
 
 _LIBCPP_BEGIN_NAMESPACE_STD
 
+// Implement partial sort base on selection sort. Rearranges elements such that the range [__first, __middle) contains
+// the sorted __middle - __first smallest elements in the range [__first, __last).
+// For selection sort, after each round, the position of element that __first represents is determined. That is, if we
+// end up when __first == __middle, [__first, __middle) contains the sorted __middle - __first smallest elements in the
+// range [__first, __last) already.
+template <class _AlgPolicy, class _Compare, class _BidirectionalIterator>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void __selection_partial_sort(
+    _BidirectionalIterator __first, _BidirectionalIterator __middle, _BidirectionalIterator __last, _Compare __comp) {
+  // if [__first, __last - 1) is determined, then [__first, __last) is determined.
+  if (__middle == __last)
+    --__middle;
+
+  for (; __first != __middle; ++__first) {
+    _BidirectionalIterator __i = std::__min_element(__first, __last, __comp);
+    if (__i != __first)
+      _IterOps<_AlgPolicy>::iter_swap(__first, __i);
+  }
+}
+
 template<class _Compare, class _RandomAccessIterator>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 bool
 __nth_element_find_guard(_RandomAccessIterator& __i, _RandomAccessIterator& __j,
                          _RandomAccessIterator __m, _Compare __comp)
 {
     // manually guard downward moving __j against __i
     while (true) {
         if (__i == --__j) {
while (true)
             {
             _RandomAccessIterator __m = __first;
             std::__sort3<_AlgPolicy, _Compare>(__first, ++__m, --__last, __comp);
             return;
             }
         }
         if (__len <= __limit)
         {
-            std::__selection_sort<_AlgPolicy, _Compare>(__first, __last, __comp);
+            std::__selection_partial_sort<_AlgPolicy, _Compare>(__first, __nth + difference_type(1), __last, __comp);
+            // [__first, __nth) <= *__nth and *__nth <= (__nth, __last)
             return;
         }
         // __len > __limit >= 3
         _RandomAccessIterator __m = __first + __len/2;
         _RandomAccessIterator __lm1 = __last;
         unsigned __n_swaps = std::__sort3<_AlgPolicy, _Compare>(__first, __m, --__lm1, __comp);
         // *__m is median
         // partition [__first, __m) < *__m and *__m <= [__m, __last)