diff --git a/devito/core/autotuning.py b/devito/core/autotuning.py index fa833b84a1..c2dfa65a3b 100644 --- a/devito/core/autotuning.py +++ b/devito/core/autotuning.py @@ -60,7 +60,7 @@ def autotune(operator, args, level, mode): for p in operator.parameters: if isinstance(p, MPINeighborhood): at_args.update( - MPINeighborhood(p.neighborhood)._arg_values() + MPINeighborhood(p.neighborhood, name=p.name)._arg_values() ) for i in p.fields: setattr(at_args[p.name]._obj, i, MPI.PROC_NULL) diff --git a/devito/data/decomposition.py b/devito/data/decomposition.py index ce14dfb742..f39126f764 100644 --- a/devito/data/decomposition.py +++ b/devito/data/decomposition.py @@ -511,6 +511,11 @@ def reshape(self, *args): >>> d.reshape(slice(4)) Decomposition([0,2], [3,3], <<[]>>, []) + + Extensions are attached to the outermost non-empty subdomains + + >>> Decomposition([[], [0, 1], [2, 3], []], 1).reshape(1, 1) + Decomposition([], <<[0,2]>>, [3,5], []) """ if len(args) == 1: @@ -544,7 +549,9 @@ def reshape(self, *args): # Handle left extension/reduction if nleft > 0: - items = [np.concatenate([np.arange(-nleft, 0), items[0]])] + items[1:] + # Extend the first non-empty subdomain, as empty ones own no halo + i = next((i for i, sd in enumerate(items) if sd.size), 0) + items[i] = np.concatenate([np.arange(-nleft, 0), items[i]]) elif nleft < 0: n = 0 for i, sd in enumerate(list(items)): @@ -555,8 +562,10 @@ def reshape(self, *args): # Handle right extension/reduction if nright > 0: + # Extend the last non-empty subdomain, as empty ones own no halo + i = max((i for i, sd in enumerate(items) if sd.size), default=-1) extension = np.arange(self.glb_max + 1, self.glb_max + 1 + nright) - items = items[:-1] + [np.concatenate([items[-1], extension])] + items[i] = np.concatenate([items[i], extension]) elif nright < 0: n = 0 for i, sd in enumerate(reversed(list(items))): diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py index 35596c39d6..a6ba22d614 100644 --- a/devito/mpi/distributed.py +++ b/devito/mpi/distributed.py @@ -705,6 +705,14 @@ def neighborhood(self): return ret + @cached_property + def _obj_neighborhood(self): + # The neighborhood is determined by the SubDomain thicknesses, hence the + # name. E.g., `nb_x0_7` for `{x: ('middle', 0, 7), y: y}` + name = ''.join(f'_{d.root.name}{d.ltkn.value}_{d.rtkn.value}' + for d in self.dimensions if d.is_Sub) + return MPINeighborhood(self.neighborhood, name=f'nb{name}') + @cached_property def rank_populated(self): """Constant symbol for a switch indicating that data is allocated on this rank""" @@ -852,14 +860,15 @@ def _arg_values(self, *args, **kwargs): class MPINeighborhood(CompositeObject): __rargs__ = ('neighborhood',) + __rkwargs__ = ('name',) - def __init__(self, neighborhood): + def __init__(self, neighborhood, name='nb'): self._neighborhood = neighborhood self._entries = [i for i in neighborhood if isinstance(i, tuple)] fields = [(''.join(j.name[0] for j in i), c_int) for i in self.entries] - super().__init__('nb', 'neighborhood', fields) + super().__init__(name, 'neighborhood', fields) @property def entries(self): @@ -894,12 +903,11 @@ def _arg_defaults(self): return values def _arg_values(self, *args, **kwargs): - grid = kwargs.get('grid') - # Update `nb` based on object attached to `grid` - if grid is not None: - return grid.distributor._obj_neighborhood._arg_defaults() - else: - return self._arg_defaults() + # Use `nb` as provided by the (possibly overriding) Grid or SubDomain. + # E.g., an override may be defined on a new Grid + if self.name in kwargs.get('args', {}): + return {} + return self._arg_defaults() class CustomTopology(DimensionTuple): diff --git a/devito/mpi/routines.py b/devito/mpi/routines.py index 3c446b24ba..62af6653e0 100644 --- a/devito/mpi/routines.py +++ b/devito/mpi/routines.py @@ -1286,7 +1286,7 @@ def _arg_defaults(self, alias=None, args=None): super()._arg_defaults(alias, args=args) f = alias or self.target.c0 - neighborhood = args.grid.distributor.neighborhood + neighborhood = f.grid.distributor.neighborhood for i, halo in enumerate(self.halos): entry = self.value[i] diff --git a/devito/types/grid.py b/devito/types/grid.py index 4568b852bf..3721f7e631 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -402,10 +402,10 @@ def _arg_defaults(self): args.update({k.name: v for k, v in self.origin_map.items()}) # MPI-related objects - if self.distributor.is_parallel: + if self.distributor.comm: distributor = self.distributor args[distributor._obj_comm.name] = distributor._obj_comm.value - args[distributor._obj_neighborhood.name] = distributor._obj_neighborhood.value + args.update(distributor._obj_neighborhood._arg_defaults()) return args @@ -721,12 +721,18 @@ def comm(self): def _arg_values(self, **kwargs): try: - return self.grid._arg_values(**kwargs) + values = self.grid._arg_values(**kwargs) except AttributeError as e: raise AttributeError( f'{self} is not attached to a Grid and has no _arg_values' ) from e + # The SubDomain neighborhood, used by the Functions defined on it + if self.distributor.comm: + values.update(self.distributor._obj_neighborhood._arg_values(**kwargs)) + + return values + class SubDomain(AbstractSubDomain): diff --git a/tests/test_pickle.py b/tests/test_pickle.py index 838dd2ea80..544600329a 100644 --- a/tests/test_pickle.py +++ b/tests/test_pickle.py @@ -935,6 +935,10 @@ def test_mpi_objects(self, pickle, mode): assert obj.pname == new_obj.pname assert obj.pfields == new_obj.pfields + obj = obj._rebuild(name='nbsd') + new_obj = pickle.loads(pickle.dumps(obj)) + assert new_obj.name == 'nbsd' + # Communicator obj = grid.distributor._obj_comm pkl_obj = pickle.dumps(obj) diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index eff6758636..59b695c906 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -1792,6 +1792,31 @@ def define(self, dimensions): assert np.all(h.data == h_check[slices]) assert np.all(i.data == i_check) + @pytest.mark.parallel(mode=[(2, 'basic'), (2, 'diag'), (2, 'diag2')]) + def test_halo_empty_neighbor_mpi(self, mode): + """ + Check that halo exchanges of Functions on SubDomains skip neighboring + ranks holding none of the SubDomain, also alongside Grid Functions. + """ + grid = Grid(shape=(11, 11), extent=(10., 10.), topology=(2, 1)) + # Only held by rank 0 + reduced_domain = ReducedDomain(('middle', 0, 7), None, grid=grid) + + f = Function(name='f', grid=reduced_domain, space_order=2) + g = Function(name='g', grid=grid, space_order=2) + fx = Function(name='fx', grid=grid) + gx = Function(name='gx', grid=grid) + + # Sentinel, as nothing should be exchanged into the halo of `f` + f.data_with_halo[:] = -1 + + Operator([Eq(f, 1), Eq(g, 1), + Eq(fx, f.dx, subdomain=reduced_domain), Eq(gx, g.dx)]).apply() + + assert (np.count_nonzero(f.data_with_halo == -1) == + f.data_with_halo.size - f.data.size) + assert np.all(gx.data[1:-1] == 0) + def set_indices(data): """ Set up individual indices for indexing/slicing check.