ENH: iter: Move code out of NpyIter_MultiNew into functions

1 files changed, 430 insertions, 365 deletions

diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index 06daed3ba..ee1d0b439 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -31,7 +31,7 @@
 #define NPY_ITFLAG_NOINNER     0x020
 /* The iterator is buffered */
 #define NPY_ITFLAG_BUFFER      0x040
-/* The iterator is buffered */
+/* The iterator should grow the buffered inner loop when possible */
 #define NPY_ITFLAG_GROWINNER   0x080
 
 /* Internal iterator per-operand iterator flags */
@@ -51,10 +51,6 @@
 /* The operand is aligned */
 #define NPY_OP_ITFLAG_ALIGNED      0x40
 
-/* Internal flag, for the type of operands */
-#define NPY_ITER_OP_ARRAY         0
-#define NPY_ITER_OP_NULL          1
-
 /*
  * The data layout of the iterator is fully specified by
  * a triple (itflags, ndim, niter).  These three variables
@@ -212,12 +208,19 @@ struct NpyIter_AD {
 
 /* Internal helper functions */
 static int
-pyiter_check_global_flags(npy_uint32 flags, npy_uint32* itflags);
+npyiter_check_global_flags(npy_uint32 flags, npy_uint32* itflags);
+static int
+npyiter_check_op_axes(npy_intp niter, npy_intp oa_ndim, npy_intp **op_axes);
 static int
 npyiter_check_per_op_flags(npy_uint32 flags, char *op_itflags);
 static int
-pyiter_prepare_operand(PyArrayObject **op, PyArray_Descr *op_request_dtype,
-                       PyArray_Descr** op_dtype, int* op_type,
+npyiter_prepare_operands(npy_intp niter, npy_intp *ndim, PyArrayObject **op_in,
+                    PyArrayObject **op, PyArray_Descr **op_request_dtypes,
+                    PyArray_Descr **op_dtype,
+                    npy_intp *op_ndim, npy_uint32 *op_flags, char *op_itflags);
+static int
+npyiter_prepare_one_operand(PyArrayObject **op, PyArray_Descr *op_request_dtype,
+                       PyArray_Descr** op_dtype,
                        npy_intp* op_ndim,
                        npy_uint32 op_flags, char *op_itflags);
 static int
@@ -263,6 +266,15 @@ static PyArrayObject *
 npyiter_new_temp_array(NpyIter *iter, PyTypeObject *subtype,
                 npy_intp op_ndim, npy_intp *shape,
                 PyArray_Descr *op_dtype, npy_intp *op_axes);
+static int
+npyiter_allocate_arrays(NpyIter *iter, PyArrayObject **op,
+                        PyArray_Descr **op_dtype, PyTypeObject *subtype,
+                        npy_uint32 *op_flags, char *op_itflags,
+                        npy_intp *op_ndim, npy_intp **op_axes);
+static void
+npyiter_get_priority_subtype(PyArrayObject **op, char *op_itflags,
+                            npy_intp niter, double *subtype_priority,
+                            PyTypeObject **subtype);
 
 static int
 npyiter_allocate_buffers(NpyIter *iter);
@@ -291,16 +303,15 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
     /* Per-operand values */
     PyArrayObject *op[NPY_MAXARGS];
     PyArray_Descr *op_dtype[NPY_MAXARGS];
-    int op_type[NPY_MAXARGS];
     char op_itflags[NPY_MAXARGS];
     npy_intp op_ndim[NPY_MAXARGS];
     char **op_dataptr;
 
     npy_intp *perm;
     NpyIter_BufferData *bufferdata = NULL;
-    char axes_dupcheck[NPY_MAXDIMS];
-    int any_allocate_if_null = 0, any_missing_dtypes = 0,
-            allocate_output_scalars = 0;
+    int any_allocate = 0, any_missing_dtypes = 0,
+            allocate_output_scalars = 0, need_subtype = 0;
+
     /* The subtype for automatically allocated outputs */
     double subtype_priority = NPY_PRIORITY;
     PyTypeObject *subtype = &PyArray_Type;
@@ -313,59 +324,12 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
     }
 
     /* Error check 'oa_ndim' and 'op_axes', which must be used together */
-    if (oa_ndim == 0 && op_axes != NULL) {
-        PyErr_Format(PyExc_ValueError,
-                "If 'op_axes' is not NULL in the iterator constructor, "
-                "'oa_ndim' must be greater than zero");
+    if (!npyiter_check_op_axes(niter, oa_ndim, op_axes)) {
         return NULL;
     }
-    else if (oa_ndim > 0) {
-        if (oa_ndim > NPY_MAXDIMS) {
-            PyErr_Format(PyExc_ValueError,
-                "Cannot construct an iterator with more than %d dimensions "
-                "(%d were requested for op_axes)",
-                (int)NPY_MAXDIMS, (int)oa_ndim);
-            return NULL;
-        }
-        else if (op_axes == NULL) {
-            PyErr_Format(PyExc_ValueError,
-                    "If 'oa_ndim' is greater than zero in the iterator "
-                    "constructor, then op_axes cannot be NULL");
-            return NULL;
-        }
 
-        /* Check that there are no duplicates in op_axes */
-        for (iiter = 0; iiter < niter; ++iiter) {
-            npy_intp *axes = op_axes[iiter];
-            if (axes != NULL) {
-                memset(axes_dupcheck, 0, NPY_MAXDIMS);
-                for (idim = 0; idim < oa_ndim; ++idim) {
-                    npy_intp i = axes[idim];
-                    if (i >= 0) {
-                        if (i >= NPY_MAXDIMS) {
-                            PyErr_Format(PyExc_ValueError,
-                                    "The 'op_axes' provided to the iterator "
-                                    "constructor contained invalid "
-                                    "values %d", (int)i);
-                            return NULL;
-                        } else if(axes_dupcheck[i] == 1) {
-                            PyErr_Format(PyExc_ValueError,
-                                    "The 'op_axes' provided to the iterator "
-                                    "constructor contained duplicate "
-                                    "value %d", (int)i);
-                            return NULL;
-                        }
-                        else {
-                            axes_dupcheck[i] = 1;
-                        }
-                    }
-                }
-            }
-        }
-    }
-
-    /* Checks the global iterator flags */
-    if (!pyiter_check_global_flags(flags, &itflags)) {
+    /* Check the global iterator flags */
+    if (!npyiter_check_global_flags(flags, &itflags)) {
         return NULL;
     }
 
@@ -374,76 +338,15 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
      * chosen to be big enough to get some amortization benefits, but
      * small enough to be cache-friendly.
      */
-    if (itflags&NPY_ITFLAG_BUFFER && buffersize <= 0) {
+    if ((itflags&NPY_ITFLAG_BUFFER) && buffersize <= 0) {
         buffersize = 1 << 12;
     }
 
     /* Prepare all the operands */
-    for (iiter = 0; iiter < niter; ++iiter) {
-        /*
-         * Make a copy of the input operands so we can substitute
-         * new values in place when necessary without affecting
-         * the caller's array.
-         */
-        op[iiter] = op_in[iiter];
-        Py_XINCREF(op[iiter]);
-        op_dtype[iiter] = NULL;
-
-        /* Check the readonly/writeonly flags, and fill in op_itflags */
-        if (!npyiter_check_per_op_flags(op_flags[iiter], &op_itflags[iiter])) {
-            npy_intp i;
-
-            for (i = 0; i <= iiter; ++i) {
-                Py_XDECREF(op[i]);
-                Py_XDECREF(op_dtype[i]);
-            }
-            return NULL;
-        }
-
-        /*
-         * Prepare the operand.  This produces an op_dtype[iiter] reference
-         * on success.
-         */
-        if (!pyiter_prepare_operand(&op[iiter],
-                        op_request_dtypes ? op_request_dtypes[iiter] : NULL,
-                        &op_dtype[iiter], &op_type[iiter],
-                        &op_ndim[iiter],
-                        op_flags[iiter], &op_itflags[iiter])) {
-            npy_intp i;
-
-            for (i = 0; i <= iiter; ++i) {
-                Py_XDECREF(op[i]);
-                Py_XDECREF(op_dtype[i]);
-            }
-            return NULL;
-        }
-        /* The iterator dimensions is the maximum of all the inputs */
-        if (op_ndim[iiter] > ndim) {
-            ndim = op_ndim[iiter];
-        }
-    }
-
-
-    /* If all the operands were NULL, it's an error */
-    if (op_type[0] == NPY_ITER_OP_NULL) {
-        int all_null = 1;
-        for (iiter = 1; iiter < niter; ++iiter) {
-            if (op_type[iiter] != NPY_ITER_OP_NULL) {
-                all_null = 0;
-                break;
-            }
-        }
-        if (all_null) {
-            npy_intp i;
-
-            for (i = 0; i < niter; ++i) {
-                Py_XDECREF(op[i]);
-                Py_XDECREF(op_dtype[i]);
-            }
-            PyErr_SetString(PyExc_ValueError,
-                    "At least one iterator input must be non-NULL");
-            return NULL;
-        }
+    if (!npyiter_prepare_operands(niter, &ndim, op_in, op,
+                        op_request_dtypes, op_dtype,
+                        op_ndim, op_flags, op_itflags)) {
+        return NULL;
     }
 
     /* If 'op_axes' is being used, force 'ndim' */
@@ -474,36 +377,37 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
         NIT_BASEOFFSETS(iter)[iiter] = 0;
 
         /* Get the data pointer for this operand */
-        switch (op_type[iiter]) {
-            case NPY_ITER_OP_ARRAY:
-                /*
-                 * Array casting/copying is handled later, once the
-                 * iteration order is finalized.  Here, we
-                 * optimistically assume the array will be used
-                 * as is.
-                 */
-                op_dataptr[iiter] = PyArray_DATA(op[iiter]);
-                break;
-
-            case NPY_ITER_OP_NULL:
-                op_dataptr[iiter] = NULL;
-                /* Now that ndim is fixed, outputs get the full ndim */
-                if (allocate_output_scalars) {
-                    op_ndim[iiter] = 0;
-                }
-                else {
-                    op_ndim[iiter] = ndim;
-                }
-                /* Flag this so later we can avoid flipping axes */
-                any_allocate_if_null = 1;
-                /*
-                 * If the data type wasn't provided, will need to
-                 * calculate it later.
-                 */
-                if (op_dtype[iiter] == NULL) {
-                    any_missing_dtypes = 1;
-                }
-                break;
+        if (op[iiter] != NULL) {
+            /*
+             * Array casting/copying is handled later, once the
+             * iteration order is finalized.  Here, we
+             * optimistically assume the array will be used
+             * as is.
+             */
+            op_dataptr[iiter] = PyArray_DATA(op[iiter]);
+        }
+        else {
+            op_dataptr[iiter] = NULL;
+            /* Now that ndim is fixed, outputs get the full ndim */
+            if (allocate_output_scalars) {
+                op_ndim[iiter] = 0;
+            }
+            else {
+                op_ndim[iiter] = ndim;
+            }
+            /* Flag this so later we can avoid flipping axes */
+            any_allocate = 1;
+            /* If a subtype may be used, indicate so */
+            if (!(op_flags[iiter]&NPY_ITER_NO_SUBTYPE)) {
+                need_subtype = 1;
+            }
+            /*
+             * If the data type wasn't provided, will need to
+             * calculate it later.
+             */
+            if (op_dtype[iiter] == NULL) {
+                any_missing_dtypes = 1;
+            }
         }
     }
     /* Set resetindex to zero as well (it's just after the resetdataptr) */
@@ -569,28 +473,15 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
          * If there's an output being allocated, we must not negate
          * any strides.
          */
-        if (!any_allocate_if_null) {
+        if (!any_allocate) {
             npyiter_flip_negative_strides(iter);
         }
         itflags = NIT_ITFLAGS(iter);
     }
 
-    if (any_allocate_if_null) {
-        /*
-         * The __array_priority__ attribute of the inputs determines
-         * the subtype of any output arrays.  Take the subtype
-         * with highest priority.
-         */
-        for (iiter = 0; iiter < niter; ++iiter) {
-            if (op_itflags[iiter]&NPY_OP_ITFLAG_READ) {
-                double priority =
-                            PyArray_GetPriority((PyObject *)op[iiter], 0.0);
-                if (priority > subtype_priority) {
-                    subtype_priority = priority;
-                    subtype = Py_TYPE(op[iiter]);
-                }
-            }
-        }
+    if (need_subtype) {
+        npyiter_get_priority_subtype(op, op_itflags, niter,
+                                     &subtype_priority, &subtype);
     }
 
     /*
@@ -646,189 +537,10 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
      * copying due to casting/byte order/alignment can be
      * done now using a memory layout matching the iterator.
      */
-    for (iiter = 0; iiter < niter; ++iiter) {
-        if (op_type[iiter] == NPY_ITER_OP_NULL) {
-            PyArrayObject *out;
-            PyTypeObject *op_subtype;
-
-            /* Check whether the subtype was disabled */
-            if (op_flags[iiter]&NPY_ITER_NO_SUBTYPE) {
-                op_subtype = &PyArray_Type;
-            }
-            else {
-                op_subtype = subtype;
-            }
-
-            /* Allocate the output array, if possible */
-            out = npyiter_new_temp_array(iter, op_subtype,
-                                        op_ndim[iiter], NULL,
-                                        op_dtype[iiter],
-                                        op_axes ? op_axes[iiter] : NULL);
-            if (out == NULL) {
-                NpyIter_Deallocate(iter);
-                return NULL;
-            }
-
-            op[iiter] = out;
-            NIT_OBJECTS(iter)[iiter] = out;
-
-            /*
-             * Now we need to replace the pointers and strides with values
-             * from the new array.
-             */
-            npyiter_replace_axisdata(iter, iiter, op[iiter], op_ndim[iiter],
-                    PyArray_DATA(op[iiter]), op_axes ? op_axes[iiter] : NULL);
-
-            /* New arrays are aligned and need no swapping or casting */
-            op_itflags[iiter] |= NPY_OP_ITFLAG_ALIGNED;
-            op_itflags[iiter] &= ~(NPY_OP_ITFLAG_COPYSWAP|NPY_OP_ITFLAG_CAST);
-        }
-        else if ((op_itflags[iiter]&
-                        (NPY_OP_ITFLAG_CAST|NPY_OP_ITFLAG_COPYSWAP)) &&
-                              (op_itflags[iiter]&NPY_OP_ITFLAG_COPY)) {
-            PyArrayObject *temp;
-
-            /* Allocate the temporary array, if possible */
-            temp = npyiter_new_temp_array(iter, &PyArray_Type,
-                                        PyArray_NDIM(op[iiter]),
-                                        PyArray_DIMS(op[iiter]),
-                                        op_dtype[iiter],
-                                        op_axes ? op_axes[iiter] : NULL);
-            if (temp == NULL) {
-                NpyIter_Deallocate(iter);
-                return NULL;
-            }
-
-            /* If the data will be read, copy it into temp */
-            if (op_itflags[iiter]&NPY_OP_ITFLAG_READ) {
-                if (PyArray_CopyInto(temp, op[iiter]) != 0) {
-                    Py_DECREF(temp);
-                    NpyIter_Deallocate(iter);
-                    return NULL;
-                }
-            }
-            /* If the data will be written to, set UPDATEIFCOPY */
-            if (op_itflags[iiter]&NPY_OP_ITFLAG_WRITE) {
-                PyArray_FLAGS(temp) |= NPY_UPDATEIFCOPY;
-                PyArray_FLAGS(op[iiter]) &= ~NPY_WRITEABLE;
-                Py_INCREF(op[iiter]);
-                temp->base = (PyObject *)op[iiter];
-            }
-
-            Py_DECREF(op[iiter]);
-            op[iiter] = temp;
-            NIT_OBJECTS(iter)[iiter] = temp;
-
-            /*
-             * Now we need to replace the pointers and strides with values
-             * from the temporary array.
-             */
-            npyiter_replace_axisdata(iter, iiter, op[iiter], op_ndim[iiter],
-                    PyArray_DATA(op[iiter]), op_axes ? op_axes[iiter] : NULL);
-
-            /* Now it is aligned, and no longer needs a swap or cast */
-            op_itflags[iiter] |= NPY_OP_ITFLAG_ALIGNED;
-            op_itflags[iiter] &= ~(NPY_OP_ITFLAG_COPYSWAP|NPY_OP_ITFLAG_CAST);
-        }
-        else {
-            /*
-             * Buffering must be enabled for casting/conversion if copy
-             * wasn't specified.
-             */
-            if (op_itflags[iiter]&
-                        (NPY_OP_ITFLAG_CAST|NPY_OP_ITFLAG_COPYSWAP) &&
-                          !(itflags&NPY_ITFLAG_BUFFER)) {
-                PyErr_SetString(PyExc_TypeError,
-                        "Iterator input required copying or buffering, "
-                        "but neither copying nor buffering was enabled");
-                NpyIter_Deallocate(iter);
-                return NULL;
-            }
-
-            /*
-             * If the operand is aligned, any buffering can use aligned
-             * optimizations.
-             */
-            if (PyArray_ISALIGNED(op[iiter])) {
-                op_itflags[iiter] |= NPY_OP_ITFLAG_ALIGNED;
-            }
-        }
-
-        /*
-         * If no alignment, byte swap, or casting is needed, and
-         * the inner stride of this operand works for the whole
-         * array, we can set NPY_OP_ITFLAG_BUFNEVER.
-         * But, if buffering is enabled, write-buffering must be
-         * one-to-one, because the buffering write back won't combine
-         * values correctly. This test doesn't catch everything, but it will
-         * catch the most common case of a broadcasting a write-buffered
-         * dimension.
-         */
-        if ((itflags&NPY_ITFLAG_BUFFER) &&
-                        (!(op_itflags[iiter]&(NPY_OP_ITFLAG_CAST|
-                                             NPY_OP_ITFLAG_COPYSWAP)) ||
-                          (op_itflags[iiter]&NPY_OP_ITFLAG_WRITE))) {
-            int is_one_to_one = 1;
-            npy_intp stride, shape, innerstride = 0, innershape;
-            NpyIter_AxisData *axisdata = NIT_AXISDATA(iter);
-            npy_intp sizeof_axisdata =
-                                NIT_SIZEOF_AXISDATA(itflags, ndim, niter);
-            /* Find stride of the first non-empty shape */
-            for (idim = 0; idim < ndim; ++idim) {
-                innershape = NAD_SHAPE(axisdata);
-                if (innershape != 1) {
-                    innerstride = NAD_STRIDES(axisdata)[iiter];
-                    if (innerstride == 0) {
-                        is_one_to_one = 0;
-                    }
-                    break;
-                }
-                NIT_ADVANCE_AXISDATA(axisdata, 1);
-            }
-            ++idim;
-            NIT_ADVANCE_AXISDATA(axisdata, 1);
-            /* Check that everything could have coalesced together */
-            for (; idim < ndim; ++idim) {
-                stride = NAD_STRIDES(axisdata)[iiter];
-                shape = NAD_SHAPE(axisdata);
-                if (shape != 1) {
-                    if (stride == 0) {
-                        is_one_to_one = 0;
-                    }
-                    /*
-                     * If N times the inner stride doesn't equal this
-                     * stride, the multi-dimensionality is needed.
-                     */
-                    if (innerstride*innershape != stride) {
-                        break;
-                    }
-                    else {
-                        innershape *= shape;
-                    }
-                }
-                NIT_ADVANCE_AXISDATA(axisdata, 1);
-            }
-            /*
-             * If we looped all the way to the end, one stride works.
-             * Set that stride, because it may not belong to the first
-             * dimension.
-             */
-            if (idim == ndim &&
-                        !(op_itflags[iiter]&(NPY_OP_ITFLAG_CAST|
-                                             NPY_OP_ITFLAG_COPYSWAP))) {
-                op_itflags[iiter] |= NPY_OP_ITFLAG_BUFNEVER;
-                NBF_STRIDES(bufferdata)[iiter] = innerstride;
-            }
-            else if (!is_one_to_one &&
-                        (op_itflags[iiter]&NPY_OP_ITFLAG_WRITE)) {
-                PyErr_SetString(PyExc_ValueError,
-                        "Iterator operand requires write buffering, "
-                        "but has dimensions which have been broadcasted "
-                        "and would be combined incorrectly");
-                NpyIter_Deallocate(iter);
-                return NULL;
-            }
-        }
+    if (!npyiter_allocate_arrays(iter, op, op_dtype, subtype, op_flags,
+                            op_itflags, op_ndim, op_axes)) {
+        NpyIter_Deallocate(iter);
+        return NULL;
     }
 
     /*
@@ -1895,7 +1607,7 @@ npy_intp* NpyIter_GetInnerLoopSizePtr(NpyIter *iter)
  * Returns 1 on success, 0 on error.
  */
 static int
-pyiter_check_global_flags(npy_uint32 flags, npy_uint32* itflags)
+npyiter_check_global_flags(npy_uint32 flags, npy_uint32* itflags)
 {
     if ((flags&NPY_ITER_PER_OP_FLAGS) != 0) {
         PyErr_SetString(PyExc_ValueError,
@@ -1944,6 +1656,68 @@ pyiter_check_global_flags(npy_uint32 flags, npy_uint32* itflags)
     return 1;
 }
 
+static int
+npyiter_check_op_axes(npy_intp niter, npy_intp oa_ndim, npy_intp **op_axes)
+{
+    char axes_dupcheck[NPY_MAXDIMS];
+    npy_intp iiter, idim;
+
+    if (oa_ndim == 0 && op_axes != NULL) {
+        PyErr_Format(PyExc_ValueError,
+                "If 'op_axes' is not NULL in the iterator constructor, "
+                "'oa_ndim' must be greater than zero");
+        return 0;
+    }
+    else if (oa_ndim > 0) {
+        if (oa_ndim > NPY_MAXDIMS) {
+            PyErr_Format(PyExc_ValueError,
+                "Cannot construct an iterator with more than %d dimensions "
+                "(%d were requested for op_axes)",
+                (int)NPY_MAXDIMS, (int)oa_ndim);
+            return 0;
+        }
+        else if (op_axes == NULL) {
+            PyErr_Format(PyExc_ValueError,
+                    "If 'oa_ndim' is greater than zero in the iterator "
+                    "constructor, then op_axes cannot be NULL");
+            return 0;
+        }
+
+        /* Check that there are no duplicates in op_axes */
+        for (iiter = 0; iiter < niter; ++iiter) {
+            npy_intp *axes = op_axes[iiter];
+            if (axes != NULL) {
+                memset(axes_dupcheck, 0, NPY_MAXDIMS);
+                for (idim = 0; idim < oa_ndim; ++idim) {
+                    npy_intp i = axes[idim];
+                    if (i >= 0) {
+                        if (i >= NPY_MAXDIMS) {
+                            PyErr_Format(PyExc_ValueError,
+                                    "The 'op_axes' provided to the iterator "
+                                    "constructor for operand %d "
+                                    "contained invalid "
+                                    "values %d", (int)iiter, (int)i);
+                            return 0;
+                        } else if(axes_dupcheck[i] == 1) {
+                            PyErr_Format(PyExc_ValueError,
+                                    "The 'op_axes' provided to the iterator "
+                                    "constructor for operand %d "
+                                    "contained duplicate "
+                                    "value %d", (int)iiter, (int)i);
+                            return 0;
+                        }
+                        else {
+                            axes_dupcheck[i] = 1;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    return 1;
+}
+
 /*
  * Checks the per-operand input flags, and fills in op_itflags.
  *
@@ -2011,6 +1785,85 @@ npyiter_check_per_op_flags(npy_uint32 op_flags, char *op_itflags)
 }
 
 /*
+ * Process all the operands, copying new references so further processing
+ * can replace the arrays if copying is necessary.  Fill in the iterator's
+ * natural ndim.
+ */
+static int
+npyiter_prepare_operands(npy_intp niter, npy_intp *ndim, PyArrayObject **op_in,
+                    PyArrayObject **op, PyArray_Descr **op_request_dtypes,
+                    PyArray_Descr **op_dtype,
+                    npy_intp *op_ndim, npy_uint32 *op_flags, char *op_itflags)
+{
+    npy_intp iiter;
+
+    *ndim = 0;
+
+    for (iiter = 0; iiter < niter; ++iiter) {
+        op[iiter] = op_in[iiter];
+        Py_XINCREF(op[iiter]);
+        op_dtype[iiter] = NULL;
+
+        /* Check the readonly/writeonly flags, and fill in op_itflags */
+        if (!npyiter_check_per_op_flags(op_flags[iiter], &op_itflags[iiter])) {
+            npy_intp i;
+
+            for (i = 0; i <= iiter; ++i) {
+                Py_XDECREF(op[i]);
+                Py_XDECREF(op_dtype[i]);
+            }
+            return 0;
+        }
+
+        /*
+         * Prepare the operand.  This produces an op_dtype[iiter] reference
+         * on success.
+         */
+        if (!npyiter_prepare_one_operand(&op[iiter],
+                        op_request_dtypes ? op_request_dtypes[iiter] : NULL,
+                        &op_dtype[iiter], &op_ndim[iiter],
+                        op_flags[iiter], &op_itflags[iiter])) {
+            npy_intp i;
+
+            for (i = 0; i <= iiter; ++i) {
+                Py_XDECREF(op[i]);
+                Py_XDECREF(op_dtype[i]);
+            }
+            return 0;
+        }
+        /* The iterator dimensions is the maximum of all the inputs */
+        if (op_ndim[iiter] > *ndim) {
+            *ndim = op_ndim[iiter];
+        }
+    }
+
+
+    /* If all the operands were NULL, it's an error */
+    if (op[0] == NULL) {
+        int all_null = 1;
+        for (iiter = 1; iiter < niter; ++iiter) {
+            if (op[iiter] != NULL) {
+                all_null = 0;
+                break;
+            }
+        }
+        if (all_null) {
+            npy_intp i;
+
+            for (i = 0; i < niter; ++i) {
+                Py_XDECREF(op[i]);
+                Py_XDECREF(op_dtype[i]);
+            }
+            PyErr_SetString(PyExc_ValueError,
+                    "At least one iterator input must be non-NULL");
+            return 0;
+        }
+    }
+
+    return 1;
+}
+
+/*
  * Returns 1 if the from -> to cast can be done, based on the casting
  * flags provided in op_flags, and 0 otherwise.
  *
@@ -2086,14 +1939,14 @@ npyiter_can_cast(PyArray_Descr *from, PyArray_Descr *to, NPY_CASTING casting)
 
 /*
  * Prepares a a constructor operand.  Assumes a reference to 'op'
- * is owned, and that 'op' may be replaced.  Fills in 'op_dtype',
- * 'op_type' and 'ndim'.
+ * is owned, and that 'op' may be replaced.  Fills in 'op_dtype'
+ * and 'ndim'.
  *
  * Returns 1 on success, 0 on failure.
  */
 static int
-pyiter_prepare_operand(PyArrayObject **op, PyArray_Descr *op_request_dtype,
-                       PyArray_Descr **op_dtype, int* op_type,
+npyiter_prepare_one_operand(PyArrayObject **op, PyArray_Descr *op_request_dtype,
+                       PyArray_Descr **op_dtype,
                        npy_intp* op_ndim,
                        npy_uint32 op_flags, char *op_itflags)
 {
@@ -2116,7 +1969,6 @@ pyiter_prepare_operand(PyArrayObject **op, PyArray_Descr *op_request_dtype,
         /* If a requested dtype was provided, use it, otherwise NULL */
         Py_XINCREF(op_request_dtype);
         *op_dtype = op_request_dtype;
-        *op_type = NPY_ITER_OP_NULL;
         *op_ndim = 0;
         /* No copying of NULL operands */
         *op_itflags &= ~NPY_OP_ITFLAG_COPY;
@@ -2132,7 +1984,6 @@ pyiter_prepare_operand(PyArrayObject **op, PyArray_Descr *op_request_dtype,
             return 0;
 
         }
-        *op_type = NPY_ITER_OP_ARRAY;
         *op_ndim = PyArray_NDIM(*op);
         /* PyArray_DESCR does not give us a reference */
         *op_dtype = PyArray_DESCR(*op);
@@ -3267,6 +3118,220 @@ npyiter_new_temp_array(NpyIter *iter, PyTypeObject *subtype,
     return ret;
 }
 
+static int
+npyiter_allocate_arrays(NpyIter *iter, PyArrayObject **op,
+                        PyArray_Descr **op_dtype, PyTypeObject *subtype,
+                        npy_uint32 *op_flags, char *op_itflags,
+                        npy_intp *op_ndim, npy_intp **op_axes)
+{
+    npy_uint32 itflags = NIT_ITFLAGS(iter);
+    npy_intp idim, ndim = NIT_NDIM(iter);
+    npy_intp iiter, niter = NIT_NITER(iter);
+
+    NpyIter_BufferData *bufferdata = NIT_BUFFERDATA(iter);
+
+    for (iiter = 0; iiter < niter; ++iiter) {
+        if (op[iiter] == NULL) {
+            PyArrayObject *out;
+            PyTypeObject *op_subtype;
+
+            /* Check whether the subtype was disabled */
+            op_subtype = (op_flags[iiter]&NPY_ITER_NO_SUBTYPE) ?
+                                                &PyArray_Type : subtype;
+
+            /* Allocate the output array, if possible */
+            out = npyiter_new_temp_array(iter, op_subtype,
+                                        op_ndim[iiter], NULL,
+                                        op_dtype[iiter],
+                                        op_axes ? op_axes[iiter] : NULL);
+            if (out == NULL) {
+                return 0;
+            }
+
+            op[iiter] = out;
+            NIT_OBJECTS(iter)[iiter] = out;
+
+            /*
+             * Now we need to replace the pointers and strides with values
+             * from the new array.
+             */
+            npyiter_replace_axisdata(iter, iiter, op[iiter], op_ndim[iiter],
+                    PyArray_DATA(op[iiter]), op_axes ? op_axes[iiter] : NULL);
+
+            /* New arrays are aligned and need no swapping or casting */
+            op_itflags[iiter] |= NPY_OP_ITFLAG_ALIGNED;
+            op_itflags[iiter] &= ~(NPY_OP_ITFLAG_COPYSWAP|NPY_OP_ITFLAG_CAST);
+        }
+        else if ((op_itflags[iiter]&
+                        (NPY_OP_ITFLAG_CAST|NPY_OP_ITFLAG_COPYSWAP)) &&
+                              (op_itflags[iiter]&NPY_OP_ITFLAG_COPY)) {
+            PyArrayObject *temp;
+
+            /* Allocate the temporary array, if possible */
+            temp = npyiter_new_temp_array(iter, &PyArray_Type,
+                                        PyArray_NDIM(op[iiter]),
+                                        PyArray_DIMS(op[iiter]),
+                                        op_dtype[iiter],
+                                        op_axes ? op_axes[iiter] : NULL);
+            if (temp == NULL) {
+                return 0;
+            }
+
+            /* If the data will be read, copy it into temp */
+            if (op_itflags[iiter]&NPY_OP_ITFLAG_READ) {
+                if (PyArray_CopyInto(temp, op[iiter]) != 0) {
+                    Py_DECREF(temp);
+                    return 0;
+                }
+            }
+            /* If the data will be written to, set UPDATEIFCOPY */
+            if (op_itflags[iiter]&NPY_OP_ITFLAG_WRITE) {
+                PyArray_FLAGS(temp) |= NPY_UPDATEIFCOPY;
+                PyArray_FLAGS(op[iiter]) &= ~NPY_WRITEABLE;
+                Py_INCREF(op[iiter]);
+                temp->base = (PyObject *)op[iiter];
+            }
+
+            Py_DECREF(op[iiter]);
+            op[iiter] = temp;
+            NIT_OBJECTS(iter)[iiter] = temp;
+
+            /*
+             * Now we need to replace the pointers and strides with values
+             * from the temporary array.
+             */
+            npyiter_replace_axisdata(iter, iiter, op[iiter], op_ndim[iiter],
+                    PyArray_DATA(op[iiter]), op_axes ? op_axes[iiter] : NULL);
+
+            /* The temporary copy is aligned and needs no swap or cast */
+            op_itflags[iiter] |= NPY_OP_ITFLAG_ALIGNED;
+            op_itflags[iiter] &= ~(NPY_OP_ITFLAG_COPYSWAP|NPY_OP_ITFLAG_CAST);
+        }
+        else {
+            /*
+             * Buffering must be enabled for casting/conversion if copy
+             * wasn't specified.
+             */
+            if (op_itflags[iiter]&
+                        (NPY_OP_ITFLAG_CAST|NPY_OP_ITFLAG_COPYSWAP) &&
+                          !(itflags&NPY_ITFLAG_BUFFER)) {
+                PyErr_SetString(PyExc_TypeError,
+                        "Iterator operand required copying or buffering, "
+                        "but neither copying nor buffering was enabled");
+                return 0;
+            }
+
+            /*
+             * If the operand is aligned, any buffering can use aligned
+             * optimizations.
+             */
+            if (PyArray_ISALIGNED(op[iiter])) {
+                op_itflags[iiter] |= NPY_OP_ITFLAG_ALIGNED;
+            }
+        }
+
+        /*
+         * If no alignment, byte swap, or casting is needed, and
+         * the inner stride of this operand works for the whole
+         * array, we can set NPY_OP_ITFLAG_BUFNEVER.
+         * But, if buffering is enabled, write-buffering must be
+         * one-to-one, because the buffering write back won't combine
+         * values correctly. This test doesn't catch everything, but it will
+         * catch the most common case of a broadcasting a write-buffered
+         * dimension.
+         */
+        if ((itflags&NPY_ITFLAG_BUFFER) &&
+                        (!(op_itflags[iiter]&(NPY_OP_ITFLAG_CAST|
+                                             NPY_OP_ITFLAG_COPYSWAP)) ||
+                          (op_itflags[iiter]&NPY_OP_ITFLAG_WRITE))) {
+            int is_one_to_one = 1;
+            npy_intp stride, shape, innerstride = 0, innershape;
+            NpyIter_AxisData *axisdata = NIT_AXISDATA(iter);
+            npy_intp sizeof_axisdata =
+                                NIT_SIZEOF_AXISDATA(itflags, ndim, niter);
+            /* Find stride of the first non-empty shape */
+            for (idim = 0; idim < ndim; ++idim) {
+                innershape = NAD_SHAPE(axisdata);
+                if (innershape != 1) {
+                    innerstride = NAD_STRIDES(axisdata)[iiter];
+                    if (innerstride == 0) {
+                        is_one_to_one = 0;
+                    }
+                    break;
+                }
+                NIT_ADVANCE_AXISDATA(axisdata, 1);
+            }
+            ++idim;
+            NIT_ADVANCE_AXISDATA(axisdata, 1);
+            /* Check that everything could have coalesced together */
+            for (; idim < ndim; ++idim) {
+                stride = NAD_STRIDES(axisdata)[iiter];
+                shape = NAD_SHAPE(axisdata);
+                if (shape != 1) {
+                    if (stride == 0) {
+                        is_one_to_one = 0;
+                    }
+                    /*
+                     * If N times the inner stride doesn't equal this
+                     * stride, the multi-dimensionality is needed.
+                     */
+                    if (innerstride*innershape != stride) {
+                        break;
+                    }
+                    else {
+                        innershape *= shape;
+                    }
+                }
+                NIT_ADVANCE_AXISDATA(axisdata, 1);
+            }
+            /*
+             * If we looped all the way to the end, one stride works.
+             * Set that stride, because it may not belong to the first
+             * dimension.
+             */
+            if (idim == ndim &&
+                        !(op_itflags[iiter]&(NPY_OP_ITFLAG_CAST|
+                                             NPY_OP_ITFLAG_COPYSWAP))) {
+                op_itflags[iiter] |= NPY_OP_ITFLAG_BUFNEVER;
+                NBF_STRIDES(bufferdata)[iiter] = innerstride;
+            }
+            else if (!is_one_to_one &&
+                        (op_itflags[iiter]&NPY_OP_ITFLAG_WRITE)) {
+                PyErr_SetString(PyExc_ValueError,
+                        "Iterator operand requires write buffering, "
+                        "but has dimensions which have been broadcasted "
+                        "and would be combined incorrectly");
+                return 0;
+            }
+        }
+    }
+
+    return 1;
+}
+
+/*
+ * The __array_priority__ attribute of the inputs determines
+ * the subtype of any output arrays.  This function finds the
+ * subtype of the input array with highest priority.
+ */
+static void
+npyiter_get_priority_subtype(PyArrayObject **op, char *op_itflags,
+                            npy_intp niter, double *subtype_priority,
+                            PyTypeObject **subtype)
+{
+    npy_intp iiter;
+
+    for (iiter = 0; iiter < niter; ++iiter) {
+        if (op_itflags[iiter]&NPY_OP_ITFLAG_READ) {
+            double priority = PyArray_GetPriority((PyObject *)op[iiter], 0.0);
+            if (priority > *subtype_priority) {
+                *subtype_priority = priority;
+                *subtype = Py_TYPE(op[iiter]);
+            }
+        }
+    }
+}
+
 /*
  * Calculates a dtype that all the types can be promoted to, using the
  * ufunc rules.  If only_inputs is 1, it leaves any operands that