prange generates unnecessary python interactions when indexing a 2d view

Indexing a 2d view in a prange will generate some unnecessary python interactions related to the GIL.

Here is a reproducing example. LMK if I can help with anything.

#cython: boundscheck=False
#cython: language_level=3
from cython.parallel import prange

# Indexing a 2d view in a prange loop: generates lots of python interactions.
def sum_2d_a(float[:, :] view_2d):

    cdef:
        float out = 0
        int row_idx = 0

    for row_idx in prange(view_2d.shape[0], nogil=True):
        out += sum_1d_a(view_2d[row_idx, :])
        # lots of python interactions at the end of the loop

    return out

cdef float sum_1d_a(float[:] view_1d) nogil:
    return 3.4  # irrelevant code


# workaround: pass the whole 2d view and the row index. No interactions in this case.
def sum_2d_b(float[:, :] view_2d):

    cdef:
        float out = 0
        int row_idx = 0

    for row_idx in prange(view_2d.shape[0], nogil=True):
        out += sum_1d_b(view_2d, row_idx)
        # no python interaction

    return out

cdef float sum_1d_b(float[:, :] view_2d, int row_idx) nogil:
    return 3.4  # irrelevant code

version : 0.29.6

[scoder]

I can well imagine that that’s a use case for something like scikit-learn, thanks for the insights.

Atomic ints are not costly per se, they are often implemented in hardware and thus fairly cheap in comparison to many other forms of locking mechanisms. But they do represent a form of locking, and that’s still not entirely for free, and (as with all such mechanisms) gets worse under congestion. With very small work packets like in your case, they’re probably getting close to saturation.

I’m wondering, #2227 has been pending for a while now. Would it be a good idea to tweak the semantics of for-loop iteration by axis to return borrowed slices, with their lifetime bound to the loop iteration? I think that would also help your use case.

[NicolasHug]

yes, I’ll keep an eye on #2227 and the PR

Thanks a lot for your time and help @scoder