mpi: Honor MultiSubDomain thicknesses in overlap modes

devito/mpi/halo_scheme.py

@@ -142,7 +142,7 @@ def __init__(self, exprs, ispace):
         self._honored = {}
         dims = set().union(*[d._defines for d in ispace.dimensions
                              if d._defines & self.dimensions])
-        subdims = [d for d in dims if d.is_Sub and not d.local]
+        subdims = [d for d in dims if d.is_AbstractSub and not d.local]
         for i in subdims:
             ltk, rtk = i.tkns
             self._honored[i.root] = frozenset([(ltk, rtk)])

devito/mpi/routines.py

@@ -16,8 +16,8 @@
 )
 from devito.mpi import MPI
 from devito.symbolics import (
-    Byref, CondNe, FieldFromComposite, FieldFromPointer, IndexedPointer, Macro, cast,
-    subs_op_args
+    Byref, CondNe, DefFunction, FieldFromComposite, FieldFromPointer, IndexedPointer,
+    Macro, cast, subs_op_args
 )
 from devito.tools import (
     as_mapper, dtype_len, dtype_to_mpitype, flatten, generator, infer_datasize, is_integer
@@ -791,7 +791,10 @@ class Overlap2HaloExchangeBuilder(OverlapHaloExchangeBuilder):
     """
 
     def _make_region(self, hs, key):
-        return MPIRegion('reg', key, hs.arguments, hs.omapper.owned)
+        # Only the OWNED split of the root Dimensions, as the SubDimension
+        # thicknesses are applied in the remainder
+        owned = hs.halo_scheme._rebuild(honored={}).omapper.owned
+        return MPIRegion('reg', key, hs.dimensions, owned)
 
     def _make_msg(self, f, hse, key):
         # Only retain the halos required by the Diag scheme
@@ -912,12 +915,17 @@ def _make_remainder(self, hs, key, callcompute, region):
         region_i = IndexedPointer(region, dim)
 
         dynamic_args_mapper = {}
-        for i in hs.arguments:
-            if i.is_Dimension:
-                dynamic_args_mapper[i] = (FieldFromComposite(i.min_name, region_i),
-                                          FieldFromComposite(i.max_name, region_i))
-            else:
-                dynamic_args_mapper[i] = (FieldFromComposite(i.name, region_i),)
+        for d in hs.dimensions:
+            m = FieldFromComposite(d.min_name, region_i)
+            M = FieldFromComposite(d.max_name, region_i)
+            dynamic_args_mapper[d] = (m, M)
+
+            # Intersect the region with the SubDimension bounds, which may only
+            # be known within the Operator (e.g., MultiSubDimension thicknesses)
+            for ltkn, rtkn in hs.halo_scheme.honored.get(d, []):
+                ltkn1 = DefFunction('MAX', [0, d.symbolic_min + ltkn - m])
+                rtkn1 = DefFunction('MAX', [0, M - d.symbolic_max + rtkn])
+                dynamic_args_mapper.update({ltkn: (ltkn1,), rtkn: (rtkn1,)})
 
         iet = callcompute._rebuild(dynamic_args_mapper=dynamic_args_mapper)
         # The -1 below is because an Iteration, by default, generates <=
@@ -1348,12 +1356,9 @@ def __init__(self, prefix, key, arguments, owned):
         pname = f'region{key}'
 
         fields = []
-        for i in self.arguments:
-            if i.is_Dimension:
-                fields.append((i.min_name, c_int))
-                fields.append((i.max_name, c_int))
-            else:
-                fields.append((i.name, c_int))
+        for d in self.arguments:
+            fields.append((d.min_name, c_int))
+            fields.append((d.max_name, c_int))
 
         super().__init__(name, pname, fields)
 
@@ -1387,16 +1392,10 @@ def _arg_values(self, args=None, **kwargs):
         values = self._arg_defaults()
         for i, (_, mapper) in enumerate(self.owned):
             entry = values[self.name][i]
-            for a in self.arguments:
-                if a.is_Dimension:
-                    a_m, a_M = mapper[a]
-                    setattr(entry, a.min_name, subs_op_args(a_m, args))
-                    setattr(entry, a.max_name, subs_op_args(a_M, args))
-                else:
-                    try:
-                        setattr(entry, a.name, subs_op_args(mapper[a][0], args))
-                    except AttributeError:
-                        setattr(entry, a.name, mapper[a][0])
+            for d in self.arguments:
+                d_m, d_M = mapper[d]
+                setattr(entry, d.min_name, subs_op_args(d_m, args))
+                setattr(entry, d.max_name, subs_op_args(d_M, args))
         return values
 
 

devito/types/dimension.py

@@ -639,6 +639,9 @@ class AbstractSubDimension(DerivedDimension):
 
     is_AbstractSub = True
 
+    # Whether the iteration fits entirely within a single MPI domain
+    local = False
+
     __rargs__ = DerivedDimension.__rargs__ + ('thickness',)
     __rkwargs__ = ('separated',)
 

tests/test_subdomains.py

@@ -932,6 +932,36 @@ def test_subdomainset_w_sparse(self):
         check[0, 0] = 1
         assert np.all(np.isclose(f.data, check))
 
+    @pytest.mark.parallel(mode=[(4, 'basic'), (4, 'diag2'), (4, 'overlap'),
+                                (4, 'overlap2'), (4, 'full')])
+    def test_subdomainset_stencil_mpi(self, mode):
+        """
+        A stencil update restricted to a set of boxes straddling the rank
+        boundaries, as DevitoPRO's ABox does around the sources and receivers
+        (devitopro/subdomains/abox.py). Under comp/comm overlap, the CORE and
+        OWNED regions must be intersected with each box.
+        """
+        class Boxes(SubDomainSet):
+            name = 'boxes'
+
+        grid = Grid(shape=(12, 12))
+        bounds = tuple(np.array(i, dtype=np.int32)
+                       for i in ([2, 7], [3, 1], [1, 6], [4, 1]))
+        boxes = Boxes(N=2, bounds=bounds, grid=grid)
+
+        u = TimeFunction(name='u', grid=grid, space_order=2)
+
+        op = Operator(Eq(u.forward, u.laplace + 1, subdomain=boxes))
+        op.apply(time_M=0)
+
+        check = np.zeros((12, 12))
+        check[2:9, 1:8] = 1
+        check[7:11, 6:11] = 1
+
+        data = u.data_gather()
+        if grid.distributor.myrank == 0:
+            assert np.all(data[1] == check)
+
 
 class TestBorder:
     # Note: This class is partially covered by doctests