diff --git a/devito/mpi/halo_scheme.py b/devito/mpi/halo_scheme.py index 3b4bcb5bc8..ca56d89c86 100644 --- a/devito/mpi/halo_scheme.py +++ b/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)]) diff --git a/devito/mpi/routines.py b/devito/mpi/routines.py index 3c446b24ba..42051091a5 100644 --- a/devito/mpi/routines.py +++ b/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 diff --git a/devito/types/dimension.py b/devito/types/dimension.py index c81eb90d37..02e58df77c 100644 --- a/devito/types/dimension.py +++ b/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',) diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index eff6758636..1a4886e583 100644 --- a/tests/test_subdomains.py +++ b/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