Merge pull request #2671 from devitocodes/muli-allred-v3

compiler: fix conditional reductions

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -457,8 +457,8 @@ def reduction_comms(clusters):
         # if `c`'s IterationSpace is such that the reduction can be carried out
         found, fifo = split(fifo, lambda dr: dr.ispace.is_subset(c.ispace))
         for ispace, reds in groupby(found, key=lambda r: r.ispace):
-            exprs = [Eq(dr.var, dr) for dr in reds]
-            processed.append(Cluster(exprs=exprs, ispace=ispace))
+            exprs = flatten([dr.exprs for dr in reds])
+            processed.append(c.rebuild(exprs=exprs, ispace=ispace))
 
         # Detect the global distributed reductions in `c`
         for e in c.exprs:
@@ -487,15 +487,16 @@ def reduction_comms(clusters):
             # The IterationSpace within which the global distributed reduction
             # must be carried out
             ispace = c.ispace.prefix(lambda d: d in var.free_symbols)
-
-            fifo.append(DistReduce(var, op=op, grid=grid, ispace=ispace))
+            expr = [Eq(var, DistReduce(var, op=op, grid=grid, ispace=ispace))]
+            fifo.append(c.rebuild(exprs=expr, ispace=ispace))
 
         processed.append(c)
 
     # Leftover reductions are placed at the very end
     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))
+        reds = list(reds)
+        exprs = flatten([dr.exprs for dr in reds])
+        processed.append(reds[0].rebuild(exprs=exprs, ispace=ispace))
 
     return processed
 

tests/test_mpi.py

@@ -2068,6 +2068,39 @@ def test_multi_allreduce_time(self, mode):
         assert np.isclose(np.max(g.data), 4356.0)
         assert np.isclose(np.max(h.data), 4356.0)
 
+    @pytest.mark.parallel(mode=1)
+    def test_multi_allreduce_time_cond(self, mode):
+        space_order = 8
+        nx, ny = 11, 11
+
+        grid = Grid(shape=(nx, ny))
+        tt = grid.time_dim
+        nt = 20
+        ct = ConditionalDimension(name="ct", parent=tt, factor=2)
+
+        ux = TimeFunction(name="ux", grid=grid, time_order=1, space_order=space_order)
+        g = TimeFunction(name="g", grid=grid, dimensions=(ct, ), shape=(int(nt/2),),
+                         time_dim=ct)
+        h = TimeFunction(name="h", grid=grid, dimensions=(ct, ), shape=(int(nt/2),),
+                         time_dim=ct)
+
+        op = Operator([Eq(g, 0), Eq(ux.forward, tt), Inc(g, ux), Inc(h, ux)], name="Op")
+        assert_structure(op, ['t', 't,x,y', 't,x,y'], 'txyxy')
+
+        # Make sure the two allreduce calls are in the time the loop
+        iters = FindNodes(Iteration).visit(op)
+        for i in iters:
+            if i.dim.is_Time:
+                assert len(FindNodes(Call).visit(i)) == 2  # Two allreduce
+            else:
+                assert len(FindNodes(Call).visit(i)) == 0
+
+        op.apply(time_m=0, time_M=nt-1)
+
+        expected = [nx * ny * max(t-1, 0) for t in range(0, nt, 2)]
+        assert np.allclose(g.data, expected)
+        assert np.allclose(h.data, expected)
+
 
 class TestOperatorAdvanced: