mpi: fix allreduce iteration space

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -8,7 +8,7 @@
 from devito.exceptions import CompilationError
 from devito.finite_differences.elementary import Max, Min
 from devito.ir.support import (Any, Backward, Forward, IterationSpace, erange,
-                               pull_dims, null_ispace)
+                               pull_dims)
 from devito.ir.equations import OpMin, OpMax, identity_mapper
 from devito.ir.clusters.analysis import analyze
 from devito.ir.clusters.cluster import Cluster, ClusterGroup
@@ -493,9 +493,9 @@ def reduction_comms(clusters):
         processed.append(c)
 
     # Leftover reductions are placed at the very end
-    if fifo:
-        exprs = [Eq(dr.var, dr) for dr in fifo]
-        processed.append(Cluster(exprs=exprs, ispace=null_ispace))
+    for ispace, reds in groupby(fifo, key=lambda r: r.ispace):
+        exprs = [Eq(dr.var, dr) for dr in reds]
+        processed.append(Cluster(exprs=exprs, ispace=ispace))
 
     return processed
 

devito/ir/iet/visitors.py

@@ -1292,6 +1292,20 @@ def visit_tuple(self, o, **kwargs):
 
     visit_list = visit_tuple
 
+    def visit_ExpressionBundle(self, o, **kwargs):
+        children = [self._visit(i) for i in o.children]
+        children = [i for i in children if i]
+        if not children:
+            return None
+        return o._rebuild(*children, **o.args_frozen)
+
+    def visit_Iteration(self, o, **kwargs):
+        children = [self._visit(i) for i in o.children]
+        children = [i for i in children if i]
+        if not children:
+            return None
+        return o._rebuild(*children, **o.args_frozen)
+
     def visit_Node(self, o, **kwargs):
         if o in self.mapper:
             handle = self.mapper[o]

tests/test_mpi.py

@@ -2014,6 +2014,24 @@ def test_merge_smart_if_within_conditional(self, mode):
         for n in FindNodes(Conditional).visit(op1):
             assert len(FindNodes(HaloUpdateCall).visit(n)) == 0
 
+    @pytest.mark.parallel(mode=2)
+    def test_allreduce_time(self, mode):
+        space_order = 8
+        nx, ny = 11, 11
+
+        grid = Grid(shape=(nx, ny))
+        tt = grid.time_dim
+        nt = 10
+
+        ux = TimeFunction(name="ux", grid=grid, time_order=1, space_order=space_order)
+        g = TimeFunction(name="g", grid=grid, dimensions=(tt, ), shape=(nt,))
+
+        op = Operator([Eq(ux.forward, ux + tt), Inc(g, ux)], name="Op")
+        assert_structure(op, ['t,x,y', 't'], 'txy')
+
+        op.apply(time_m=0, time_M=nt-1)
+        assert np.isclose(np.max(g.data), 4356.0)
+
 
 class TestOperatorAdvanced: