data: Fix strided slice with start inside a subdomain

index_glb_to_loc mis-converted a strided global slice when its start (or,
for a negative step, its high end) fell strictly inside the owning
subdomain: the stride phase was computed from the subdomain boundary
instead of the slice start, so the local block picked spurious leading or
trailing elements (e.g. global 2:8:2 gave rank 0 local [0, 2] instead of
[2]). This disagreed with the induced decomposition and made data_gather
on such slices raise a shape mismatch.

Begin the stride phase at the slice start when it lies within the
subdomain, for both the positive (loc_min) and negative (loc_max) cases.

Add a serial regression test for the conversion and exercise the failing
gather slices (2:8:2, 7:0:-2) in the 2D/3D sliced-gather tests.

Author: mloubout

devito/data/decomposition.py

@@ -296,7 +296,14 @@ def index_glb_to_loc(self, *args, rel=True):
                         loc_min = loc_abs_min - base \
                             + np.mod(glb_idx.step - np.mod(base, glb_idx.step),
                                      glb_idx.step)
-                    else:  # glb start is given explicitly
+                    elif glb_idx_min >= loc_abs_min:
+                        # The slice starts within this subdomain, so its first
+                        # selected index is the start itself (the stride phase
+                        # begins here, not at loc_abs_min).
+                        loc_min = glb_idx_min - base
+                    else:  # the slice started in an earlier subdomain
+                        # Advance to the first in-stride index at or after this
+                        # subdomain's start.
                         loc_min = loc_abs_min - base \
                             + np.mod(glb_idx.step - np.mod(base - glb_idx.start,
                                                            glb_idx.step), glb_idx.step)
@@ -319,7 +326,14 @@ def index_glb_to_loc(self, *args, rel=True):
                         loc_max = top - base \
                             + np.mod(glb_idx.step - np.mod(top - glb_max,
                                                            glb_idx.step), glb_idx.step)
-                    else:
+                    elif glb_idx_max <= loc_abs_max:
+                        # The slice's high end (its start) is within this
+                        # subdomain, so the first selected index is the start
+                        # itself (the stride phase begins here, not at loc_abs_max).
+                        loc_max = glb_idx_max - base
+                    else:  # the slice's high end is in a later subdomain
+                        # Step back to the first in-stride index at or below this
+                        # subdomain's end.
                         loc_max = top - base \
                             + np.mod(glb_idx.step - np.mod(top - glb_idx.start,
                                                            glb_idx.step), glb_idx.step)

tests/test_data.py

@@ -451,6 +451,25 @@ def test_glb_to_loc_index_conversions(self):
         assert d.index_glb_to_loc((1, 6), rel=False) == (5, 6)
         assert d.index_glb_to_loc((None, None), rel=False) == (5, 7)
 
+    def test_glb_to_loc_strided_start_in_subdomain(self):
+        """
+        Regression: a strided slice whose start (or, for a negative step, its
+        high end) falls strictly inside the owning subdomain must begin the
+        stride phase at the start, not at the subdomain boundary. Previously the
+        local slice picked spurious leading/trailing elements (e.g. global
+        ``2:8:2`` gave rank 0 local ``[0, 2]`` instead of ``[2]``).
+        """
+        d0 = Decomposition([[0, 1, 2], [3, 4], [5, 6, 7], [8, 9, 10, 11]], 0)
+        d2 = Decomposition([[0, 1, 2], [3, 4], [5, 6, 7], [8, 9, 10, 11]], 2)
+
+        # Positive step, start strictly inside the owning subdomain
+        assert d0.index_glb_to_loc(slice(2, 12, 2)) == slice(2, 3, 2)   # -> [2]
+        assert d2.index_glb_to_loc(slice(6, 12, 2)) == slice(1, 3, 2)   # -> [6]
+
+        # Negative step, high end (the start) strictly inside the subdomain
+        assert d2.index_glb_to_loc(slice(5, 0, -2)) == slice(0, None, -2)   # -> [5]
+        assert d0.index_glb_to_loc(slice(2, None, -2)) == slice(2, None, -2)
+
     def test_glb_to_loc_w_side(self):
         d = Decomposition([[0, 1, 2], [3, 4], [5, 6, 7], [8, 9, 10, 11]], 2)
 
@@ -1551,6 +1570,10 @@ def test_simple_gather(self, rank, mode):
         (None, None, -1),
         (None, None, -2),
         (1, 8, 3),
+        # Strided slices whose start falls inside a subdomain (see
+        # TestDecomposition.test_glb_to_loc_strided_start_in_subdomain)
+        (2, 8, 2),
+        (7, 0, -2),
         ((0, 4), None, (2, 1))])
     def test_sliced_gather_2D(self, start, stop, step, mode):
         """ Test gather for various 2D slices."""
@@ -1585,6 +1608,8 @@ def test_sliced_gather_2D(self, start, stop, step, mode):
         (None, None, -1),
         (None, None, -2),
         (1, 8, 3),
+        (2, 8, 2),
+        (7, 0, -2),
         ((0, 4, 4), None, (2, 1, 1))])
     def test_sliced_gather_3D(self, start, stop, step, mode):
         """ Test gather for various 3D slices."""