Merge pull request #4497 from juliantaylor/index-gil-release

ENH: release gil for indexing operations

3 files changed, 55 insertions, 6 deletions

diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
index 71c17c9c9..e4126109e 100644
--- a/numpy/core/src/multiarray/item_selection.c
+++ b/numpy/core/src/multiarray/item_selection.c
@@ -350,12 +350,14 @@ PyArray_PutTo(PyArrayObject *self, PyObject* values0, PyObject *indices0,
         }
     }
     else {
+        NPY_BEGIN_THREADS_DEF;
+        NPY_BEGIN_THREADS_THRESHOLDED(ni);
         switch(clipmode) {
         case NPY_RAISE:
             for (i = 0; i < ni; i++) {
                 src = PyArray_BYTES(values) + chunk * (i % nv);
                 tmp = ((npy_intp *)(PyArray_DATA(indices)))[i];
-                if (check_and_adjust_index(&tmp, max_item, 0, NULL) < 0) {
+                if (check_and_adjust_index(&tmp, max_item, 0, _save) < 0) {
                     goto fail;
                 }
                 memmove(dest + tmp * chunk, src, chunk);
@@ -392,6 +394,7 @@ PyArray_PutTo(PyArrayObject *self, PyObject* values0, PyObject *indices0,
             }
             break;
         }
+        NPY_END_THREADS;
     }
 
  finish:
@@ -487,6 +490,8 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
         }
     }
     else {
+        NPY_BEGIN_THREADS_DEF;
+        NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(self));
         func = PyArray_DESCR(self)->f->fastputmask;
         if (func == NULL) {
             for (i = 0; i < ni; i++) {
@@ -500,6 +505,7 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
         else {
             func(dest, PyArray_DATA(mask), ni, PyArray_DATA(values), nv);
         }
+        NPY_END_THREADS;
     }
 
     Py_XDECREF(values);
@@ -2279,6 +2285,7 @@ count_boolean_trues(int ndim, char *data, npy_intp *ashape, npy_intp *astrides)
     npy_intp shape[NPY_MAXDIMS], strides[NPY_MAXDIMS];
     npy_intp i, coord[NPY_MAXDIMS];
     npy_intp count = 0;
+    NPY_BEGIN_THREADS_DEF;
 
     /* Use raw iteration with no heap memory allocation */
     if (PyArray_PrepareOneRawArrayIter(
@@ -2294,6 +2301,8 @@ count_boolean_trues(int ndim, char *data, npy_intp *ashape, npy_intp *astrides)
         return 0;
     }
 
+    NPY_BEGIN_THREADS_THRESHOLDED(shape[0]);
+
     /* Special case for contiguous inner loop */
     if (strides[0] == 1) {
         NPY_RAW_ITER_START(idim, ndim, coord, shape) {
@@ -2323,6 +2332,8 @@ count_boolean_trues(int ndim, char *data, npy_intp *ashape, npy_intp *astrides)
         } NPY_RAW_ITER_ONE_NEXT(idim, ndim, coord, shape, data, strides);
     }
 
+    NPY_END_THREADS;
+
     return count;
 }
 
@@ -2343,6 +2354,7 @@ PyArray_CountNonzero(PyArrayObject *self)
     NpyIter_IterNextFunc *iternext;
     char **dataptr;
     npy_intp *strideptr, *innersizeptr;
+    NPY_BEGIN_THREADS_DEF;
 
     /* Special low-overhead version specific to the boolean type */
     if (PyArray_DESCR(self)->type_num == NPY_BOOL) {
@@ -2392,6 +2404,9 @@ PyArray_CountNonzero(PyArrayObject *self)
         NpyIter_Deallocate(iter);
         return -1;
     }
+
+    NPY_BEGIN_THREADS_NDITER(iter);
+
     dataptr = NpyIter_GetDataPtrArray(iter);
     strideptr = NpyIter_GetInnerStrideArray(iter);
     innersizeptr = NpyIter_GetInnerLoopSizePtr(iter);
@@ -2411,6 +2426,8 @@ PyArray_CountNonzero(PyArrayObject *self)
 
     } while(iternext(iter));
 
+    NPY_END_THREADS;
+
     NpyIter_Deallocate(iter);
 
     return PyErr_Occurred() ? -1 : nonzero_count;
diff --git a/numpy/core/src/multiarray/lowlevel_strided_loops.c.src b/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
index bd69c5ca9..d6a0113f8 100644
--- a/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
+++ b/numpy/core/src/multiarray/lowlevel_strided_loops.c.src
@@ -1384,6 +1384,7 @@ mapiter_trivial_@name@(PyArrayObject *self, PyArrayObject *ind,
     int needs_api = PyDataType_REFCHK(PyArray_DESCR(self));
 
     PyArray_CopySwapFunc *copyswap = PyArray_DESCR(self)->f->copyswap;
+    NPY_BEGIN_THREADS_DEF;
 
     base_ptr = PyArray_BYTES(self);
     self_stride = PyArray_STRIDE(self, 0);
@@ -1392,11 +1393,14 @@ mapiter_trivial_@name@(PyArrayObject *self, PyArrayObject *ind,
                                            ind_ptr, result_ptr,
                                            ind_stride, result_stride)
 
+    if (!needs_api) {
+        NPY_BEGIN_THREADS_THRESHOLDED(PyArray_SIZE(ind));
+    }
 #if !@isget@
     /* Check the indices beforehand */
     while (itersize--) {
         indval = *((npy_intp*)ind_ptr);
-        if (check_and_adjust_index(&indval, fancy_dim, 1, NULL) < 0 ) {
+        if (check_and_adjust_index(&indval, fancy_dim, 1, _save) < 0 ) {
             return -1;
         }
         ind_ptr += ind_stride;
@@ -1426,7 +1430,7 @@ mapiter_trivial_@name@(PyArrayObject *self, PyArrayObject *ind,
         while (itersize--) {
             indval = *((npy_intp*)ind_ptr);
 #if @isget@
-            if (check_and_adjust_index(&indval, fancy_dim, 1, NULL) < 0 ) {
+            if (check_and_adjust_index(&indval, fancy_dim, 1, _save) < 0 ) {
                 return -1;
             }
 #else
@@ -1459,6 +1463,8 @@ mapiter_trivial_@name@(PyArrayObject *self, PyArrayObject *ind,
 /**end repeat1**/
     }
 
+    NPY_END_THREADS;
+
     return 0;
 }
 
@@ -1527,6 +1533,10 @@ mapiter_@name@(PyArrayMapIterObject *mit)
 #else
         else {
 #endif
+            NPY_BEGIN_THREADS_DEF;
+            if (!needs_api) {
+                NPY_BEGIN_THREADS;
+            }
 
             /* Optimization for aligned types that do not need the api */
             switch ((is_aligned && !needs_api) ? PyArray_ITEMSIZE(array) : 0) {
@@ -1560,7 +1570,7 @@ mapiter_@name@(PyArrayMapIterObject *mit)
 
 #if @isget@ && @one_iter@
                             if (check_and_adjust_index(&indval, fancy_dims[i],
-                                                       iteraxis, NULL) < 0 ) {
+                                                       iteraxis, _save) < 0 ) {
                                 return -1;
                             }
 #else
@@ -1596,6 +1606,7 @@ mapiter_@name@(PyArrayMapIterObject *mit)
 
 /**end repeat2**/
             }
+            NPY_END_THREADS;
         }
 /**end repeat1**/
     }
@@ -1653,6 +1664,10 @@ mapiter_@name@(PyArrayMapIterObject *mit)
 #else
         else {
 #endif
+            NPY_BEGIN_THREADS_DEF;
+            if (!needs_api) {
+                NPY_BEGIN_THREADS;
+            }
 
             /* Outer iteration (safe because mit->size != 0) */
             do {
@@ -1662,7 +1677,7 @@ mapiter_@name@(PyArrayMapIterObject *mit)
 
 #if @isget@ && @one_iter@
                     if (check_and_adjust_index(&indval, fancy_dims[i],
-                                               iteraxis, NULL) < 0 ) {
+                                               iteraxis, _save) < 0 ) {
                         NPY_AUXDATA_FREE(transferdata);
                         return -1;
                     }
@@ -1680,12 +1695,16 @@ mapiter_@name@(PyArrayMapIterObject *mit)
                  * only a single inner loop.
                  */
                 if (!skip) {
+                    char * errmsg = NULL;
                     subspace_baseptrs[0] = self_ptr;
                     subspace_baseptrs[1] = mit->extra_op_ptrs[0];
 
                     /* (can't really fail, since no buffering necessary) */
                     if (!NpyIter_ResetBasePointers(mit->subspace_iter,
-                                                   subspace_baseptrs, NULL)) {
+                                                   subspace_baseptrs,
+                                                   &errmsg)) {
+                        NPY_END_THREADS;
+                        PyErr_SetString(PyExc_ValueError, errmsg);
                         NPY_AUXDATA_FREE(transferdata);
                         return -1;
                     }
@@ -1722,6 +1741,7 @@ mapiter_@name@(PyArrayMapIterObject *mit)
 
                 mit->extra_op_next(mit->extra_op_iter);
             } while (mit->outer_next(mit->outer));
+            NPY_END_THREADS;
         }
 /**end repeat1**/
 
diff --git a/numpy/core/src/multiarray/mapping.c b/numpy/core/src/multiarray/mapping.c
index cafec7689..1e30c66cf 100644
--- a/numpy/core/src/multiarray/mapping.c
+++ b/numpy/core/src/multiarray/mapping.c
@@ -970,6 +970,7 @@ array_boolean_subscript(PyArrayObject *self,
         npy_intp self_stride, bmask_stride, subloopsize;
         char *self_data;
         char *bmask_data;
+        NPY_BEGIN_THREADS_DEF;
 
         /* Set up the iterator */
         flags = NPY_ITER_EXTERNAL_LOOP | NPY_ITER_REFS_OK;
@@ -1004,6 +1005,9 @@ array_boolean_subscript(PyArrayObject *self,
             NPY_AUXDATA_FREE(transferdata);
             return NULL;
         }
+
+        NPY_BEGIN_THREADS_NDITER(iter);
+
         innerstrides = NpyIter_GetInnerStrideArray(iter);
         dataptrs = NpyIter_GetDataPtrArray(iter);
 
@@ -1031,6 +1035,8 @@ array_boolean_subscript(PyArrayObject *self,
             }
         } while (iternext(iter));
 
+        NPY_END_THREADS;
+
         NpyIter_Deallocate(iter);
         NPY_AUXDATA_FREE(transferdata);
     }
@@ -1143,6 +1149,7 @@ array_assign_boolean_subscript(PyArrayObject *self,
         npy_intp self_stride, bmask_stride, subloopsize;
         char *self_data;
         char *bmask_data;
+        NPY_BEGIN_THREADS_DEF;
 
         /* Set up the iterator */
         flags = NPY_ITER_EXTERNAL_LOOP | NPY_ITER_REFS_OK;
@@ -1161,6 +1168,9 @@ array_assign_boolean_subscript(PyArrayObject *self,
             NpyIter_Deallocate(iter);
             return -1;
         }
+
+        NPY_BEGIN_THREADS_NDITER(iter);
+
         innerstrides = NpyIter_GetInnerStrideArray(iter);
         dataptrs = NpyIter_GetDataPtrArray(iter);
 
@@ -1202,6 +1212,8 @@ array_assign_boolean_subscript(PyArrayObject *self,
             }
         } while (iternext(iter));
 
+        NPY_END_THREADS;
+
         NPY_AUXDATA_FREE(transferdata);
         NpyIter_Deallocate(iter);
     }