Split highway and x86-simd-sort dispatch to separate files

Author: Raghuveer Devulapalli

numpy/_core/meson.build

@@ -764,23 +764,29 @@ endforeach
 # -----------------------------------
 foreach gen_mtargets : [
   [
-    'simd_argsort.dispatch.h',
-    'src/npysort/simd_argsort.dispatch.cpp',
+    'x86_simd_argsort.dispatch.h',
+    'src/npysort/x86_simd_argsort.dispatch.cpp',
     [AVX512_SKX]
   ],
   [
-    'simd_qsort.dispatch.h',
-    'src/npysort/simd_qsort.dispatch.cpp',
+    'x86_simd_qsort.dispatch.h',
+    'src/npysort/x86_simd_qsort.dispatch.cpp',
     [
       AVX512_SKX, AVX2,
-      ASIMD,
     ]
   ],
   [
-    'simd_qsort_16bit.dispatch.h',
-    'src/npysort/simd_qsort_16bit.dispatch.cpp',
+    'x86_simd_qsort_16bit.dispatch.h',
+    'src/npysort/x86_simd_qsort_16bit.dispatch.cpp',
     [AVX512_SPR, AVX512_ICL]
   ],
+  [
+    'highway_qsort.dispatch.h',
+    'src/npysort/highway_qsort.dispatch.cpp',
+    [
+      ASIMD,
+    ]
+  ],
 ]
   mtargets = mod_features.multi_targets(
     gen_mtargets[0], multiarray_gen_headers + gen_mtargets[1],

numpy/_core/src/npysort/highway_qsort.dispatch.cpp

@@ -0,0 +1,48 @@
+/*@targets
+ * $maxopt $keep_baseline
+ * asimd
+ */
+// policy $keep_baseline is used to avoid skip building avx512_skx
+// when its part of baseline features (--cpu-baseline), since
+// 'baseline' option isn't specified within targets.
+
+#include "highway_qsort.hpp"
+#ifndef __CYGWIN__
+
+#if NPY_HAVE_ASIMD
+    #define VQSORT_ONLY_STATIC 1
+    #include "hwy/contrib/sort/vqsort-inl.h"
+#endif
+
+namespace np { namespace qsort_simd {
+
+#if NPY_HAVE_ASIMD
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(int32_t *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
+}
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(uint32_t *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
+}
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(int64_t *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
+}
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(uint64_t *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
+}
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(float *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
+}
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(double *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
+}
+#endif // NPY_HAVE_ASIMD
+
+}} // namespace np::simd
+
+#endif // __CYGWIN__

numpy/_core/src/npysort/highway_qsort.hpp

@@ -0,0 +1,16 @@
+#ifndef NUMPY_SRC_COMMON_NPYSORT_HWY_SIMD_QSORT_HPP
+#define NUMPY_SRC_COMMON_NPYSORT_HWY_SIMD_QSORT_HPP
+
+#include "common.hpp"
+
+namespace np { namespace qsort_simd {
+
+#ifndef NPY_DISABLE_OPTIMIZATION
+    #include "highway_qsort.dispatch.h"
+#endif
+NPY_CPU_DISPATCH_DECLARE(template <typename T> void QSort, (T *arr, npy_intp size))
+NPY_CPU_DISPATCH_DECLARE(template <typename T> void QSelect, (T* arr, npy_intp num, npy_intp kth))
+
+} } // np::qsort_simd
+
+#endif // NUMPY_SRC_COMMON_NPYSORT_HWY_SIMD_QSORT_HPP

numpy/_core/src/npysort/quicksort.cpp

@@ -54,13 +54,13 @@
 #include "npysort_common.h"
 #include "npysort_heapsort.h"
 #include "numpy_tag.h"
-#include "simd_qsort.hpp"
+#include "x86_simd_qsort.hpp"
+#include "highway_qsort.hpp"
 
 #include <cstdlib>
 #include <utility>
 
 #define NOT_USED NPY_UNUSED(unused)
-#define DISABLE_HIGHWAY_OPTIMIZATION defined(__arm__)
 
 /*
  * pushing largest partition has upper bound of log2(n) space
@@ -81,18 +81,23 @@ inline bool quicksort_dispatch(T *start, npy_intp num)
     void (*dispfunc)(TF*, intptr_t) = nullptr;
     if (sizeof(T) == sizeof(uint16_t)) {
         #ifndef NPY_DISABLE_OPTIMIZATION
-            #include "simd_qsort_16bit.dispatch.h"
+            #include "x86_simd_qsort_16bit.dispatch.h"
         #endif
         NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSort, <TF>);
     }
-    #if !DISABLE_HIGHWAY_OPTIMIZATION
     else if (sizeof(T) == sizeof(uint32_t) || sizeof(T) == sizeof(uint64_t)) {
         #ifndef NPY_DISABLE_OPTIMIZATION
-            #include "simd_qsort.dispatch.h"
+            #include "x86_simd_qsort.dispatch.h"
         #endif
         NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSort, <TF>);
+        // If not dispatched for avx-512 or avx2:
+        if (dispfunc == nullptr) {
+            #ifndef NPY_DISABLE_OPTIMIZATION
+                #include "highway_qsort.dispatch.h"
+            #endif
+            NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSort, <TF>);
+        }
     }
-    #endif
     if (dispfunc) {
         (*dispfunc)(reinterpret_cast<TF*>(start), static_cast<intptr_t>(num));
         return true;
@@ -109,7 +114,7 @@ inline bool aquicksort_dispatch(T *start, npy_intp* arg, npy_intp num)
     using TF = typename np::meta::FixedWidth<T>::Type;
     void (*dispfunc)(TF*, npy_intp*, npy_intp) = nullptr;
     #ifndef NPY_DISABLE_OPTIMIZATION
-        #include "simd_argsort.dispatch.h"
+        #include "x86_simd_argsort.dispatch.h"
     #endif
     NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template ArgQSort, <TF>);
     if (dispfunc) {

numpy/_core/src/npysort/selection.cpp

@@ -25,7 +25,7 @@
 #include <array>
 #include <cstdlib>
 #include <utility>
-#include "simd_qsort.hpp"
+#include "x86_simd_qsort.hpp"
 
 #define NOT_USED NPY_UNUSED(unused)
 #define DISABLE_HIGHWAY_OPTIMIZATION (defined(__arm__) || defined(__aarch64__))
@@ -45,14 +45,14 @@ inline bool quickselect_dispatch(T* v, npy_intp num, npy_intp kth)
         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"
+                #include "x86_simd_qsort_16bit.dispatch.h"
             #endif
             NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSelect, <TF>);
         }
         #if !DISABLE_HIGHWAY_OPTIMIZATION
         else if constexpr (sizeof(T) == sizeof(uint32_t) || sizeof(T) == sizeof(uint64_t)) {
             #ifndef NPY_DISABLE_OPTIMIZATION
-                #include "simd_qsort.dispatch.h"
+                #include "x86_simd_qsort.dispatch.h"
             #endif
             NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSelect, <TF>);
         }
@@ -79,7 +79,7 @@ inline bool argquickselect_dispatch(T* v, npy_intp* arg, npy_intp num, npy_intp
         (sizeof(T) == sizeof(uint32_t) || sizeof(T) == sizeof(uint64_t))) {
         using TF = typename np::meta::FixedWidth<T>::Type;
         #ifndef NPY_DISABLE_OPTIMIZATION
-            #include "simd_argsort.dispatch.h"
+            #include "x86_simd_argsort.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>);

numpy/_core/src/npysort/simd_argsort.dispatch.cpp renamed to numpy/_core/src/npysort/x86_simd_argsort.dispatch.cpp

@@ -6,7 +6,7 @@
 // when its part of baseline features (--cpu-baseline), since
 // 'baseline' option isn't specified within targets.
 
-#include "simd_qsort.hpp"
+#include "x86_simd_qsort.hpp"
 #ifndef __CYGWIN__
 
 #if defined(NPY_HAVE_AVX512_SKX)

numpy/_core/src/npysort/simd_qsort.dispatch.cpp renamed to numpy/_core/src/npysort/x86_simd_qsort.dispatch.cpp

@@ -1,27 +1,21 @@
 /*@targets
  * $maxopt $keep_baseline
  * avx512_skx avx2
- * asimd
  */
 // policy $keep_baseline is used to avoid skip building avx512_skx
 // when its part of baseline features (--cpu-baseline), since
 // 'baseline' option isn't specified within targets.
 
-#include "simd_qsort.hpp"
+#include "x86_simd_qsort.hpp"
 #ifndef __CYGWIN__
 
-#define USE_HIGHWAY defined(__aarch64__)
-
 #if defined(NPY_HAVE_AVX512_SKX)
     #include "x86-simd-sort/src/avx512-32bit-qsort.hpp"
     #include "x86-simd-sort/src/avx512-64bit-qsort.hpp"
     #include "x86-simd-sort/src/avx512-64bit-argsort.hpp"
 #elif defined(NPY_HAVE_AVX2)
     #include "x86-simd-sort/src/avx2-32bit-qsort.hpp"
     #include "x86-simd-sort/src/avx2-64bit-qsort.hpp"
-#elif USE_HIGHWAY
-    #define VQSORT_ONLY_STATIC 1
-    #include "hwy/contrib/sort/vqsort-inl.h"
 #endif
 
 namespace np { namespace qsort_simd {
@@ -123,31 +117,6 @@ template<> void NPY_CPU_DISPATCH_CURFX(QSort)(double *arr, npy_intp num)
     avx2_qsort(arr, num, true);
 #endif
 }
-#elif USE_HIGHWAY
-template<> void NPY_CPU_DISPATCH_CURFX(QSort)(int32_t *arr, intptr_t size)
-{
-    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
-}
-template<> void NPY_CPU_DISPATCH_CURFX(QSort)(uint32_t *arr, intptr_t size)
-{
-    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
-}
-template<> void NPY_CPU_DISPATCH_CURFX(QSort)(int64_t *arr, intptr_t size)
-{
-    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
-}
-template<> void NPY_CPU_DISPATCH_CURFX(QSort)(uint64_t *arr, intptr_t size)
-{
-    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
-}
-template<> void NPY_CPU_DISPATCH_CURFX(QSort)(float *arr, intptr_t size)
-{
-    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
-}
-template<> void NPY_CPU_DISPATCH_CURFX(QSort)(double *arr, intptr_t size)
-{
-    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
-}
 #endif // NPY_HAVE_AVX512_SKX || NPY_HAVE_AVX2
 
 }} // namespace np::simd

numpy/_core/src/npysort/simd_qsort.hpp renamed to numpy/_core/src/npysort/x86_simd_qsort.hpp

@@ -1,34 +1,28 @@
-#ifndef NUMPY_SRC_COMMON_NPYSORT_SIMD_QSORT_HPP
-#define NUMPY_SRC_COMMON_NPYSORT_SIMD_QSORT_HPP
+#ifndef NUMPY_SRC_COMMON_NPYSORT_X86_SIMD_QSORT_HPP
+#define NUMPY_SRC_COMMON_NPYSORT_X86_SIMD_QSORT_HPP
 
 #include "common.hpp"
 
-#define DISABLE_HIGHWAY_OPTIMIZATION defined(__arm__)
-
 namespace np { namespace qsort_simd {
 
-#if !DISABLE_HIGHWAY_OPTIMIZATION
 #ifndef NPY_DISABLE_OPTIMIZATION
-    #include "simd_qsort.dispatch.h"
+    #include "x86_simd_qsort.dispatch.h"
 #endif
 NPY_CPU_DISPATCH_DECLARE(template <typename T> void QSort, (T *arr, npy_intp size))
 NPY_CPU_DISPATCH_DECLARE(template <typename T> void QSelect, (T* arr, npy_intp num, npy_intp kth))
-#endif // DISABLE_HIGHWAY_OPTIMIZATION
 
 #ifndef NPY_DISABLE_OPTIMIZATION
-    #include "simd_argsort.dispatch.h"
+    #include "x86_simd_argsort.dispatch.h"
 #endif
 NPY_CPU_DISPATCH_DECLARE(template <typename T> void ArgQSort, (T *arr, npy_intp* arg, npy_intp size))
 NPY_CPU_DISPATCH_DECLARE(template <typename T> void ArgQSelect, (T *arr, npy_intp* arg, npy_intp kth, npy_intp size))
 
 #ifndef NPY_DISABLE_OPTIMIZATION
-    #include "simd_qsort_16bit.dispatch.h"
+    #include "x86_simd_qsort_16bit.dispatch.h"
 #endif
 NPY_CPU_DISPATCH_DECLARE(template <typename T> void QSort, (T *arr, npy_intp size))
 NPY_CPU_DISPATCH_DECLARE(template <typename T> void QSelect, (T* arr, npy_intp num, npy_intp kth))
 
 } } // np::qsort_simd
 
-#undef DISABLE_HIGHWAY_OPTIMIZATION
-
-#endif // NUMPY_SRC_COMMON_NPYSORT_SIMD_QSORT_HPP
+#endif // NUMPY_SRC_COMMON_NPYSORT_X86_SIMD_QSORT_HPP

numpy/_core/src/npysort/simd_qsort_16bit.dispatch.cpp renamed to numpy/_core/src/npysort/x86_simd_qsort_16bit.dispatch.cpp

@@ -5,7 +5,7 @@
 // when its part of baseline features (--cpu-baseline), since
 // 'baseline' option isn't specified within targets.
 
-#include "simd_qsort.hpp"
+#include "x86_simd_qsort.hpp"
 #ifndef __CYGWIN__
 
 #if defined(NPY_HAVE_AVX512_SPR)