compiler: Improve DDA involving SubDimensions

Author: FabioLuporini

devito/ir/support/basic.py

@@ -369,6 +369,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]>`
@@ -1527,6 +1535,92 @@ 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):
+    """
+    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.
+    """
+    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
+
+    if not is_partition:
+        return MAYBE_OVERLAP
+
+    if not q_affine(e0, d0) or not q_affine(e1, d1):
+        return MAYBE_OVERLAP
+
+    if it0.offsets != it1.offsets or it0.direction is not it1.direction:
+        return MAYBE_OVERLAP
+
+    e0 = e0._subs(d0, d0.root)
+    e1 = e1._subs(d1, d1.root)
+
+    if e0 - e1 == 0:
+        return DISJOINT
+    else:
+        return MAYBE_OVERLAP
+
+
 def disjoint_test(e0, e1, d, it):
     """
     A rudimentary test to check if two accesses `e0` and `e1` along `d` within

tests/test_ir.py

@@ -357,6 +357,87 @@ def test_timed_access_cmp(self, ta_literal):
         assert tcyx_irr0 >= tcyx_irr0
         assert tcyx_irr0 == tcyx_irr0
 
+    def test_timed_access_distance_subdimensions(self):
+        grid = Grid(shape=(24, 24))
+        x, y = grid.dimensions
+
+        xl = SubDimension.left('xl', x, 4)
+        xm = SubDimension.middle('xm', x, 4, 4)
+        xr = SubDimension.right('xr', x, 4)
+        xm_bad = SubDimension.middle('xm_bad', x, 3, 4)
+        yl = SubDimension.left('yl', y, 4)
+        ym = SubDimension.middle('ym', y, 4, 4)
+        xl_overlap = SubDimension.left('xl_overlap', x, 4)
+        xm_overlap = SubDimension.middle('xm_overlap', x, 4, 4)
+        xr_overlap = SubDimension.right('xr_overlap', x, 4)
+
+        f = Function(name='f', grid=grid)
+
+        left = TimedAccess(
+            f[xl, y], 'W', 0, IterationSpace([Interval(xl), Interval(y)])
+        )
+        middle = TimedAccess(
+            f[xm, y], 'R', 1, IterationSpace([Interval(xm), Interval(y)])
+        )
+        right = TimedAccess(
+            f[xr, y], 'R', 1, IterationSpace([Interval(xr), Interval(y)])
+        )
+        left_overlap = TimedAccess(
+            f[xl_overlap, y], 'R', 1,
+            IterationSpace([Interval(xl_overlap), Interval(y)])
+        )
+        middle_overlap = TimedAccess(
+            f[xm_overlap, y], 'R', 1,
+            IterationSpace([Interval(xm_overlap), Interval(y)])
+        )
+        right_overlap = TimedAccess(
+            f[xr_overlap, y], 'R', 1,
+            IterationSpace([Interval(xr_overlap), Interval(y)])
+        )
+        bad = TimedAccess(
+            f[xm_bad, y], 'R', 1,
+            IterationSpace([Interval(xm_bad), Interval(y)])
+        )
+        shifted = TimedAccess(
+            f[xm + 1, y], 'R', 1,
+            IterationSpace([Interval(xm), Interval(y)])
+        )
+        shifted_range = TimedAccess(
+            f[xm, y], 'R', 1,
+            IterationSpace([Interval(xm, 1, 1), Interval(y)])
+        )
+        left_nonlinear = TimedAccess(
+            f[xl % 2, y], 'W', 0, IterationSpace([Interval(xl), Interval(y)])
+        )
+        middle_nonlinear = TimedAccess(
+            f[xm % 2, y], 'R', 1, IterationSpace([Interval(xm), Interval(y)])
+        )
+        orthogonal = TimedAccess(
+            f[x, yl], 'R', 1, IterationSpace([Interval(x), Interval(yl)])
+        )
+        corner_left = TimedAccess(
+            f[xl, yl], 'W', 0, IterationSpace([Interval(xl), Interval(yl)])
+        )
+        corner_middle = TimedAccess(
+            f[xl_overlap, ym], 'R', 1,
+            IterationSpace([Interval(xl_overlap), Interval(ym)])
+        )
+
+        assert left.distance(middle) == (S.ImaginaryUnit,)
+        assert middle.distance(right) == (S.ImaginaryUnit,)
+        assert left.distance(right) == (S.ImaginaryUnit,)
+        assert right.distance(left) == (S.ImaginaryUnit,)
+        assert corner_left.distance(corner_middle) == (S.ImaginaryUnit,)
+
+        assert left.distance(left_overlap) == (S.Infinity, 0)
+        assert middle.distance(middle_overlap) == (S.Infinity, 0)
+        assert right.distance(right_overlap) == (S.Infinity, 0)
+        assert left.distance(bad) == (S.Infinity, 0)
+        assert left.distance(shifted) == (S.Infinity, 0)
+        assert left.distance(shifted_range) == (S.Infinity, 0)
+        assert left_nonlinear.distance(middle_nonlinear) == (S.Infinity, 0)
+        assert left.distance(orthogonal) == (S.Infinity,)
+
 
 class TestSpace: