[libcxx] nth_element use introselect to avoid quadratic behavior

[Backl1ght]

Summary

Current implementation of nth_element will be quadratic under some cases, use introselect to avoid quadratic behavior.

I also add a benchmark for nth_element.

related issue https://github.com/llvm/llvm-project/issues/52747

[Backl1ght]

Here is the benchmark result for this patch.

-------------------------------------------------------------------------------------------
Benchmark                                                 Time             CPU   Iterations
-------------------------------------------------------------------------------------------
BM_NthElement_uint32_QuickSelectAdversary_4            7.97 ns         7.96 ns     87818240
BM_NthElement_uint32_QuickSelectAdversary_16           64.9 ns         64.8 ns     10928128
BM_NthElement_uint32_QuickSelectAdversary_64            712 ns          711 ns       991232
BM_NthElement_uint32_QuickSelectAdversary_256          9542 ns         9537 ns        73728
BM_NthElement_uint32_QuickSelectAdversary_1024       132177 ns       132090 ns         5376
BM_NthElement_uint32_QuickSelectAdversary_16384    32366447 ns     32343198 ns           16
BM_NthElement_uint32_QuickSelectAdversary_262144 8268646976 ns   8260294657 ns            1
BM_NthElement_uint64_QuickSelectAdversary_4            7.98 ns         7.97 ns     88211456
BM_NthElement_uint64_QuickSelectAdversary_16           73.5 ns         73.3 ns      9814016
BM_NthElement_uint64_QuickSelectAdversary_64            977 ns          976 ns       737280
BM_NthElement_uint64_QuickSelectAdversary_256         13829 ns        13817 ns        51200
BM_NthElement_uint64_QuickSelectAdversary_1024       195468 ns       195371 ns         3584
BM_NthElement_uint64_QuickSelectAdversary_16384    47547973 ns     47521170 ns           16
BM_NthElement_uint64_QuickSelectAdversary_262144 12350804915 ns   12340965027 ns            1
-------------------------------------------------------------------------------------------
Benchmark                                                 Time             CPU   Iterations
-------------------------------------------------------------------------------------------
BM_NthElement_uint32_QuickSelectAdversary_4            9.70 ns         9.69 ns     72351744
BM_NthElement_uint32_QuickSelectAdversary_16           66.9 ns         66.8 ns     10387456
BM_NthElement_uint32_QuickSelectAdversary_64            676 ns          676 ns      1028096
BM_NthElement_uint32_QuickSelectAdversary_256          3950 ns         3947 ns       178176
BM_NthElement_uint32_QuickSelectAdversary_1024        19214 ns        19203 ns        36608
BM_NthElement_uint32_QuickSelectAdversary_16384      417974 ns       417712 ns         1680
BM_NthElement_uint32_QuickSelectAdversary_262144    8901424 ns      8895905 ns           79
BM_NthElement_uint64_QuickSelectAdversary_4            10.5 ns         10.5 ns     68026368
BM_NthElement_uint64_QuickSelectAdversary_16           71.9 ns         71.9 ns      9732096
BM_NthElement_uint64_QuickSelectAdversary_64            835 ns          834 ns       839680
BM_NthElement_uint64_QuickSelectAdversary_256          5169 ns         5166 ns       136192
BM_NthElement_uint64_QuickSelectAdversary_1024        25318 ns        25302 ns        27904
BM_NthElement_uint64_QuickSelectAdversary_16384      565411 ns       565056 ns         1248
BM_NthElement_uint64_QuickSelectAdversary_262144   11918315 ns     11911538 ns           59
RUNNING: ./nth_element.libcxx.before.out --benchmark_filter=BM_NthElement_(uint32|uint64)_QuickSelectAdversary.* --benchmark_out=/tmp/tmpm2fwl5zr
RUNNING: ./nth_element.libcxx.after.out --benchmark_filter=BM_NthElement_(uint32|uint64)_QuickSelectAdversary.* --benchmark_out=/tmp/tmpeoej8bia
Comparing ./nth_element.libcxx.before.out to ./nth_element.libcxx.after.out
Benchmark                                                          Time             CPU      Time Old      Time New       CPU Old       CPU New
-----------------------------------------------------------------------------------------------------------------------------------------------
BM_NthElement_uint32_QuickSelectAdversary_4                     +0.2172         +0.2170             8            10             8            10
BM_NthElement_uint32_QuickSelectAdversary_16                    +0.0304         +0.0308            65            67            65            67
BM_NthElement_uint32_QuickSelectAdversary_64                    -0.0500         -0.0500           712           676           711           676
BM_NthElement_uint32_QuickSelectAdversary_256                   -0.5861         -0.5861          9542          3950          9537          3947
BM_NthElement_uint32_QuickSelectAdversary_1024                  -0.8546         -0.8546        132177         19214        132090         19203
BM_NthElement_uint32_QuickSelectAdversary_16384                 -0.9871         -0.9871      32366447        417974      32343198        417712
BM_NthElement_uint32_QuickSelectAdversary_262144                -0.9989         -0.9989    8268646976       8901424    8260294657       8895905
BM_NthElement_uint64_QuickSelectAdversary_4                     +0.3182         +0.3182             8            11             8            11
BM_NthElement_uint64_QuickSelectAdversary_16                    -0.0206         -0.0192            73            72            73            72
BM_NthElement_uint64_QuickSelectAdversary_64                    -0.1454         -0.1451           977           835           976           834
BM_NthElement_uint64_QuickSelectAdversary_256                   -0.6262         -0.6261         13829          5169         13817          5166
BM_NthElement_uint64_QuickSelectAdversary_1024                  -0.8705         -0.8705        195468         25318        195371         25302
BM_NthElement_uint64_QuickSelectAdversary_16384                 -0.9881         -0.9881      47547973        565411      47521170        565056
BM_NthElement_uint64_QuickSelectAdversary_262144                -0.9990         -0.9990   12350804915      11918315   12340965027      11911538
OVERALL_GEOMEAN                                                 -0.8659         -0.8658             0             0             0             0

[Backl1ght]

There are some other select algorithms outperform introselect, ie pdqselect, adqselect. But I think we could prevent quadratic behavior first, and then pursue extreme performance.

[Backl1ght]

fix error: ADL lookup

[philnik]

In D141205#4033709, @Backl1ght wrote:

There are some other select algorithms outperform introselect, ie pdqselect, adqselect. But I think we could prevent quadratic behavior first, and then pursue extreme performance.

Do you plan to make a patch for one of them? If yes, it's probably easier to just go for it directly, since it's probably a re-implementation anyways?

[Backl1ght]

In D141205#4033996, @philnik wrote:

In D141205#4033709, @Backl1ght wrote:

There are some other select algorithms outperform introselect, ie pdqselect, adqselect. But I think we could prevent quadratic behavior first, and then pursue extreme performance.

Do you plan to make a patch for one of them? If yes, it's probably easier to just go for it directly, since it's probably a re-implementation anyways?

I do plan to make one, and I agree with that maybe go for one of them directly will be better.

[Mordante]

Thanks for working on this. Can you also provide benchmarks to see the effect on the other kinds of data?

libcxx/include/__algorithm/nth_element.h

104

Can you add comment why this is done? Currently the comment and the code tell the same thing.
For maintainability it would be very good to know why we switch the algorithm.

[danlark]

libcxx/include/__algorithm/nth_element.h

104

Even though I am in favor of this patch, this does not fully fix the issue. With the comparator that changes the inputs on the fly, it's possible to get into a O(n log n) average case.

Otherwise LGTM, thanks!

[Backl1ght]

libcxx/include/__algorithm/nth_element.h

104

I think the average time complexity of introselect should be O(n), since the average time complexity of quickselect is O(n), and according to the original paper of introselect:

If as in introsort we limit the depth to a logarithmic bound the average time remains linear but the worst case is O(N log N).

[danlark]

libcxx/include/__algorithm/nth_element.h

104

The example from the link https://github.com/llvm/llvm-project/issues/52747#issuecomment-996916909 will still be of average case O(N log N)

[huixie90]

libcxx/include/__algorithm/nth_element.h

52

perhaps use typename _IterOps<_AlgPolicy>::__diference_type

104

I think that example should be UB. I can't remember the legacy version's wording, but I am pretty sure the ranges version requires the predicate to be "equality-preserving". The same inputs to the predicate should return the same output. But that example's predicate seems to modify the states in such way that same inputs can result in different output

105

nit: _Compare can be deduced

libcxx/benchmarks/CMakeLists.txt

 # Register Benchmark tests
 #==============================================================================
 set(BENCHMARK_TESTS
     algorithms.partition_point.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/common.h

 enum class Order {
   Random,
   Ascending,
   Descending,
   SingleElement,
   PipeOrgan,
   Heap,
   QuickSortAdversary,
+  QuickSelectAdversary,
 };
-struct AllOrders : EnumValuesAsTuple<AllOrders, Order, 7> {
+struct AllOrders : EnumValuesAsTuple<AllOrders, Order, 8> {
   static constexpr const char* Names[] = {"Random",     "Ascending",
                                           "Descending", "SingleElement",
                                           "PipeOrgan",  "Heap",
-                                          "QuickSortAdversary"};
+                                          "QuickSortAdversary", "QuickSelectAdversary"};
 };
 
 // fillAdversarialQuickSortInput fills the input vector with N int-like values.
 // These values are arranged in such a way that they would invoke O(N^2)
 // behavior on any quick sort implementation that satisifies certain conditions.
 // Details are available in the following paper:
 // "A Killer Adversary for Quicksort", M. D. McIlroy, Software-Practice &
 // Experience Volume 29 Issue 4 April 10, 1999 pp 341-344.
if (V[x] == gas) {
       candidate = x;
     } else if (V[y] == gas) {
       candidate = y;
     }
     return V[x] < V[y];
   });
 }
 
+// fillAdversarialQuickSelectInput fills the input vector with N int-like values.
+// These values are arranged in such a way that they will make quickselect using median
+// of 3 quadratic.
+template <class T>
+void fillAdversarialQuickSelectInput(T& V, size_t N) {
+  assert(N > 0);
+
+  V.reserve(N);
+  size_t K = N / 2;
+  for (size_t i = 0; i < K; ++i) {
+    if (i & 1) {
+      V.push_back(K + i - 1);
+    } else {
+      V.push_back(i);
+    }
+  }
+  for (size_t i = 0; i < K; ++i)
+    V.push_back(2 * i);
+}
+
 template <typename T>
 void fillValues(std::vector<T>& V, size_t N, Order O) {
   if (O == Order::SingleElement) {
     V.resize(N, 0);
   } else if (O == Order::QuickSortAdversary) {
     fillAdversarialQuickSortInput(V, N);
+  } else if (O == Order::QuickSelectAdversary) {
+    fillAdversarialQuickSelectInput(V, N);
   } else {
     while (V.size() < N)
       V.push_back(V.size());
   }
 }
 
 template <typename T>
 void fillValues(std::vector<std::pair<T, T> >& V, size_t N, Order O) {
void sortValues(T& V, Order O) {
   case Order::PipeOrgan:
     std::sort(V.begin(), V.end());
     std::reverse(V.begin() + V.size() / 2, V.end());
     break;
   case Order::Heap:
     std::make_heap(V.begin(), V.end());
     break;
   case Order::QuickSortAdversary:
+  case Order::QuickSelectAdversary:
     // Nothing to do
     break;
   }
 }
 
 constexpr size_t TestSetElements =
 #if !TEST_HAS_FEATURE(memory_sanitizer)
     1 << 18;

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 <random>
+
+#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::CountBatch,
+        [](auto& Copy) {
+            size_t k = Copy.size() / 2;
+            std::nth_element(Copy.begin(), Copy.begin() + k, Copy.end());
+        });
+  }
+
+  bool skip() const { return Order() == ::Order::QuickSelectAdversary && Quantity == 1; }
+
+  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();
+}

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/make_heap.h>
+#include <__algorithm/sift_down.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)
while (true) {
         if (__comp(*__j, *__m)) {
             return true;  // found guard for downward moving __j, now use unguarded partition
         }
     }
 }
 
 template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
+__heap_select(_RandomAccessIterator __first, _RandomAccessIterator __middle, _RandomAccessIterator __last, _Compare __comp)
+{
+    std::__make_heap<_AlgPolicy>(__first, __middle, __comp);
+
+    typename iterator_traits<_RandomAccessIterator>::difference_type __len = __middle - __first;
+    _RandomAccessIterator __i = __middle;
+    for (; __i != __last; ++__i)
+    {
+        if (__comp(*__i, *__first))
+        {
+            _IterOps<_AlgPolicy>::iter_swap(__i, __first);
+            std::__sift_down<_AlgPolicy>(__first, __comp, __len, __first);
+        }
+    }
+
+    _IterOps<_AlgPolicy>::iter_swap(__first, __middle);
+}
+
+template <class _AlgPolicy, class _Compare, class _RandomAccessIterator>
+_LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 void
 __nth_element(_RandomAccessIterator __first, _RandomAccessIterator __nth, _RandomAccessIterator __last, _Compare __comp)
 {
     using _Ops = _IterOps<_AlgPolicy>;
 
     // _Compare is known to be a reference type
     typedef typename iterator_traits<_RandomAccessIterator>::difference_type difference_type;
     const difference_type __limit = 7;
+    difference_type __depth = 2 * std::__log2i(__last - __first);
     while (true)
     {
         if (__nth == __last)
             return;
         difference_type __len = __last - __first;
         switch (__len)
         {
         case 0:
while (true)
             return;
             }
         }
         if (__len <= __limit)
         {
             std::__selection_sort<_AlgPolicy, _Compare>(__first, __last, __comp);
             return;
         }
+        if (__depth == 0)
+        {
+            // fall back to heap select.
+            __heap_select<_AlgPolicy, _Compare>(__first, __nth, __last, __comp);
+            return;
+        }
+        --__depth;
         // __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)
         // (this inhibits tossing elements equivalent to __m around unnecessarily)
         _RandomAccessIterator __i = __first;

libcxx/include/__algorithm/sort.h

while (true) {
     std::__introsort<_AlgPolicy, _Compare, _RandomAccessIterator, _UseBitSetPartition>(
         __first, __i, __comp, __depth, __leftmost);
     __leftmost = false;
     __first    = ++__i;
   }
 }
 
 template <typename _Number>
-inline _LIBCPP_HIDE_FROM_ABI _Number __log2i(_Number __n) {
+inline _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX14 _Number __log2i(_Number __n) {
   if (__n == 0)
     return 0;
   if (sizeof(__n) <= sizeof(unsigned))
     return sizeof(unsigned) * CHAR_BIT - 1 - __libcpp_clz(static_cast<unsigned>(__n));
   if (sizeof(__n) <= sizeof(unsigned long))
     return sizeof(unsigned long) * CHAR_BIT - 1 - __libcpp_clz(static_cast<unsigned long>(__n));
   if (sizeof(__n) <= sizeof(unsigned long long))
     return sizeof(unsigned long long) * CHAR_BIT - 1 - __libcpp_clz(static_cast<unsigned long long>(__n));