compiler: Fixup halo exchange with sequentially-iterated, distributed dims

Author: FabioLuporini

devito/ir/clusters/algorithms.py

@@ -471,9 +471,11 @@ def callback(self, clusters, prefix, seen=None):
             # Construct the HaloTouch Cluster
             expr = Eq(self.B, HaloTouch(*points, halo_scheme=hs))
 
-            key = lambda i: i in prefix[:-1] or i in hs.loc_indices  # noqa: B023
+            key0 = lambda i: i in prefix[:-1] or i in hs.loc_indices  # noqa: B023
+            key1 = lambda i: i not in hs.distributed_defined  # noqa: B023
+            key = lambda i: key0(i) and key1(i)  # noqa: B023
             ispace = c.ispace.project(key)
-            # HaloTouches are not parallel
+
             properties = c.properties.sequentialize()
 
             halo_touch = c.rebuild(exprs=expr, ispace=ispace, properties=properties)

devito/mpi/halo_scheme.py

@@ -429,6 +429,10 @@ def distributed(self):
     def distributed_aindices(self):
         return set().union(*[i.dims for i in self.fmapper.values()])
 
+    @cached_property
+    def distributed_defined(self):
+        return set().union(*[i._defines for i in self.distributed])
+
     @cached_property
     def loc_indices(self):
         return set().union(*[i.loc_indices.keys() for i in self.fmapper.values()])

tests/test_mpi.py

@@ -2181,6 +2181,28 @@ def test_allreduce_multicond(self, mode):
         assert np.allclose(g.data, expected)
         assert np.allclose(h.data, expected)
 
+    @pytest.mark.parallel(mode=1)
+    def test_lift_halo_update_outside_distributed(self, mode):
+        grid = Grid(shape=(28, 28, 28))
+        x, _, _ = grid.dimensions
+        h_x = x.spacing
+
+        vx = TimeFunction(name="vx", grid=grid, space_order=4, time_order=1)
+        vy = TimeFunction(name="vy", grid=grid, space_order=4, time_order=1)
+        vz = TimeFunction(name="vz", grid=grid, space_order=4, time_order=1)
+
+        eqn = Eq(vz, vz.dy + vz.dx(x0=x - 5.*h_x) + (vx.dy + vy.dx).dx.dy + vz)
+
+        # Ensure the Operator can be constructed and jit-compiled correctly
+        op = Operator(eqn)
+        _ = op.cfunction
+
+        # Check generated code -- expected one halo exchange exactly at the top
+        # of the time loop
+        tloop = get_time_loop(op)
+        halo_update = tloop.nodes[0].body[0].body[0].body[0]
+        assert isinstance(halo_update, HaloUpdateList)
+
 
 class TestOperatorAdvanced: