8000 ENH: Vectorize np.sort and np.partition with AVX2 by r-devulap · Pull Request #25045 · numpy/numpy · GitHub
[go: up one dir, main page]
Skip to content

ENH: Vectorize np.sort and np.partition with AVX2 #25045

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 22 commits into from
Dec 4, 2023
Merged

ENH: Vectorize np.sort and np.partition with AVX2 #25045

Changes from 1 commit
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
95f6158
Update x86-simd-sort to latest
Oct 13, 2023
125017e
Dont include arg methods on 32-bit platforms
Oct 17, 2023
65d6506
Update x86-simd-sort to latest
Oct 17, 2023
2766e9b
Revert "Dont include arg methods on 32-bit platforms"
Oct 18, 2023
77278fe
Enable np.partition and np.argpartition on 32-bit
Oct 18, 2023
d729a49
Update x86-simd-sort to latest
Oct 18, 2023
05000e1
Update x86-simd-sort to latest
Oct 18, 2023
c2bdc0e
Update x86-simd-sort to latest
Oct 19, 2023
a303128
Update x86-simd-sort to latest
Oct 26, 2023
9a7e109
avx512 qsort for fp16, fp32 and fp64 require an explicit flag to sort…
Oct 26, 2023
d4169da
Dispatch AVX2 qsort and partition
Oct 26, 2023
9f1faa1 Separate testing np.partition and np.argpartition
Nov 2, 2023
7c39472
update x86-simd-sort to latest
Nov 2, 2023
3b6643d
Disable test_shape_base.py::TestTakeAlongAxis::test_argequivalent for…
Nov 2, 2023
fb69b5b
Linter fixes
Nov 2, 2023
190e80e
Fix rebase bug in meson.build
Nov 27, 2023
c9588ca
Split highway and x86-simd-sort dispatch to separate files
Nov 27, 2023
0cb29a2
Update x86-simd-sort to latest
Nov 27, 2023
a3ca84b
Include highway/x86-simd-sort at compile time
Nov 30, 2023
9bde4ae
Remove distutils related dispatch code and create new namespace for h…
Dec 1, 2023
b28ed78
Add optional hasnan arguement to avx sorting
Dec 1, 2023
675cb07
Add x86_qsort/qselect functions inside anonymous namespace
Dec 1, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Dont include arg methods on 32-bit platforms
  • Loading branch information
Raghuveer Devulapalli committed Nov 30, 2023
commit 125017eaa05232a67e51d0079a2fa70ff63cd1d8
40 changes: 40 additions & 0 deletions numpy/_core/src/npysort/simd_qsort.dispatch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,10 @@
#if defined(NPY_HAVE_AVX512_SKX)
#include "x86-simd-sort/src/avx512-32bit-qsort.hpp"
#include "x86-simd-sort/src/avx512-64bit-qsort.hpp"
/* arg methods are only intended to 64-bit platforms */
#ifdef NPY_CPU_AMD64
#include "x86-simd-sort/src/avx512-64bit-argsort.hpp"
#endif
#elif USE_HIGHWAY
#define VQSORT_ONLY_STATIC 1
#include "hwy/contrib/sort/vqsort-inl.h"
Expand All @@ -23,6 +27,42 @@
namespace np { namespace qsort_simd {

#if defined(NPY_HAVE_AVX512_SKX)
template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(int32_t *arr, npy_intp* arg, npy_intp num, npy_intp kth)
{
#ifdef NPY_CPU_AMD64
avx512_argselect(arr, reinterpret_cast<int64_t*>(arg), kth, num);
#endif
}
template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(uint32_t *arr, npy_intp* arg, npy_intp num, npy_intp kth)
{
#ifdef NPY_CPU_AMD64
avx512_argselect(arr, reinterpret_cast<int64_t*>(arg), kth, num);
#endif
}
template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(int64_t*arr, npy_intp* arg, npy_intp num, npy_intp kth)
{
#ifdef NPY_CPU_AMD64
avx512_argselect(arr, reinterpret_cast<int64_t*>(arg), kth, num);
#endif
}
template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(uint64_t*arr, npy_intp* arg, npy_intp num, npy_intp kth)
{
#ifdef NPY_CPU_AMD64
avx512_argselect(arr, reinterpret_cast<int64_t*>(arg), kth, num);
#endif
}
template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(float *arr, npy_intp* arg, npy_intp num, npy_intp kth)
{
#ifdef NPY_CPU_AMD64
avx512_argselect(arr, reinterpret_cast<int64_t*>(arg), kth, num);
#endif
}
template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(double *arr, npy_intp* arg, npy_intp num, npy_intp kth)
{
#ifdef NPY_CPU_AMD64
avx512_argselect(arr, reinterpret_cast<int64_t*>(arg), kth, num);
#endif
}
template<> void NPY_CPU_DISPATCH_CURFX(QSelect)(int32_t *arr, npy_intp num, npy_intp kth)
{
avx512_qselect(arr, kth, num, true);
Expand Down
0