mpi: Prevent parallel halo exchanges

devito/ir/stree/algorithms.py

@@ -218,10 +218,10 @@ def preprocess(clusters, options=None, **kwargs):
                 # dimension requiring a loc-index in the HaloScheme. The compiler
                 # will generate the non-perfect loop nest `t,f ; t,x,y,z,f`, with
                 # the first nest triggering all necessary halo exchanges along `f`
-                mapper = as_mapper(found, lambda c1: c1.ispace)
-                for k, v in mapper.items():
+                mapper = as_mapper(found, lambda c1: (c1.ispace, c1.properties))
+                for v in mapper.values():
                     hs = HaloScheme.union([c1.halo_scheme for c1 in v])
-                    processed.append(c.rebuild(exprs=[], ispace=k, halo_scheme=hs))
+                    processed.append(v[0].rebuild(exprs=[], halo_scheme=hs))
                 processed.append(c)
             else:
                 # Avoid ugly empty loops

tests/test_mpi.py

@@ -2242,6 +2242,38 @@ 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_ensure_halo_updates_within_seq_loops(self, mode):
+        n = 30
+        grid = Grid(shape=(n, n))
+
+        x, y = grid.dimensions
+        d = Dimension(name="d")
+
+        a = Function(name="a", grid=grid, space_order=8)
+        b = a.func(name='b')
+
+        g = Function(name="g", grid=grid, space_order=8, dimensions=(x, y, d),
+                     shape=(n, n, 2))
+        h = g.func(name='h')
+
+        eqns = [Eq(a, 0.0),
+                Eq(h, b * g),
+                Inc(a, b * (h.dx + h.dy))]
+
+        op = Operator(eqns, opt=('advanced', {'openmp': True}))
+        _ = op.cfunction
+
+        # Check generated code -- expected one halo exchange within a sequential
+        # `d` loop
+        assert_structure(op, ['x,y', 'x,y,d', 'd', 'x,y,d'], 'xyddxyd')
+        iters = FindNodes(Iteration).visit(op)
+        assert iters[2].dim is d and iters[2].is_Parallel
+        assert iters[3].dim is d and iters[3].is_Sequential
+
+        # Probably unnecessary, but ensures there's no hanging
+        op.apply()
+
     @pytest.mark.parallel(mode=4)
     def test_halo_inner_dim(self, mode):
         grid = Grid((11, 11, 11))