mpi: Honor MultiSubDomain thicknesses in overlap modes

devito/mpi/halo_scheme.py

@@ -139,14 +139,20 @@ def __init__(self, exprs, ispace):
 
         # Track the IterationSpace offsets induced by SubDomains/SubDimensions,
         # which are honored in the derivation of the `omapper`
+        # The thicknesses of a MultiSubDimension are only known at runtime, as
+        # they are set inside the Operator, one SubDomain at a time
         self._honored = {}
+        runtime = set()
         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_Sub and not d.local) or d.is_MultiSub]
         for i in subdims:
             ltk, rtk = i.tkns
             self._honored[i.root] = frozenset([(ltk, rtk)])
+            if i.is_MultiSub:
+                runtime.update(i.tkns)
         self._honored = frozendict(self._honored)
+        self._runtime_tkns = frozenset(runtime)
 
         # Further constraints on the `omapper` derivation. At construction time
         # there's none, but lowering passes may change this
@@ -169,7 +175,7 @@ def __len__(self):
     def __hash__(self):
         return hash((self._mapper.__hash__(), self.honored.__hash__()))
 
-    def _rebuild(self, fmapper=None, honored=None, alignment=None):
+    def _rebuild(self, fmapper=None, honored=None, alignment=None, runtime_tkns=None):
         """
         Rebuild a HaloScheme from the provided `fmapper` and `honored`. Reuse
         `self`'s values for the missing arguments.
@@ -180,9 +186,12 @@ def _rebuild(self, fmapper=None, honored=None, alignment=None):
             fmapper = self._mapper
         if honored is None:
             honored = self._honored
+        if runtime_tkns is None:
+            runtime_tkns = self._runtime_tkns
 
         obj._mapper = frozendict(fmapper)
         obj._honored = frozendict(honored)
+        obj._runtime_tkns = frozenset(runtime_tkns)
         obj._alignment = alignment or self._alignment
 
         return obj
@@ -201,6 +210,7 @@ def union(self, halo_schemes):
 
         fmapper = {}
         honored = {}
+        runtime_tkns = set()
         for i in as_tuple(halo_schemes):
             # Compute the `fmapper `union`
             for k, v in i.fmapper.items():
@@ -237,13 +247,18 @@ def union(self, halo_schemes):
             # Compute the `honored` union
             for d, v in i.honored.items():
                 honored[d] = honored.get(d, frozenset()) | v
+            runtime_tkns.update(i.runtime_tkns)
 
-        return i._rebuild(fmapper=fmapper, honored=honored)
+        return i._rebuild(fmapper=fmapper, honored=honored, runtime_tkns=runtime_tkns)
 
     @property
     def honored(self):
         return self._honored
 
+    @property
+    def runtime_tkns(self):
+        return self._runtime_tkns
+
     @cached_property
     def fmapper(self):
         return OrderedDict([(i, self._mapper[i]) for i in

devito/mpi/routines.py

@@ -5,7 +5,7 @@
 from itertools import product
 from operator import mul
 
-from sympy import Integer
+from sympy import Integer, sympify
 
 from devito.data import CENTER, HALO, LEFT, NOPAD, OWNED, RIGHT
 from devito.ir.equations import DummyEq, OpInc, OpMax, OpMin
@@ -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,15 @@ class Overlap2HaloExchangeBuilder(OverlapHaloExchangeBuilder):
     """
 
     def _make_region(self, hs, key):
-        return MPIRegion('reg', key, hs.arguments, hs.omapper.owned)
+        # Runtime thicknesses can't be stored in the region, so they are
+        # zeroed here and applied in `_make_remainder` instead
+        runtime = hs.halo_scheme.runtime_tkns
+        arguments = [i for i in hs.arguments if i not in runtime]
+        zeros = dict.fromkeys(runtime, 0)
+        owned = [(w, {k: tuple(sympify(i).xreplace(zeros) for i in v)
+                      for k, v in mapper.items()})
+                 for w, mapper in hs.omapper.owned]
+        return MPIRegion('reg', key, arguments, owned)
 
     def _make_msg(self, f, hse, key):
         # Only retain the halos required by the Diag scheme
@@ -911,11 +919,26 @@ def _make_remainder(self, hs, key, callcompute, region):
         dim = Dimension(name='i')
         region_i = IndexedPointer(region, dim)
 
+        # A runtime thickness is applied relative to the OWNED region bounds
+        sides = {}
+        for d, v in hs.halo_scheme.honored.items():
+            for ltkn, rtkn in v:
+                sides.update({ltkn: (d, LEFT), rtkn: (d, RIGHT)})
+
         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))
+            elif i in hs.halo_scheme.runtime_tkns:
+                d, side = sides[i]
+                if side is LEFT:
+                    m = FieldFromComposite(d.min_name, region_i)
+                    v = DefFunction('MAX', [0, d.symbolic_min + i - m])
+                else:
+                    M = FieldFromComposite(d.max_name, region_i)
+                    v = DefFunction('MAX', [0, M - d.symbolic_max + i])
+                dynamic_args_mapper[i] = (v,)
             else:
                 dynamic_args_mapper[i] = (FieldFromComposite(i.name, region_i),)
 

tests/test_subdomains.py

@@ -647,6 +647,49 @@ class Inner(SubDomainSet):
 
         assert((np.array(result) == np.array(fex.data[:])).all())
 
+    @pytest.mark.parallel(mode=[(4, 'basic'), (4, 'diag2'), (4, 'overlap'),
+                                (4, 'overlap2'), (4, 'full')])
+    def test_subdomainset_stencil_mpi(self, mode):
+        """
+        Test a wave equation updated only within a box, defined via a
+        SubDomainSet, as done by DevitoPRO's ABox to restrict the computation
+        to the region around the sources and receivers. The box straddles the
+        rank boundaries, so the CORE and OWNED regions of the `overlap`,
+        `overlap2` and `full` MPI modes must honor its runtime thicknesses.
+        """
+        tkns = (3, 11, 9, 2)
+
+        class Box(SubDomainSet):
+            name = 'box'
+
+        class Reference(SubDomain):
+            name = 'reference'
+
+            def define(self, dimensions):
+                x, y = dimensions
+                return {x: ('middle', tkns[0], tkns[1]),
+                        y: ('middle', tkns[2], tkns[3])}
+
+        grid = Grid(shape=(40, 40), extent=(39., 39.))
+        bounds = tuple(np.array([i], dtype=np.int32) for i in tkns)
+        box = Box(N=1, bounds=bounds, grid=grid)
+        reference = Reference(grid=grid)
+
+        u0 = TimeFunction(name='u0', grid=grid, time_order=2, space_order=8)
+        u1 = TimeFunction(name='u1', grid=grid, time_order=2, space_order=8)
+        for u in (u0, u1):
+            u.data[:, 20, 20] = 1.
+
+        op0 = Operator(Eq(u0.forward, solve(u0.dt2 - u0.laplace, u0.forward),
+                          subdomain=reference))
+        op1 = Operator(Eq(u1.forward, solve(u1.dt2 - u1.laplace, u1.forward),
+                          subdomain=box))
+        op0(time_M=20, dt=0.3)
+        op1(time_M=20, dt=0.3)
+
+        assert np.linalg.norm(u0.data[:]) > 0
+        assert np.allclose(u0.data[:], u1.data[:], rtol=0, atol=1e-6)
+
     def test_multi_sets_eq(self):
         """
         Check functionality for when multiple subdomain sets are present, each