diff --git a/devito/passes/iet/mpi.py b/devito/passes/iet/mpi.py index db835400ae..5d79e6ca3d 100644 --- a/devito/passes/iet/mpi.py +++ b/devito/passes/iet/mpi.py @@ -1,4 +1,5 @@ from collections import defaultdict +from functools import cache, partial from itertools import combinations from sympy import S @@ -147,12 +148,13 @@ def _merge_halospots(iet): if _check_control_flow(hs0, hs1, cond_mapper): continue - scope = _derive_scope(it, hs0, hs1) + scope = _lazy_scope(it, hs0, hs1) for f in hs1.fmapper: hsf0 = mapper.get(hs0).halo_scheme hsf1 = mapper.get(hs1).halo_scheme.project(f) - if not _is_mergeable(hsf0, hsf1, scope): + if not _is_compatible(hsf0, hsf1) or \ + not _is_iter_carried(hsf1, scope()): continue # All good -- `hsf1` can be merged within `hs0` @@ -197,7 +199,7 @@ def _hoist_invariant(iet): if _check_control_flow(hs0, hs1, cond_mapper): continue - scope = _derive_scope(it, hs0, hs1) + scope = _lazy_scope(it, hs0, hs1) for f in hs1.fmapper: hsf0 = hs0.halo_scheme.project(f) @@ -208,7 +210,7 @@ def _hoist_invariant(iet): # Ensure there's another HaloScheme that could cover for # us should we get hoisted while still satisfying the # data dependences - if hsf1.issubset(hsf0) and _is_iter_carried(hsf1, scope): + if hsf1.issubset(hsf0) and _is_iter_carried(hsf1, scope()): hs, hsf = hs1, hsf1 elif hsf0.issubset(hsf1) and hs0 is halo_spots[0]: # Special case @@ -476,6 +478,16 @@ def _derive_scope(it, hs0, hs1): return Scope(e.expr for e in expressions) +def _lazy_scope(it, hs0, hs1): + """ + As `_derive_scope`, but the Scope is only derived upon the first call, + and then reused. Deriving a Scope requires a visit of `it`, which is + wasted on the many HaloSpot pairs that turn out incompatible regardless + of the data dependences. + """ + return cache(partial(_derive_scope, it, hs0, hs1)) + + def _check_control_flow(hs0, hs1, cond_mapper): """ If there are Conditionals involved, both `hs0` and `hs1` must be @@ -513,22 +525,24 @@ def rule1(dep, loc_indices): return True -def _is_mergeable(hsf0, hsf1, scope): +def _is_compatible(hsf0, hsf1): """ - True if `hsf1` can be merged into `hsf0`, i.e., if they are compatible - and the data dependences would be satisfied, False otherwise. + True if `hsf1` is non-empty and compatible with `hsf0`, False otherwise. """ # If `hsf1` is empty there's nothing to merge if hsf1.is_void: return False - # Ensure `hsf0` and `hsf1` are compatible - if not hsf0.dimensions.issubset(hsf1.dimensions) or \ - not hsf0.functions & hsf1.functions: - return False + return (hsf0.dimensions.issubset(hsf1.dimensions) and + bool(hsf0.functions & hsf1.functions)) + - # Finally, check the data dependences would be satisfied - return _is_iter_carried(hsf1, scope) +def _is_mergeable(hsf0, hsf1, scope): + """ + True if `hsf1` can be merged into `hsf0`, i.e., if they are compatible + and the data dependences would be satisfied, False otherwise. + """ + return _is_compatible(hsf0, hsf1) and _is_iter_carried(hsf1, scope) def _semantical_eq_loc_indices(hsf0, hsf1): diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 7081be3beb..cc9e180a58 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -1,4 +1,5 @@ from functools import cached_property +from unittest.mock import patch import numpy as np import pytest @@ -23,6 +24,7 @@ from devito.mpi import MPI from devito.mpi.distributed import CustomTopology from devito.mpi.routines import ComputeCall, HaloUpdateCall, HaloUpdateList, MPICall +from devito.passes.iet import mpi as mpi_passes from devito.tools import Bunch from devito.types import Bundle from devito.types.dimension import ModuloDimension @@ -2453,6 +2455,54 @@ def define(self, dimensions): # to Function on SubDomain check_halo_exchanges(op, 0, 0) + def test_merge_many_subdomains_linear(self): + """ + Each Eq over its own row SubDomain carries its own HaloSpot. Merging + and hoisting them must not derive a Scope for each of the quadratically + many HaloSpot pairs, most of which are incompatible. + """ + + class Row(SubDomain): + + def __init__(self, i, n, **kwargs): + self.i, self.n = i, n + self.name = f'row{i}' + super().__init__(**kwargs) + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), y: ('middle', self.i, self.n - 1 - self.i)} + + n = 8 + grid = Grid(shape=(8, n)) + rows = [Row(i, n, grid=grid) for i in range(n)] + x, y = grid.dimensions + t = grid.stepping_dim + + u = TimeFunction(name='u', grid=grid, space_order=2) + v = TimeFunction(name='v', grid=grid, space_order=2) + ws = [TimeFunction(name=f'w{i}', grid=grid, space_order=2) for i in range(n)] + f = Function(name='f', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) + + # The HaloSpots of the first `n` Eqs share `v`, hence they may merge, + # while those of the last `n` Eqs are over distinct Functions + eqns = [Eq(f, (i + 1)*v[t, x, y + 1] + v[t, x + 1, y], + subdomain=r, implicit_dims=t) for i, r in enumerate(rows)] + eqns += [Eq(u.forward, u + f)] + eqns += [Eq(g, w[t, x, y + 1] + w[t, x - 1, y] + u[t + 1, x, y], + subdomain=r, implicit_dims=t) for w, r in zip(ws, rows, strict=True)] + eqns += [Eq(v.forward, v + g)] + + with patch('devito.passes.iet.mpi._derive_scope', + wraps=mpi_passes._derive_scope) as derive_scope: + Operator(eqns) + + # One HaloSpot per Eq over a SubDomain, i.e. `2*n`, whose pairs are + # quadratically many. Only the pairs that may actually merge or hoist, + # i.e. those exchanging the same Function, require a Scope + assert derive_scope.call_count < 2*n + class TestOperatorAdvanced: