Merge pull request numpy#26972 from yairchu/overlap-off-by-one

BUG: Off by one in memory overlap check

Author: seberg

numpy/_core/src/umath/fast_loop_macros.h

@@ -341,15 +341,15 @@ abs_ptrdiff(char *a, char *b)
     ((labs(steps[0]) < MAX_STEP_SIZE)  && \
      (labs(steps[1]) < MAX_STEP_SIZE)  && \
      (labs(steps[2]) < MAX_STEP_SIZE)  && \
-     (nomemoverlap(args[0], steps[0] * dimensions[0], args[2], steps[2] * dimensions[0])) && \
-     (nomemoverlap(args[1], steps[1] * dimensions[0], args[2], steps[2] * dimensions[0])))
+     (nomemoverlap(args[0], steps[0], args[2], steps[2], dimensions[0])) && \
+     (nomemoverlap(args[1], steps[1], args[2], steps[2], dimensions[0])))
 
 #define IS_UNARY_TWO_OUT_SMALL_STEPS_AND_NOMEMOVERLAP \
     ((labs(steps[0]) < MAX_STEP_SIZE)  && \
      (labs(steps[1]) < MAX_STEP_SIZE)  && \
      (labs(steps[2]) < MAX_STEP_SIZE)  && \
-     (nomemoverlap(args[0], steps[0] * dimensions[0], args[2], steps[2] * dimensions[0])) && \
-     (nomemoverlap(args[0], steps[0] * dimensions[0], args[1], steps[1] * dimensions[0])))
+     (nomemoverlap(args[0], steps[0], args[2], steps[2], dimensions[0])) && \
+     (nomemoverlap(args[0], steps[0], args[1], steps[1], dimensions[0])))
 
 /*
  * 1) Output should be contiguous, can handle strided input data
@@ -359,7 +359,7 @@ abs_ptrdiff(char *a, char *b)
 #define IS_OUTPUT_BLOCKABLE_UNARY(esizein, esizeout, vsize) \
     ((steps[0] & (esizein-1)) == 0 && \
      steps[1] == (esizeout) && llabs(steps[0]) < MAX_STEP_SIZE && \
-     (nomemoverlap(args[1], steps[1] * dimensions[0], args[0], steps[0] * dimensions[0])))
+     (nomemoverlap(args[1], steps[1], args[0], steps[0], dimensions[0])))
 
 #define IS_BLOCKABLE_REDUCE(esize, vsize) \
     (steps[1] == (esize) && abs_ptrdiff(args[1], args[0]) >= (vsize) && \

numpy/_core/src/umath/loops_utils.h.src

@@ -16,28 +16,31 @@
 #endif
 /*
  * nomemoverlap - returns false if two strided arrays have an overlapping
- * region in memory. ip_size/op_size = size of the arrays which can be negative
- * indicating negative steps.
+ * region in memory.
  */
 NPY_FINLINE npy_bool
-nomemoverlap(char *ip, npy_intp ip_size, char *op, npy_intp op_size)
+nomemoverlap(char *ip, npy_intp ip_step, char *op, npy_intp op_step, npy_intp len)
 {
+    // Calculate inclusive ranges for offsets of items in arrays.
+    // The end pointer points to address of the last item.
+    const npy_intp ip_offset = ip_step * (len - 1);
+    const npy_intp op_offset = op_step * (len - 1);
     char *ip_start, *ip_end, *op_start, *op_end;
-    if (ip_size < 0) {
-        ip_start = ip + ip_size;
+    if (ip_step < 0) {
+        ip_start = ip + ip_offset;
         ip_end = ip;
     }
     else {
         ip_start = ip;
-        ip_end = ip + ip_size;
+        ip_end = ip + ip_offset;
     }
-    if (op_size < 0) {
-        op_start = op + op_size;
+    if (op_step < 0) {
+        op_start = op + op_offset;
         op_end = op;
     }
     else {
         op_start = op;
-        op_end = op + op_size;
+        op_end = op + op_offset;
     }
     return (ip_start == op_start && op_end == ip_end) ||
            (ip_start > op_end) || (op_start > ip_end);
@@ -48,7 +51,7 @@ nomemoverlap(char *ip, npy_intp ip_size, char *op, npy_intp op_size)
 NPY_FINLINE npy_bool
 is_mem_overlap(const void *src, npy_intp src_step, const void *dst, npy_intp dst_step, npy_intp len)
 {
-    return !(nomemoverlap((char*)src, src_step*len, (char*)dst, dst_step*len));
+    return !(nomemoverlap((char*)src, src_step, (char*)dst, dst_step, len));
 }
 
 /*

numpy/_core/tests/test_regression.py

@@ -2622,3 +2622,17 @@ def test_vectorize_fixed_width_string(self):
         f = str.casefold
         res = np.vectorize(f, otypes=[arr.dtype])(arr)
         assert res.dtype == "U30"
+
+    def test_repeated_square_consistency(self):
+        # gh-26940
+        buf = np.array([-5.171866611150749e-07 + 2.5618634555957426e-07j,
+                        0, 0, 0, 0, 0])
+        # Test buffer with regular and reverse strides
+        for in_vec in [buf[:3], buf[:3][::-1]]:
+            expected_res = np.square(in_vec)
+            # Output vector immediately follows input vector
+            # to reproduce off-by-one in nomemoverlap check.
+            for res in [buf[3:], buf[3:][::-1]]:
+                res = buf[3:]
+                np.square(in_vec, out=res)
+                assert_equal(res, expected_res)