compiler: fix halo placement for non out dimm exchange

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -443,7 +443,7 @@ def callback(self, clusters, prefix, seen=None):
 
         # Construct a representation of the halo accesses
         processed = []
-        for c in clusters:
+        for i, c in enumerate(clusters):
             if c.properties.is_sequential(d) or \
                c in seen:
                 continue
@@ -453,6 +453,10 @@ def callback(self, clusters, prefix, seen=None):
                not d._defines & hs.distributed_aindices:
                 continue
 
+            if any(halo_write(ci, hs) for ci in clusters[:i]):
+                # If there's a halo write before `c`, then we cannot inject the HaloTouch
+                continue
+
             points = set()
             for f in hs.fmapper:
                 for a in c.scope.getreads(f):
@@ -781,3 +785,14 @@ def normalize_reductions_sparse(cluster, sregistry):
             processed.append(e)
 
     return cluster.rebuild(processed)
+
+
+def halo_write(c, hs):
+    loc_vals = hs.loc_values
+
+    for f in hs.fmapper:
+        for a in c.scope.getwrites(f):
+            if set(a.access.indices) & loc_vals:
+                return True
+
+    return False

devito/mpi/halo_scheme.py

@@ -5,7 +5,7 @@
 from operator import attrgetter
 
 import sympy
-from sympy import Max, Min
+from sympy import Max, Min, S
 
 from devito import configuration
 from devito.data import CENTER, CORE, LEFT, OWNED, RIGHT
@@ -514,6 +514,31 @@ def merge(self, hs):
             fmapper[f] = fmapper.get(f, hse).merge(hse)
         return HaloScheme.build(fmapper, self.honored)
 
+    def _is_iter_carried(self, scope):
+        """
+        True if the HaloScheme is iteration-carried, i.e., it induces
+        a halo exchange that requires values from the previous iteration(s); False
+        otherwise.
+        """
+
+        def rule0(dep):
+            # E.g., `dep=W<f,[t1, x]> -> R<f,[t0, x-1]>`, `d=t` => OK
+            return not any(dep.distance_mapper[d] is S.Infinity for d in dep.cause)
+
+        def rule1(dep, loc_indices):
+            # E.g., `dep=W<f,[t1, x+1]> -> R<f,[t1, xl+1]>`, `loc_indices={t: t0}` => OK
+            return any(dep.distance_mapper[d] == 0 and
+                       dep.source[d] is not v and
+                       dep.sink[d] is not v
+                       for d, v in loc_indices.items())
+
+        for f, v in self.fmapper.items():
+            for dep in scope.d_flow.project(f):
+                if not rule0(dep) and not rule1(dep, v.loc_indices):
+                    return False
+
+        return True
+
 
 def classify(exprs, ispace):
     """

devito/passes/iet/mpi.py

@@ -206,7 +206,7 @@ def _hoist_invariant(iet):
                 # Ensure there's another HaloScheme that could cover for
                 # us should we get hoisted while still satisfying the
                 # data dependences
-                if hsf1.issubset(hsf0) and _is_iter_carried(hsf1, scope):
+                if hsf1.issubset(hsf0) and hsf1._is_iter_carried(scope):
                     hs, hsf = hs1, hsf1
                 elif hsf0.issubset(hsf1) and hs0 is halo_spots[0]:
                     # Special case
@@ -474,32 +474,6 @@ def _check_control_flow(hs0, hs1, cond_mapper):
     return cond0 != cond1
 
 
-def _is_iter_carried(hsf, scope):
-    """
-    True if the provided HaloScheme `hsf` is iteration-carried, i.e., it induces
-    a halo exchange that requires values from the previous iteration(s); False
-    otherwise.
-    """
-
-    def rule0(dep):
-        # E.g., `dep=W<f,[t1, x]> -> R<f,[t0, x-1]>`, `d=t` => OK
-        return not any(dep.distance_mapper[d] is S.Infinity for d in dep.cause)
-
-    def rule1(dep, loc_indices):
-        # E.g., `dep=W<f,[t1, x+1]> -> R<f,[t1, xl+1]>`, `loc_indices={t: t0}` => OK
-        return any(dep.distance_mapper[d] == 0 and
-                   dep.source[d] is not v and
-                   dep.sink[d] is not v
-                   for d, v in loc_indices.items())
-
-    for f, v in hsf.fmapper.items():
-        for dep in scope.d_flow.project(f):
-            if not rule0(dep) and not rule1(dep, v.loc_indices):
-                return False
-
-    return True
-
-
 def _is_mergeable(hsf0, hsf1, scope):
     """
     True if `hsf1` can be merged into `hsf0`, i.e., if they are compatible
@@ -515,7 +489,7 @@ def _is_mergeable(hsf0, hsf1, scope):
         return False
 
     # Finally, check the data dependences would be satisfied
-    return _is_iter_carried(hsf1, scope)
+    return hsf1._is_iter_carried(scope)
 
 
 def _semantical_eq_loc_indices(hsf0, hsf1):

tests/test_mpi.py

@@ -2203,6 +2203,24 @@ def test_lift_halo_update_outside_distributed(self, mode):
         halo_update = tloop.nodes[0].body[0].body[0].body[0]
         assert isinstance(halo_update, HaloUpdateList)
 
+    @pytest.mark.parallel(mode=4)
+    def test_halo_inner_dim(self, mode):
+        grid = Grid((11, 11, 11))
+
+        np.random.seed(0)
+        v = TimeFunction(name="v", grid=grid, space_order=4,
+                         time_order=1, save=Buffer(1))
+        v.data[:] = np.random.randn(*grid.shape)
+        e = TimeFunction(name="dummy", grid=grid, space_order=4, time_order=0)
+
+        eq = [Eq(v.forward, v + 1), Eq(e, v.forward.dydz)]
+
+        op = Operator(eq, opt=('advanced', {'blocklevels': 0}))
+        assert_structure(op, ['txyz', 't', 'txyz', 'txyz'], 'txyzxyzz')
+        op(time=100)
+
+        assert np.isclose(norm(e), 23484.863, rtol=0, atol=1e-1)
+
 
 class TestOperatorAdvanced:
 
@@ -2736,7 +2754,7 @@ def test_haloupdate_same_timestep_v2(self, mode):
         assert titer.dim is grid.time_dim
         assert titer.nodes[0].body[0].body[0].is_List
         assert len(titer.nodes[0].body[0].body[0].body[0].body) == 1
-        assert titer.nodes[0].body[0].body[0].body[0].body[0].is_Call
+        assert not titer.nodes[0].body[0].body[0].body[0].body[0].is_Call
 
         op.apply(time=0)
 
@@ -3138,8 +3156,8 @@ def test_fission_due_to_antidep(self, mode):
         # First, check the generated code
         assert_structure(op1, ['t',
                                't,x0_blk0,y0_blk0,x,y,z',
-                               't,x0_blk0,y0_blk0,x,y,z'],
-                         't,x0_blk0,y0_blk0,x,y,z,z')
+                               't,x1_blk0,y1_blk0,x,y,z'],
+                         'tx0_blk0y0_blk0xyzx1_blk0y1_blk0xyz')
 
         def init(f, v=1):
             f.data[:] = np.indices(grid.shape).sum(axis=0) % (.004*v) + .01