MAINT: Reinstante the CONTIG flag (as buggy as it was)

seberg · seberg · commit 92ecbaf29b28 · 2024-12-05T19:39:30.000+01:00
diff --git a/numpy/_core/src/multiarray/nditer_api.c b/numpy/_core/src/multiarray/nditer_api.c
@@ -1515,6 +1515,8 @@ NpyIter_DebugPrint(NpyIter *iter)
             printf("WRITEMASKED ");
         if ((NIT_OPITFLAGS(iter)[iop])&NPY_OP_ITFLAG_BUF_SINGLESTRIDE)
             printf("BUF_SINGLESTRIDE ");
+        if ((NIT_OPITFLAGS(iter)[iop])&NPY_OP_ITFLAG_CONTIG)
+            printf("CONTIG ");
         printf("\n");
     }
     printf("|\n");
@@ -1532,9 +1534,9 @@ NpyIter_DebugPrint(NpyIter *iter)
         if (itflags&NPY_ITFLAG_REDUCE) {
             printf("|   REDUCE Pos: %d\n",
                         (int)NBF_REDUCE_POS(bufferdata));
-            printf("|   BUFFER Reduce outersize: %d\n",
+            printf("|   REDUCE OuterSize: %d\n",
                         (int)NBF_REDUCE_OUTERSIZE(bufferdata));
-            printf("|   BUFFER OuterDim: %d\n",
+            printf("|   REDUCE OuterDim: %d\n",
                         (int)NBF_OUTERDIM(bufferdata));
         }
         printf("|   Strides: ");
@@ -1894,12 +1896,17 @@ npyiter_fill_buffercopy_params(
     }
 
     if (opitflags & NPY_OP_ITFLAG_BUF_SINGLESTRIDE) {
-        /* Flatten the copy into a single stride. */
         *ndim_transfer = 1;
         *op_shape = op_transfersize;
         assert(**op_coords == 0);  /* initialized by caller currently */
         *op_strides = &NAD_STRIDES(axisdata)[iop];
-        if ((*op_strides)[0] == 0) {
+        if ((*op_strides)[0] == 0 && (
+                !(opitflags & NPY_OP_ITFLAG_CONTIG) ||
+                (opitflags & NPY_OP_ITFLAG_WRITE))) {
+            /*
+             * If the user didn't force contig, optimize single element.
+             * (Unless CONTIG was requested and this is not a write/reduce!)
+             */
             *op_transfersize = 1;
             *buf_stride = 0;
         }
diff --git a/numpy/_core/src/multiarray/nditer_constr.c b/numpy/_core/src/multiarray/nditer_constr.c
@@ -1010,6 +1010,10 @@ npyiter_check_per_op_flags(npy_uint32 op_flags, npyiter_opitflags *op_itflags)
         *op_itflags |= NPY_OP_ITFLAG_VIRTUAL;
     }
 
+    if (op_flags & NPY_ITER_CONTIG) {
+        *op_itflags |= NPY_OP_ITFLAG_CONTIG;
+    }
+
     return 1;
 }
 
@@ -2175,6 +2179,11 @@ npyiter_find_buffering_setup(NpyIter *iter)
         npy_bool is_reduce_op;
         npy_bool op_is_buffered = (op_itflags[iop]&NPY_OP_ITFLAG_CAST) != 0;
 
+        /* If contig was requested and this is not writeable avoid zero strides */
+        npy_bool avoid_zero_strides = (
+                (op_itflags[iop] & NPY_OP_ITFLAG_CONTIG) != 0
+                && !(op_itflags[iop] & NPY_OP_ITFLAG_WRITE));
+
         /*
          * Figure out if this is iterated as a reduce op.  Even one marked
          * for reduction may not be iterated as one.
@@ -2189,16 +2198,19 @@ npyiter_find_buffering_setup(NpyIter *iter)
         else if (op_single_stride_dims[iop] == best_dim && !op_is_buffered) {
             /*
              * Optimization: This operand is not buffered and we might as well
-             * iterate it as an unbuffered reduce operand (if not yet buffered).
+             * iterate it as an unbuffered reduce operand.
              */
             is_reduce_op = 1;
         }
         else if (NAD_STRIDES(reduce_axisdata)[iop] == 0
-                    && op_single_stride_dims[iop] <= best_dim) {
+                    && op_single_stride_dims[iop] <= best_dim
+                    && !avoid_zero_strides) {
             /*
              * Optimization: If the outer (reduce) stride is 0 on the operand
              * then we can iterate this in a reduce way: buffer the core only
              * and repeat it in the "outer" dimension.
+             * If user requested contig, we may have to avoid 0 strides, this
+             * is incompatible with the reduce path.
              */
             is_reduce_op = 1;
         }
@@ -2229,7 +2241,8 @@ npyiter_find_buffering_setup(NpyIter *iter)
              */
             if (!is_reduce_op
                     && NIT_OPITFLAGS(iter)[iop] & NPY_OP_ITFLAG_BUF_SINGLESTRIDE
-                    && NAD_STRIDES(axisdata)[iop] == 0) {
+                    && NAD_STRIDES(axisdata)[iop] == 0
+                    && !avoid_zero_strides) {
                 /* This op is always 0 strides, so even the buffer is that. */
                 inner_stride = 0;
                 reduce_outer_stride = 0;
@@ -3310,11 +3323,16 @@ npyiter_allocate_arrays(NpyIter *iter,
         }
 
         /* Here we can finally check for contiguous iteration */
-        if (op_flags[iop] & NPY_ITER_CONTIG) {
+        if (op_itflags[iop] & NPY_OP_ITFLAG_CONTIG) {
             NpyIter_AxisData *axisdata = NIT_AXISDATA(iter);
             npy_intp stride = NAD_STRIDES(axisdata)[iop];
 
             if (stride != op_dtype[iop]->elsize) {
+                /*
+                 * Need to copy to buffer (cast) to ensure contiguous
+                 * NOTE: This is the wrong place in case of axes reorder
+                 *       (there is an xfailing test for this).
+                 */
                 NPY_IT_DBG_PRINT("Iterator: Setting NPY_OP_ITFLAG_CAST "
                                     "because of NPY_ITER_CONTIG\n");
                 op_itflags[iop] |= NPY_OP_ITFLAG_CAST;
diff --git a/numpy/_core/src/multiarray/nditer_impl.h b/numpy/_core/src/multiarray/nditer_impl.h
@@ -139,6 +139,8 @@
 #define NPY_OP_ITFLAG_FORCECOPY    0x0200
 /* The operand has temporary data, write it back at dealloc */
 #define NPY_OP_ITFLAG_HAS_WRITEBACK 0x0400
+/* Whether the user request a contiguous operand */
+#define NPY_OP_ITFLAG_CONTIG 0x0800
 
 /*
  * The data layout of the iterator is fully specified by
diff --git a/numpy/_core/tests/test_nditer.py b/numpy/_core/tests/test_nditer.py
@@ -2320,6 +2320,79 @@ def test_iter_buffering_growinner():
     assert_equal(i[0].size, a.size)
 
 
+@pytest.mark.parametrize("read_or_readwrite", ["readonly", "readwrite"])
+def test_iter_contig_flag_reduce_error(read_or_readwrite):
+    # Test that a non-contiguous operand is rejected without buffering.
+    # NOTE: This is true even for a reduction, where we return a 0-stride
+    #       below!
+    with pytest.raises(TypeError, match="Iterator operand required buffering"):
+        it = np.nditer(
+            (np.zeros(()),), flags=["external_loop", "reduce_ok"],
+            op_flags=[(read_or_readwrite, "contig"),], itershape=(10,))
+
+
+@pytest.mark.parametrize("arr", [
+        lambda: np.zeros(()),
+        lambda: np.zeros((20, 1))[::20],
+        lambda: np.zeros((1, 20))[:, ::20]
+    ])
+def test_iter_contig_flag_single_operand_strides(arr):
+    """
+    Tests the strides with the contig flag for both broadcast and non-broadcast
+    operands in 3 cases where the logic is needed:
+    1. When everything has a zero stride, the broadcast op needs to repeated
+    2. When the reduce axis is the last axis (first to iterate).
+    3. When the reduce axis is the first axis (last to iterate).
+
+    NOTE: The semantics of the cast flag are not clearly defined when
+          it comes to reduction.  It is unclear that there are any users.
+    """
+    first_op = np.ones((10, 10))
+    broadcast_op = arr()
+    red_op = arr()
+    # Add a first operand to ensure no axis-reordering and the result shape.
+    iterator = np.nditer(
+        (first_op, broadcast_op, red_op),
+        flags=["external_loop", "reduce_ok", "buffered", "delay_bufalloc"],
+        op_flags=[("readonly", "contig")] * 2 + [("readwrite", "contig")])
+
+    with iterator:
+        iterator.reset()
+        try:
+            for f, b, r in iterator:
+                # The first operand is contigouos, we should have a view
+                assert np.shares_memory(f, first_op)
+                # Although broadcast, the second op always has a contiguous stride
+                assert b.strides[0] == 8
+                assert not np.shares_memory(b, broadcast_op)
+                # The reduction has a contiguous stride or a 0 stride
+                if red_op.ndim == 0 or red_op.shape[-1] == 1:
+                    assert r.strides[0] == 0
+                else:
+                    # The stride is 8, although it was not originally:
+                    assert r.strides[0] == 8
+                # If the reduce stride is 0, buffering makes no difference, but we
+                # do it anyway right now:
+                assert not np.shares_memory(r, red_op)
+        finally:
+            iterator.debug_print()
+
+
+@pytest.mark.xfail(reason="The contig flag was always buggy.")
+def test_iter_contig_flag_incorrect():
+    # This case does the wrong thing...
+    iterator = np.nditer(
+        (np.ones((10, 10)).T, np.ones((1, 10))),
+        flags=["external_loop", "reduce_ok", "buffered", "delay_bufalloc"],
+        op_flags=[("readonly", "contig")] * 2)
+
+    with iterator:
+        iterator.reset()
+        for a, b in iterator:
+            assert a.strides == 8
+            assert b.strides == 8  # should be 8 but is 0 due to axis reorder
+
+
 @pytest.mark.slow
 def test_iter_buffered_reduce_reuse():
     # large enough array for all views, including negative strides.
@@ -3324,8 +3397,8 @@ def test_debug_print(capfd):
     |   BufIterEnd: 5
     |   BUFFER CoreSize: 5
     |   REDUCE Pos: 0
-    |   BUFFER Reduce outersize: 10
-    |   BUFFER OuterDim: 1
+    |   REDUCE OuterSize: 10
+    |   REDUCE OuterDim: 1
     |   Strides: 8 4
     |   Ptrs:
     |   REDUCE Outer Strides: 40 0
