MPI: clamp localised SubDimension thickness to the rank's local size

devito/data/decomposition.py

@@ -146,9 +146,13 @@ def index_glb_to_loc(self, *args, rel=True):
               relative local index if ``I`` belongs to the local subdomain,
               ``None`` otherwise.
             * int, DataSide. Given ``O`` and ``S``, respectively a global
-              offset and a side, return the relative local offset. This
-              can be ``None`` if the local subdomain doesn't intersect with the
-              region defined by the given global offset.
+              offset and a side, return the corresponding relative local index
+              (``rel=True``, the default). If the offset covers the local
+              subdomain entirely, this is the subdomain's last index -- never a
+              count. This can be ``None`` if the local subdomain doesn't
+              intersect with the region defined by the given global offset.
+              The ``rel=False`` behaviour of this two-argument form is not
+              well-defined and no caller relies on it.
             * (int, int).  Given global ``(min, max)``, return ``(min', max')``
               representing the corresponding relative local min/max. If the
               input doesn't intersect with the local subdomain, then ``min'``
@@ -353,20 +357,32 @@ def index_glb_to_loc(self, *args, rel=True):
             if self.loc_empty:
                 return None
             abs_ofs, side = args
+            # Both branches below return an *index* in the same frame as `base`
+            # -- relative when `rel`, absolute otherwise. In particular the
+            # saturating case, where the offset covers this subdomain entirely,
+            # returns the last index of the subdomain and NOT a count: callers
+            # convert to a count themselves.
+            #
+            # The old returns, `top + 1` and `glb_max - base + 1`, were not
+            # indices at all: they sat one past the end, outside the valid local
+            # index range, on every subdomain including the first. On the LEFT
+            # they also mixed frames, `top` being absolute while `rel_ofs` is
+            # relative to `base`, so the overshoot grew with the subdomain's
+            # offset rather than staying at one.
             if side == LEFT:
                 rel_ofs = glb_min + abs_ofs - base
                 if abs_ofs >= base and abs_ofs <= top:
                     return rel_ofs
                 elif abs_ofs > top:
-                    return top + 1
+                    return top - base
                 else:
                     return None
             else:
                 rel_ofs = abs_ofs - (glb_max - top)
                 if abs_ofs >= glb_max - top and abs_ofs <= glb_max - base:
                     return rel_ofs
                 elif abs_ofs > glb_max - base:
-                    return glb_max - base + 1
+                    return top - base
                 else:
                     return None
         else:

devito/types/dimension.py

@@ -605,8 +605,17 @@ def _arg_values(self, grid=None, **kwargs):
                 else:
                     tkn = 0
             else:
-                # Dimension is of type `middle`
-                tkn = grid.distributor.glb_to_loc(self.root, rtkn, self.side) or 0
+                # Dimension is of type `middle`. Convert as an INDEX -- hence
+                # `rtkn-1` -- and add one to get a count, the same convention as
+                # the left/right branch above. Passing `rtkn` and using the result
+                # directly as a count only worked while the saturating branch of
+                # `index_glb_to_loc` returned a count; it returns an index now, so
+                # both branches must convert the same way.
+                if rtkn:
+                    tkn = grid.distributor.glb_to_loc(self.root, rtkn-1, self.side)
+                    tkn = tkn+1 if tkn is not None else 0
+                else:
+                    tkn = 0
         else:
             tkn = rtkn or 0
 

tests/test_data.py

@@ -489,6 +489,41 @@ def test_glb_to_loc_w_side(self):
         assert d.index_glb_to_loc(2, LEFT) is None
         assert d.index_glb_to_loc(3, RIGHT) is None
 
+        # Saturating case: the offset covers this subdomain entirely. The result
+        # is still an index in the same frame as the in-range results above --
+        # here the subdomain's last relative index, 2 -- and not a count.
+        assert d.index_glb_to_loc(8, LEFT) == 2
+        assert d.index_glb_to_loc(11, LEFT) == 2
+        assert d.index_glb_to_loc(7, RIGHT) == 2
+        assert d.index_glb_to_loc(11, RIGHT) == 2
+
+    def test_glb_to_loc_w_side_is_always_an_index(self):
+        """
+        `index_glb_to_loc(offset, side)` must return an index, never a count, on
+        every subdomain and for both sides.
+
+        The saturating branch used to return `loc_abs_max + 1` (LEFT) and
+        `glb_max - loc_abs_min + 1` (RIGHT). Neither is an index: both sit one
+        past the end of the valid local range, on every subdomain including the
+        first. The LEFT form additionally mixes frames -- `loc_abs_max` is
+        absolute while the in-range result is relative to `loc_abs_min` -- so
+        its overshoot grows with the subdomain's offset instead of staying at
+        one. It reached `Thickness._arg_values`, which adds 1 to turn the index
+        into a count, as a boundary loop running past the end of the rank's data.
+        """
+        parts = [list(range(i, i + 16)) for i in range(0, 64, 16)]
+        for r in range(len(parts)):
+            d = Decomposition(parts, r)
+            extent = d.loc_abs_max - d.loc_abs_min + 1
+            for side in (LEFT, RIGHT):
+                for ofs in range(64):
+                    v = d.index_glb_to_loc(ofs, side)
+                    if v is None:
+                        continue
+                    assert 0 <= v < extent, \
+                        f"rank {r} {side} offset {ofs}: {v} not an index in " \
+                        f"[0, {extent})"
+
     def test_loc_to_glb_index_conversions(self):
         d = Decomposition([[0, 1, 2], [3, 4], [5, 6, 7], [8, 9, 10, 11]], 2)
 

tests/test_mpi.py

@@ -1235,6 +1235,57 @@ def test_avoid_haloupdate_with_local_subdims(self, mode):
         calls = FindNodes(Call).visit(op)
         assert len(calls) == 1
 
+    @pytest.mark.parallel(mode=4)
+    def test_subdimension_thickness_localisation(self, mode):
+        # A rank lying entirely inside the boundary layer used to be handed a
+        # localised thickness of `local_extent + 1`, because
+        # `index_glb_to_loc(offset, side)` saturates to a *count* while
+        # `_arg_values` unconditionally added 1 to turn an *index* into a count.
+        # The generated boundary loop then ran one point past the rank's extent.
+        grid = Grid(shape=(64, 64, 64), topology=(4, 1, 1))
+        x = grid.dimensions[0]
+        decomp = grid.distributor.decomposition[x]
+        local_extent = decomp.loc_abs_max - decomp.loc_abs_min + 1
+
+        # thickness 20 > local extent 16, so ranks 0 and 3 sit wholly inside
+        for name, ctor in [('xl', SubDimension.left), ('xr', SubDimension.right)]:
+            sd = ctor(name=name, parent=x, thickness=20)
+            for tkn in sd.thickness:
+                for v in tkn._arg_values(grid=grid).values():
+                    assert v <= local_extent
+
+    @pytest.mark.parallel(mode=4)
+    def test_subdimension_middle_thickness_localisation(self, mode):
+        # Companion to the test above, and the regression it guards against.
+        # `middle` localises through the same saturating conversion, but consumed
+        # the result directly as a count while `left`/`right` converted an index
+        # and added one. Once the saturating branch was made to return an index,
+        # the two conventions disagreed and a wholly-covered rank lost a point
+        # off each end.
+        grid = Grid(shape=(64, 64, 64), topology=(4, 1, 1))
+        x = grid.dimensions[0]
+        decomp = grid.distributor.decomposition[x]
+        local_extent = decomp.loc_abs_max - decomp.loc_abs_min + 1
+
+        # 20 either side of a 64-point dimension split 4 ways. Rank 0 lies wholly
+        # inside the left band and rank 3 wholly inside the right, so each takes
+        # its full extent; ranks 1 and 2 keep the 4-point remainder on one side
+        # and nothing on the other. The regression showed up as 15 for the
+        # wholly-covered ranks.
+        sd = SubDimension.middle(name='xm', parent=x, thickness_left=20,
+                                 thickness_right=20)
+        got = {}
+        for tkn in sd.thickness:
+            got.update(tkn._arg_values(grid=grid))
+        left, right = (got[t.name] for t in sd.thickness)
+
+        expected = {0: (local_extent, 0), 1: (4, 0),
+                    2: (0, 4), 3: (0, local_extent)}
+        assert (left, right) == expected[grid.distributor.myrank], \
+            f"rank {grid.distributor.myrank}: localised middle thicknesses " \
+            f"{(left, right)}, expected {expected[grid.distributor.myrank]}"
+        assert left + right <= local_extent
+
     @pytest.mark.parallel(mode=1)
     def test_avoid_haloupdate_with_local_customdim(self, mode):
         grid = Grid(shape=(10, 10))