numpy · mattip · Oct 1, 2023 · Jul 13, 2023 · Jun 29, 2023 · Jul 13, 2023
diff --git a/benchmarks/benchmarks/bench_function_base.py b/benchmarks/benchmarks/bench_function_base.py
@@ -132,7 +132,7 @@ def memoize(f):
     def wrapped(*args):
         if args not in _memoized:
             _memoized[args] = f(*args)
-        
+
         return _memoized[args].copy()
 
     return f
@@ -154,7 +154,7 @@ def random(size, dtype):
         arr = np.arange(size, dtype=dtype)
         np.random.shuffle(arr)
         return arr
-    
+
     @staticmethod
     @memoize
     def ordered(size, dtype):
@@ -237,7 +237,6 @@ def random_bubble(cls, size, dtype, bubble_num, bubble_size=None):
 
         return cls.random_unsorted_area(size, dtype, frac, bubble_size)
 
-
 class Sort(Benchmark):
     """
     This benchmark tests sorting performance with several
@@ -288,6 +287,37 @@ def time_argsort(self, kind, dtype, array_type):
         np.argsort(self.arr, kind=kind)
 
 
+class Partition(Benchmark):
+    params = [
+        ['float64', 'int64', 'float32', 'int32', 'int16', 'float16'],
+        [
+            ('random',),
+            ('ordered',),
+            ('reversed',),
+            ('uniform',),
+            ('sorted_block', 10),
+            ('sorted_block', 100),
+            ('sorted_block', 1000),
+        ],
+        [10, 100, 1000],
+    ]
+    param_names = ['dtype', 'array_type', 'k']
+
+    # The size of the benchmarked arrays.
+    ARRAY_SIZE = 100000
+
+    def setup(self, dtype, array_type, k):
+        np.random.seed(1234)
+        array_class = array_type[0]
+        self.arr = getattr(SortGenerator, array_class)(self.ARRAY_SIZE,
+                dtype, *array_type[1:])
+
+    def time_partition(self, dtype, array_type, k):
+        temp = np.partition(self.arr, k)
+
+    def time_argpartition(self, dtype, array_type, k):
+        temp = np.argpartition(self.arr, k)
+
 class SortWorst(Benchmark):
     def setup(self):
         # quicksort median of 3 worst case

diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py
@@ -4239,7 +4239,7 @@
     >>> a = np.array([3, 4, 2, 1])
     >>> a.partition(3)
     >>> a
-    array([2, 1, 3, 4])
+    array([2, 1, 3, 4]) # may vary
 
     >>> a.partition((1, 3))
     >>> a

diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
@@ -740,7 +740,7 @@ def partition(a, kth, axis=-1, kind='introselect', order=None):
     >>> a = np.array([7, 1, 7, 7, 1, 5, 7, 2, 3, 2, 6, 2, 3, 0])
     >>> p = np.partition(a, 4)
     >>> p
-    array([0, 1, 2, 1, 2, 5, 2, 3, 3, 6, 7, 7, 7, 7])
+    array([0, 1, 2, 1, 2, 5, 2, 3, 3, 6, 7, 7, 7, 7]) # may vary
 
     ``p[4]`` is 2;  all elements in ``p[:4]`` are less than or equal
     to ``p[4]``, and all elements in ``p[5:]`` are greater than or
@@ -838,13 +838,13 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
 
     >>> x = np.array([3, 4, 2, 1])
     >>> x[np.argpartition(x, 3)]
-    array([2, 1, 3, 4])
+    array([2, 1, 3, 4]) # may vary
     >>> x[np.argpartition(x, (1, 3))]
-    array([1, 2, 3, 4])
+    array([1, 2, 3, 4]) # may vary
 
     >>> x = [3, 4, 2, 1]
     >>> np.array(x)[np.argpartition(x, 3)]
-    array([2, 1, 3, 4])
+    array([2, 1, 3, 4]) # may vary
 
     Multi-dimensional array:
 

diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
@@ -400,12 +400,6 @@ typedef int (PyArray_FillFunc)(void *, npy_intp, void *);
 
 typedef int (PyArray_SortFunc)(void *, npy_intp, void *);
 typedef int (PyArray_ArgSortFunc)(void *, npy_intp *, npy_intp, void *);
-typedef int (PyArray_PartitionFunc)(void *, npy_intp, npy_intp,
-                                    npy_intp *, npy_intp *,
-                                    void *);
-typedef int (PyArray_ArgPartitionFunc)(void *, npy_intp *, npy_intp, npy_intp,
-                                       npy_intp *, npy_intp *,
-                                       void *);
 
 typedef int (PyArray_FillWithScalarFunc)(void *, npy_intp, void *, void *);
 

diff --git a/numpy/core/src/common/npy_partition.h b/numpy/core/src/common/npy_partition.h
@@ -11,12 +11,19 @@
 
 #define NPY_MAX_PIVOT_STACK 50
 
+typedef int (PyArray_PartitionFunc)(void *, npy_intp, npy_intp,
+                                    npy_intp *, npy_intp *, npy_intp,
+                                    void *);
+typedef int (PyArray_ArgPartitionFunc)(void *, npy_intp *, npy_intp, npy_intp,
+                                       npy_intp *, npy_intp *, npy_intp,
+                                       void *);
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 NPY_NO_EXPORT PyArray_PartitionFunc *
 get_partition_func(int type, NPY_SELECTKIND which);
+
 NPY_NO_EXPORT PyArray_ArgPartitionFunc *
 get_argpartition_func(int type, NPY_SELECTKIND which);
 

diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
@@ -832,7 +832,7 @@ npy_fastrepeat(
 
     return npy_fastrepeat_impl(
         n_outer, n, nel, chunk, broadcast, counts, new_data, old_data, elsize,
-        cast_info, needs_refcounting);    
+        cast_info, needs_refcounting);
 }
 
 
@@ -1255,7 +1255,7 @@ _new_sortlike(PyArrayObject *op, int axis, PyArray_SortFunc *sort,
             npy_intp npiv = 0;
             npy_intp i;
             for (i = 0; i < nkth; ++i) {
-                ret = part(bufptr, N, kth[i], pivots, &npiv, op);
+                ret = part(bufptr, N, kth[i], pivots, &npiv, nkth, op);
                 if (needs_api && PyErr_Occurred()) {
                     ret = -1;
                 }
@@ -1439,7 +1439,7 @@ _new_argsortlike(PyArrayObject *op, int axis, PyArray_ArgSortFunc *argsort,
             npy_intp npiv = 0;
 
             for (i = 0; i < nkth; ++i) {
-                ret = argpart(valptr, idxptr, N, kth[i], pivots, &npiv, op);
+                ret = argpart(valptr, idxptr, N, kth[i], pivots, &npiv, nkth, op);
                 /* Object comparisons may raise an exception in Python 3 */
                 if (needs_api && PyErr_Occurred()) {
                     ret = -1;

diff --git a/numpy/core/src/npysort/quicksort.cpp b/numpy/core/src/npysort/quicksort.cpp
@@ -69,23 +69,12 @@
 #define SMALL_MERGESORT 20
 #define SMALL_STRING 16
 
-// Temporarily disable AVX512 sorting on WIN32 and CYGWIN until we can figure
+// Disable AVX512 sorting on CYGWIN until we can figure
 // out why it has test failures
-#if defined(_MSC_VER) || defined(__CYGWIN__)
-template<typename T>
-inline bool quicksort_dispatch(T*, npy_intp)
-{
-    return false;
-}
-template<typename T>
-inline bool aquicksort_dispatch(T*, npy_intp*, npy_intp)
-{
-    return false;
-}
-#else
 template<typename T>
 inline bool quicksort_dispatch(T *start, npy_intp num)
 {
+#if !defined(__CYGWIN__)
     using TF = typename np::meta::FixedWidth<T>::Type;
     void (*dispfunc)(TF*, intptr_t) = nullptr;
     if (sizeof(T) == sizeof(uint16_t)) {
@@ -104,12 +93,15 @@ inline bool quicksort_dispatch(T *start, npy_intp num)
         (*dispfunc)(reinterpret_cast<TF*>(start), static_cast<intptr_t>(num));
         return true;
     }
+#endif // __CYGWIN__
+    (void)start; (void)num; // to avoid unused arg warn
     return false;
 }
 
 template<typename T>
 inline bool aquicksort_dispatch(T *start, npy_intp* arg, npy_intp num)
 {
+#if !defined(__CYGWIN__)
     using TF = typename np::meta::FixedWidth<T>::Type;
     void (*dispfunc)(TF*, npy_intp*, npy_intp) = nullptr;
     #ifndef NPY_DISABLE_OPTIMIZATION
@@ -124,9 +116,10 @@ inline bool aquicksort_dispatch(T *start, npy_intp* arg, npy_intp num)
         (*dispfunc)(reinterpret_cast<TF*>(start), arg, num);
         return true;
     }
+#endif // __CYGWIN__
+    (void)start; (void)arg; (void)num; // to avoid unused arg warn
     return false;
 }
-#endif // _MSC_VER || CYGWIN
 
 /*
  *****************************************************************************

diff --git a/numpy/core/src/npysort/selection.cpp b/numpy/core/src/npysort/selection.cpp
@@ -25,13 +25,83 @@
 #include <array>
 #include <cstdlib>
 #include <utility>
+#include "simd_qsort.hpp"
 
 #define NOT_USED NPY_UNUSED(unused)
 
+template<typename T>
+inline bool quickselect_dispatch(T* v, npy_intp num, npy_intp kth)
+{
+#ifndef __CYGWIN__
+    /*
+     * Only defined for int16_t, uint16_t, float16, int32_t, uint32_t, float32,
+     * int64_t, uint64_t, double
+     * TODO: Enable quickselect for 32-bit. np.argpartition is only vectorized
+     * for 64-bit when npy_intp is 8 bytes, which means np.partition and
+     * np.argpartition will produces different results on 32-bit systems.
+     * Several tests in test_multiarray.py rely on these results being
+     * identical. We should get rid of this constraint and re-write
+     * these tests once we enable this for 32-bit.
+     */
+    if constexpr (
+        (std::is_integral_v<T> || std::is_floating_point_v<T> || std::is_same_v<T, np::Half>) &&
+        (sizeof(T) == sizeof(uint16_t) || sizeof(T) == sizeof(uint32_t) || sizeof(T) == sizeof(uint64_t)) &&
+        sizeof(npy_intp) == sizeof(int64_t)) {
+        using TF = typename np::meta::FixedWidth<T>::Type;
+        void (*dispfunc)(TF*, npy_intp, npy_intp) = nullptr;
+        if constexpr (sizeof(T) == sizeof(uint16_t)) {
+            #ifndef NPY_DISABLE_OPTIMIZATION
+                #include "simd_qsort_16bit.dispatch.h"
+            #endif
+            NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSelect, <TF>);
+        }
+        else if constexpr (sizeof(T) == sizeof(uint32_t) || sizeof(T) == sizeof(uint64_t)) {
+            #ifndef NPY_DISABLE_OPTIMIZATION
+                #include "simd_qsort.dispatch.h"
+            #endif
+            NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSelect, <TF>);
+        }
+        if (dispfunc) {
+            (*dispfunc)(reinterpret_cast<TF*>(v), num, kth);
+            return true;
+        }
+    }
+#endif
+    (void)v; (void)num; (void)kth; // to avoid unused arg warn
+    return false;
+}
+
+template<typename T>
+inline bool argquickselect_dispatch(T* v, npy_intp* arg, npy_intp num, npy_intp kth)
+{
+#ifndef __CYGWIN__
+    /*
+     * Only defined for int32_t, uint32_t, float32, int64_t, uint64_t, double
+     */
+    if constexpr (
+        (std::is_integral_v<T> || std::is_floating_point_v<T>) &&
+        (sizeof(T) == sizeof(uint32_t) || sizeof(T) == sizeof(uint64_t)) &&
+        // x86-simd-sort uses 8-byte int to store arg values, npy_intp is 4 bytes in 32-bit
+        sizeof(npy_intp) == sizeof(int64_t)) {
+        using TF = typename np::meta::FixedWidth<T>::Type;
+        #ifndef NPY_DISABLE_OPTIMIZATION
+            #include "simd_qsort.dispatch.h"
+        #endif
+        void (*dispfunc)(TF*, npy_intp*, npy_intp, npy_intp) = nullptr;
+        NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template ArgQSelect, <TF>);
+        if (dispfunc) {
+            (*dispfunc)(reinterpret_cast<TF*>(v), arg, num, kth);
+            return true;
+        }
+    }
+#endif
+    (void)v; (void)arg; (void)num; (void)kth; // to avoid unused arg warn
+    return false;
+}
+
 template <typename Tag, bool arg, typename type>
 NPY_NO_EXPORT int
-introselect_(type *v, npy_intp *tosort, npy_intp num, npy_intp kth,
-             npy_intp *pivots, npy_intp *npiv);
+introselect_(type *v, npy_intp *tosort, npy_intp num, npy_intp kth, npy_intp *pivots, npy_intp *npiv);
 
 /*
  *****************************************************************************
@@ -397,17 +467,25 @@ introselect_(type *v, npy_intp *tosort, npy_intp num, npy_intp kth,
 template <typename Tag>
 static int
 introselect_noarg(void *v, npy_intp num, npy_intp kth, npy_intp *pivots,
-                  npy_intp *npiv, void *)
+                  npy_intp *npiv, npy_intp nkth, void *)
 {
+    using T = typename std::conditional<std::is_same_v<Tag, npy::half_tag>, np::Half, typename Tag::type>::type;
+    if ((nkth == 1) && (quickselect_dispatch((T *)v, num, kth))) {
+        return 0;
+    }
     return introselect_<Tag, false>((typename Tag::type *)v, nullptr, num, kth,
                                     pivots, npiv);
 }
 
 template <typename Tag>
 static int
 introselect_arg(void *v, npy_intp *tosort, npy_intp num, npy_intp kth,
-                npy_intp *pivots, npy_intp *npiv, void *)
+                npy_intp *pivots, npy_intp *npiv, npy_intp nkth, void *)
 {
+    using T = typename Tag::type;
+    if ((nkth == 1) && (argquickselect_dispatch((T *)v, tosort, num, kth))) {
+        return 0;
+    }
     return introselect_<Tag, true>((typename Tag::type *)v, tosort, num, kth,
                                    pivots, npiv);
 }