Merge pull request #5501 from jaimefrio/faster_fastputmask

charris · charris · commit 2e016ac65ace · 2015-02-24T16:48:26.000-05:00
ENH: speed up putmask avoiding % in inner loop
diff --git a/numpy/core/src/multiarray/arraytypes.c.src b/numpy/core/src/multiarray/arraytypes.c.src
@@ -3680,21 +3680,23 @@ static void
 @name@_fastputmask(@type@ *in, npy_bool *mask, npy_intp ni, @type@ *vals,
         npy_intp nv)
 {
-    npy_intp i;
-    @type@ s_val;
+    npy_intp i, j;
 
     if (nv == 1) {
-        s_val = *vals;
+        @type@ s_val = *vals;
         for (i = 0; i < ni; i++) {
             if (mask[i]) {
                 in[i] = s_val;
             }
         }
     }
     else {
-        for (i = 0; i < ni; i++) {
+        for (i = 0, j = 0; i < ni; i++, j++) {
+            if (j >= nv) {
+                j = 0;
+            }
             if (mask[i]) {
-                in[i] = vals[i%nv];
+                in[i] = vals[j];
             }
         }
     }
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
@@ -427,10 +427,11 @@ NPY_NO_EXPORT PyObject *
 PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
 {
     PyArray_FastPutmaskFunc *func;
-    PyArrayObject  *mask, *values;
+    PyArrayObject *mask, *values;
     PyArray_Descr *dtype;
-    npy_intp i, chunk, ni, max_item, nv, tmp;
+    npy_intp i, j, chunk, ni, max_item, nv;
     char *src, *dest;
+    npy_bool *mask_data;
     int copied = 0;
 
     mask = NULL;
@@ -469,6 +470,7 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
                         "the same size");
         goto fail;
     }
+    mask_data = PyArray_DATA(mask);
     dtype = PyArray_DESCR(self);
     Py_INCREF(dtype);
     values = (PyArrayObject *)PyArray_FromAny(values0, dtype,
@@ -483,14 +485,20 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
         Py_INCREF(Py_None);
         return Py_None;
     }
+    src = PyArray_DATA(values);
+
     if (PyDataType_REFCHK(PyArray_DESCR(self))) {
-        for (i = 0; i < ni; i++) {
-            tmp = ((npy_bool *)(PyArray_DATA(mask)))[i];
-            if (tmp) {
-                src = PyArray_BYTES(values) + chunk * (i % nv);
-                PyArray_Item_INCREF(src, PyArray_DESCR(self));
-                PyArray_Item_XDECREF(dest+i*chunk, PyArray_DESCR(self));
-                memmove(dest + i * chunk, src, chunk);
+        for (i = 0, j = 0; i < ni; i++, j++) {
+            if (j >= nv) {
+                j = 0;
+            }
+            if (mask_data[i]) {
+                char *src_ptr = src + j*chunk;
+                char *dest_ptr = dest + i*chunk;
+
+                PyArray_Item_INCREF(src_ptr, PyArray_DESCR(self));
+                PyArray_Item_XDECREF(dest_ptr, PyArray_DESCR(self));
+                memmove(dest_ptr, src_ptr, chunk);
             }
         }
     }
@@ -499,16 +507,17 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
         NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(self));
         func = PyArray_DESCR(self)->f->fastputmask;
         if (func == NULL) {
-            for (i = 0; i < ni; i++) {
-                tmp = ((npy_bool *)(PyArray_DATA(mask)))[i];
-                if (tmp) {
-                    src = PyArray_BYTES(values) + chunk*(i % nv);
-                    memmove(dest + i*chunk, src, chunk);
+            for (i = 0, j = 0; i < ni; i++, j++) {
+                if (j >= nv) {
+                    j = 0;
+                }
+                if (mask_data[i]) {
+                    memmove(dest + i*chunk, src + j*chunk, chunk);
                 }
             }
         }
         else {
-            func(dest, PyArray_DATA(mask), ni, PyArray_DATA(values), nv);
+            func(dest, mask_data, ni, src, nv);
         }
         NPY_END_THREADS;
     }

Original file line number	Diff line number	Diff line change
`@@ -3680,21 +3680,23 @@ static void`
`3680`	`3680`	`@name@_fastputmask(@type@ in, npy_bool mask, npy_intp ni, @type@ *vals,`
`3681`	`3681`	`npy_intp nv)`
`3682`	`3682`	`{`
`3683`		`- npy_intp i;`
`3684`		`- @type@ s_val;`
	`3683`	`+ npy_intp i, j;`
`3685`	`3684`
`3686`	`3685`	`if (nv == 1) {`
`3687`		`- s_val = *vals;`
	`3686`	`+ @type@ s_val = *vals;`
`3688`	`3687`	`for (i = 0; i < ni; i++) {`
`3689`	`3688`	`if (mask[i]) {`
`3690`	`3689`	`in[i] = s_val;`
`3691`	`3690`	`}`
`3692`	`3691`	`}`
`3693`	`3692`	`}`
`3694`	`3693`	`else {`
`3695`		`- for (i = 0; i < ni; i++) {`
	`3694`	`+ for (i = 0, j = 0; i < ni; i++, j++) {`
	`3695`	`+ if (j >= nv) {`
	`3696`	`+ j = 0;`
	`3697`	`+ }`
`3696`	`3698`	`if (mask[i]) {`
`3697`		`- in[i] = vals[i%nv];`
	`3699`	`+ in[i] = vals[j];`
`3698`	`3700`	`}`
`3699`	`3701`	`}`
`3700`	`3702`	`}`