compiler: Revamp SubDimension DDA

Author: FabioLuporini

devito/ir/support/basic.py

@@ -7,7 +7,7 @@
 from sympy import Expr, S
 
 from devito.ir.support.space import Backward, null_ispace
-from devito.ir.support.utils import AccessMode, extrema
+from devito.ir.support.utils import AccessMode, erange, extrema
 from devito.ir.support.vector import LabeledVector, Vector
 from devito.symbolics import (
     compare_ops, q_affine, q_comp_acc, q_constant, retrieve_indexed, search
@@ -358,6 +358,9 @@ def distance(self, other, logical=False):
             # E.g., `uv(x).x` and `uv(x).y` -- not a real dependence!
             return Vector(S.ImaginaryUnit)
 
+        if disjoint_subdims(self, other):
+            return Vector(S.ImaginaryUnit)
+
         ret = []
         for sit, oit in zip(self.itintervals, other.itintervals, strict=False):
             n = len(ret)
@@ -369,20 +372,14 @@ def distance(self, other, logical=False):
                 # E.g., `self=R<f,[x]>` and `self.itintervals=(x, i)`
                 break
 
-            # If over SubDimensions, check disjointness
-            test = disjoint_subdims(self[n], other[n], sai, oai, sit, oit)
-            if test == DISJOINT:
-                return Vector(S.ImaginaryUnit)
-            elif test == MAYBE_OVERLAP:
-                ret.append(S.Infinity)
-                continue
-
             try:
                 if not (sit == oit and sai.root is oai.root):
                     # E.g., `self=R<f,[x + 2]>` and `other=W<f,[i + 1]>`
                     # E.g., `self=R<f,[x]>`, `other=W<f,[x + 1]>`,
                     #       `self.itintervals=(x<0>,)`, `other.itintervals=(x<1>,)`
-                    return vinf(ret)
+                    # Keep looking: a later axis may prove disjointness
+                    ret.append(S.Infinity)
+                    continue
             except AttributeError:
                 # E.g., `self=R<f,[cy]>` and `self.itintervals=(y,)` => `sai=None`
                 pass
@@ -1152,7 +1149,29 @@ def reads_smart_gen(self, f):
         """
         Generate all read accesses to a given function.
 
-        StencilDimensions, if any, are replaced with their extrema.
+        StencilDimensions, if any, are replaced with:
+
+            * in presence of SubDimensions: the range of points they span;
+            * in all other cases: just their extrema, since it suffices to
+              capture all possible dependencies.
+
+        The reason SubDimensions must be treated specially -- with a full set
+        of TimedAccess objects getting generated -- is to handle the special
+        case of slabs thinner than the stencil’s reach. For example, consider
+        the following scenario:
+
+        * A SubDimension with just two points, 10 and 11;
+        * One equation writes `F[10]` and `F[11]`;
+        * Another equation runs over the same SubDimension reading the stencil
+          `F[x-4] ... F[x+4]`.
+
+        If we examine only the two extreme stencil offsets:
+
+        * `F[x-4]` reads points 6–7: no overlap.
+        * `F[x+4]` reads points 14–15: no overlap.
+
+        But interior offsets certainly overlap -- for instance, `F[x-1]` reads
+        9–10, which includes the producer’s point 10.
 
         Notes
         -----
@@ -1163,9 +1182,13 @@ def reads_smart_gen(self, f):
         be found. For example, a DiscreteFunction would never appear among
         the iteration symbols.
         """
+        uses_subdims = lambda i: any(d.is_Sub for d in i.ispace.dimensions)
+
         if isinstance(f, (Function, Temp, TempArray, TBArray)):
             for i in self.getreads(f):
-                for j in extrema(i.access):
+                expand = erange if uses_subdims(i) else extrema
+
+                for j in expand(i.access):
                     yield TimedAccess(j, i.mode, i.timestamp, i.ispace)
 
         else:
@@ -1581,90 +1604,95 @@ def skippable_interval(d, ispace, it):
     return d is None or (d in ispace and not d._defines & it.dim._defines)
 
 
-# Possible return values for `disjoint_subdims`
-INAPPLICABLE = 0
-DISJOINT = 1
-MAYBE_OVERLAP = 2
-
-
-def disjoint_subdims(e0, e1, d0, d1, it0, it1):
+def disjoint_subdims(a0, a1):
     """
-    Determine whether two accesses span distinct pieces of the same
-    SubDimension decomposition.
-
-    Consider a root Dimension `x` with bounds `x_m` and `x_M`. A valid
-    left/middle/right decomposition with thicknesses `L` and `R` is::
-
-        xl = [x_m,         x_m + L - 1]
-        xm = [x_m + L,     x_M - R]
-        xr = [x_M - R + 1, x_M]
-
-    These intervals are pairwise disjoint. Replacing `xl`, `xm`, or `xr`
-    with `x` in an affine access removes the choice of partition piece while
-    retaining the relative access. If two such normalized accesses have zero
-    distance, they apply the same affine map to disjoint intervals and therefore
-    cannot refer to the same data point. The apparent dependence is imaginary.
-
-    For example, `f[xl]` and `f[xm]` normalize to `f[x]` and `f[x]`;
-    they are independent. The same holds for `f[xl + 1]` and `f[xm + 1]`
-    when their iteration intervals have equal offsets. By contrast, `f[xl]`
-    and `f[xm - 1]` normalize to different accesses, and the latter may reach
-    into the left piece, so they must be treated conservatively.
-
-    This proof requires distinct pieces of the same root, compatible declared
-    thicknesses, affine accesses, and iteration intervals with equal offsets and
-    directions. Runtime bounds are assumed to preserve the declared partition.
-    Return DISJOINT if disjointness is proven, and MAYBE_OVERLAP if the
-    intervals are aligned SubDimensions but are not proven disjoint. In
-    particular, two declarations of the same left, right, or middle piece
-    overlap along this Dimension. MAYBE_OVERLAP lets the caller record an
-    infinite distance and inspect later Dimensions, which may still prove the
-    multidimensional accesses disjoint. Return INAPPLICABLE if this test does not
-    apply, so that the general distance analysis can classify the dependence.
+    Determine whether two TimedAccesses touch disjoint SubDimension regions
+    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.
+
+    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.
+
+    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.
+
+    Match data axes independently of the iteration nests. Return True if any
+    axis proves separation, False otherwise. Accesses over the same interval
+    use the general distance analysis.
     """
-    try:
-        # E.g., `f[xl]` over `(xl,)` and `f[xm]` over `(xm,)` need this
-        # special test, while accesses over the same `(xl,)` should use general
-        # distance analysis, so we can return immediately in such a case
-        if not (d0.is_Sub and
-                d1.is_Sub and
-                d0.root is d1.root and
-                it0.dim.root is d0.root and
-                it1.dim.root is d1.root and
+    for e0, e1, d0, d1 in zip(a0, a1, a0.aindices, a1.aindices, strict=False):
+        it0 = a0.intervals[d0]
+        it1 = a1.intervals[d1]
+        if it0.is_Null or it1.is_Null:
+            continue
+
+        it0 = it0.promote(lambda d: d.is_Incr)
+        it1 = it1.promote(lambda d: d.is_Incr)
+        if not (it0.dim.is_Sub and
+                it1.dim.is_Sub and
+                it0.dim.root is it1.dim.root and
                 it0 != it1):
-            return INAPPLICABLE
-    except AttributeError:
-        return INAPPLICABLE
-
-    if (d0.is_left and d1.is_middle) or \
-       (d0.is_middle and d1.is_left):
-        is_partition = d0.ltkn.value == d1.ltkn.value
-    elif (d0.is_middle and d1.is_right) or \
-         (d0.is_right and d1.is_middle):
-        is_partition = d0.rtkn.value == d1.rtkn.value
-    elif d0.is_left and d1.is_right:
-        is_partition = d0.ltkn.value is not None and d1.rtkn.value is not None
-    elif d0.is_right and d1.is_left:
-        is_partition = d0.rtkn.value is not None and d1.ltkn.value is not None
-    else:
-        is_partition = False
-
-    if not is_partition:
-        return MAYBE_OVERLAP
-
-    if not q_affine(e0, d0) or not q_affine(e1, d1):
-        return MAYBE_OVERLAP
+            continue
 
-    if it0.offsets != it1.offsets or it0.direction is not it1.direction:
-        return MAYBE_OVERLAP
+        bounds = []
+        for e, d, it in ((e0, d0, it0), (e1, d1, it1)):
+            if not q_affine(e, d):
+                break
 
-    e0 = e0._subs(d0, d0.root)
-    e1 = e1._subs(d1, d1.root)
+            lower, upper = [], []
+            for v in erange(e):
+                slope = v.diff(d)
+                if slope.is_nonnegative:
+                    m, M = it.symbolic_min, it.symbolic_max
+                elif slope.is_nonpositive:
+                    M, m = it.symbolic_min, it.symbolic_max
+                else:
+                    break
+                lower.append(v._subs(d, m))
+                upper.append(v._subs(d, M))
+            else:
+                bounds.append((sympy.Min(*lower), sympy.Max(*upper)))
+
+        if len(bounds) == 2:
+            (m0, M0), (m1, M1) = bounds
+            mapper = {}
+
+            dl, dr = (it0.dim, it1.dim) if it0.dim.is_left else (it1.dim, it0.dim)
+            dlp, drp = dl.parent, dr.parent
+
+            if dl.is_left and dr.is_right and dlp is drp:
+                # A valid partition satisfies L + R <= N, where N is the parent
+                # extent; an explicit middle SubDomain checks this at construction.
+                # Further, for stencils, we require that:
+                # `N - L - R >= the combined inward reach`
+                # so accesses from opposite slabs cannot meet. Explicit middle
+                # SubDimensions check for at least space_order interior points
+                # at *op.apply time*, accounting for runtime overrides.  Without
+                # an explicit middle, the gap assumption is unchecked
+                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()
+                    if is_integer(reach):
+                        gap += max(0, reach)
+
+                mapper[dlp.symbolic_max] = dlp.symbolic_min + dl.ltkn + dr.rtkn + gap - 1
+
+            if (M0 - m1).subs(mapper).is_negative or \
+               (M1 - m0).subs(mapper).is_negative:
+                return True
 
-    if e0 - e1 == 0:
-        return DISJOINT
-    else:
-        return MAYBE_OVERLAP
+    return False
 
 
 def disjoint_test(e0, e1, d, it):

devito/operator/operator.py

@@ -716,7 +716,7 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs):
             except AttributeError:
                 pass
             if d.is_Derived:
-                d._arg_check(args)
+                d._arg_check(args, **kwargs)
 
         # Turn arguments into a format suitable for the generated code
         # E.g., instead of NumPy arrays for Functions, the generated code expects

devito/passes/iet/orchestration.py

@@ -125,9 +125,9 @@ def _make_syncarray(self, iet, sync_ops, layer):
     def _make_prefetchupdate(self, iet, sync_ops, layer, wrap=True):
         return self._make_async_task(prefetchupdate, iet, sync_ops, layer, wrap)
 
-    @iet_pass
-    def process(self, iet):
-        callbacks = {
+    @property
+    def _callbacks(self):
+        return {
             WaitLock: self._make_waitlock,
             WithLock: self._make_withlock,
             SyncArray: self._make_syncarray,
@@ -139,13 +139,15 @@ def process(self, iet):
             AsyncCallable: self._make_async_callable
         }
 
+    @iet_pass
+    def process(self, iet):
         # Collect the compatible asynchronous task groups, if any
         task_groups = TaskGroups()
         if self.npthreads:
             CollectTasks(task_groups).visit(iet)
 
         # Lower the SyncSpots in a single bottom-up traversal, atomically lowering
-        lowerer = LowerSyncSpots(callbacks, task_groups, self.sregistry)
+        lowerer = LowerSyncSpots(self._callbacks, task_groups, self.sregistry)
         iet = lowerer.visit(iet)
 
         return iet, {'efuncs': lowerer.efuncs}

devito/types/dimension.py

@@ -818,6 +818,46 @@ def _arg_values(self, interval, grid=None, **kwargs):
         # themselves
         return {}
 
+    def _arg_check(self, args, *_args, **kwargs):
+        # These modules depend on Dimension, so importing them above would cycle
+        from devito.mpi import mpi_raise  # noqa: PLC0415
+        from devito.symbolics import subs_op_args  # noqa: PLC0415
+
+        if not self.is_middle:
+            return
+
+        # Function._arg_check visits original axes (e.g. x_ltkn), whereas `args`
+        # contains the concretized thicknesses (x_ltkn0, ...). The Operator checks
+        # the matching concrete SubDimensions separately in its dimension loop
+        if self not in args.op.dimensions:
+            return
+
+        d = self.root
+        if args.grid is not None and args.grid.is_distributed(d):
+            # Check the global runtime region: its MPI-local slices may be empty
+            # or smaller than space_order even for a non-degenerate global interior
+            size = args.grid.size_map[d].glb
+            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}}
+        else:
+            values = args
+        size = int(subs_op_args(self.symbolic_size, values))
+
+        # Runtime overrides do not change the compiled stencil order
+        items = [f.space_order for f in args.op.input if f.is_DiscreteFunction]
+        space_order = max(items, default=0)
+
+        if size < space_order:
+            error = (f"Expected at least {space_order} interior points along "
+                     f"`{self.parent}` (space_order), but runtime arguments leave {size}")
+        else:
+            error = None
+
+        comm = args.comm if args.options['mpi'] else None
+        mpi_raise(error, InvalidArgument, comm=comm)
+
 
 class MultiSubDimension(AbstractSubDimension):
 
@@ -1436,7 +1476,7 @@ def _arg_values(self, interval, grid=None, args=None, **kwargs):
                 # Avoid OOB (will end up here only in case of tiny iteration spaces)
                 return {name: 1}
 
-    def _arg_check(self, args, *_args):
+    def _arg_check(self, args, *_args, **kwargs):
         try:
             name = self.step.name
         except AttributeError:

devito/types/grid.py

@@ -641,12 +641,6 @@ def __subdomain_finalize_legacy__(self, grid):
                 try:
                     # Case ('middle', int, int)
                     side, ltkn, rtkn = v
-                    if side != 'middle':
-                        raise ValueError(f"Expected side 'middle', not `{side}`")
-                    sub_dimensions.append(SubDimension.middle(f'i{k.name}',
-                                                              k, ltkn, rtkn))
-                    thickness = s-ltkn-rtkn
-                    sdshape.append(thickness)
                 except ValueError:
                     side, thickness = v
                     constructor = {'left': SubDimension.left,
@@ -663,6 +657,23 @@ def __subdomain_finalize_legacy__(self, grid):
                         ) from None
                     sub_dimensions.append(constructor(f'i{k.name}', k, thickness))
                     sdshape.append(thickness)
+                else:
+                    if side != 'middle':
+                        raise ValueError(f"Expected side 'middle', not `{side}`")
+
+                    # A `middle` region expects `ltkn + rtkn <= s` in the global Grid.
+                    # This ensures that the left and right regions won't overlap
+                    thickness = s-ltkn-rtkn
+                    if thickness < 0:
+                        raise ValueError(
+                            f"SubDomain `{self.name}` has combined thickness "
+                            f"{ltkn + rtkn} along `{k}`, exceeding the Grid size {s}"
+                        )
+
+                    sub_dimensions.append(
+                        SubDimension.middle(f'i{k.name}', k, ltkn, rtkn)
+                    )
+                    sdshape.append(thickness)
 
         self._shape = tuple(sdshape)
         self._dimensions = tuple(sub_dimensions)