10000 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 8000
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
Separate testing np.partition and np.argpartition
  • Loading branch information
Raghuveer Devulapalli committed Nov 30, 2023
commit 9f1faa191cc927bd229207b25d9edbfe24b2f14d
67 changes: 36 additions & 31 deletions numpy/_core/tests/test_multiarray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2817,10 +2817,10 @@ def test_partition(self):
tgt = np.sort(d)
assert_array_equal(np.partition(d, 0, kind=k)[0], tgt[0])
assert_array_equal(np.partition(d, 1, kind=k)[1], tgt[1])
assert_array_equal(d[np.argpartition(d, 0, kind=k)],
np.partition(d, 0, kind=k))
assert_array_equal(d[np.argpartition(d, 1, kind=k)],
np.partition(d, 1, kind=k))
self.assert_partitioned(np.partition(d, 0, kind=k), [0])
self.assert_partitioned(d[np.argpartition(d, 0, kind=k)], [0])
self.assert_partitioned(np.partition(d, 1, kind=k), [1])
self.assert_partitioned(d[np.argpartition(d, 1, kind=k)], [1])
for i in range(d.size):
d[i:].partition(0, kind=k)
assert_array_equal(d, tgt)
Expand All @@ -2832,12 +2832,12 @@ def test_partition(self):
assert_array_equal(np.partition(d, 0, kind=k)[0], tgt[0])
assert_array_equal(np.partition(d, 1, kind=k)[1], tgt[1])
assert_array_equal(np.partition(d, 2, kind=k)[2], tgt[2])
assert_array_equal(d[np.argpartition(d, 0, kind=k)],
np.partition(d, 0, kind=k))
assert_array_equal(d[np.argpartition(d, 1, kind=k)],
np.partition(d, 1, kind=k))
assert_array_equal(d[np.argpartition(d, 2, kind=k)],
np.partition(d, 2, kind=k))
self.assert_partitioned(np.partition(d, 0, kind=k), [0])
self.assert_partitioned(d[np.argpartition(d, 0, kind=k)], [0])
self.assert_partitioned(np.partition(d, 1, kind=k), [1])
self.assert_partitioned(d[np.argpartition(d, 1, kind=k)], [1])
self.assert_partitioned(np.partition(d, 2, kind=k), [2])
self.assert_partitioned(d[np.argpartition(d, 2, kind=k)], [2])
for i in range(d.size):
d[i:].partition(0, kind=k)
assert_array_equal(d, tgt)
Expand All @@ -2851,26 +2851,26 @@ def test_partition(self):
d = np.arange(49)
assert_equal(np.partition(d, 5, kind=k)[5], 5)
assert_equal(np.partition(d, 15, kind=k)[15], 15)
assert_array_equal(d[np.argpartition(d, 5, kind=k)],
np.partition(d, 5, kind=k))
assert_array_equal(d[np.argpartition(d, 15, kind=k)],
np.partition(d, 15, kind=k))
self.assert_partitioned(np.partition(d, 5, kind=k), [5])
self.assert_partitioned(d[np.argpartition(d, 5, kind=k)], [5])
self.assert_partitioned(np.partition(d, 15, kind=k), [15])
self.assert_partitioned(d[np.argpartition(d, 15, kind=k)], [15])

# rsorted
d = np.arange(47)[::-1]
assert_equal(np.partition(d, 6, kind=k)[6], 6)
assert_equal(np.partition(d, 16, kind=k)[16], 16)
assert_array_equal(d[np.argpartition(d, 6, kind=k)],
np.partition(d, 6, kind=k))
assert_array_equal(d[np.argpartition(d, 16, kind=k)],
np.partition(d, 16, kind=k))
self.assert_partitioned(np.partition(d, 6, kind=k), [6])
self.assert_partitioned(d[np.argpartition(d, 6, kind=k)], [6])
self.assert_partitioned(np.partition(d, 16, kind=k), [16])
self.assert_partitioned(d[np.argpartition(d, 16, kind=k)], [16])

assert_array_equal(np.partition(d, -6, kind=k),
np.partition(d, 41, kind=k))
assert_array_equal(np.partition(d, -16, kind=k),
np.partition(d, 31, kind=k))
assert_array_equal(d[np.argpartition(d, -6, kind=k)],
np.partition(d, 41, kind=k))
self.assert_partitioned(np.partition(d, 41, kind=k), [41])
self.assert_partitioned(d[np.argpartition(d, -6, kind=k)], [41])

# median of 3 killer, O(n^2) on pure median 3 pivot quickselect
# exercises the median of median of 5 code used to keep O(n)
Expand Down Expand Up @@ -2900,10 +2900,10 @@ def test_partition(self):
np.random.shuffle(d)
for i in range(d.size):
assert_equal(np.partition(d, i, kind=k)[i], tgt[i])
assert_array_equal(d[np.argpartition(d, 6, kind=k)],
np.partition(d, 6, kind=k))
assert_array_equal(d[np.argpartition(d, 16, kind=k)],
np.partition(d, 16, kind=k))
self.assert_partitioned(np.partition(d, 6, kind=k), [6])
self.assert_partitioned(d[np.argpartition(d, 6, kind=k)], [6])
self.assert_partitioned(np.partition(d, 16, kind=k), [16])
self.assert_partitioned(d[np.argpartition(d, 16, kind=k)], [16])
for i in range(d.size):
d[i:].partition(0, kind=k)
assert_array_equal(d, tgt)
Expand Down Expand Up @@ -2965,27 +2965,32 @@ def test_partition(self):
assert_array_less(p[:i], p[i])
# all after are larger
assert_array_less(p[i], p[i + 1:])
aae(p, d[np.argpartition(d, i, kind=k)])
self.assert_partitioned(p, [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)]
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(),
msg="%d: %r <= %r" % (i, p[:, i], p[:, :i].T))
at((p[:, i + 1:].T > p[:, i]).all(),
msg="%d: %r < %r" % (i, p[:, i], p[:, i + 1:].T))
aae(p, d1[np.arange(d1.shape[0])[:, None],
np.argpartition(d1, i, axis=1, kind=k)])
for row in range(p.shape[0]):
self.assert_partitioned(p[row], [i])
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, :]]
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:]))
aae(p, d0[np.argpartition(d0, i, axis=0, kind=k),
np.arange(d0.shape[1])[None, :]])
for col in range(p.shape[1]):
self.assert_partitioned(p[:,col], [i])
self.assert_partitioned(parg[:,col], [i])

# check inplace
dc = d.copy()
Expand All @@ -3001,9 +3006,9 @@ def test_partition(self):
def assert_partitioned(self, d, kth):
prev = 0
for k in np.sort(kth):
assert_array_less(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 %d" % (k, d[k:], d[k]))
msg="kth %d, %r not greater equal %r" % (k, d[k:], d[k]))
prev = k + 1

def test_partition_iterative(self):
Expand Down
0