compiler: Improve SubDimension DDA

Author: FabioLuporini

devito/ir/support/basic.py

@@ -1610,21 +1610,20 @@ def disjoint_subdims(a0, a1):
     of the same Function.
 
     Compare symbolic accessed bounds, including shifts and stencil points.
-    Block intervals are promoted to their logical SubDimensions. Bounds and
-    thicknesses remain symbolic: MPI decomposition and runtime overrides can
-    change their values independently of the defaults.
+    Block intervals are promoted to their logical SubDimensions. Declared
+    thicknesses determine the global regions: explicit overrides are forbidden,
+    while MPI clips these regions to each rank. Parent bounds and access offsets
+    remain symbolic; only iteration bounds use the declared thicknesses.
 
-    For example, `xl = [m, m + L - 1]` and `xm = [m + L, M - R]` are
-    disjoint when they share the symbol `L`, whatever its runtime value.
-    Equal default thicknesses alone do not establish that relationship.
+    For example, a left slab of thickness 4 ends before a middle excluding 4
+    points, even when the two thickness symbols are distinct.
 
     Opposite left/right slabs are assumed to form a valid partition: their
     thicknesses satisfy `L + R <= N`. For translated stencil accesses, the
     interior must also accommodate their combined inward reach. For example,
     a pointwise left write and a right read at offset -4 require four interior
     points. Runtime space_order checks cover explicit middle SubDimensions,
-    not arbitrary left/right pairs; no concrete domain size or thickness is
-    used here.
+    not arbitrary left/right pairs; no concrete domain size is used here.
 
     Match data axes independently of the iteration nests. Return True if any
     axis proves separation, False otherwise. Accesses over the same interval
@@ -1644,6 +1643,8 @@ def disjoint_subdims(a0, a1):
                 it0 != it1):
             continue
 
+        thicknesses = {t: t.value for it in (it0, it1)
+                       for t in it.dim.thickness if t.value is not None}
         bounds = []
         for e, d, it in ((e0, d0, it0), (e1, d1, it1)):
             if not q_affine(e, d):
@@ -1658,8 +1659,8 @@ def disjoint_subdims(a0, a1):
                     M, m = it.symbolic_min, it.symbolic_max
                 else:
                     break
-                lower.append(v._subs(d, m))
-                upper.append(v._subs(d, M))
+                lower.append(v._subs(d, m.xreplace(thicknesses)))
+                upper.append(v._subs(d, M.xreplace(thicknesses)))
             else:
                 bounds.append((sympy.Min(*lower), sympy.Max(*upper)))
 
@@ -1682,11 +1683,13 @@ def disjoint_subdims(a0, a1):
                 gap = sympy.Dummy(nonnegative=True)
                 if e0.diff(d0) == e1.diff(d1) == 1:
                     M, m = (M0, m1) if it0.dim.is_left else (M1, m0)
-                    reach = (M - dl.symbolic_max - m + dr.symbolic_min).expand()
+                    reach = (M - dl.symbolic_max.xreplace(thicknesses) - m +
+                             dr.symbolic_min.xreplace(thicknesses)).expand()
                     if is_integer(reach):
                         gap += max(0, reach)
 
-                mapper[dlp.symbolic_max] = dlp.symbolic_min + dl.ltkn + dr.rtkn + gap - 1
+                mapper[dlp.symbolic_max] = (dlp.symbolic_min + dl.ltkn.value +
+                                            dr.rtkn.value + gap - 1)
 
             if (M0 - m1).subs(mapper).is_negative or \
                (M1 - m0).subs(mapper).is_negative:

devito/types/dimension.py

@@ -554,7 +554,12 @@ def _arg_check(self, *args, **kwargs):
 # the user
 
 class Thickness(DataSymbol):
-    """A DataSymbol to represent a thickness of a SubDimension"""
+    """
+    A SubDimension thickness, fixed at construction and localized by MPI.
+
+    Explicit runtime overrides are not supported: dependence analysis uses the
+    declared thickness to determine the global region before MPI decomposition.
+    """
 
     __rkwargs__ = DataSymbol.__rkwargs__ + ('root', 'side', 'local', 'value')
 
@@ -589,11 +594,7 @@ def value(self):
         return self._value
 
     def _arg_values(self, grid=None, **kwargs):
-        # Allow override of thickness values to disable BCs
-        # However, arguments from the user are considered global
-        # So overriding the thickness to a nonzero value should not cause
-        # boundaries to exist between ranks where they did not before
-        rtkn = kwargs.get(self.name, self.value)
+        rtkn = self.value
         if grid is not None and grid.is_distributed(self.root):
             # Get local thickness
             if self.local:
@@ -612,6 +613,15 @@ def _arg_values(self, grid=None, **kwargs):
 
         return {self.name: tkn}
 
+    def _arg_check(self, args, *_args, **kwargs):
+        # This module depends on Dimension, so importing it above would cycle
+        from devito.mpi import mpi_raise  # noqa: PLC0415
+
+        error = (f"Cannot override SubDimension thickness `{self.name}`"
+                 if self.name in kwargs else None)
+        comm = args.comm if args.options['mpi'] else None
+        mpi_raise(error, InvalidArgument, comm=comm)
+
 
 class AbstractSubDimension(DerivedDimension):
 
@@ -840,7 +850,7 @@ def _arg_check(self, args, *_args, **kwargs):
             values = {**args,
                       d.min_name: kwargs.get(d.min_name, 0),
                       d.max_name: kwargs.get(d.max_name, kwargs.get(d.name, size - 1)),
-                      **{t.name: kwargs.get(t.name, t.value) for t in self.thickness}}
+                      **{t.name: t.value for t in self.thickness}}
         else:
             values = args
         size = int(subs_op_args(self.symbolic_size, values))

tests/test_dimension.py

@@ -794,17 +794,18 @@ def test_expandingbox_like(self, opt):
         grid = Grid(shape=(8, 8))
         x, y = grid.dimensions
 
+        # Declare the widest box; runtime root bounds control the active box
+        xi = SubDimension.middle(name='xi', parent=x, thickness_left=0, thickness_right=0)
+        yi = SubDimension.middle(name='yi', parent=y, thickness_left=0, thickness_right=0)
+
         u = TimeFunction(name='u', grid=grid)
-        xi = SubDimension.middle(name='xi', parent=x, thickness_left=2, thickness_right=2)
-        yi = SubDimension.middle(name='yi', parent=y, thickness_left=2, thickness_right=2)
 
         eqn = Eq(u.forward, u + 1)
         eqn = eqn.subs({x: xi, y: yi})
 
         op = Operator(eqn, opt=opt)
 
-        op.apply(time=3, x_m=2, x_M=5, y_m=2, y_M=5,
-                 x_ltkn0=0, x_rtkn0=0, y_ltkn0=0, y_rtkn0=0)
+        op.apply(time=3, x_m=2, x_M=5, y_m=2, y_M=5)
 
         assert np.all(u.data[0, 2:-2, 2:-2] == 4.)
         assert np.all(u.data[1, 2:-2, 2:-2] == 3.)

tests/test_ir.py

@@ -562,8 +562,6 @@ def test_subdimension_affine_bounds(self, shared_boundary, symbolic, slope, offs
         xm = SubDimension.middle('xm', x, 8, 40)
         if shared_boundary:
             xm = xm._rebuild(thickness=(xl.ltkn, xm.rtkn))
-        else:
-            expected = S.Infinity
         if symbolic:
             f = Array(name='f', dimensions=(x,))
         else:
@@ -1305,8 +1303,8 @@ def test_bundle_components(self):
 
     @pytest.mark.parametrize('symbolic', [False, True])
     @pytest.mark.parametrize('lower,upper,ndeps', [
-        (0, 31, (8, 32)), (0, 7, (0, 8)), (16, 31, (0, 16)),
-        (7, 8, (1, 2)), (15, 16, (1, 2))
+        (0, 31, (8, 24)), (0, 7, (0, 0)), (16, 31, (0, 16)),
+        (7, 8, (1, 1)), (15, 16, (1, 2))
     ])
     def test_stencil_contains_producer(self, symbolic, lower, upper, ndeps):
         grid = Grid(shape=(32,))
@@ -1317,6 +1315,8 @@ def test_stencil_contains_producer(self, symbolic, lower, upper, ndeps):
         f = Function(name='f', grid=grid)
         g = Function(name='g', grid=grid)
         if symbolic:
+            # Thickness 8 fixes the left boundary; the parent upper bound remains
+            # symbolic, so stencil offsets >= 8 can still touch the producer
             pi, ci = Interval(xm), Interval(xl)
         else:
             # Encode actual fixed iteration bounds, not runtime defaults

tests/test_subdomains.py

@@ -1937,14 +1937,16 @@ def define(self, dimensions):
 class TestSubDomainArguments:
 
     @staticmethod
-    def _make_operator(left_shift=0, right_shift=0, grid=None, middle=False):
+    def _make_operator(left_shift=0, right_shift=0, grid=None, middle=False,
+                       thickness=(8, 8)):
         grid = grid or Grid(shape=(8, 32))
         y = grid.dimensions[-1]
 
-        yl = SubDimension.left('yl', y, 8)
-        yr = SubDimension.right('yr', y, 8)
+        left, right = thickness
+        yl = SubDimension.left('yl', y, left)
+        yr = SubDimension.right('yr', y, right)
         if middle:
-            yl = yr = SubDimension.middle('ym', y, 8, 8)
+            yl = yr = SubDimension.middle('ym', y, left, right)
 
         u = TimeFunction(name='u', grid=grid, space_order=8)
         v = TimeFunction(name='v', grid=grid, space_order=8)
@@ -1979,7 +1981,11 @@ def test_stencil_gap(self, middle, left_shift, right_shift, margin, override):
             dl, = [d for d in op.dimensions if d.is_Sub and not d.is_right]
             kwargs = {dl.ltkn.name: 32 - 8 - 8 - margin}
 
-        if middle and margin < 0:
+        if override == 'thickness':
+            with pytest.raises(InvalidArgument,
+                               match='Cannot override SubDimension thickness'):
+                op.arguments(time_M=0, **kwargs)
+        elif middle and margin < 0:
             with pytest.raises(InvalidArgument, match='at least 8 interior points'):
                 op.arguments(time_M=0, **kwargs)
         else:
@@ -2007,13 +2013,11 @@ def test_runtime_space_order(self, middle, space_order, margin):
 
     @pytest.mark.parametrize('side', ['left', 'right'])
     def test_empty_slab(self, side):
-        op, _ = self._make_operator(right_shift=-4)
-
-        d, = [d for d in op.dimensions if d.is_Sub and getattr(d, f'is_{side}')]
-        thickness = d.ltkn if side == 'left' else d.rtkn
+        thickness = (0, 8) if side == 'left' else (8, 0)
+        op, _ = self._make_operator(right_shift=-4, thickness=thickness)
 
         # The active slab fills the local domain; the opposite slab is absent
-        op.arguments(time_M=0, y_M=7, **{thickness.name: 0})
+        op.arguments(time_M=0, y_M=7)
 
     def test_before_autotuning(self):
         op, _ = self._make_operator(right_shift=-4, middle=True)
@@ -2032,32 +2036,27 @@ def test_before_autotuning(self):
     def test_distributed_middle(self, left, mode):
         grid = Grid(shape=(16, 32), topology=(1, 2))
 
-        op, _ = self._make_operator(grid=grid, middle=True)
-
-        d, = [d for d in op.dimensions if d.is_Sub]
+        op, _ = self._make_operator(grid=grid, middle=True, thickness=(left, 0))
 
         # Rank 0 has one point or an empty middle (possibly with inverted bounds).
         # Only the global size determines whether the middle is large enough
-        kwargs = {d.ltkn.name: left, d.rtkn.name: 0}
-
         if left == 25:
             with pytest.raises(InvalidArgument, match='at least 8 interior points'):
-                op.arguments(time_M=0, **kwargs)
+                op.arguments(time_M=0)
         else:
-            op.arguments(time_M=0, **kwargs)
+            op.arguments(time_M=0)
 
     @pytest.mark.parallel(mode=[(2, 'basic')])
     def test_collective_rejection(self, mode):
         grid = Grid(shape=(32, 32), topology=(2, 1))
 
         op, _ = self._make_operator(right_shift=-4, grid=grid, middle=True)
 
-        dl, = [d for d in op.dimensions if d.is_Sub and d.is_middle]
-        left = 24 if grid.distributor.myrank == 0 else 0
+        upper = 22 if grid.distributor.myrank == 0 else 31
 
-        # Only rank 0 has an insufficient middle; its peer has 24 interior points
+        # Only rank 0 has an insufficient middle; its peer has 16 interior points
         with pytest.raises(InvalidArgument, match='interior points'):
-            op.arguments(time_M=0, **{dl.ltkn.name: left})
+            op.arguments(time_M=0, y_M=upper)
 
     def test_function_on_subdomain(self):
         class Interior(SubDomain):
@@ -2086,4 +2085,52 @@ def define(self, dimensions):
         assert concrete.ltkn.name in args
 
         with pytest.raises(InvalidArgument, match='at least 8 interior points'):
+            op.arguments(y_M=22)
+
+        with pytest.raises(InvalidArgument,
+                           match='Cannot override SubDimension thickness'):
             op.arguments(**{concrete.ltkn.name: 17})
+
+    def test_thickness_overrides(self):
+        grid = Grid(shape=(32,))
+        x, = grid.dimensions
+        xi = SubDimension.left('xi', x, 4)
+
+        f = Function(name='f', grid=grid, space_order=0)
+
+        eq = Eq(f[xi], 1)
+
+        op = Operator(eq, name='thickness_overrides')
+
+        d, = [d for d in op.dimensions if d.is_Sub]
+
+        with pytest.raises(InvalidArgument,
+                           match='Cannot override SubDimension thickness'):
+            op.apply(**{d.ltkn.name: 2})
+
+        assert np.all(f.data == 0)
+
+    @pytest.mark.parallel(mode=[(2, 'basic')])
+    def test_collective_thickness_rejection(self, mode):
+        grid = Grid(shape=(32,))
+        x, = grid.dimensions
+        xi = SubDimension.middle('xi', x, 4, 4)
+
+        f = Function(name='f', grid=grid, space_order=0)
+
+        eq = Eq(f[xi], 1)
+
+        op = Operator(eq, name='collective_thickness_rejection')
+
+        d, = [d for d in op.dimensions if d.is_Sub]
+        t = d.ltkn
+
+        # MPI clips the declared thicknesses without any explicit overrides
+        args = op.arguments()
+        assert set(grid.distributor.comm.allgather(args[t.name])) == {0, 4}
+
+        # Only rank 0 supplies an override; all ranks must reject it
+        kwargs = {t.name: 4} if grid.distributor.myrank == 0 else {}
+        with pytest.raises(InvalidArgument,
+                           match='Cannot override SubDimension thickness'):
+            op.arguments(**kwargs)