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
8000 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
Linter fixes
  • Loading branch information
Raghuveer Devulapalli committed Nov 30, 2023
commit fb69b5bd1d1f6d170d669edc42199c8182cb247f
16 changes: 10 additions & 6 deletions numpy/_core/tests/test_multiarray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2966,10 +2966,12 @@ def test_partition(self):
# all after are larger
assert_array_less(p[i], p[i + 1:])
self.assert_partitioned(p, [i])
self.assert_partitioned(d[np.argpartition(d, i, kind=k)], [i])
self.assert_partitioned(
d[np.argpartition(d, i, kind=k)], [i])

p = np.partition(d1, i, axis=1, kind=k)
parg = d1[np.arange(d1.shape[0])[:, None], np.argpartition(d1, i, axis=1, kind=k)]
parg = d1[np.arange(d1.shape[0])[:, None],
np.argpartition(d1, i, axis=1, kind=k)]
aae(p[:, i], np.array([i] * d1.shape[0], dtype=dt))
# array_less does not seem to work right
at((p[:, :i].T <= p[:, i]).all(),
Expand All @@ -2981,16 +2983,17 @@ def test_partition(self):
self.assert_partitioned(parg[row], [i])

p = np.partition(d0, i, axis=0, kind=k)
parg = d0[np.argpartition(d0, i, axis=0, kind=k), np.arange(d0.shape[1])[None, :]]
parg = d0[np.argpartition(d0, i, axis=0, kind=k),
np.arange(d0.shape[1])[None, :]]
aae(p[i, :], np.array([i] * d1.shape[0], dtype=dt))
# array_less does not seem to work right
at((p[:i, :] <= p[i, :]).all(),
msg="%d: %r <= %r" % (i, p[i, :], p[:i, :]))
at((p[i + 1:, :] > p[i, :]).all(),
msg="%d: %r < %r" % (i, p[i, :], p[:, i + 1:]))
for col in range(p.shape[1]):
self.assert_partitioned(p[:,col], [i])
self.assert_partitioned(parg[:,col], [i])
self.assert_partitioned(p[:, col], [i])
self.assert_partitioned(parg[:, col], [i])

# check inplace
dc = d.copy()
Expand All @@ -3006,7 +3009,8 @@ def test_partition(self):
def assert_partitioned(self, d, kth):
prev = 0
for k in np.sort(kth):
assert_array_compare(operator.__le__, d[prev:k], d[k], err_msg='kth %d' % k)
assert_array_compare(operator.__le__, d[prev:k], d[k],
err_msg='kth %d' % k)
assert_((d[k:] >= d[k]).all(),
msg="kth %d, %r not greater equal %r" % (k, d[k:], d[k]))
prev = k + 1
Expand Down
0