From 61b3374dd8625096cd42185051b034a57cb8e992 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Tue, 3 Nov 2026 11:04:37 +0000 Subject: [PATCH 1/6] mpi: Fix halos of SubDomain Functions next to empty ranks --- devito/core/autotuning.py | 2 +- devito/data/decomposition.py | 11 +++++++++-- devito/mpi/distributed.py | 32 +++++++++++++++++++++++++------- devito/mpi/routines.py | 2 +- tests/test_pickle.py | 4 ++++ tests/test_subdomains.py | 25 +++++++++++++++++++++++++ 6 files changed, 65 insertions(+), 11 deletions(-) diff --git a/devito/core/autotuning.py b/devito/core/autotuning.py index fa833b84a16..b6f7e899f33 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() + p._rebuild()._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 ce14dfb7427..fc9a001bc33 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,8 @@ def reshape(self, *args): # Handle left extension/reduction if nleft > 0: - items = [np.concatenate([np.arange(-nleft, 0), items[0]])] + items[1:] + 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 +561,9 @@ def reshape(self, *args): # Handle right extension/reduction if nright > 0: + 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 35596c39d64..5edaaee0648 100644 --- a/devito/mpi/distributed.py +++ b/devito/mpi/distributed.py @@ -1,5 +1,6 @@ import atexit from abc import ABC, abstractmethod +from contextlib import suppress from ctypes import c_int, c_void_p, sizeof from functools import cached_property from itertools import groupby, product @@ -705,6 +706,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 +861,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,13 +904,21 @@ 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: + try: + inputs = kwargs['args'].op.input + except KeyError: return self._arg_defaults() + # Update `nb` based on the distributor of the (possibly overriding) + # Functions using it. E.g., an override may be defined on a new Grid + for p in inputs: + # Skip, e.g., Dimensions, SparseFunctions, or array overrides + with suppress(AttributeError): + if p._distributor._obj_neighborhood == self: + nb = kwargs.get(p.name, p)._distributor._obj_neighborhood + return {self.name: nb._arg_defaults()[nb.name]} + return self._arg_defaults() + class CustomTopology(DimensionTuple): diff --git a/devito/mpi/routines.py b/devito/mpi/routines.py index 3c446b24ba4..62af6653e02 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/tests/test_pickle.py b/tests/test_pickle.py index 838dd2ea800..544600329a8 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 eff67586367..59b695c9067 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. From b428b0e0df58072e55630d9704a2fed7ee690022 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Tue, 3 Nov 2026 16:53:40 +0000 Subject: [PATCH 2/6] misc: Revert to explicit copy --- devito/core/autotuning.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/core/autotuning.py b/devito/core/autotuning.py index b6f7e899f33..c2dfa65a3b4 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( - p._rebuild()._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) From 9366d40dbb2d73ebd9e16af15023b913ab5b4779 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Thu, 5 Nov 2026 15:03:50 +0000 Subject: [PATCH 3/6] compiler: Remove input search for MPINeighborhood --- devito/mpi/distributed.py | 15 +++------------ devito/types/grid.py | 12 +++++++++--- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py index 5edaaee0648..9fa3ffaf9b1 100644 --- a/devito/mpi/distributed.py +++ b/devito/mpi/distributed.py @@ -1,6 +1,5 @@ import atexit from abc import ABC, abstractmethod -from contextlib import suppress from ctypes import c_int, c_void_p, sizeof from functools import cached_property from itertools import groupby, product @@ -904,21 +903,13 @@ def _arg_defaults(self): return values def _arg_values(self, *args, **kwargs): + # Use `nb` as provided by the (possibly overriding) Grid or SubDomain. + # E.g., an override may be defined on a new Grid try: - inputs = kwargs['args'].op.input + return {self.name: kwargs['args'][self.name]} except KeyError: return self._arg_defaults() - # Update `nb` based on the distributor of the (possibly overriding) - # Functions using it. E.g., an override may be defined on a new Grid - for p in inputs: - # Skip, e.g., Dimensions, SparseFunctions, or array overrides - with suppress(AttributeError): - if p._distributor._obj_neighborhood == self: - nb = kwargs.get(p.name, p)._distributor._obj_neighborhood - return {self.name: nb._arg_defaults()[nb.name]} - return self._arg_defaults() - class CustomTopology(DimensionTuple): diff --git a/devito/types/grid.py b/devito/types/grid.py index 4568b852bf1..7658f0acc6f 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_defaults()) + + return values + class SubDomain(AbstractSubDomain): From 2ed4bc6e2570291dc95b4e1cf01cefb481db454a Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Thu, 5 Nov 2026 15:25:56 +0000 Subject: [PATCH 4/6] compiler: Refactor MPINeighborhood._arg_values --- devito/mpi/distributed.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py index 9fa3ffaf9b1..a6ba22d614a 100644 --- a/devito/mpi/distributed.py +++ b/devito/mpi/distributed.py @@ -905,10 +905,9 @@ def _arg_defaults(self): def _arg_values(self, *args, **kwargs): # Use `nb` as provided by the (possibly overriding) Grid or SubDomain. # E.g., an override may be defined on a new Grid - try: - return {self.name: kwargs['args'][self.name]} - except KeyError: - return self._arg_defaults() + if self.name in kwargs.get('args', {}): + return {} + return self._arg_defaults() class CustomTopology(DimensionTuple): From ba9d60d0ef1ad5d6b69dba1ce5f49a80226a2704 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Thu, 5 Nov 2026 15:39:23 +0000 Subject: [PATCH 5/6] misc: Improve comments --- devito/data/decomposition.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/devito/data/decomposition.py b/devito/data/decomposition.py index fc9a001bc33..f39126f7641 100644 --- a/devito/data/decomposition.py +++ b/devito/data/decomposition.py @@ -549,6 +549,7 @@ def reshape(self, *args): # Handle left extension/reduction if nleft > 0: + # 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: @@ -561,6 +562,7 @@ 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[i] = np.concatenate([items[i], extension]) From 0be3f4761bdd53e1bab18754fef89252ed3940ee Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Fri, 6 Nov 2026 16:07:39 +0000 Subject: [PATCH 6/6] misc: Fixup arg_values oversight for distributor --- devito/types/grid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devito/types/grid.py b/devito/types/grid.py index 7658f0acc6f..3721f7e6311 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -729,7 +729,7 @@ def _arg_values(self, **kwargs): # The SubDomain neighborhood, used by the Functions defined on it if self.distributor.comm: - values.update(self.distributor._obj_neighborhood._arg_defaults()) + values.update(self.distributor._obj_neighborhood._arg_values(**kwargs)) return values