mpi: Fix halo=0 SubDomain mask at rank edges

devito/types/dimension.py

@@ -464,6 +464,16 @@ class SpaceDimension(BasicDimension):
 
     is_Space = True
 
+    @cached_property
+    def rank_offsets(self):
+        """
+        The number of points of `root` owned by the MPI ranks before (`side=LEFT`)
+        or after (`side=RIGHT`) this one, that is the offset of this rank from that
+        edge of the Grid. Zero without domain decomposition.
+        """
+        return tuple(RankOffset(name=f"{self.name}_{s}_ofs", side=s, root=self)
+                     for s in (LEFT, RIGHT))
+
 
 class TimeDimension(BasicDimension):
 
@@ -566,6 +576,7 @@ class Thickness(DataSymbol):
     __rkwargs__ = DataSymbol.__rkwargs__ + ('root', 'side', 'local', 'value', 'separated')
 
     def __new__(cls, *args, root=None, side=None, local=False, separated=False, **kwargs):
+        kwargs.setdefault('nonnegative', True)
         newobj = super().__new__(cls, *args, **kwargs)
         newobj._root = root
         newobj._side = side
@@ -625,6 +636,31 @@ def _arg_check(self, args, *_args, **kwargs):
         mpi_raise(error, InvalidArgument, comm=comm)
 
 
+class RankOffset(Thickness):
+    """
+    The number of points of `root` owned by the MPI ranks before (`side=LEFT`)
+    or after (`side=RIGHT`) this one, that is the offset of this rank from that
+    edge of the Grid. Zero without domain decomposition.
+    """
+
+    def _arg_values(self, grid=None, **kwargs):
+        root = self.root
+        if grid is not None:
+            root = next((d for d in grid.dimensions if d.name == root.name), root)
+
+        if grid is None or not grid.is_distributed(root):
+            return {self.name: 0}
+
+        dec = grid.distributor.decomposition[root]
+        if dec.loc_empty:
+            # No iterations run on an empty rank
+            return {self.name: 0}
+        elif self.side == LEFT:
+            return {self.name: dec.loc_abs_min - dec.glb_min}
+        else:
+            return {self.name: dec.glb_max - dec.loc_abs_max}
+
+
 class AbstractSubDimension(DerivedDimension):
 
     """
@@ -670,11 +706,9 @@ def _interval(self):
 
     @memoized_meth
     def _symbolic_thickness(self, **kwargs):
-        kwargs = {'dtype': np.int32, 'is_const': True, 'nonnegative': True,
-                  'separated': self.separated}
-
         names = [f"{self.parent.name}_{s}tkn" for s in ('l', 'r')]
-        return SubDimensionThickness(*[Thickness(name=n, **kwargs) for n in names])
+        return SubDimensionThickness(*[Thickness(name=n, separated=self.separated)
+                                       for n in names])
 
     @cached_property
     def symbolic_min(self):
@@ -800,9 +834,7 @@ def middle(cls, name, parent, thickness_left, thickness_right, local=False,
 
     @memoized_meth
     def _symbolic_thickness(self, thickness=None):
-        kwargs = {'dtype': np.int32, 'is_const': True, 'nonnegative': True,
-                  'root': self.root, 'local': self.local,
-                  'separated': self.separated}
+        kwargs = {'root': self.root, 'local': self.local, 'separated': self.separated}
 
         names = [f"{self.parent.name}_{s}tkn" for s in ('l', 'r')]
         sides = [LEFT, RIGHT]
@@ -824,6 +856,27 @@ def _interval(self):
 
         return sympy.Interval(left, right)
 
+    @cached_property
+    def _glb_interval(self):
+        """
+        The `(min, max)` bounds of this SubDimension in the global Grid, in this
+        rank's local indices. Unlike `_interval`, these are not clipped to this
+        rank, so with domain decomposition they may lie in the halo or beyond.
+        """
+        lofs, rofs = self.root.rank_offsets
+
+        # The global Grid edges
+        left = self.parent.symbolic_min - lofs
+        right = self.parent.symbolic_max + rofs
+
+        # A tuple, not a `sympy.Interval`, which would collapse with literal bounds
+        if self.is_left:
+            return (left, left + self.ltkn.value - 1)
+        elif self.is_right:
+            return (right - self.rtkn.value + 1, right)
+        else:
+            return (left + self.ltkn.value, right - self.rtkn.value)
+
     @property
     def local(self):
         return self._local

devito/types/grid.py

@@ -902,7 +902,8 @@ def _indicator(self, dim, offset):
         """
         1 if the point `offset` points away from the current one along `dim` lies
         in this SubDomain, 0 otherwise, computed without branching. Always 1 if
-        this SubDomain spans all of `dim`.
+        this SubDomain spans all of `dim`. With domain decomposition, the test
+        is against the global bounds, which may lie beyond this rank.
 
         Parameters
         ----------
@@ -915,8 +916,9 @@ def _indicator(self, dim, offset):
         if not maybe_subdim.is_Sub:
             return sympy.S.One
         point = dim + offset
-        return (Max(0, Min(1, point - maybe_subdim.symbolic_min + 1)) *
-                Max(0, Min(1, maybe_subdim.symbolic_max - point + 1)))
+        lower, upper = maybe_subdim._glb_interval
+        return (Max(0, Min(1, point - lower + 1)) *
+                Max(0, Min(1, upper - point + 1)))
 
     def _restrict(self, expr, exclude=()):
         """

tests/test_derivatives.py

@@ -1895,7 +1895,7 @@ def test_default_unchanged(self):
         """
         grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64)
         subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid)
-        kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64}
+        kwargs = {'grid': grid, 'space_order': 4}
         g_bar = Function(name='g_bar', **kwargs)
         ref = Function(name='ref', **kwargs)
         out_bar = Function(name='out_bar', **kwargs)
@@ -1927,7 +1927,7 @@ def test_transpose_keeps_halo(self):
     def test_no_subdomain(self):
         """Without a SubDomain, halo=0 has no effect."""
         grid = Grid(shape=(16,), extent=(15.,), dtype=np.float64)
-        kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64}
+        kwargs = {'grid': grid, 'space_order': 4}
         f = Function(name='f', **kwargs)
         g = Function(name='g', **kwargs)
         ref = Function(name='ref', **kwargs)
@@ -1966,7 +1966,7 @@ def test_terms(self):
         """
         grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64)
         subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid)
-        kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64}
+        kwargs = {'grid': grid, 'space_order': 4}
         f, g, c, out, ref = (Function(name=n, **kwargs)
                              for n in ('f', 'g', 'c', 'out', 'ref'))
         rng = np.random.default_rng(8)
@@ -1987,7 +1987,7 @@ def test_sum_factor(self):
         """
         grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64)
         subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid)
-        kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64}
+        kwargs = {'grid': grid, 'space_order': 4}
         f, g, c, d, out = (Function(name=n, **kwargs)
                            for n in ('f', 'g', 'c', 'd', 'out'))
         rng = np.random.default_rng(9)
@@ -2008,7 +2008,7 @@ def test_vector(self):
         """
         grid = Grid(shape=(12, 14), extent=(11., 13.), dtype=np.float64)
         subdomain = Region('s', {'x': ('left', 5)}, grid=grid)
-        kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64}
+        kwargs = {'grid': grid, 'space_order': 4}
         f = Function(name='f', **kwargs)
         c = VectorFunction(name='c', staggered=(None, None), **kwargs)
         w = VectorFunction(name='w', staggered=(None, None), **kwargs)
@@ -2034,7 +2034,7 @@ def test_vector_staggered(self):
         grid = Grid(shape=(14, 16), extent=(13., 15.), dtype=np.float64)
         x, y = grid.dimensions
         subdomain = Region('s', {'x': ('left', 5), 'y': ('middle', 4, 5)}, grid=grid)
-        kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64}
+        kwargs = {'grid': grid, 'space_order': 4}
         p = Function(name='p', **kwargs)
         p_bar = Function(name='p_bar', **kwargs)
         v = VectorFunction(name='v', **kwargs)
@@ -2057,7 +2057,7 @@ def test_unrestricted_dimension(self):
         """halo=0 along a Dimension the SubDomain spans has no effect."""
         grid = Grid(shape=(12, 14), extent=(11., 13.), dtype=np.float64)
         subdomain = Region('s', {'x': ('left', 5)}, grid=grid)
-        kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64}
+        kwargs = {'grid': grid, 'space_order': 4}
         f = Function(name='f', **kwargs)
         g = Function(name='g', **kwargs)
         ref = Function(name='ref', **kwargs)

tests/test_mpi.py

@@ -1,3 +1,4 @@
+import pickle
 from functools import cached_property
 
 import numpy as np
@@ -25,7 +26,7 @@
 from devito.mpi.routines import ComputeCall, HaloUpdateCall, HaloUpdateList, MPICall
 from devito.tools import Bunch
 from devito.types import Bundle
-from devito.types.dimension import ModuloDimension
+from devito.types.dimension import ModuloDimension, RankOffset
 from examples.seismic.acoustic import acoustic_setup
 
 
@@ -3351,6 +3352,64 @@ def test_subdim_thickness_spanning_rank(self, mode):
             assert np.all(glb[:, 3:-3] == 0.)
             assert np.all(glb[:, -3:] == 2.)
 
+    @pytest.mark.parallel(mode=4)
+    @pytest.mark.parametrize('expand', [True, False])
+    @pytest.mark.parametrize('region', [('middle', 6, 6), ('middle', 16, 16),
+                                        ('middle', 15, 17), ('left', 17),
+                                        ('right', 18)])
+    def test_halo_zero_subdomain(self, region, expand, mode):
+        """
+        Derivatives with `halo=0` treat their argument as zero outside the global
+        SubDomain, not outside its part on each rank. The rank edges lie at 16,
+        32 and 48, within the stencil radius of some of the SubDomain edges.
+        """
+        class Region(SubDomain):
+            name = 'region'
+
+            def define(self, dimensions):
+                x, = dimensions
+                return {x: region}
+
+        serial = Grid(shape=(64,), comm=MPI.COMM_SELF)
+        grid = Grid(shape=(64,))
+        subdomain = Region(grid=serial)
+
+        def functions(grid):
+            kwargs = {'grid': grid, 'space_order': 4}
+            return [Function(name=n, **kwargs) for n in ('f', 'o', 'g')]
+
+        f0, o0, g0 = functions(serial)
+        f1, o1, g1 = functions(grid)
+        f0.data[:] = f1.data[:] = np.random.default_rng(0).normal(size=grid.shape)
+
+        # Compiled for the serial Grid, applied to both
+        op = Operator([Eq(o0, f0.dx(halo=0), subdomain=subdomain),
+                       Inc(g0, f0.dx(halo=0).T, subdomain=subdomain)],
+                      opt=('advanced', {'expand': expand}))
+        op.apply()
+        op.apply(f=f1, o=o1, g=g1)
+
+        glb = grid.distributor.glb_slices[grid.dimensions[0]]
+        assert np.allclose(o1.data, o0.data[glb], rtol=1e-12)
+        assert np.allclose(g1.data, g0.data[glb], rtol=1e-12)
+
+    @pytest.mark.parallel(mode=4)
+    def test_rank_offset_empty(self, mode):
+        """
+        The offsets of a rank from the Grid edges, including an empty rank, also
+        once unpickled.
+        """
+        grid = Grid(shape=(3,))
+        x = grid.dimensions[0]
+
+        lofs = pickle.loads(pickle.dumps(RankOffset(name='x_lofs', side=LEFT, root=x)))
+        rofs = pickle.loads(pickle.dumps(RankOffset(name='x_rofs', side=RIGHT, root=x)))
+        values = (lofs._arg_values(grid=grid)['x_lofs'],
+                  rofs._arg_values(grid=grid)['x_rofs'])
+
+        # Rank 3 is empty
+        assert values == [(0, 2), (1, 1), (2, 0), (0, 0)][grid.distributor.myrank]
+
     @pytest.mark.parallel(mode=[(4, 'full')])
     def test_custom_subdomain(self, mode):
         """