From ab62c888dac9ae307d409e273f3259fff26faec6 Mon Sep 17 00:00:00 2001 From: Gerard Gorman Date: Mon, 26 Oct 2026 17:46:10 +0000 Subject: [PATCH] mpi: Honor MultiSubDomain thicknesses in overlap modes Under DEVITO_MPI=overlap/overlap2/full an equation over a SubDomainSet (or DevitoPRO's ABox) skipped points near the box edges: the HaloScheme only honored SubDimension thicknesses, so a MultiSubDimension's thickness was applied on top of the CORE/OWNED bounds rather than intersected with them. Honor the thicknesses of every non-local AbstractSubDimension. The overlap2 regions now describe the OWNED split of the root Dimensions only, and the remainder intersects them with the iteration bounds, as these may only be known within the Operator. --- devito/mpi/halo_scheme.py | 2 +- devito/mpi/routines.py | 49 +++++++++++++++++++-------------------- devito/types/dimension.py | 3 +++ tests/test_subdomains.py | 30 ++++++++++++++++++++++++ 4 files changed, 58 insertions(+), 26 deletions(-) diff --git a/devito/mpi/halo_scheme.py b/devito/mpi/halo_scheme.py index 3b4bcb5bc80..ca56d89c866 100644 --- a/devito/mpi/halo_scheme.py +++ b/devito/mpi/halo_scheme.py @@ -142,7 +142,7 @@ def __init__(self, exprs, ispace): self._honored = {} 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_AbstractSub and not d.local] for i in subdims: ltk, rtk = i.tkns self._honored[i.root] = frozenset([(ltk, rtk)]) diff --git a/devito/mpi/routines.py b/devito/mpi/routines.py index 3c446b24ba4..42051091a51 100644 --- a/devito/mpi/routines.py +++ b/devito/mpi/routines.py @@ -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,10 @@ class Overlap2HaloExchangeBuilder(OverlapHaloExchangeBuilder): """ def _make_region(self, hs, key): - return MPIRegion('reg', key, hs.arguments, hs.omapper.owned) + # Only the OWNED split of the root Dimensions, as the SubDimension + # thicknesses are applied in the remainder + owned = hs.halo_scheme._rebuild(honored={}).omapper.owned + return MPIRegion('reg', key, hs.dimensions, owned) def _make_msg(self, f, hse, key): # Only retain the halos required by the Diag scheme @@ -912,12 +915,17 @@ def _make_remainder(self, hs, key, callcompute, region): region_i = IndexedPointer(region, dim) 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)) - else: - dynamic_args_mapper[i] = (FieldFromComposite(i.name, region_i),) + for d in hs.dimensions: + m = FieldFromComposite(d.min_name, region_i) + M = FieldFromComposite(d.max_name, region_i) + dynamic_args_mapper[d] = (m, M) + + # Intersect the region with the SubDimension bounds, which may only + # be known within the Operator (e.g., MultiSubDimension thicknesses) + for ltkn, rtkn in hs.halo_scheme.honored.get(d, []): + ltkn1 = DefFunction('MAX', [0, d.symbolic_min + ltkn - m]) + rtkn1 = DefFunction('MAX', [0, M - d.symbolic_max + rtkn]) + dynamic_args_mapper.update({ltkn: (ltkn1,), rtkn: (rtkn1,)}) iet = callcompute._rebuild(dynamic_args_mapper=dynamic_args_mapper) # The -1 below is because an Iteration, by default, generates <= @@ -1348,12 +1356,9 @@ def __init__(self, prefix, key, arguments, owned): pname = f'region{key}' fields = [] - for i in self.arguments: - if i.is_Dimension: - fields.append((i.min_name, c_int)) - fields.append((i.max_name, c_int)) - else: - fields.append((i.name, c_int)) + for d in self.arguments: + fields.append((d.min_name, c_int)) + fields.append((d.max_name, c_int)) super().__init__(name, pname, fields) @@ -1387,16 +1392,10 @@ def _arg_values(self, args=None, **kwargs): values = self._arg_defaults() for i, (_, mapper) in enumerate(self.owned): entry = values[self.name][i] - for a in self.arguments: - if a.is_Dimension: - a_m, a_M = mapper[a] - setattr(entry, a.min_name, subs_op_args(a_m, args)) - setattr(entry, a.max_name, subs_op_args(a_M, args)) - else: - try: - setattr(entry, a.name, subs_op_args(mapper[a][0], args)) - except AttributeError: - setattr(entry, a.name, mapper[a][0]) + for d in self.arguments: + d_m, d_M = mapper[d] + setattr(entry, d.min_name, subs_op_args(d_m, args)) + setattr(entry, d.max_name, subs_op_args(d_M, args)) return values diff --git a/devito/types/dimension.py b/devito/types/dimension.py index c81eb90d37d..02e58df77c1 100644 --- a/devito/types/dimension.py +++ b/devito/types/dimension.py @@ -639,6 +639,9 @@ class AbstractSubDimension(DerivedDimension): is_AbstractSub = True + # Whether the iteration fits entirely within a single MPI domain + local = False + __rargs__ = DerivedDimension.__rargs__ + ('thickness',) __rkwargs__ = ('separated',) diff --git a/tests/test_subdomains.py b/tests/test_subdomains.py index eff67586367..1a4886e5837 100644 --- a/tests/test_subdomains.py +++ b/tests/test_subdomains.py @@ -932,6 +932,36 @@ def test_subdomainset_w_sparse(self): check[0, 0] = 1 assert np.all(np.isclose(f.data, check)) + @pytest.mark.parallel(mode=[(4, 'basic'), (4, 'diag2'), (4, 'overlap'), + (4, 'overlap2'), (4, 'full')]) + def test_subdomainset_stencil_mpi(self, mode): + """ + A stencil update restricted to a set of boxes straddling the rank + boundaries, as DevitoPRO's ABox does around the sources and receivers + (devitopro/subdomains/abox.py). Under comp/comm overlap, the CORE and + OWNED regions must be intersected with each box. + """ + class Boxes(SubDomainSet): + name = 'boxes' + + grid = Grid(shape=(12, 12)) + bounds = tuple(np.array(i, dtype=np.int32) + for i in ([2, 7], [3, 1], [1, 6], [4, 1])) + boxes = Boxes(N=2, bounds=bounds, grid=grid) + + u = TimeFunction(name='u', grid=grid, space_order=2) + + op = Operator(Eq(u.forward, u.laplace + 1, subdomain=boxes)) + op.apply(time_M=0) + + check = np.zeros((12, 12)) + check[2:9, 1:8] = 1 + check[7:11, 6:11] = 1 + + data = u.data_gather() + if grid.distributor.myrank == 0: + assert np.all(data[1] == check) + class TestBorder: # Note: This class is partially covered by doctests