compiler: Tweak preprocess to avoid erroneous matching of subdomains when inserting halo exchanges

devito/ir/stree/algorithms.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from itertools import groupby
 
 from anytree import findall
@@ -206,6 +207,16 @@ def preprocess(clusters, options=None, **kwargs):
                 diff = dims - distributed_aindices
                 intersection = dims & distributed_aindices
 
+                # TODO: Can check intersection length here and short-circuit early
+
+                # A non-empty `intersection` doesn't guarantee `c1` belongs to
+                # `c`: two SubDomains sharing an axis's (side, thickness) alias
+                # to the same cached SubDimension, so the overlap may be one
+                # incidental axis while the rest belongs to an unrelated
+                # SubDomain -- see `is_halo_scheme_conflicting`.
+                if is_halo_scheme_conflicting(dims, distributed_aindices):
+                    continue
+
                 if all(c1.guards.get(d) == c.guards.get(d) for d in diff) and \
                    len(intersection) > 0:
                     found.append(c1)
@@ -243,6 +254,25 @@ def preprocess(clusters, options=None, **kwargs):
     return processed
 
 
+def is_halo_scheme_conflicting(dims, distributed_aindices):
+    """
+    True if `distributed_aindices` cannot safely be attached to a Cluster
+    whose block-promoted itintervals are `dims`, because some Dimension in
+    `distributed_aindices` shares a root with a Dimension in `dims` without
+    being identical to it (a root missing from `dims` is still fine, e.g.
+    `t, f` triggering a halo for `t, x, y, z, f`). `dims` may hold several
+    Dimensions per root at once (e.g. a point Dimension plus a derived
+    radius CustomDimension), hence a set per root below.
+    """
+    d_by_root = defaultdict(set)
+    for d in dims:
+        d_by_root[d.root].add(d)
+    return any(
+        e.root in d_by_root and e not in d_by_root[e.root]
+        for e in distributed_aindices
+    )
+
+
 def reuse_partial_subtree(c0, c1, d=None):
     return c0.guards.get(d) == c1.guards.get(d)
 

tests/test_ir.py

@@ -4,15 +4,16 @@
 
 from conftest import EVAL, skipif  # noqa
 from devito import (  # noqa
-    Constant, Dimension, Eq, Function, Grid, Inc, Operator, SubDimension, TimeFunction,
-    switchconfig
+    Constant, CustomDimension, Dimension, Eq, Function, Grid, Inc, Operator,
+    SubDimension, TimeFunction, switchconfig
 )
 from devito.ir.cgen import ccode
 from devito.ir.clusters import Cluster, ClusterGroup
 from devito.ir.equations import LoweredEq
 from devito.ir.equations.algorithms import dimension_sort
 from devito.ir.iet import FindNodes, Iteration
 from devito.ir.stree import stree_build
+from devito.ir.stree.algorithms import is_halo_scheme_conflicting
 from devito.ir.support.basic import (
     AFFINE, IRREGULAR, REGULAR, IterationInstance, Scope, TimedAccess, Vector, mocksym0,
     mocksym1
@@ -1188,6 +1189,63 @@ def test_from_clusters_mixed_dtypes(self):
         assert len([i for i in stree.visit() if i.is_Iteration]) == 1
 
 
+class TestSubdomainHaloMatching:
+
+    """
+    Tests for `devito.ir.stree.algorithms.is_halo_scheme_conflicting`, the
+    logic that decides whether a "wild" HaloTouch Cluster's HaloScheme may be
+    attached to a given computational Cluster during `preprocess()`.
+    """
+
+    def test_cross_subdomain_aliasing_conflict(self):
+        """
+        Two different SubDomains sharing one axis's (side, thickness) alias
+        to the same cached SubDimension on that axis; a mismatch on another
+        shared-root axis must still be flagged as conflicting.
+        """
+        x, y = Dimension('x'), Dimension('y')
+        ix_shared = SubDimension.left('ix', x, 3)
+        iy_a = SubDimension.right('iy', y, 3)  # domain A's own y
+        iy_b = SubDimension.left('iy', y, 3)  # domain B's y -- different!
+
+        dims = {ix_shared, iy_a}
+        distributed_aindices = {ix_shared, iy_b}
+
+        assert is_halo_scheme_conflicting(dims, distributed_aindices)
+
+    def test_bare_dimension_vs_subdimension_conflict(self):
+        """
+        A bare (unrestricted) Dimension in `dims` colliding with an
+        unrelated SubDomain's SubDimension for the same root.
+        """
+        x = Dimension('x')
+        xi = SubDimension.left('xi', x, 3)
+
+        assert is_halo_scheme_conflicting({x}, {xi})
+
+    def test_no_conflict_when_root_missing_from_dims(self):
+        """
+        A root present in `distributed_aindices` but absent from `dims`
+        entirely (e.g. an outer `t, f` Cluster triggering a halo for an
+        inner `t, x, y, z, f` Cluster) is not a conflict.
+        """
+        x, f = Dimension('x'), Dimension('f')
+
+        assert not is_halo_scheme_conflicting({x}, {x, f})
+
+    def test_multiple_dimensions_per_root_not_collapsed(self):
+        """
+        `dims` may legitimately hold several Dimensions sharing one root at
+        once (e.g. a sparse-function point Dimension plus a CustomDimension
+        derived from it for an interpolation radius); this must not be
+        collapsed to a single arbitrary match.
+        """
+        p = Dimension('p')
+        r = CustomDimension(name='r', symbolic_min=0, symbolic_max=1, parent=p)
+
+        assert not is_halo_scheme_conflicting({p, r}, {p})
+
+
 class TestClusterGroup:
 
     def test_eq_hash_include_ispace(self):

tests/test_mpi.py

@@ -14,7 +14,8 @@
 from devito.arch.compiler import OneapiCompiler
 from devito.data import LEFT, RIGHT
 from devito.ir.iet import (
-    Call, Conditional, FindNodes, FindSymbols, Iteration, retrieve_iteration_tree
+    Call, Conditional, FindNodes, FindSymbols, HaloSpot, Iteration,
+    retrieve_iteration_tree
 )
 from devito.ir.support.space import Backward, Forward
 from devito.mpi import MPI
@@ -1107,6 +1108,117 @@ def check_halo_exchanges(op, exp0, exp1):
     return calls, tloop
 
 
+def check_cpml_no_misplaced_halo(specs):
+    """
+    Build an Operator modelled on a CPML-style absorbing-boundary formulation:
+    several overlapping-thickness SubDomains, each carrying a pair of auxiliary
+    VectorTimeFunctions updated via a "diagonal" derivative pattern (component
+    `i` derived along dimension `i`). This is the minimal pattern that exposed
+    a bug in `devito.ir.stree.algorithms.preprocess`, whereby the HaloScheme of
+    one SubDomain could be erroneously attached to a *different* SubDomain's
+    Cluster whenever the two happened to share a single axis's SubDimension
+    (SubDimensions are cached by `(name, parent, thickness, local)`, so two
+    unrelated SubDomains restricting one axis identically alias to the same
+    object).
+
+    The Operator is built and run to completion; the values are not physically
+    meaningful and are not checked. The regression check is purely structural:
+    no halo-exchange node may end up nested *inside* a blocking (`is_Incr`)
+    Iteration.
+    """
+    class SpeccedDomain(SubDomain):
+        def __init__(self, name, thickness, spec, **kwargs):
+            self.name = name
+            self.thickness = thickness
+            self.spec = spec
+            super().__init__(**kwargs)
+
+        def define(self, dimensions):
+            retval = {}
+            for d, s in zip(dimensions, self.spec, strict=True):
+                if s == 'none':
+                    # Unrestricted -- bare Dimension, not a SubDimension
+                    retval[d] = d
+                elif s == 'middle':
+                    retval[d] = (s, self.thickness, self.thickness)
+                else:
+                    retval[d] = (s, self.thickness)
+            return retval
+
+    so, to, nb = 4, 2, 3
+
+    grid = Grid(shape=(21, 21, 21), extent=(1., 1., 1.))
+    domains = {spec: SpeccedDomain(f"x{spec[0]}_y{spec[1]}_z{spec[2]}", nb, spec,
+                                   grid=grid)
+               for spec in specs}
+
+    p = TimeFunction(name='p', grid=grid, space_order=so, time_order=to)
+
+    psi_eqs, zeta_eqs, p_eqs = [], [], []
+    for v in domains.values():
+        # Fields live on the base grid, not `v` -- `subdomain=v` is passed
+        # explicitly on each Eq instead.
+        psi = VectorTimeFunction(name=f"psi_{v.name}", grid=grid, space_order=so,
+                                 time_order=to, staggered=(None, None, None))
+        zeta = VectorTimeFunction(name=f"zeta_{v.name}", grid=grid, space_order=so,
+                                  time_order=to, staggered=(None, None, None))
+        psi_eqs.append(Eq(psi, 1, subdomain=v))
+
+        # "Diagonal" derivative pattern -- component `i` of zeta is the
+        # derivative of component `i` of psi along dimension `i`. This
+        # specific pattern is required to trigger the bug; grad()/div()/a
+        # plain vector add alone do not.
+        zeta_diag = VectorTimeFunction([getattr(psi[i], f"d{d.name}")
+                                        for i, d in enumerate(grid.dimensions)])
+        zeta_eqs.append(Eq(zeta, zeta_diag, subdomain=v))
+        p_eqs.append(Eq(p.forward, psi.div(), subdomain=v))
+
+    op = Operator(psi_eqs + zeta_eqs + p_eqs,
+                  opt=('advanced', {'blockrelax': 'device-aware'}))
+    op.apply(time_M=1)
+
+    # No misplaced halo exchanges: every halo-exchange node must sit above
+    # (never inside) any blocking Iteration.
+    incr_iterations = [i for i in FindNodes(Iteration).visit(op) if i.dim.is_Incr]
+    assert incr_iterations, "no blocking Iterations found"
+
+    # HaloSpots have already been lowered into concrete Calls by mpiize() in the MPI
+    # case, but HaloSpots remain in the IET as transparent (no-op) wrappers in the
+    # serial case.
+    halo_types = (HaloUpdateCall, HaloUpdateList) if configuration['mpi'] else (HaloSpot,)
+
+    # Assert that HaloSpots etc are present, just not located within inner loops
+    assert len(FindNodes(halo_types).visit(op)) > 0
+
+    for i in incr_iterations:
+        assert len(FindNodes(halo_types).visit(i)) == 0
+
+
+CPML_SPECS = [
+    pytest.param(
+        [('left', 'right', 'left'),
+         ('left', 'right', 'middle'),
+         ('left', 'right', 'right'),
+         ('middle', 'left', 'right')],
+        id='full-restriction'
+    ),
+    pytest.param(
+        [('left', 'right', 'left'),
+         ('none', 'none', 'middle'),
+         ('none', 'none', 'right'),
+         ('middle', 'left', 'right')],
+        id='partial-none-xy'
+    ),
+    pytest.param(
+        [('left', 'right', 'left'),
+         ('left', 'right', 'none'),
+         ('left', 'right', 'right'),
+         ('middle', 'left', 'right')],
+        id='partial-none-z'
+    ),
+]
+
+
 class TestCodeGeneration:
 
     @pytest.mark.parallel(mode=1)
@@ -1670,6 +1782,15 @@ def test_process_but_avoid_haloupdate_along_replicated(self, mode):
             assert len(calls) == 1
             assert calls[0].arguments[0] is u
 
+    @pytest.mark.parametrize('specs', CPML_SPECS)
+    def test_cpml_no_misplaced_halo(self, specs):
+        check_cpml_no_misplaced_halo(specs)
+
+    @pytest.mark.parametrize('specs', CPML_SPECS)
+    @pytest.mark.parallel(mode=1)
+    def test_cpml_no_misplaced_halo_mpi(self, specs, mode):
+        check_cpml_no_misplaced_halo(specs)
+
     @pytest.mark.parallel(mode=1)
     def test_conditional_dimension(self, mode):
         """