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,11 @@ 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]) and \
+               hs.dimensions & set(prefix.itdims):
+                # 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 +786,15 @@ 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:
+        # Does `c` write to any of the halo points of `f`?
+        w = c.scope.writes.get(f, None)
+        if w is None:
+            continue
+        elif any(set(a.indices) & loc_vals for a in [wi.access for wi in w]):
+            return True
+    return False

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: