: selection: bugfix reifying strided views

Summary:
**fix two bugs**

thanks go to pzhan9 for reporting these issues! (see D78298712)

(1) one in `reify_view` where views with strides were reified incorrectly. the old code assumed views could always be expressed as contiguous `range(start..start + len)` expressions, even when the view's stride differed from the base. this produced incorrect selections when the view was strided relative to the base layout.

the fix is to recognize when a view is layout-aligned; that is, when each of its strides is an integer multiple of the corresponding base stride, including the unitary case; and to reify such views using `Range(start, Some(end), step)` expressions that preserve both shape and layout. previously, only the unitary case (`step` = 1) was handled. this change extends support to non-unitary aligned views, such as `step` = 2, by correctly computing the step factor and corresponding end coordinate.

if any dimension is not layout-aligned — that is, if `view_stride` % `base_stride` ≠ 0 — we conservatively fall back to enumerating all selected coordinates explicitly.

(2) also fixes a bug in `Shape::select()` where the length of a strided range was computed using truncating division. this now correctly uses ceiling division to ensure all selected indices are included.

Differential Revision: D78315005

Author: shayne-fletcher

ndslice/src/selection.rs

@@ -1021,16 +1021,22 @@ impl ReifyView for Slice {
     /// matches all coordinates in the given `view`, expressed in the
     /// coordinate system of the provided `base` slice (`self`).
     ///
-    /// The resulting expression uses nested `range(start..end, ...)`
-    /// combinators to represent the rectangular region selected by
-    /// the view within the base slice.
+    /// The result is a nested sequence of `range(start..end, step)`
+    /// combinators that match the rectangular region covered by `view`
+    /// in base coordinates. This preserves geometry and layout when
+    /// `view` is *layout-aligned* — that is, each of its strides is
+    /// a multiple of the corresponding base stride.
     ///
-    /// Returns [`dsl::false_()`] for empty views.
+    /// If any dimension is not layout-aligned, the view is reified
+    /// by explicitly enumerating its coordinates.
+    ///
+    /// Returns [`dsl::false_()`] if the view is empty.
     ///
     /// # Errors
     ///
     /// Returns an error if:
-    /// - The view lies outside the bounds of the base slice
+    /// - The base is not contiguous and row-major
+    /// - The view lies outside the bounds of the base
     ///
     /// # Example
     ///
@@ -1043,6 +1049,11 @@ impl ReifyView for Slice {
     /// let selection = base.reify_view(view).unwrap();
     /// ```
     fn reify_view(&self, view: &Slice) -> Result<Selection, SliceError> {
+        // Precondition: the base is contiguous and row major.
+        if !self.is_contiguous() {
+            return Err(SliceError::NonContiguous);
+        }
+
         if view.is_empty() {
             return Ok(dsl::false_());
         }
@@ -1055,8 +1066,21 @@ impl ReifyView for Slice {
 
         let origin = self.coordinates(view.offset())?;
         let mut acc = dsl::true_();
-        for (&start, &len) in origin.iter().zip(view.sizes()).rev() {
-            acc = dsl::range(start..start + len, acc);
+        for ((&start, &len), (&view_stride, &base_stride)) in origin
+            .iter()
+            .zip(view.sizes())
+            .zip(view.strides().iter().zip(self.strides()))
+            .rev()
+        {
+            if view_stride % base_stride == 0 {
+                // Layout-aligned with base.
+                let step = view_stride / base_stride;
+                let end = start + step * len;
+                acc = dsl::range(crate::shape::Range(start, Some(end), step), acc);
+            } else {
+                // Irregular layout; fallback to explicit enumeration.
+                return Selection::of_ranks(self, &view.iter().collect::<BTreeSet<_>>());
+            }
         }
 
         Ok(acc)
@@ -1338,6 +1362,7 @@ mod tests {
     use super::Selection;
     use super::dsl::*;
     use super::is_equivalent_true;
+    use crate::Range;
     use crate::Slice;
     use crate::assert_structurally_eq;
     use crate::select;
@@ -2134,6 +2159,52 @@ mod tests {
         );
     }
 
+    #[test]
+    #[allow(clippy::identity_op)]
+    fn test_reify_view_1d_with_stride() {
+        let shape = shape!(x = 7); // 1D shape with 7 elements
+        let selected = shape.select("x", Range(0, None, 2)).unwrap();
+        let view = selected.slice();
+        assert_eq!(view, &Slice::new(0, vec![4], vec![1 * 2]).unwrap());
+
+        let base = shape.slice();
+        let selection = base.reify_view(view).unwrap();
+        let expected = range(Range(0, Some(8), 2), true_());
+        assert_structurally_eq!(&selection, expected);
+
+        let flat: Vec<_> = selection.eval(&EvalOpts::strict(), base).unwrap().collect();
+        assert_eq!(flat, vec![0, 2, 4, 6]);
+    }
+
+    #[test]
+    #[allow(clippy::identity_op)]
+    fn test_reify_view_2d_with_stride() {
+        // 4 x 4: x = 4, y = 4.
+        let base = shape!(x = 4, y = 4);
+        // Step 1: select odd rows (x = 1..4 step 2)
+        let shape = base.select("x", Range(1, Some(4), 2)).unwrap();
+        // Step 2: then select odd columns (y = 1..4 step 2)
+        let shape = shape.select("y", Range(1, Some(4), 2)).unwrap();
+        let view = shape.slice();
+        assert_eq!(
+            view,
+            &Slice::new(5, vec![2, 2], vec![4 * 2, 1 * 2]).unwrap()
+        );
+
+        let base = base.slice();
+        let selection = base.reify_view(view).unwrap();
+        // We use `end = start + step * len` to reify the selection.
+        // Note: This may yield `end > original_end` (e.g., 5 instead of 4)
+        // when the selection length was computed via ceiling division.
+        // This is safe: the resulting range will still select the correct
+        // indices (e.g., 1 and 3 for Range(1, Some(5), 2)).
+        let expected = range(Range(1, Some(5), 2), range(Range(1, Some(5), 2), true_()));
+        assert_structurally_eq!(&selection, expected);
+
+        let flat: Vec<_> = selection.eval(&EvalOpts::strict(), base).unwrap().collect();
+        assert_eq!(flat, vec![5, 7, 13, 15]);
+    }
+
     #[test]
     fn test_reify_view_selects_column_across_rows() {
         let shape = shape!(host = 2, gpu = 4); // shape [2, 4]

ndslice/src/shape.rs

@@ -133,7 +133,10 @@ impl Shape {
         }
 
         offset += begin * strides[dim];
-        sizes[dim] = (end - begin) / stride;
+        // The # of elems in `begin..end` with step `stride`. This is
+        // ⌈(end - begin) / stride⌉ — the number of stride steps that
+        // fit in the half-open interval.
+        sizes[dim] = (end - begin).div_ceil(stride);
         strides[dim] *= stride;
 
         Ok(Self {
@@ -626,4 +629,18 @@ mod tests {
                 .is_err()
         );
     }
+
+    #[test]
+    fn test_shape_select_stride_rounding() {
+        let shape = shape!(x = 10);
+        // Select x = 0..10 step 3 → expect indices [0, 3, 6, 9]
+        let sub = shape.select("x", Range(0, Some(10), 3)).unwrap();
+        let slice = sub.slice();
+        // 10 / 3 = 3.33..., so ceil(10 / 3) = 4
+        assert_eq!(
+            slice,
+            &Slice::new(0, vec![4], vec![3]).unwrap(),
+            "Expected offset 0, size 4, stride 3"
+        );
+    }
 }

ndslice/src/slice.rs

@@ -35,6 +35,9 @@ pub enum SliceError {
 
     #[error("incompatible view: {reason}")]
     IncompatibleView { reason: String },
+
+    #[error("noncontiguous shape")]
+    NonContiguous,
 }
 
 /// Slice is a compact representation of indices into the flat