diff --git a/devito/builtins/initializers.py b/devito/builtins/initializers.py index d8c6055e20..19bbeaa3b2 100644 --- a/devito/builtins/initializers.py +++ b/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: diff --git a/devito/builtins/utils.py b/devito/builtins/utils.py index e271b86a64..908207f0fb 100644 --- a/devito/builtins/utils.py +++ b/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 diff --git a/devito/data/data.py b/devito/data/data.py index b41fc0eed3..38bbc80b04 100644 --- a/devito/data/data.py +++ b/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) diff --git a/devito/types/dimension.py b/devito/types/dimension.py index fec9fb78a3..45ecd5ed9a 100644 --- a/devito/types/dimension.py +++ b/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: diff --git a/tests/test_data.py b/tests/test_data.py index 20563b8398..40a2d403b7 100644 --- a/tests/test_data.py +++ b/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 diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 6081fa627e..de115d3a5d 100644 --- a/tests/test_mpi.py +++ b/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): """