Merge pull request numpy#25346 from Mousius/highway-vqsort-float16

ENH: Enable 16-bit VQSort routines on AArch64

numpy/_core/meson.build

@@ -787,6 +787,13 @@ foreach gen_mtargets : [
       SVE, ASIMD,
     ]
   ],
+  [
+    'highway_qsort_16bit.dispatch.h',
+    'src/npysort/highway_qsort_16bit.dispatch.cpp',
+    [
+      ASIMDHP,
+    ]
+  ],
 ]
   mtargets = mod_features.multi_targets(
     gen_mtargets[0], multiarray_gen_headers + gen_mtargets[1],

numpy/_core/src/npysort/highway_qsort.hpp

@@ -11,6 +11,13 @@ namespace np { namespace highway { namespace qsort_simd {
 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))
 
+
+#ifndef NPY_DISABLE_OPTIMIZATION
+    #include "highway_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::highway::qsort_simd
 
 #endif // NUMPY_SRC_COMMON_NPYSORT_HWY_SIMD_QSORT_HPP

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

@@ -0,0 +1,20 @@
+#include "highway_qsort.hpp"
+#define VQSORT_ONLY_STATIC 1
+#include "hwy/contrib/sort/vqsort-inl.h"
+
+namespace np { namespace highway { namespace qsort_simd {
+
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(Half *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(reinterpret_cast<hwy::float16_t*>(arr), size, hwy::SortAscending());
+}
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(uint16_t *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
+}
+template<> void NPY_CPU_DISPATCH_CURFX(QSort)(int16_t *arr, intptr_t size)
+{
+    hwy::HWY_NAMESPACE::VQSortStatic(arr, size, hwy::SortAscending());
+}
+
+} } } // np::highway::qsort_simd

numpy/_core/src/npysort/quicksort.cpp

@@ -81,9 +81,14 @@ 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 "x86_simd_qsort_16bit.dispatch.h"
+            #if defined(NPY_CPU_AMD64) || defined(NPY_CPU_X86) // x86 32-bit and 64-bit
+                #include "x86_simd_qsort_16bit.dispatch.h"
+                NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSort, <TF>);
+            #else
+                #include "highway_qsort_16bit.dispatch.h"
+                NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::highway::qsort_simd::template QSort, <TF>);
+            #endif
         #endif
-        NPY_CPU_DISPATCH_CALL_XB(dispfunc = np::qsort_simd::template QSort, <TF>);
     }
     else if (sizeof(T) == sizeof(uint32_t) || sizeof(T) == sizeof(uint64_t)) {
         #ifndef NPY_DISABLE_OPTIMIZATION