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

@@ -1295,6 +1295,28 @@ def __init__(self, mapper, nested=False):
         self.mapper = mapper
         self.nested = nested
 
+    def transform(self, o, handle, **kwargs):
+        if handle is None:
+            # None -> drop `o`
+            return None
+        elif isinstance(handle, Iterable):
+            # Iterable -> inject `handle` into `o`'s children
+            if not o.children:
+                raise CompilationError("Cannot inject nodes in a leaf node")
+            if self.nested:
+                children = [self._visit(i, **kwargs) for i in o.children]
+            else:
+                children = o.children
+            children = (tuple(handle) + children[0],) + tuple(children[1:])
+            return o._rebuild(*children, **o.args_frozen)
+        else:
+            # Replace `o` with `handle`
+            if self.nested:
+                children = [self._visit(i, **kwargs) for i in handle.children]
+                return handle._rebuild(*children, **handle.args_frozen)
+            else:
+                return handle
+
     def visit_object(self, o, **kwargs):
         return o
 
@@ -1304,32 +1326,30 @@ def visit_tuple(self, o, **kwargs):
 
     visit_list = visit_tuple
 
+    def visit_ExpressionBundle(self, o, **kwargs):
+        if o in self.mapper:
+            handle = self.mapper[o]
+            return self.transform(o, handle, **kwargs)
+        children = [self._visit(i) for i in o.children]
+        if not [i for i in children if i]:
+            return None
+        return o._rebuild(*children, **o.args_frozen)
+
+    def visit_Iteration(self, o, **kwargs):
+        if o in self.mapper:
+            handle = self.mapper[o]
+            return self.transform(o, handle, **kwargs)
+        children = [self._visit(i) for i in o.children]
+        if not [i for i in children if i]:
+            return None
+        return o._rebuild(*children, **o.args_frozen)
+
     def visit_Node(self, o, **kwargs):
         if o in self.mapper:
             handle = self.mapper[o]
-            if handle is None:
-                # None -> drop `o`
-                return None
-            elif isinstance(handle, Iterable):
-                # Iterable -> inject `handle` into `o`'s children
-                if not o.children:
-                    raise CompilationError("Cannot inject nodes in a leaf node")
-                if self.nested:
-                    children = [self._visit(i, **kwargs) for i in o.children]
-                else:
-                    children = o.children
-                children = (tuple(handle) + children[0],) + tuple(children[1:])
-                return o._rebuild(*children, **o.args_frozen)
-            else:
-                # Replace `o` with `handle`
-                if self.nested:
-                    children = [self._visit(i, **kwargs) for i in handle.children]
-                    return handle._rebuild(*children, **handle.args_frozen)
-                else:
-                    return handle
-        else:
-            children = [self._visit(i, **kwargs) for i in o.children]
-            return o._rebuild(*children, **o.args_frozen)
+            return self.transform(o, handle, **kwargs)
+        children = [self._visit(i, **kwargs) for i in o.children]
+        return o._rebuild(*children, **o.args_frozen)
 
     def visit_Operator(self, o, **kwargs):
         raise ValueError("Cannot apply a Transformer visitor to an Operator directly")

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: