Merge pull request #3033 from devitocodes/tweak-graph-construction-final

compiler: Revamp DDA for SubDimensions

Author: FabioLuporini

devito/exceptions.py

@@ -54,3 +54,19 @@ class ExecutionError(DevitoError):
         * Device shared memory or registers (e.g., too many threads per block);
     * etc.
     """
+
+
+def mpi_raise(error, exception=ValueError, comm=None):
+    """
+    Raise `exception` with the first non-None error message in rank order.
+
+    All ranks in `comm` must call this routine, including those with no local
+    error (`error=None`). This prevents a rank-local exception from stranding
+    peers in subsequent MPI calls. With no communicator or `MPI.COMM_NULL`,
+    only the local error is checked.
+    """
+    # A null MPI communicator is false, like None
+    if comm:
+        error = next((i for i in comm.allgather(error) if i is not None), None)
+    if error is not None:
+        raise exception(error)

devito/finite_differences/differentiable.py

@@ -563,7 +563,8 @@ def _gather_for_diff(self):
 
     # Bypass useless expensive SymPy _eval_ methods, for which we either already
     # know or don't care about the answer, because it'd have ~zero impact on our
-    # average expressions
+    # average expressions. Sign inference may also call `diff`, which here
+    # constructs finite differences rather than symbolic derivatives
 
     def _eval_is_even(self):
         return None
@@ -580,12 +581,18 @@ def _eval_is_negative(self):
     def _eval_is_extended_negative(self):
         return None
 
+    def _eval_is_extended_nonpositive(self):
+        return None
+
     def _eval_is_positive(self):
         return None
 
     def _eval_is_extended_positive(self):
         return None
 
+    def _eval_is_extended_nonnegative(self):
+        return None
+
     def _eval_is_zero(self):
         return None
 

devito/ir/cgen/printer.py

@@ -16,6 +16,7 @@
 
 from devito import configuration
 from devito.arch.compiler import AOMPCompiler
+from devito.symbolics.extended_sympy import BitwiseAnd
 from devito.symbolics.inspection import has_integer_args, sympy_dtype
 from devito.symbolics.queries import q_leaf
 from devito.tools import (
@@ -255,7 +256,26 @@ def _print_RoundUp(self, expr):
         return f'ROUND_UP({value}, {step})'
 
     def _print_Mod(self, expr):
-        """Print a Mod as a C-like %-based operation."""
+        """
+        Print a Mod as an integer remainder or a power-of-two mask.
+
+        Python's `%` and SymPy's `Mod` give a nonnegative result for a positive
+        divisor, whereas C's `%` can be negative when the dividend is negative.
+        For example, `Mod(-1, 4) == 3`, but C's `-1 % 4 == -1`. For integer
+        operands and a positive power-of-two divisor `b`, emit `a & (b - 1)`
+        unless `a` is known nonnegative, preserving the Python/SymPy semantics.
+        """
+        a, b = expr.args
+
+        # Unlike C's signed remainder, a mask preserves `Mod`'s nonnegative
+        # result for positive power-of-two divisors
+        if b.is_Integer and \
+           b > 0 and \
+           not (int(b) & (int(b) - 1)) and \
+           has_integer_args(a, b) and \
+           a.is_nonnegative is not True:
+            return f'({self._print(BitwiseAnd(a, b - 1))})'
+
         args = [f'({self._print(a)})' for a in expr.args]
         return '%'.join(args)
 

devito/ir/iet/nodes.py

@@ -1078,7 +1078,7 @@ def expr_symbols(self):
             with suppress(AttributeError):
                 ret.update(f.initvalue.free_symbols)
             return tuple(ret)
-        elif f.is_Array and f.initvalue is not None:
+        elif f.is_ArrayLike and f.initvalue is not None:
             # These are just a handful of values so it's OK to iterate them over
             ret = set()
             for i in f.initvalue:

devito/ir/iet/visitors.py

@@ -360,11 +360,11 @@ def _gen_value(self, obj, mode=1, masked=()):
         except AttributeError:
             pass
 
-        if obj.is_Array and obj.initvalue is not None and mode == 1:
+        if obj.is_ArrayLike and obj.initvalue is not None and mode == 1:
             init = ListInitializer(obj.initvalue)
             if not obj._mem_constant or init.is_numeric:
                 # printed at the Array's own precision, not the Operator's
-                value = c.Initializer(value, self.ccode(init, dtype=obj.dtype))
+                value = c.Initializer(value, self.ccode(init, dtype=obj.c0.dtype))
         elif obj.is_LocalObject and obj.initvalue is not None and mode == 1:
             value = c.Initializer(value, self.ccode(obj.initvalue))
 

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 SubDimensions 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,103 @@ 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(a0, a1):
+    """
+    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. 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, a left SubDimension of thickness 4 ends before a middle
+    SubDimension excluding 4 points, even when the two thickness symbols are distinct.
+
+    Left/right SubDimensions of the same parent with `separated=True` satisfy
+    `L + R + space_order <= N`, checked against the full global parent extent
+    at `Operator.apply`. Their gap therefore accommodates stencil accesses;
+    larger shifts are still compared explicitly. If either SubDimension has
+    `separated=False`, no minimum separation is assumed.
+
+    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.
+    """
+    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):
+            continue
+
+        f = a0.function.c0
+        space_order = f.space_order if isinstance(f, Function) else 0
+        if disjoint_subdims_axis(e0, e1, d0, d1, it0, it1, space_order):
+            return True
+
+    return False
 
 
-def disjoint_subdims(e0, e1, d0, d1, it0, it1):
+@memoized_func(scope='build')
+def disjoint_subdims_axis(e0, e1, d0, d1, it0, it1, space_order):
     """
-    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.
+    Test separation along one data axis of two SubDimension accesses.
+
+    The result depends on the indices, intervals and stencil order, rather than
+    the access timestamps, modes or other axes. Cache it across the many
+    TimedAccess pairs and Scopes that reuse the same one-dimensional regions.
+    The cache is cleared at the end of Operator construction.
     """
-    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
-                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
+    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):
+            break
+
+        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.xreplace(thicknesses)))
+            upper.append(v._subs(d, M.xreplace(thicknesses)))
+        else:
+            bounds.append((sympy.Min(*lower), sympy.Max(*upper)))
 
-    if not is_partition:
-        return MAYBE_OVERLAP
+    if len(bounds) == 2:
+        (m0, M0), (m1, M1) = bounds
+        mapper = {}
 
-    if not q_affine(e0, d0) or not q_affine(e1, d1):
-        return MAYBE_OVERLAP
+        dl, dr = (it0.dim, it1.dim) if it0.dim.is_left else (it1.dim, it0.dim)
+        dlp, drp = dl.parent, dr.parent
 
-    if it0.offsets != it1.offsets or it0.direction is not it1.direction:
-        return MAYBE_OVERLAP
+        if dl.is_left and dr.is_right and \
+           dl.separated and dr.separated and \
+           dlp is drp:
+            # Runtime validation guarantees N - L - R >= space_order
+            gap = sympy.Dummy(nonnegative=True)
+            mapper[dlp.symbolic_max] = (dlp.symbolic_min + dl.ltkn.value +
+                                        dr.rtkn.value + space_order + gap - 1)
 
-    e0 = e0._subs(d0, d0.root)
-    e1 = e1._subs(d1, d1.root)
+        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):