Disable avx512_qselect on 32-bit systems

r-devulap · r-devulap · commit 97e6714c85ff · 2023-09-08T13:05:26.000-07:00
diff --git a/numpy/core/src/npysort/selection.cpp b/numpy/core/src/npysort/selection.cpp
@@ -36,10 +36,17 @@ inline bool quickselect_dispatch(T* v, npy_intp num, npy_intp kth)
     /*
      * 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(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)) {
diff --git a/numpy/core/src/npysort/simd_qsort.dispatch.cpp b/numpy/core/src/npysort/simd_qsort.dispatch.cpp
@@ -6,6 +6,7 @@
 // 'baseline' option isn't specified within targets.
 
 #include "simd_qsort.hpp"
+#ifndef __CYGWIN__
 
 #if defined(NPY_HAVE_AVX512_SKX)
     #include "x86-simd-sort/src/avx512-32bit-qsort.hpp"
@@ -16,7 +17,6 @@
 namespace np { namespace qsort_simd {
 
 #if defined(NPY_HAVE_AVX512_SKX)
-#ifndef __CYGWIN__
 template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(int32_t *arr, npy_intp* arg, npy_intp num, npy_intp kth)
 {
     avx512_argselect(arr, reinterpret_cast<int64_t*>(arg), kth, num);
@@ -113,7 +113,8 @@ template<> void NPY_CPU_DISPATCH_CURFX(ArgQSort)(double *arr, npy_intp *arg, npy
 {
     avx512_argsort(arr, reinterpret_cast<int64_t*>(arg), size);
 }
-#endif // __CYGWIN__
 #endif  // NPY_HAVE_AVX512_SKX
 
 }} // namespace np::simd
+
+#endif // __CYGWIN__
diff --git a/numpy/core/src/npysort/simd_qsort_16bit.dispatch.cpp b/numpy/core/src/npysort/simd_qsort_16bit.dispatch.cpp
@@ -6,6 +6,7 @@
 // 'baseline' option isn't specified within targets.
 
 #include "simd_qsort.hpp"
+#ifndef __CYGWIN__
 
 #if defined(NPY_HAVE_AVX512_SPR)
     #include "x86-simd-sort/src/avx512fp16-16bit-qsort.hpp"
@@ -15,13 +16,12 @@
  */
 void avx512_qsort_uint16(uint16_t*, intptr_t);
 void avx512_qsort_int16(int16_t*, intptr_t);
-void avx512_qselect_uint16(uint16_t*, npy_intp, npy_intp)
-void avx512_qselect_int16(int16_t*, npy_intp, npy_intp)
+void avx512_qselect_uint16(uint16_t*, npy_intp, npy_intp);
+void avx512_qselect_int16(int16_t*, npy_intp, npy_intp);
 
 #elif defined(NPY_HAVE_AVX512_ICL)
     #include "x86-simd-sort/src/avx512-16bit-qsort.hpp"
 /* Wrapper function defintions here: */
-#ifndef __CYGWIN__
 void avx512_qsort_uint16(uint16_t* arr, intptr_t size)
 {
     avx512_qsort(arr, size);
@@ -38,7 +38,6 @@ void avx512_qselect_int16(int16_t* arr, npy_intp kth, npy_intp size)
 {
     avx512_qselect(arr, kth, size, true);
 }
-#endif // __CYGWIN__
 #endif
 
 namespace np { namespace qsort_simd {
@@ -47,7 +46,6 @@ namespace np { namespace qsort_simd {
  * QSelect dispatch functions:
  */
 #if defined(NPY_HAVE_AVX512_ICL) || defined(NPY_HAVE_AVX512_SPR)
-#ifndef __CYGWIN__
 template<> void NPY_CPU_DISPATCH_CURFX(QSelect)(Half *arr, npy_intp num, npy_intp kth)
 {
 #if defined(NPY_HAVE_AVX512_SPR)
@@ -102,7 +100,8 @@ template<> void NPY_CPU_DISPATCH_CURFX(QSort)(int16_t *arr, intptr_t size)
     avx512_qsort(arr, size);
 #endif
 }
-#endif // __CYGWIN__
 #endif // NPY_HAVE_AVX512_ICL || SPR
 
 }} // namespace np::qsort_simd
+
+#endif // __CYGWIN__
