api: Fix initializers with mpi

devito/builtins/initializers.py

@@ -1,7 +1,9 @@
 import numpy as np
 
 import devito as dv
-from devito.builtins.utils import check_builtins_args, nbl_to_padsize, pad_outhalo
+from devito.builtins.utils import (
+    axis_slice, check_builtins_args, nbl_to_padsize, pad_outhalo
+)
 from devito.tools import as_list, as_tuple
 
 __all__ = ['assign', 'gaussian_smooth', 'initialize_function', 'smooth']
@@ -217,15 +219,95 @@ def fset(f, g):
     return f
 
 
-def _initialize_function(function, data, nbl, mapper=None, mode='constant'):
+def _rank_without_interior(function, nbl):
     """
-    Construct the symbolic objects for `initialize_function`.
+    Whether some rank owns no part of the interior along a padded Dimension.
+
+    The padding is extended by reading the boundary plane at a fixed global
+    index, which only works while every rank owns interior to read from.
+
+    Collective: the answer is reduced over all ranks, since they must then all
+    extend the padding the same way.
     """
-    nbl, slices = nbl_to_padsize(nbl, function.ndim)
+    distributor = function.grid.distributor
+    if not distributor.is_parallel:
+        return False
+
+    owns = True
+    for i, (d, (nl, nr)) in enumerate(
+        zip(function.space_dimensions, as_tuple(nbl), strict=True)
+    ):
+        if function.grid.is_distributed(d):
+            # The interior spans `[nl, glb_max - nr]`
+            span = distributor.glb_slices[d]
+            owns &= not (span.stop <= nl or
+                         span.start > distributor.glb_shape[i] - 1 - nr)
+
+    return not distributor.comm.allreduce(owns, op=dv.mpi.MPI.LAND)
+
+
+def _global_plane(function, axis, index):
+    """
+    The plane of `function` at global position `index` along `axis`, assembled
+    identically on every rank and keeping `axis` as a length-1 dimension.
+    """
+    distributor = function.grid.distributor
+
+    local = np.asarray(function.data[axis_slice(function.ndim, axis, index,
+                                                index + 1)])
+    plane = np.zeros(tuple(distributor.glb_shape[i] if i != axis else 1
+                           for i in range(function.ndim)),
+                     dtype=function.dtype)
+
+    # `glb_slices` tile the domain, so each slot gets a single contribution
+    if local.size:
+        slot = tuple(distributor.glb_slices[dim] if i != axis else slice(None)
+                     for i, dim in enumerate(function.dimensions))
+        plane[slot] = local.reshape(plane[slot].shape)
+    distributor.comm.Allreduce(dv.mpi.MPI.IN_PLACE, plane, op=dv.mpi.MPI.SUM)
+
+    return plane
+
+
+def _extend_padding_mpi(function, nbl):
+    """
+    Replicate the boundary planes of `function` outwards into its padding.
+
+    Used in place of the symbolic extension when a rank owns no interior to read
+    from, which the halo exchanges cannot express. Writing through the
+    distributed `Data` instead resolves global indices on whichever rank owns
+    them.
+    """
+    distributor = function.grid.distributor
+
+    # `TimeFunction`s are rejected upstream, so `axis` indexes `glb_shape` too
+    for axis, (nl, nr) in enumerate(as_tuple(nbl)):
+        glb_max = distributor.glb_shape[axis] - 1
+
+        # One Dimension at a time, each reading the planes left by the previous
+        # ones, so that corners come out filled
+        for nb, src, lo in ((nl, nl, 0), (nr, glb_max - nr, glb_max - nr + 1)):
+            if nb <= 0:
+                continue
+            # `plane` is size 1 along `axis`, so it broadcasts over the slot
+            plane = _global_plane(function, axis, src)
+            function.data[axis_slice(function.ndim, axis, lo, lo + nb)] = plane
+
+
+def _write_interior(function, data, slices):
+    """Fill the interior of `function`, that is everything but the padding."""
     if isinstance(data, dv.Function):
         function.data[slices] = data.data[:]
     else:
         function.data[slices] = data
+
+
+def _initialize_function(function, data, nbl, mapper=None, mode='constant'):
+    """
+    Construct the symbolic objects for `initialize_function`.
+    """
+    nbl, slices = nbl_to_padsize(nbl, function.ndim)
+    _write_interior(function, data, slices)
     lhs = []
     rhs = []
     options = []
@@ -385,17 +467,24 @@ def initialize_function(function, data, nbl, mapper=None, mode='constant',
     else:
         lhss, rhss, optionss = [], [], []
         for f, data in zip(functions, datas, strict=True):
+            padsizes, slices = nbl_to_padsize(nbl, f.ndim)
 
-            lhs, rhs, options = _initialize_function(f, data, nbl, mapper, mode)
+            # `reflect` has its own, stricter check on the halo size
+            if mode == 'constant' and _rank_without_interior(f, padsizes):
+                _write_interior(f, data, slices)
+                _extend_padding_mpi(f, padsizes)
+                continue
 
+            lhs, rhs, options = _initialize_function(f, data, nbl, mapper, mode)
             lhss.extend(lhs)
             rhss.extend(rhs)
             optionss.extend(options)
 
         assert len(lhss) == len(rhss) == len(optionss)
 
-        name = name or f'initialize_{"_".join(f.name for f in functions)}'
-        assign(lhss, rhss, options=optionss, name=name, **kwargs)
+        if lhss:
+            name = name or f'initialize_{"_".join(f.name for f in functions)}'
+            assign(lhss, rhss, options=optionss, name=name, **kwargs)
 
     if pad_halo:
         for f in functions:

devito/builtins/utils.py

@@ -9,6 +9,7 @@
 
 __all__ = [
     'abstract_args',
+    'axis_slice',
     'check_builtins_args',
     'make_retval',
     'nbl_to_padsize',
@@ -90,6 +91,12 @@ def nbl_to_padsize(nbl, ndim):
     return tuple(nb_pad), tuple(slices)
 
 
+def axis_slice(ndim, axis, start, stop):
+    """An index tuple selecting `[start, stop)` along `axis`."""
+    return tuple(slice(start, stop) if i == axis else slice(None)
+                 for i in range(ndim))
+
+
 def pad_outhalo(function):
     """ Pad outer halo with edge values."""
     h_shape = function._data_with_outhalo.shape

devito/data/data.py

@@ -347,7 +347,11 @@ def __setitem__(self, glb_idx, val):
                         processed.append(slice(j.start, j.stop, 1))
                     else:
                         processed.append(j)
-                val_idx = as_tuple(processed)
+                # An axis `val` broadcasts along carries the same value for every
+                # rank, so it must be taken whole rather than cut down to this
+                # rank's share of the destination
+                val_idx = tuple(slice(None) if n == 1 else j
+                                for j, n in zip(processed, val.shape, strict=True))
             elif is_windowed(val):
                 # Read all of it, there is no decomposition to restrict it to
                 val_idx = tuple(slice(0, s) for s in val.shape)

devito/types/dimension.py

@@ -597,11 +597,19 @@ def _arg_values(self, grid=None, **kwargs):
         if grid is not None and grid.is_distributed(self.root):
             # Get local thickness
             if self.local:
-                # Dimension is of type `left`/`right` - compute the offset
-                # and then add 1 to get the appropriate thickness
+                # Dimension is of type `left`/`right` - compute the offset and
+                # then add 1 to get the appropriate thickness. `glb_to_loc`
+                # saturates to the local size when the layer covers the whole
+                # subdomain, which is already a thickness, hence the clamp
                 if self.value is not None:
-                    tkn = grid.distributor.glb_to_loc(self.root, rtkn-1, self.side)
-                    tkn = tkn+1 if tkn is not None else 0
+                    distributor = grid.distributor
+                    tkn = distributor.glb_to_loc(self.root, rtkn-1, self.side)
+                    if tkn is None:
+                        tkn = 0
+                    else:
+                        loc_size = distributor.shape[
+                            distributor.dimensions.index(self.root)]
+                        tkn = min(tkn + 1, loc_size)
                 else:
                     tkn = 0
             else:

tests/test_data.py

@@ -1217,6 +1217,26 @@ def test_big_steps(self, mode):
                        (slice(6, None, -3), slice(6, None, -3))]:
             self._assert_induced(f.data[gslice], a[gslice])
 
+    @pytest.mark.parallel(mode=4)
+    def test_setitem_broadcast(self, mode):
+        """
+        An RHS that is size 1 along an axis broadcasts over that axis, rather
+        than being cut down to each rank's share of the destination.
+        """
+        grid = Grid(shape=(4, 8), topology=(1, 4))
+        f = Function(name='f', grid=grid, space_order=0, dtype=np.int32)
+        f.data[:] = 0
+
+        col = np.array([10, 20, 30, 40], dtype=np.int32)[:, None]
+        f.data[:, 0:3] = col
+
+        expected = np.zeros((4, 8), dtype=np.int32)
+        expected[:, 0:3] = col
+
+        glb = f.data_gather(rank=0)
+        if grid.distributor.myrank == 0:
+            assert np.array_equal(np.asarray(glb), expected)
+
     @pytest.mark.parallel(mode=4)
     def test_setitem(self, mode):
         # __setitem__ mpi slicing tests

tests/test_mpi.py

@@ -12,6 +12,7 @@
     inner, norm, solve, switchconfig
 )
 from devito.arch.compiler import OneapiCompiler
+from devito.builtins import initialize_function
 from devito.data import LEFT, RIGHT
 from devito.ir import Cluster, Interval, IterationSpace
 from devito.ir.clusters.algorithms import check_halo_writes
@@ -3293,6 +3294,63 @@ def test_misc_subdims_3D(self, mode):
         assert np.all(u.data[1, :, :, 0:2] == 1.)
         assert np.all(u.data[1, :, :, -2:] == 1.)
 
+    @pytest.mark.parallel(mode=[(4, 'basic')])
+    def test_subdim_thickness_spanning_rank(self, mode):
+        """
+        A `left`/`right` SubDimension whose thickness covers a whole rank must
+        not produce a local thickness larger than that rank's extent, or the
+        generated loop runs off the ends of its temporaries.
+        """
+        # 4 ranks along `y`, so each holds 2 points against 3-point layers
+        grid = Grid(shape=(4, 8), topology=(1, 4))
+        y = grid.dimensions[1]
+
+        f = Function(name='f', grid=grid)
+
+        yl = SubDimension.left(name='yl', parent=y, thickness=3)
+        yr = SubDimension.right(name='yr', parent=y, thickness=3)
+
+        op = Operator([Eq(f.subs(y, yl), 1), Eq(f.subs(y, yr), 2)])
+
+        args = op.arguments(f=f)
+        loc_size = grid.distributor.shape[1]
+        assert args['y_ltkn0'] <= loc_size
+        assert args['y_rtkn1'] <= loc_size
+
+        op.apply(f=f)
+
+        glb = f.data_gather(rank=0)
+        if grid.distributor.myrank == 0:
+            assert np.all(glb[:, :3] == 1.)
+            assert np.all(glb[:, 3:-3] == 0.)
+            assert np.all(glb[:, -3:] == 2.)
+
+    @pytest.mark.parallel(mode=[(3, 'basic'), (4, 'basic')])
+    def test_initialize_function_pad_over_rank(self, mode):
+        """
+        `initialize_function` must pad correctly even when a rank owns no part
+        of the interior, since the boundary plane it reads is then on a
+        different rank.
+        """
+        nbl = 6
+        inner = (7, 5, 6)
+        shape = tuple(s + 2*nbl for s in inner)
+
+        # Decompose along `y` only, so some rank holds nothing but padding
+        grid = Grid(shape=shape, topology=(1, '*', 1))
+
+        f = Function(name='f', grid=grid, space_order=(8, 4, 4), dtype=np.float32)
+
+        # Asymmetric, so that a misplaced plane cannot go unnoticed
+        data = np.arange(np.prod(inner), dtype=np.float32).reshape(inner) + 1.
+
+        initialize_function(f, data, [(nbl, nbl)]*3)
+
+        glb = f.data_gather(rank=0)
+        if grid.distributor.myrank == 0:
+            expected = np.pad(data, nbl, mode='edge')
+            assert np.array_equal(np.asarray(glb), expected)
+
     @pytest.mark.parallel(mode=[(4, 'full')])
     def test_custom_subdomain(self, mode):
         """