diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/src/multiarray/item_selection.c | 30 |
1 files changed, 25 insertions, 5 deletions
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index 086b674c8..9fad153a3 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -2641,13 +2641,33 @@ PyArray_Nonzero(PyArrayObject *self) *multi_index++ = j++; } } + /* + * Fallback to a branchless strategy to avoid branch misprediction + * stalls that are very expensive on most modern processors. + */ else { - npy_intp j; - for (j = 0; j < count; ++j) { - if (*data != 0) { - *multi_index++ = j; - } + npy_intp *multi_index_end = multi_index + nonzero_count; + npy_intp j = 0; + + /* Manually unroll for GCC and maybe other compilers */ + while (multi_index + 4 < multi_index_end) { + *multi_index = j; + multi_index += data[0] != 0; + *multi_index = j + 1; + multi_index += data[stride] != 0; + *multi_index = j + 2; + multi_index += data[stride * 2] != 0; + *multi_index = j + 3; + multi_index += data[stride * 3] != 0; + data += stride * 4; + j += 4; + } + + while (multi_index < multi_index_end) { + *multi_index = j; + multi_index += *data != 0; data += stride; + ++j; } } } |