From 0a1835a558202d5dd90141e8cc1d0db5be2a2f17 Mon Sep 17 00:00:00 2001 From: Gerard Gorman Date: Mon, 19 Oct 2026 18:31:12 +0000 Subject: [PATCH] mpi: Honor MultiSubDomain thicknesses in overlap modes Under DEVITO_MPI=overlap/overlap2/full the CORE and OWNED regions ignored the thicknesses of MultiSubDimensions, so points near the edges of a SubDomainSet (or DevitoPRO ABox) were skipped. These thicknesses are only known at runtime, so overlap2/full apply them to the OWNED region bounds in the remainder rather than storing them in the region struct. --- devito/mpi/halo_scheme.py | 21 ++++++++++++++++--- devito/mpi/routines.py | 31 ++++++++++++++++++++++++---- tests/test_subdomains.py | 43 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 88 insertions(+), 7 deletions(-) diff --git a/devito/mpi/halo_scheme.py b/devito/mpi/halo_scheme.py index 3b4bcb5bc80..b5132fc4aa7 100644 --- a/devito/mpi/halo_scheme.py +++ b/devito/mpi/halo_scheme.py @@ -139,14 +139,20 @@ def __init__(self, exprs, ispace): # Track the IterationSpace offsets induced by SubDomains/SubDimensions, # which are honored in the derivation of the `omapper` + # The thicknesses of a MultiSubDimension are only known at runtime, as + # they are set inside the Operator, one SubDomain at a time self._honored = {} + runtime = set() dims = set().union(*[d._defines for d in ispace.dimensions if d._defines & self.dimensions]) - subdims = [d for d in dims if d.is_Sub and not d.local] + subdims = [d for d in dims if (d.is_Sub and not d.local) or d.is_MultiSub] for i in subdims: ltk, rtk = i.tkns self._honored[i.root] = frozenset([(ltk, rtk)]) + if i.is_MultiSub: + runtime.update(i.tkns) self._honored = frozendict(self._honored) + self._runtime_tkns = frozenset(runtime) # Further constraints on the `omapper` derivation. At construction time # there's none, but lowering passes may change this @@ -169,7 +175,7 @@ def __len__(self): def __hash__(self): return hash((self._mapper.__hash__(), self.honored.__hash__())) - def _rebuild(self, fmapper=None, honored=None, alignment=None): + def _rebuild(self, fmapper=None, honored=None, alignment=None, runtime_tkns=None): """ Rebuild a HaloScheme from the provided `fmapper` and `honored`. Reuse `self`'s values for the missing arguments. @@ -180,9 +186,12 @@ def _rebuild(self, fmapper=None, honored=None, alignment=None): fmapper = self._mapper if honored is None: honored = self._honored + if runtime_tkns is None: + runtime_tkns = self._runtime_tkns obj._mapper = frozendict(fmapper) obj._honored = frozendict(honored) + obj._runtime_tkns = frozenset(runtime_tkns) obj._alignment = alignment or self._alignment return obj @@ -201,6 +210,7 @@ def union(self, halo_schemes): fmapper = {} honored = {} + runtime_tkns = set() for i in as_tuple(halo_schemes): # Compute the `fmapper `union` for k, v in i.fmapper.items(): @@ -237,13 +247,18 @@ def union(self, halo_schemes): # Compute the `honored` union for d, v in i.honored.items(): honored[d] = honored.get(d, frozenset()) | v + runtime_tkns.update(i.runtime_tkns) - return i._rebuild(fmapper=fmapper, honored=honored) + return i._rebuild(fmapper=fmapper, honored=honored, runtime_tkns=runtime_tkns) @property def honored(self): return self._honored + @property + def runtime_tkns(self): + return self._runtime_tkns + @cached_property def fmapper(self): return OrderedDict([(i, self._mapper[i]) for i in diff --git a/devito/mpi/routines.py b/devito/mpi/routines.py index 3c446b24ba4..57b3ea3346a 100644 --- a/devito/mpi/routines.py +++ b/devito/mpi/routines.py @@ -5,7 +5,7 @@ from itertools import product from operator import mul -from sympy import Integer +from sympy import Integer, sympify from devito.data import CENTER, HALO, LEFT, NOPAD, OWNED, RIGHT from devito.ir.equations import DummyEq, OpInc, OpMax, OpMin @@ -16,8 +16,8 @@ ) from devito.mpi import MPI from devito.symbolics import ( - Byref, CondNe, FieldFromComposite, FieldFromPointer, IndexedPointer, Macro, cast, - subs_op_args + Byref, CondNe, DefFunction, FieldFromComposite, FieldFromPointer, IndexedPointer, + Macro, cast, subs_op_args ) from devito.tools import ( as_mapper, dtype_len, dtype_to_mpitype, flatten, generator, infer_datasize, is_integer @@ -791,7 +791,15 @@ class Overlap2HaloExchangeBuilder(OverlapHaloExchangeBuilder): """ def _make_region(self, hs, key): - return MPIRegion('reg', key, hs.arguments, hs.omapper.owned) + # Runtime thicknesses can't be stored in the region, so they are + # zeroed here and applied in `_make_remainder` instead + runtime = hs.halo_scheme.runtime_tkns + arguments = [i for i in hs.arguments if i not in runtime] + zeros = dict.fromkeys(runtime, 0) + owned = [(w, {k: tuple(sympify(i).xreplace(zeros) for i in v) + for k, v in mapper.items()}) + for w, mapper in hs.omapper.owned] + return MPIRegion('reg', key, arguments, owned) def _make_msg(self, f, hse, key): # Only retain the halos required by the Diag scheme @@ -911,11 +919,26 @@ def _make_remainder(self, hs, key, callcompute, region): dim = Dimension(name='i') region_i = IndexedPointer(region, dim) + # A runtime thickness is applied relative to the OWNED region bounds + sides = {} + for d, v in hs.halo_scheme.honored.items(): + for ltkn, rtkn in v: + sides.update({ltkn: (d, LEFT), rtkn: (d, RIGHT)}) + dynamic_args_mapper = {} for i in hs.arguments: if i.is_Dimension: dynamic_args_mapper[i] = (FieldFromComposite(i.min_name, region_i), FieldFromComposite(i.max_name, region_i)) + elif i in hs.halo_scheme.runtime_tkns: + d, side = sides[i] + if side is LEFT: + m = FieldFromComposite(d.min_name, region_i) + v = DefFunction('MAX', [0, d.symbolic_min + i - m]) + else: + M = FieldFromComposite(d.max_name, region_i) + v = DefFunction('MAX', [0, M - d.symbolic_max + i]) + dynamic_args_mapper[i] = (v,) else: dynamic_args_mapper[i] = (FieldFromComposite(i.name, region_i),) diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index eff67586367..58e92e9c5db 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -647,6 +647,49 @@ class Inner(SubDomainSet): assert((np.array(result) == np.array(fex.data[:])).all()) + @pytest.mark.parallel(mode=[(4, 'basic'), (4, 'diag2'), (4, 'overlap'), + (4, 'overlap2'), (4, 'full')]) + def test_subdomainset_stencil_mpi(self, mode): + """ + Test a wave equation updated only within a box, defined via a + SubDomainSet, as done by DevitoPRO's ABox to restrict the computation + to the region around the sources and receivers. The box straddles the + rank boundaries, so the CORE and OWNED regions of the `overlap`, + `overlap2` and `full` MPI modes must honor its runtime thicknesses. + """ + tkns = (3, 11, 9, 2) + + class Box(SubDomainSet): + name = 'box' + + class Reference(SubDomain): + name = 'reference' + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', tkns[0], tkns[1]), + y: ('middle', tkns[2], tkns[3])} + + grid = Grid(shape=(40, 40), extent=(39., 39.)) + bounds = tuple(np.array([i], dtype=np.int32) for i in tkns) + box = Box(N=1, bounds=bounds, grid=grid) + reference = Reference(grid=grid) + + u0 = TimeFunction(name='u0', grid=grid, time_order=2, space_order=8) + u1 = TimeFunction(name='u1', grid=grid, time_order=2, space_order=8) + for u in (u0, u1): + u.data[:, 20, 20] = 1. + + op0 = Operator(Eq(u0.forward, solve(u0.dt2 - u0.laplace, u0.forward), + subdomain=reference)) + op1 = Operator(Eq(u1.forward, solve(u1.dt2 - u1.laplace, u1.forward), + subdomain=box)) + op0(time_M=20, dt=0.3) + op1(time_M=20, dt=0.3) + + assert np.linalg.norm(u0.data[:]) > 0 + assert np.allclose(u0.data[:], u1.data[:], rtol=0, atol=1e-6) + def test_multi_sets_eq(self): """ Check functionality for when multiple subdomain sets are present, each