Merge pull request #2869 from devitocodes/patch-degenerating-indices

compiler: Patch degenerating_dimensions to catch SteppingDimensions over size 1 buffer

Author: EdCaunt

devito/ir/support/basic.py

@@ -392,7 +392,7 @@ def distance(self, other):
                 # objects falls back to zero, as any other value would be
                 # nonsensical
                 ret.append(S.Zero)
-            elif degenerating_dimensions(sai, oai):
+            elif degenerating_indices(self[n], other[n], self.function):
                 # Special case: `sai` and `oai` may be different symbolic objects
                 # but they can be proved to systematically generate the same value
                 ret.append(S.Zero)
@@ -1425,17 +1425,32 @@ def disjoint_test(e0, e1, d, it):
     return not bool(i0.intersect(i1))
 
 
-def degenerating_dimensions(d0, d1):
+def degenerating_indices(i0, i1, function):
     """
-    True if `d0` and `d1` are Dimensions that are possibly symbolically
+    True if `i0` and `i1` are indices that are possibly symbolically
     different, but they can be proved to systematically degenerate to the
     same value, False otherwise.
     """
     # Case 1: ModuloDimensions of size 1
     try:
-        if d0.is_Modulo and d1.is_Modulo and d0.modulo == d1.modulo == 1:
+        if i0.is_Modulo and i1.is_Modulo and i0.modulo == i1.modulo == 1:
             return True
     except AttributeError:
         pass
 
+    # Case 2: SteppingDimension corresponding to buffer of size 1
+    # Extract dimension from both IndexAccessFunctions -> d0, d1
+    try:
+        d0 = i0.d
+    except AttributeError:
+        d0 = i0
+    try:
+        d1 = i1.d
+    except AttributeError:
+        d1 = i1
+
+    with suppress(AttributeError):
+        if d0 is d1 and d0.is_Stepping and function._size_domain[d0] == 1:
+            return True
+
     return False

tests/test_fission.py

@@ -2,8 +2,11 @@
 
 from conftest import assert_structure
 from devito import (
-    Eq, Function, Grid, Inc, Operator, SubDimension, SubDomain, TimeFunction, solve
+    Buffer, Eq, Function, Grid, Inc, Operator, SubDimension, SubDomain, TimeFunction,
+    solve
 )
+from devito.ir.iet import retrieve_iteration_tree
+from devito.ir.support.properties import PARALLEL
 
 
 def test_issue_1725():
@@ -126,3 +129,70 @@ def test_issue_1921():
     op1.apply(time_m=1, time_M=5, g=g1)
 
     assert np.all(g.data == g1.data)
+
+
+def test_buffer1_fissioning():
+    """
+    Tests an edge case whereby inability to spot the equivalence of
+    `f.forward`/`backward` and `f` when using `Buffer(1)` would cause
+    data dependance analysis to incorrectly determine that `x` and `y`
+    loops should be sequential.
+    """
+    so = 2
+
+    class SD0(SubDomain):
+        name = 'sd0'
+
+        def __init__(self, so, **kwargs):
+            self.so = so
+            super().__init__(**kwargs)
+
+        def define(self, dimensions):
+            map_d = {d: d for d in dimensions}
+            map_d[dimensions[-1]] = ('left', self.so)
+            return map_d
+
+    grid = Grid(shape=(101, 101, 101))
+    z = grid.dimensions[-1]
+
+    sd0 = SD0(so, grid=grid)
+    iz = sd0.dimensions[-1]
+
+    f0 = TimeFunction(name='f0', grid=grid,
+                      space_order=so, save=Buffer(1))
+    f1 = TimeFunction(name='f1', grid=grid,
+                      space_order=so, save=Buffer(1))
+
+    f2 = TimeFunction(name='f2', grid=grid,
+                      space_order=so, save=Buffer(1))
+    f3 = TimeFunction(name='f3', grid=grid,
+                      space_order=so, save=Buffer(1))
+
+    # Free-surface like conditions plus derivative creates a dependence
+    eq0 = Eq(f0.backward.subs(z, iz-so), -f0.backward.subs(z, so-iz))
+    eq1 = Eq(f1.backward.subs(z, iz-so), -f1.backward.subs(z, so-iz))
+
+    eq2 = Eq(f2, f0.dx + f2)
+    eq3 = Eq(f3, f1.dy + f3)
+
+    eqs = [eq0, eq1, eq2, eq3]
+
+    op = Operator(eqs)
+
+    trees = retrieve_iteration_tree(op)
+
+    # If the equivalence of f0.backward and f0 are not spotted by the
+    # compiler, data dependence analysis will determine x and y loops
+    # to be sequential.
+    for tree in trees:
+        for i in range(1, 3):
+            assert PARALLEL in tree[i].properties
+
+    # If parallelisation failed, then there will also be no blocking on
+    # x and y
+    assert 'x0_blk0_size' in str(op.parameters)
+    assert 'y0_blk0_size' in str(op.parameters)
+
+    # Two loop nests: free-surface-like and update-like
+    assert_structure(op, ['t,x,y,z', 't,x0_blk0,y0_blk0,x,y,z'],
+                     't,x,y,z,x0_blk0,y0_blk0,x,y,z')

tests/test_mpi.py

@@ -1887,10 +1887,12 @@ def test_enforce_haloupdate_if_unwritten_function(self, mode):
     @pytest.mark.parallel(mode=1)
     def test_haloupdate_buffer1(self, mode):
         grid = Grid(shape=(4, 4))
-        x, y = grid.dimensions
 
-        u = TimeFunction(name='u', grid=grid, time_order=1, save=Buffer(1))
-        v = TimeFunction(name='v', grid=grid, time_order=1, save=Buffer(1))
+        # With order 2 (1) forward derivatives, the loops can (and will) be fused,
+        # removing the need for a halo update so space_order=4 is used to ensure
+        # derivatives are centred resulting in parallel loops and halo updates
+        u = TimeFunction(name='u', grid=grid, time_order=1, space_order=4, save=Buffer(1))
+        v = TimeFunction(name='v', grid=grid, time_order=1, space_order=4, save=Buffer(1))
 
         eqns = [Eq(u.forward, div(v) + 1.),
                 Eq(v.forward, div(u.forward) + 1.)]
@@ -1909,22 +1911,22 @@ def test_haloupdate_buffer1(self, mode):
     @pytest.mark.parallel(mode=1)
     @pytest.mark.parametrize('sz,fwd,expr,exp0,exp1,args', [
         (1, True, 'rec.interpolate(v2)', 3, 2, ('v1', 'v2')),
-        (1, True, 'Eq(v3.forward, v2.laplace + 1)', 1, 1, ('v2',)),
-        (1, True, 'Eq(v3.forward, v2.forward.laplace + 1)', 3, 2, ('v1', 'v2',)),
-        (2, True, 'Eq(v3.forward, v2.forward.laplace + 1)', 3, 2, ('v1', 'v2',)),
+        (1, True, 'Eq(v3.forward, v2.laplace + 1)', 3, 2, ('v1', 'v2')),
+        (1, True, 'Eq(v3.forward, v2.forward.laplace + 1)', 3, 2, ('v1', 'v2')),
+        (2, True, 'Eq(v3.forward, v2.forward.laplace + 1)', 3, 2, ('v1', 'v2')),
         (1, False, 'rec.interpolate(v2)', 3, 2, ('v1', 'v2')),
-        (1, False, 'Eq(v3.backward, v2.laplace + 1)', 1, 1, ('v2',)),
-        (1, False, 'Eq(v3.backward, v2.backward.laplace + 1)', 3, 2, ('v1', 'v2',)),
-        (2, False, 'Eq(v3.backward, v2.backward.laplace + 1)', 3, 2, ('v1', 'v2',)),
+        (1, False, 'Eq(v3.backward, v2.laplace + 1)', 3, 2, ('v1', 'v2')),
+        (1, False, 'Eq(v3.backward, v2.backward.laplace + 1)', 3, 3, ('v2', 'v1', 'v2')),
+        (2, False, 'Eq(v3.backward, v2.backward.laplace + 1)', 3, 3, ('v2', 'v1', 'v2')),
     ])
     def test_haloupdate_buffer_cases(self, sz, fwd, expr, exp0, exp1, args, mode):
         grid = Grid((65, 65, 65), topology=('*', 1, '*'))
 
-        v1 = TimeFunction(name='v1', grid=grid, space_order=2, time_order=1,
+        v1 = TimeFunction(name='v1', grid=grid, space_order=4, time_order=1,
                           save=Buffer(1))
-        v2 = TimeFunction(name='v2', grid=grid, space_order=2, time_order=1,
+        v2 = TimeFunction(name='v2', grid=grid, space_order=4, time_order=1,
                           save=Buffer(1))
-        v3 = TimeFunction(name='v3', grid=grid, space_order=2, time_order=1,  # noqa
+        v3 = TimeFunction(name='v3', grid=grid, space_order=4, time_order=1,  # noqa
                           save=Buffer(1))
 
         rec = SparseTimeFunction(name='rec', grid=grid, nt=500, npoint=65)  # noqa