Merge pull request #2673 from devitocodes/allreduce-v4

mpi: fix multi conditional for allreduce

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -453,13 +453,18 @@ def callback(self, clusters, prefix, seen=None):
 def reduction_comms(clusters):
     processed = []
     fifo = []
+
+    def _update(reductions):
+        for _, reds in groupby(reductions, key=lambda r: r.ispace):
+            reds = list(reds)
+            exprs = flatten([dr.exprs for dr in reds])
+            processed.append(reds[0].rebuild(exprs=exprs))
+
     for c in clusters:
         # Schedule the global distributed reductions encountered before `c`,
         # 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 = flatten([dr.exprs for dr in reds])
-            processed.append(c.rebuild(exprs=exprs, ispace=ispace))
+        _update(found)
 
         # Detect the global distributed reductions in `c`
         for e in c.exprs:
@@ -494,10 +499,7 @@ def reduction_comms(clusters):
         processed.append(c)
 
     # Leftover reductions are placed at the very end
-    for ispace, reds in groupby(fifo, key=lambda r: r.ispace):
-        reds = list(reds)
-        exprs = flatten([dr.exprs for dr in reds])
-        processed.append(reds[0].rebuild(exprs=exprs, ispace=ispace))
+    _update(fifo)
 
     return processed
 

tests/test_mpi.py

@@ -2068,7 +2068,7 @@ 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)
+    @pytest.mark.parallel(mode=2)
     def test_multi_allreduce_time_cond(self, mode):
         space_order = 8
         nx, ny = 11, 11
@@ -2101,6 +2101,64 @@ def test_multi_allreduce_time_cond(self, mode):
         assert np.allclose(g.data, expected)
         assert np.allclose(h.data, expected)
 
+    @pytest.mark.parallel(mode=2)
+    def test_allreduce_multicond(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)
+        ct2 = ConditionalDimension(name="ct2", parent=tt, factor=8)
+
+        ux = TimeFunction(name="ux", grid=grid, time_order=1, space_order=space_order)
+        uy = TimeFunction(name="uy", 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),
+                       Eq(uy, tt, implicit_dims=ct2)],
+                      name="Op")
+        assert_structure(op, ['t', 't,x,y', 't,x,y', 't,x,y'], 'txyxyxy')
+
+        # 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
+
+        # Check conditionals
+        conds = FindNodes(Conditional).visit(op)
+        assert len(conds) == 3
+        # First one is just g initialization
+        sym0 = FindSymbols().visit(conds[0])
+        assert set(sym0) == {ct.symbolic_factor, tt, g}
+        assert grid.distributor._obj_comm not in sym0
+        # Second one is g and h and allreduce
+        sym1 = FindSymbols().visit(conds[1])
+        assert g in sym1
+        assert h in sym1
+        assert ux in sym1
+        # The allreduce
+        assert grid.distributor._obj_comm in sym1
+        # Last one is only uy
+        sym2 = FindSymbols().visit(conds[-1])
+        assert g not in sym2
+        assert h not in sym2
+        assert uy in sym2
+        assert grid.distributor._obj_comm not in sym2
+
+        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: