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

Show file tree
Hide file tree
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
avx512 qsort for fp16, fp32 and fp64 require an explicit flag to sort… 
… NAN
  • Loading branch information
Raghuveer Devulapalli committed Nov 30, 2023
commit 9a7e109302ebdeb3b939966d64653535cb46391c
8 changes: 4 additions & 4 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 @@ -42,11 +42,11 @@ template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(uint64_t*arr, npy_intp* arg,
}
template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(float *arr, npy_intp* arg, npy_intp num, npy_intp kth)
{
avx512_argselect(arr, reinterpret_cast<size_t*>(arg), kth, num);
avx512_argselect(arr, reinterpret_cast<size_t*>(arg), kth, num, true);
}
template<> void NPY_CPU_DISPATCH_CURFX(ArgQSelect)(double *arr, npy_intp* arg, npy_intp num, npy_intp kth)
{
avx512_argselect(arr, reinterpret_cast<size_t*>(arg), kth, num);
avx512_argselect(arr, reinterpret_cast<size_t*>(arg), kth, num, true);
}
template<> void NPY_CPU_DISPATCH_CURFX(QSelect)(int32_t *arr, npy_intp num, npy_intp kth)
{
Expand Down Expand Up @@ -90,11 +90,11 @@ template<> void NPY_CPU_DISPATCH_CURFX(QSort)(uint64_t *arr, npy_intp size)
}
template<> void NPY_CPU_DISPATCH_CURFX(QSort)(float *arr, npy_intp size)
{
avx512_qsort(arr, size);
avx512_qsort(arr, size, true);
}
template<> void NPY_CPU_DISPATCH_CURFX(QSort)(double *arr, npy_intp size)
{
avx512_qsort(arr, size);
avx512_qsort(arr, size, true);
}
#elif USE_HIGHWAY
template<> void NPY_CPU_DISPATCH_CURFX(QSort)(int32_t *arr, intptr_t size)
Expand Down
4 changes: 2 additions & 2 deletions numpy/_core/src/npysort/simd_qsort_16bit.dispatch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,9 +79,9 @@ template<> void NPY_CPU_DISPATCH_CURFX(QSelect)(int16_t *arr, npy_intp num, npy_
template<> void NPY_CPU_DISPATCH_CURFX(QSort)(Half *arr, npy_intp size)
{
#if defined(NPY_HAVE_AVX512_SPR)
avx512_qsort(reinterpret_cast<_Float16*>(arr), size);
avx512_qsort(reinterpret_cast<_Float16*>(arr), size, true);
#else
avx512_qsort_fp16(reinterpret_cast<uint16_t*>(arr), size);
avx512_qsort_fp16(reinterpret_cast<uint16_t*>(arr), size, true);
#endif
}
template<> void NPY_CPU_DISPATCH_CURFX(QSort)(uint16_t *arr, npy_intp size)
Expand Down
0