From 3fd154fef6ebd43a86a03b2a9f33052bc18d2626 Mon Sep 17 00:00:00 2001 From: Gerard Gorman Date: Thu, 29 Oct 2026 21:40:40 +0000 Subject: [PATCH 1/2] compiler: Derive HaloSpot Scopes lazily in optimize_halospots `_merge_halospots` and `_hoist_invariant` derived a Scope, i.e. a full `FindWithin` visit of the enclosing Iteration, for every HaloSpot pair before checking whether the pair could merge or hoist at all. Since 0038c8cdd, an Operator with many Eqs over distinct SubDomains carries one HaloSpot per Eq, so the number of pairs, and Scope derivations, grows quadratically. Derive the Scope only when a pair passes the cheap compatibility checks; all merge and hoist decisions are unchanged. Claude-Session: https://claude.ai/code/session_01TT4Kt2oZNzhQrSw8SfVoaw --- devito/passes/iet/mpi.py | 40 +++++++++++++++++++++++----------- tests/test_mpi.py | 47 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 13 deletions(-) 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..8729b2afbc 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,51 @@ 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 + 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) + f = Function(name='f', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) + + 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, (i + 1)*u[t + 1, x, y + 1] + u[t + 1, x - 1, y], + subdomain=r, implicit_dims=t) for i, r in enumerate(rows)] + 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, i.e. + # those reading the same Function, require a Scope + assert derive_scope.call_count <= 2*n + class TestOperatorAdvanced: From 2ccfa7068216745e6a607e68a93f36d916baaf1c Mon Sep 17 00:00:00 2001 From: Gerard Gorman Date: Thu, 29 Oct 2026 22:13:53 +0000 Subject: [PATCH 2/2] tests: Cover the hoisting pass in the many-SubDomains HaloSpot test Add Eqs whose HaloSpots exchange distinct Functions, so that a Scope derived eagerly in `_hoist_invariant` alone also breaks the bound. Claude-Session: https://claude.ai/code/session_01TT4Kt2oZNzhQrSw8SfVoaw --- tests/test_mpi.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/test_mpi.py b/tests/test_mpi.py index 8729b2afbc..cc9e180a58 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -2458,8 +2458,8 @@ def define(self, dimensions): def test_merge_many_subdomains_linear(self): """ Each Eq over its own row SubDomain carries its own HaloSpot. Merging - them must not derive a Scope for each of the quadratically many - HaloSpot pairs, most of which are incompatible. + and hoisting them must not derive a Scope for each of the quadratically + many HaloSpot pairs, most of which are incompatible. """ class Row(SubDomain): @@ -2481,14 +2481,17 @@ def define(self, dimensions): 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, (i + 1)*u[t + 1, x, y + 1] + u[t + 1, x - 1, y], - subdomain=r, implicit_dims=t) for i, r in enumerate(rows)] + 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', @@ -2496,9 +2499,9 @@ def define(self, dimensions): 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, i.e. - # those reading the same Function, require a Scope - assert derive_scope.call_count <= 2*n + # 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: