ENH: Improve performance of numpy.nonzero for 1D/2D contiguous arrays #27519

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

Jump to bottom

Open

eendebakpt wants to merge 20 commits into numpy:main from eendebakpt:nonzero_optmize

Contributor

eendebakpt commented Oct 6, 2024 •

edited

Loading

Improve performance of np.non_zero for 1D and 2D contiguous arrays. The main cases improved are 2D boolean arrays and 1D/2D int and float arrays. The PR is based on work and ideas from @touqir14 in #18368

Benchmark on main:

<module 'numpy' from '/home/eendebakpt/numpy/build-install/lib/python3.12/site-packages/numpy/__init__.py'> np.__version__='2.2.0.dev0+git20241010.a5cb327'
bool 816.1 [us]
int64 2855.6 [us]
float64 5777.2 [us]
bool 2D 515.2 [us]
float64 2D 824.7 [us]
float64 2D view 149.7 [us]

Benchmark on PR:

<module 'numpy' from '/home/eendebakpt/numpy/build-install/lib/python3.12/site-packages/numpy/__init__.py'> np.__version__='2.2.0.dev0+git20241011.60c9b88'
bool 845.4 [us]
int64 2004.0 [us]
float64 3997.1 [us]
bool 2D 85.2 [us]
float64 2D 302.7 [us]
float64 2D view 150.2 [us]

Benchmark script

import sys
import timeit
import numpy as np

print(f'{np} {np.__version__=}')

b=np.arange(1_000_000).astype(bool)
x=b.astype(int)
f=b.astype(np.float64)
w=np.random.rand(600, 100)
w[w<.5]=0
ww = w[::2, ::3]
b2d=np.random.rand(600, 100) > .5

n=100
t = timeit.timeit('np.nonzero(b)', globals=globals(), number=n)
print(f'{b.dtype} {1e6*t/n:.1f} [us]')
t = timeit.timeit('np.nonzero(x)', globals=globals(), number=n)
print(f'{x.dtype} {1e6*t/n:.1f} [us]')
t = timeit.timeit('np.nonzero(f)', globals=globals(), number=n)
print(f'{f.dtype} {1e6*t/n:.1f} [us]')
t = timeit.timeit('np.nonzero(b2d)', globals=globals(), number=n)
print(f'{b2d.dtype} 2D {1e6*t/n:.1f} [us]')
t = timeit.timeit('np.nonzero(w)', globals=globals(), number=n)
print(f'{w.dtype} 2D {1e6*t/n:.1f} [us]')
t = timeit.timeit('np.nonzero(ww)', globals=globals(), number=n)
print(f'{w.dtype} 2D view {1e6*t/n:.1f} [us]')

Notes:

For float arrays performance will be improved by ENH: Improve performance of np.count_nonzero for float arrays #27523 as well


          ENH: Improve performance of numpy.nonzero for 1D contiguous arrays

d84439a

seberg marked this pull request as draft

October 9, 2024 09:17

eendebakpt added 3 commits

October 11, 2024 13:24


          refactor

c29c232


          remove assert

503ec14


          Merge branch 'main' into nonzero_optmize

60c9b88

eendebakpt changed the title ~~Draft: ENH: Improve performance of numpy.nonzero for 1D contiguous arrays~~ Draft: ENH: Improve performance of numpy.nonzero for 1D/2D contiguous arrays


          refactor

82450ae

eendebakpt changed the title ~~Draft: ENH: Improve performance of numpy.nonzero for 1D/2D contiguous arrays~~ ENH: Improve performance of numpy.nonzero for 1D/2D contiguous arrays

eendebakpt marked this pull request as ready for review

October 11, 2024 12:32

tylerjereddy added 01 - Enhancement component: numpy._core labels


          Merge branch 'main' into nonzero_optmize

5bbe2ce

eendebakpt requested a review from seberg

November 25, 2024 12:40


          Merge branch 'main' into nonzero_optmize

4ca8220

eendebakpt mentioned this pull request

ENH: Improve performance of np.count_nonzero for float arrays #27523

Merged


          use npy_bool

76e0e41

seberg reviewed

View reviewed changes

Member

seberg left a comment

A few comments. I'll ask for more style cleanup before we merge this.

numpy/_core/src/multiarray/item_selection.c Outdated

+                  PyArrayObject* original_array = self;
+                  if (PyArray_BASE(self) != NULL) {
+                      original_array = (PyArrayObject*) PyArray_BASE(self);
+                  }

Member

seberg Feb 19, 2025

This doesn't make any sense to me. What does the base have to do with anything here?

Contributor Author

eendebakpt Feb 19, 2025

I agree. Removed.

numpy/_core/src/multiarray/item_selection.c Outdated

+                  if (PyArray_BASE(self) != NULL) {
+                      original_array = (PyArrayObject*) PyArray_BASE(self);
+                  }
+                  int bytes_not_swapped = PyArray_ISNOTSWAPPED(self) && PyArray_ISNOTSWAPPED(original_array);

Member

seberg Feb 19, 2025

Checking for byte-swap is (just like in the other PR) important. Are there tests that fail without it? Because if not, there should be a new test for it (just like the one in the other PR which is currently failing).

Contributor Author

eendebakpt Feb 19, 2025

Test added

numpy/_core/src/multiarray/item_selection.c Outdated

@@ -2801,6 +2801,57 @@ PyArray_CountNonzero(PyArrayObject *self)
                   return nonzero_count;
               }
+              static inline void nonzero_idxs_1D_bool(npy_intp count, npy_intp nonzero_count, char *data, npy_intp stride, npy_intp* multi_index)

Member

seberg Feb 19, 2025

I'll have formatting nits, but this is nice to move out!

numpy/_core/src/multiarray/item_selection.c Outdated

-                                  ++j;
-                              }
-                          }
+                          nonzero_idxs_1D_bool(count, nonzero_count, data, stride, multi_index);

Member

seberg Feb 19, 2025

Doesn't this call belong to the family of fast functions in the other place? I.e. only the manual loop of nonzero should stay here?

Contributor Author

eendebakpt Feb 19, 2025

I wanted to keep the diff short, but it indeed fits a bit better in the other block. Changed.

numpy/_core/src/multiarray/item_selection.c Outdated

+                      int M_type_num = dtype->type_num;
+                      bool executed = nonzero_idxs_dispatcher((void*)data, multi_index, M_dim,
+                                                      M_shape, M_strides, M_type_num, nonzero_count);

Member

seberg Feb 19, 2025

This branch is missing the thresholded GIL release (internals of it I guess).

Contributor Author

eendebakpt Feb 19, 2025

Releasing the GIL (or running the cpython free-threading build) could result in a buffer overflow if another thread changes the number of non-zero entries in the array. I added a guard for that, but we should rerun the benchmarks (I still expect a good speedup though).

Should a test be added for this? (probably here: https://github.com/numpy/numpy/blob/main/numpy/_core/tests/test_multithreading.py)

Member

seberg Feb 20, 2025

Interesting, a case where we would actually want to lock down a whole array (but I think this is super hard with buffer exchange, so in practice...).

Would be very cool if you add a test. Larger use of free-threading is in the future after-all, and maybe not a too distant one.

Contributor Author

eendebakpt Feb 27, 2025

Tests have been added in #28361

eendebakpt added 5 commits

February 19, 2025 19:55


          review comments

53a9ef3


          refactor

5ec3c50


          fix setting executed

6c320f3


          add guard for mutating array


          fix bool conversion

8de9ff7

eendebakpt mentioned this pull request

BUG: Make np.nonzero threading safe #28361

Merged

charris mentioned this pull request

BUG: Make np.nonzero threading safe #28385

Merged

eendebakpt added 6 commits

February 23, 2025 21:44


          Merge branch 'main' into nonzero_optmize

5444f08


          Merge branch 'main' into nonzero_optmize

b539299


          add check on alignment

bc0a9ae


          lint

0b3ee4f


          fix test

f4fae21


          code style

a9905be

eendebakpt requested a review from seberg

February 27, 2025 13:13


          Merge branch 'main' into nonzero_optmize

6636bbf

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

01 - Enhancement component: numpy._core