WIP: Use the GCD of the strides to check if two array views could alias.

Author: adamreichold

Cargo.toml

@@ -17,6 +17,7 @@ license = "BSD-2-Clause"
 [dependencies]
 libc = "0.2"
 num-complex = ">= 0.2, < 0.5"
+num-integer = "0.1"
 num-traits = "0.2"
 ndarray = ">= 0.13, < 0.16"
 pyo3 = { version = "0.16", default-features = false, features = ["macros"] }

src/array.rs

@@ -340,9 +340,9 @@ impl<T, D> PyArray<T, D> {
     }
 
     /// Returns the pointer to the first element of the inner array.
-    pub(crate) unsafe fn data(&self) -> *mut T {
+    pub(crate) fn data(&self) -> *mut T {
         let ptr = self.as_array_ptr();
-        (*ptr).data as *mut _
+        unsafe { (*ptr).data as *mut _ }
     }
 }
 
@@ -381,7 +381,7 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
         let strides = self.strides();
 
         let mut new_strides = D::zeros(strides.len());
-        let mut data_ptr = unsafe { self.data() };
+        let mut data_ptr = self.data();
         let mut inverted_axes = InvertedAxes::new(strides.len());
 
         for i in 0..strides.len() {

src/borrow.rs

@@ -59,13 +59,9 @@
 //! });
 //! ```
 //!
-//! The second example shows that while non-overlapping views are supported,
-//! interleaved views which do not touch are currently not supported
-//! due to over-approximating which borrows are in conflict.
+//! The second example shows that non-overlapping and interleaved views are also supported.
 //!
 //! ```rust
-//! # use std::panic::{catch_unwind, AssertUnwindSafe};
-//! #
 //! use numpy::PyArray1;
 //! use pyo3::{types::IntoPyDict, Python};
 //!
@@ -78,16 +74,15 @@
 //!     let view3 = py.eval("array[::2]", None, Some(locals)).unwrap().downcast::<PyArray1<f64>>().unwrap();
 //!     let view4 = py.eval("array[1::2]", None, Some(locals)).unwrap().downcast::<PyArray1<f64>>().unwrap();
 //!
-//!     let _view1 = view1.readwrite();
-//!     let _view2 = view2.readwrite();
+//!     {
+//!         let _view1 = view1.readwrite();
+//!         let _view2 = view2.readwrite();
+//!     }
 //!
-//!     // Will fail at runtime even though `view3` and `view4`
-//!     // interleave as they are based on the same array.
-//!     let res = catch_unwind(AssertUnwindSafe(|| {
+//!     {
 //!         let _view3 = view3.readwrite();
 //!         let _view4 = view4.readwrite();
-//!     }));
-//!     assert!(res.is_err());
+//!     }
 //! });
 //! ```
 //!
@@ -125,6 +120,8 @@
 //!
 //! # Limitations
 //!
+//! TODO: We only leave the case of aliasing, but only out of bounds. Can this actually happen for array views?
+//!
 //! Note that the current implementation of this is an over-approximation: It will consider overlapping borrows
 //! potentially conflicting if the initial arrays have the same object at the end of their [base object chain][base].
 //! For example, creating two views of the same underlying array by slicing can yield potentially conflicting borrows
@@ -143,6 +140,7 @@ use std::collections::hash_map::{Entry, HashMap};
 use std::ops::{Deref, Range};
 
 use ndarray::{ArrayView, ArrayViewMut, Dimension, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
+use num_integer::gcd;
 use pyo3::{FromPyObject, PyAny, PyResult};
 
 use crate::array::PyArray;
@@ -155,9 +153,28 @@ use crate::npyffi::{self, PyArrayObject, NPY_ARRAY_WRITEABLE};
 #[derive(PartialEq, Eq, Hash)]
 struct BorrowKey {
     range: Range<usize>,
+    data_ptr: usize,
+    gcd_strides: isize,
 }
 
 impl BorrowKey {
+    fn from_array<T, D>(array: &PyArray<T, D>) -> Self
+    where
+        T: Element,
+        D: Dimension,
+    {
+        let range = data_range(array);
+
+        let data_ptr = array.data() as usize;
+        let gcd_strides = reduce(array.strides().iter().copied(), gcd).unwrap_or(1);
+
+        Self {
+            range,
+            data_ptr,
+            gcd_strides,
+        }
+    }
+
     fn conflicts(&self, other: &Self) -> bool {
         debug_assert!(self.range.start <= self.range.end);
         debug_assert!(other.range.start <= other.range.end);
@@ -166,6 +183,20 @@ impl BorrowKey {
             return false;
         }
 
+        // The Diophantine equation which describes whether any integers can combine the data pointers and strides of the two arrays s.t.
+        // they yield the same element has a solution if and only if the GCD of all strides divides the difference of the data pointers.
+        //
+        // That solution could be out of bounds which mean that this is still an approximation,
+        // but it seems sufficient to handle typical cases like the color channels of an image.
+        //
+        // https://users.rust-lang.org/t/math-for-borrow-checking-numpy-arrays/73303
+        let ptr_diff = abs_diff(self.data_ptr, other.data_ptr) as isize;
+        let gcd_strides = gcd(self.gcd_strides, other.gcd_strides);
+
+        if ptr_diff % gcd_strides != 0 {
+            return false;
+        }
+
         true
     }
 }
@@ -192,10 +223,7 @@ impl BorrowFlags {
         D: Dimension,
     {
         let address = base_address(array);
-
-        let key = BorrowKey {
-            range: data_range(array),
-        };
+        let key = BorrowKey::from_array(array);
 
         // SAFETY: Access to `&PyArray<T, D>` implies holding the GIL
         // and we are not calling into user code which might re-enter this function.
@@ -242,10 +270,7 @@ impl BorrowFlags {
         D: Dimension,
     {
         let address = base_address(array);
-
-        let key = BorrowKey {
-            range: data_range(array),
-        };
+        let key = BorrowKey::from_array(array);
 
         // SAFETY: Access to `&PyArray<T, D>` implies holding the GIL
         // and we are not calling into user code which might re-enter this function.
@@ -272,10 +297,7 @@ impl BorrowFlags {
         D: Dimension,
     {
         let address = base_address(array);
-
-        let key = BorrowKey {
-            range: data_range(array),
-        };
+        let key = BorrowKey::from_array(array);
 
         // SAFETY: Access to `&PyArray<T, D>` implies holding the GIL
         // and we are not calling into user code which might re-enter this function.
@@ -320,10 +342,7 @@ impl BorrowFlags {
         D: Dimension,
     {
         let address = base_address(array);
-
-        let key = BorrowKey {
-            range: data_range(array),
-        };
+        let key = BorrowKey::from_array(array);
 
         // SAFETY: Access to `&PyArray<T, D>` implies holding the GIL
         // and we are not calling into user code which might re-enter this function.
@@ -628,6 +647,25 @@ where
     Range { start, end }
 }
 
+// FIXME(adamreichold): Use `usize::abs_diff` from std when that becomes stable.
+fn abs_diff(lhs: usize, rhs: usize) -> usize {
+    if lhs >= rhs {
+        lhs - rhs
+    } else {
+        rhs - lhs
+    }
+}
+
+// FIXME(adamreichold): Use `Iterator::reduce` from std when our MSRV reaches 1.51.
+fn reduce<I, F>(mut iter: I, f: F) -> Option<I::Item>
+where
+    I: Iterator,
+    F: FnMut(I::Item, I::Item) -> I::Item,
+{
+    let first = iter.next()?;
+    Some(iter.fold(first, f))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -650,7 +688,7 @@ mod tests {
             assert_eq!(base_address, array as *const _ as usize);
 
             let data_range = data_range(array);
-            assert_eq!(data_range.start, unsafe { array.data() } as usize);
+            assert_eq!(data_range.start, array.data() as usize);
             assert_eq!(data_range.end, unsafe { array.data().add(15) } as usize);
         });
     }
@@ -668,7 +706,7 @@ mod tests {
             assert_eq!(base_address, base as usize);
 
             let data_range = data_range(array);
-            assert_eq!(data_range.start, unsafe { array.data() } as usize);
+            assert_eq!(data_range.start, array.data() as usize);
             assert_eq!(data_range.end, unsafe { array.data().add(15) } as usize);
         });
     }
@@ -694,7 +732,7 @@ mod tests {
             assert_eq!(base_address, base as usize);
 
             let data_range = data_range(view);
-            assert_eq!(data_range.start, unsafe { view.data() } as usize);
+            assert_eq!(data_range.start, view.data() as usize);
             assert_eq!(data_range.end, unsafe { view.data().add(12) } as usize);
         });
     }
@@ -724,7 +762,7 @@ mod tests {
             assert_eq!(base_address, base as usize);
 
             let data_range = data_range(view);
-            assert_eq!(data_range.start, unsafe { view.data() } as usize);
+            assert_eq!(data_range.start, view.data() as usize);
             assert_eq!(data_range.end, unsafe { view.data().add(12) } as usize);
         });
     }
@@ -763,7 +801,7 @@ mod tests {
             assert_eq!(base_address, base as usize);
 
             let data_range = data_range(view2);
-            assert_eq!(data_range.start, unsafe { view2.data() } as usize);
+            assert_eq!(data_range.start, view2.data() as usize);
             assert_eq!(data_range.end, unsafe { view2.data().add(6) } as usize);
         });
     }
@@ -806,7 +844,7 @@ mod tests {
             assert_eq!(base_address, base as usize);
 
             let data_range = data_range(view2);
-            assert_eq!(data_range.start, unsafe { view2.data() } as usize);
+            assert_eq!(data_range.start, view2.data() as usize);
             assert_eq!(data_range.end, unsafe { view2.data().add(6) } as usize);
         });
     }

tests/borrow.rs

@@ -212,28 +212,39 @@ fn conflict_due_reborrow_of_overlapping_views() {
 }
 
 #[test]
-#[should_panic(expected = "AlreadyBorrowed")]
-fn interleaved_views_conflict() {
+fn interleaved_views_do_not_conflict() {
     Python::with_gil(|py| {
-        let array = PyArray::<f64, _>::zeros(py, (1, 2, 3), false);
+        let array = PyArray::<f64, _>::zeros(py, (23, 42, 3), false);
         let locals = [("array", array)].into_py_dict(py);
 
         let view1 = py
-            .eval("array[:,:,1]", None, Some(locals))
+            .eval("array[:,:,0]", None, Some(locals))
             .unwrap()
             .downcast::<PyArray2<f64>>()
             .unwrap();
-        assert_eq!(view1.shape(), [1, 2]);
+        assert_eq!(view1.shape(), [23, 42]);
 
         let view2 = py
+            .eval("array[:,:,1]", None, Some(locals))
+            .unwrap()
+            .downcast::<PyArray2<f64>>()
+            .unwrap();
+        assert_eq!(view2.shape(), [23, 42]);
+
+        let view3 = py
             .eval("array[:,:,2]", None, Some(locals))
             .unwrap()
             .downcast::<PyArray2<f64>>()
             .unwrap();
-        assert_eq!(view2.shape(), [1, 2]);
+        assert_eq!(view2.shape(), [23, 42]);
 
-        let _exclusive1 = view1.readwrite();
-        let _exclusive2 = view2.readwrite();
+        let exclusive1 = view1.readwrite();
+        let exclusive2 = view2.readwrite();
+        let exclusive3 = view3.readwrite();
+
+        assert_eq!(exclusive3.len(), 23 * 42);
+        assert_eq!(exclusive2.len(), 23 * 42);
+        assert_eq!(exclusive1.len(), 23 * 42);
     });
 }