diff options
author | Matti Picus <matti.picus@gmail.com> | 2023-05-04 08:02:30 +0300 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-05-04 08:02:30 +0300 |
commit | 9163e933df91b516b6f0c7a9ba8ad1750e642f37 (patch) | |
tree | 51be3804e7e10b9f821243de9cf5f794d5b49ecb | |
parent | a2d21d8ac8716d0dc47aced070d67a11c36a4cc9 (diff) | |
parent | 28872618c7504fb3365be9dc91c1910e09c36496 (diff) | |
download | numpy-9163e933df91b516b6f0c7a9ba8ad1750e642f37.tar.gz |
Merge pull request #23435 from r-devulap/spr-simd-sort
ENH: Use AVX512 FP16 ISA for sorting float16 arrays
-rw-r--r-- | numpy/core/src/npysort/simd_qsort_16bit.dispatch.cpp | 16 | ||||
m--------- | numpy/core/src/npysort/x86-simd-sort | 0 |
2 files changed, 12 insertions, 4 deletions
diff --git a/numpy/core/src/npysort/simd_qsort_16bit.dispatch.cpp b/numpy/core/src/npysort/simd_qsort_16bit.dispatch.cpp index a6465a883..3f5099758 100644 --- a/numpy/core/src/npysort/simd_qsort_16bit.dispatch.cpp +++ b/numpy/core/src/npysort/simd_qsort_16bit.dispatch.cpp @@ -1,5 +1,5 @@ /*@targets - * $maxopt $keep_baseline avx512_icl + * $maxopt $keep_baseline avx512_icl avx512_spr */ // policy $keep_baseline is used to avoid skip building avx512_skx // when its part of baseline features (--cpu-baseline), since @@ -7,16 +7,23 @@ #include "simd_qsort.hpp" -#if defined(NPY_HAVE_AVX512_ICL) && !defined(_MSC_VER) +#if defined(NPY_HAVE_AVX512_SPR) && !defined(_MSC_VER) + #include "x86-simd-sort/src/avx512fp16-16bit-qsort.hpp" +#elif defined(NPY_HAVE_AVX512_ICL) && !defined(_MSC_VER) #include "x86-simd-sort/src/avx512-16bit-qsort.hpp" #endif namespace np { namespace qsort_simd { -#if defined(NPY_HAVE_AVX512_ICL) && !defined(_MSC_VER) +#if !defined(_MSC_VER) +#if defined(NPY_HAVE_AVX512_ICL) || defined(NPY_HAVE_AVX512_SPR) template<> void NPY_CPU_DISPATCH_CURFX(QSort)(Half *arr, intptr_t size) { +#if defined(NPY_HAVE_AVX512_SPR) + avx512_qsort(reinterpret_cast<_Float16*>(arr), size); +#else avx512_qsort_fp16(reinterpret_cast<uint16_t*>(arr), size); +#endif } template<> void NPY_CPU_DISPATCH_CURFX(QSort)(uint16_t *arr, intptr_t size) { @@ -26,6 +33,7 @@ template<> void NPY_CPU_DISPATCH_CURFX(QSort)(int16_t *arr, intptr_t size) { avx512_qsort(arr, size); } -#endif // NPY_HAVE_AVX512_ICL +#endif // NPY_HAVE_AVX512_ICL || SPR +#endif // _MSC_VER }} // namespace np::qsort_simd diff --git a/numpy/core/src/npysort/x86-simd-sort b/numpy/core/src/npysort/x86-simd-sort -Subproject 58501d026a390895f7fd7ebbe0fb7aea55055ad +Subproject 1735e86cda95a469357a19ab8984ad8530372e7 |