From 2dc440a217b1482a5bc3939fc089c18bd5acff29 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Tue, 3 Nov 2026 15:50:00 +0000 Subject: [PATCH 1/3] mpi: Fix halo=0 SubDomain mask at rank edges --- devito/types/dimension.py | 49 +++++++++++++++++++++++++++++++ devito/types/grid.py | 8 +++-- tests/test_mpi.py | 61 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 114 insertions(+), 4 deletions(-) diff --git a/devito/types/dimension.py b/devito/types/dimension.py index 6b21493c14..9a3ad6a1c9 100644 --- a/devito/types/dimension.py +++ b/devito/types/dimension.py @@ -625,6 +625,31 @@ def _arg_check(self, args, *_args, **kwargs): mpi_raise(error, InvalidArgument, comm=comm) +class RankOffset(Thickness): + """ + The number of points of `root` owned by the MPI ranks before (`side=LEFT`) + or after (`side=RIGHT`) this one, that is the offset of this rank from that + edge of the Grid. Zero without domain decomposition. + """ + + def _arg_values(self, grid=None, **kwargs): + root = self.root + if grid is not None: + root = next((d for d in grid.dimensions if d.name == root.name), root) + + if grid is None or not grid.is_distributed(root): + return {self.name: 0} + + dec = grid.distributor.decomposition[root] + if dec.loc_empty: + # No iterations run on an empty rank + return {self.name: 0} + elif self.side == LEFT: + return {self.name: dec.loc_abs_min - dec.glb_min} + else: + return {self.name: dec.glb_max - dec.loc_abs_max} + + class AbstractSubDimension(DerivedDimension): """ @@ -824,6 +849,30 @@ def _interval(self): return sympy.Interval(left, right) + @cached_property + def _glb_interval(self): + """ + The `(min, max)` bounds of this SubDimension in the global Grid, in this + rank's local indices. Unlike `_interval`, these are not clipped to this + rank, so with domain decomposition they may lie in the halo or beyond. + """ + kwargs = {'dtype': np.int32, 'is_const': True, 'nonnegative': True, + 'root': self.root} + lofs = RankOffset(name=f"{self.root.name}_lofs", side=LEFT, **kwargs) + rofs = RankOffset(name=f"{self.root.name}_rofs", side=RIGHT, **kwargs) + + # The global Grid edges + left = self.parent.symbolic_min - lofs + right = self.parent.symbolic_max + rofs + + # A tuple, not a `sympy.Interval`, which would collapse with literal bounds + if self.is_left: + return (left, left + self.ltkn.value - 1) + elif self.is_right: + return (right - self.rtkn.value + 1, right) + else: + return (left + self.ltkn.value, right - self.rtkn.value) + @property def local(self): return self._local diff --git a/devito/types/grid.py b/devito/types/grid.py index 4568b852bf..2ed199dd18 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -902,7 +902,8 @@ def _indicator(self, dim, offset): """ 1 if the point `offset` points away from the current one along `dim` lies in this SubDomain, 0 otherwise, computed without branching. Always 1 if - this SubDomain spans all of `dim`. + this SubDomain spans all of `dim`. With domain decomposition, the test + is against the global bounds, which may lie beyond this rank. Parameters ---------- @@ -915,8 +916,9 @@ def _indicator(self, dim, offset): if not maybe_subdim.is_Sub: return sympy.S.One point = dim + offset - return (Max(0, Min(1, point - maybe_subdim.symbolic_min + 1)) * - Max(0, Min(1, maybe_subdim.symbolic_max - point + 1))) + lower, upper = maybe_subdim._glb_interval + return (Max(0, Min(1, point - lower + 1)) * + Max(0, Min(1, upper - point + 1))) def _restrict(self, expr, exclude=()): """ diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 7081be3beb..8ab5c3575c 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -1,3 +1,4 @@ +import pickle from functools import cached_property import numpy as np @@ -25,7 +26,7 @@ from devito.mpi.routines import ComputeCall, HaloUpdateCall, HaloUpdateList, MPICall from devito.tools import Bunch from devito.types import Bundle -from devito.types.dimension import ModuloDimension +from devito.types.dimension import ModuloDimension, RankOffset from examples.seismic.acoustic import acoustic_setup @@ -3351,6 +3352,64 @@ def test_subdim_thickness_spanning_rank(self, mode): assert np.all(glb[:, 3:-3] == 0.) assert np.all(glb[:, -3:] == 2.) + @pytest.mark.parallel(mode=4) + @pytest.mark.parametrize('expand', [True, False]) + @pytest.mark.parametrize('region', [('middle', 6, 6), ('middle', 16, 16), + ('middle', 15, 17), ('left', 17), + ('right', 18)]) + def test_halo_zero_subdomain(self, region, expand, mode): + """ + Derivatives with `halo=0` treat their argument as zero outside the global + SubDomain, not outside its part on each rank. The rank edges lie at 16, + 32 and 48, within the stencil radius of some of the SubDomain edges. + """ + class Region(SubDomain): + name = 'region' + + def define(self, dimensions): + x, = dimensions + return {x: region} + + serial = Grid(shape=(64,), comm=MPI.COMM_SELF) + grid = Grid(shape=(64,)) + subdomain = Region(grid=serial) + + def functions(grid): + kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + return [Function(name=n, **kwargs) for n in ('f', 'o', 'g')] + + f0, o0, g0 = functions(serial) + f1, o1, g1 = functions(grid) + f0.data[:] = f1.data[:] = np.random.default_rng(0).normal(size=grid.shape) + + # Compiled for the serial Grid, applied to both + op = Operator([Eq(o0, f0.dx(halo=0), subdomain=subdomain), + Inc(g0, f0.dx(halo=0).T, subdomain=subdomain)], + opt=('advanced', {'expand': expand})) + op.apply() + op.apply(f=f1, o=o1, g=g1) + + glb = grid.distributor.glb_slices[grid.dimensions[0]] + assert np.allclose(o1.data, o0.data[glb], rtol=1e-12) + assert np.allclose(g1.data, g0.data[glb], rtol=1e-12) + + @pytest.mark.parallel(mode=4) + def test_rank_offset_empty(self, mode): + """ + The offsets of a rank from the Grid edges, including an empty rank, also + once unpickled. + """ + grid = Grid(shape=(3,)) + x = grid.dimensions[0] + + lofs = pickle.loads(pickle.dumps(RankOffset(name='x_lofs', side=LEFT, root=x))) + rofs = pickle.loads(pickle.dumps(RankOffset(name='x_rofs', side=RIGHT, root=x))) + values = (lofs._arg_values(grid=grid)['x_lofs'], + rofs._arg_values(grid=grid)['x_rofs']) + + # Rank 3 is empty + assert values == [(0, 2), (1, 1), (2, 0), (0, 0)][grid.distributor.myrank] + @pytest.mark.parallel(mode=[(4, 'full')]) def test_custom_subdomain(self, mode): """ From 7249fa074ae766812dea7963b394700eff902877 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Thu, 5 Nov 2026 17:21:12 +0000 Subject: [PATCH 2/3] dsl: Remove unnecessary Thickness kwargs and consolidate --- devito/types/dimension.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/devito/types/dimension.py b/devito/types/dimension.py index 9a3ad6a1c9..dfa99dbfb4 100644 --- a/devito/types/dimension.py +++ b/devito/types/dimension.py @@ -464,6 +464,16 @@ class SpaceDimension(BasicDimension): is_Space = True + @cached_property + def rank_offsets(self): + """ + The number of points of `root` owned by the MPI ranks before (`side=LEFT`) + or after (`side=RIGHT`) this one, that is the offset of this rank from that + edge of the Grid. Zero without domain decomposition. + """ + return tuple(RankOffset(name=f"{self.name}_{s}_ofs", side=s, root=self) + for s in (LEFT, RIGHT)) + class TimeDimension(BasicDimension): @@ -566,6 +576,7 @@ class Thickness(DataSymbol): __rkwargs__ = DataSymbol.__rkwargs__ + ('root', 'side', 'local', 'value', 'separated') def __new__(cls, *args, root=None, side=None, local=False, separated=False, **kwargs): + kwargs.setdefault('nonnegative', True) newobj = super().__new__(cls, *args, **kwargs) newobj._root = root newobj._side = side @@ -695,11 +706,9 @@ def _interval(self): @memoized_meth def _symbolic_thickness(self, **kwargs): - kwargs = {'dtype': np.int32, 'is_const': True, 'nonnegative': True, - 'separated': self.separated} - names = [f"{self.parent.name}_{s}tkn" for s in ('l', 'r')] - return SubDimensionThickness(*[Thickness(name=n, **kwargs) for n in names]) + return SubDimensionThickness(*[Thickness(name=n, separated=self.separated) + for n in names]) @cached_property def symbolic_min(self): @@ -825,9 +834,7 @@ def middle(cls, name, parent, thickness_left, thickness_right, local=False, @memoized_meth def _symbolic_thickness(self, thickness=None): - kwargs = {'dtype': np.int32, 'is_const': True, 'nonnegative': True, - 'root': self.root, 'local': self.local, - 'separated': self.separated} + kwargs = {'root': self.root, 'local': self.local, 'separated': self.separated} names = [f"{self.parent.name}_{s}tkn" for s in ('l', 'r')] sides = [LEFT, RIGHT] @@ -856,10 +863,7 @@ def _glb_interval(self): rank's local indices. Unlike `_interval`, these are not clipped to this rank, so with domain decomposition they may lie in the halo or beyond. """ - kwargs = {'dtype': np.int32, 'is_const': True, 'nonnegative': True, - 'root': self.root} - lofs = RankOffset(name=f"{self.root.name}_lofs", side=LEFT, **kwargs) - rofs = RankOffset(name=f"{self.root.name}_rofs", side=RIGHT, **kwargs) + lofs, rofs = self.root.rank_offsets # The global Grid edges left = self.parent.symbolic_min - lofs From 063ab3d090fd4b8d259a7ad0e2a70ed25c2660a1 Mon Sep 17 00:00:00 2001 From: Edward Caunt Date: Fri, 6 Nov 2026 13:12:25 +0000 Subject: [PATCH 3/3] tests: Remove unnecessary dtype specifications --- tests/test_derivatives.py | 14 +++++++------- tests/test_mpi.py | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py index ee93c00488..55a75e81e5 100644 --- a/tests/test_derivatives.py +++ b/tests/test_derivatives.py @@ -1895,7 +1895,7 @@ def test_default_unchanged(self): """ grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) - kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + kwargs = {'grid': grid, 'space_order': 4} g_bar = Function(name='g_bar', **kwargs) ref = Function(name='ref', **kwargs) out_bar = Function(name='out_bar', **kwargs) @@ -1927,7 +1927,7 @@ def test_transpose_keeps_halo(self): def test_no_subdomain(self): """Without a SubDomain, halo=0 has no effect.""" grid = Grid(shape=(16,), extent=(15.,), dtype=np.float64) - kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + kwargs = {'grid': grid, 'space_order': 4} f = Function(name='f', **kwargs) g = Function(name='g', **kwargs) ref = Function(name='ref', **kwargs) @@ -1966,7 +1966,7 @@ def test_terms(self): """ grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) - kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + kwargs = {'grid': grid, 'space_order': 4} f, g, c, out, ref = (Function(name=n, **kwargs) for n in ('f', 'g', 'c', 'out', 'ref')) rng = np.random.default_rng(8) @@ -1987,7 +1987,7 @@ def test_sum_factor(self): """ grid = Grid(shape=(24,), extent=(23.,), dtype=np.float64) subdomain = Region('s', {'x': ('middle', 7, 9)}, grid=grid) - kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + kwargs = {'grid': grid, 'space_order': 4} f, g, c, d, out = (Function(name=n, **kwargs) for n in ('f', 'g', 'c', 'd', 'out')) rng = np.random.default_rng(9) @@ -2008,7 +2008,7 @@ def test_vector(self): """ grid = Grid(shape=(12, 14), extent=(11., 13.), dtype=np.float64) subdomain = Region('s', {'x': ('left', 5)}, grid=grid) - kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + kwargs = {'grid': grid, 'space_order': 4} f = Function(name='f', **kwargs) c = VectorFunction(name='c', staggered=(None, None), **kwargs) w = VectorFunction(name='w', staggered=(None, None), **kwargs) @@ -2034,7 +2034,7 @@ def test_vector_staggered(self): grid = Grid(shape=(14, 16), extent=(13., 15.), dtype=np.float64) x, y = grid.dimensions subdomain = Region('s', {'x': ('left', 5), 'y': ('middle', 4, 5)}, grid=grid) - kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + kwargs = {'grid': grid, 'space_order': 4} p = Function(name='p', **kwargs) p_bar = Function(name='p_bar', **kwargs) v = VectorFunction(name='v', **kwargs) @@ -2057,7 +2057,7 @@ def test_unrestricted_dimension(self): """halo=0 along a Dimension the SubDomain spans has no effect.""" grid = Grid(shape=(12, 14), extent=(11., 13.), dtype=np.float64) subdomain = Region('s', {'x': ('left', 5)}, grid=grid) - kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + kwargs = {'grid': grid, 'space_order': 4} f = Function(name='f', **kwargs) g = Function(name='g', **kwargs) ref = Function(name='ref', **kwargs) diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 8ab5c3575c..1e2adf3dad 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -3375,7 +3375,7 @@ def define(self, dimensions): subdomain = Region(grid=serial) def functions(grid): - kwargs = {'grid': grid, 'space_order': 4, 'dtype': np.float64} + kwargs = {'grid': grid, 'space_order': 4} return [Function(name=n, **kwargs) for n in ('f', 'o', 'g')] f0, o0, g0 = functions(serial)