mpi: Fix halos of SubDomain Functions next to empty ranks

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)

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))):

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):

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]

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):
 

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)

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.