compiler: Improve halo-exchange hoisting

Author: FabioLuporini

devito/mpi/halo_scheme.py

@@ -43,6 +43,10 @@ def __hash__(self):
         return hash((self.loc_indices, self.loc_dirs, self.halos, self.dims,
                      self.bundle))
 
+    @property
+    def loc_values(self):
+        return frozenset(self.loc_indices.values())
+
     def union(self, other):
         """
         Return a new HaloSchemeEntry that is the union of this and `other`.
@@ -413,6 +417,37 @@ def loc_values(self):
     def arguments(self):
         return self.dimensions | set(flatten(self.honored.values()))
 
+    def issubset(self, other):
+        """
+        Check if `self` is a subset of `other`.
+        """
+        if not isinstance(other, HaloScheme):
+            return False
+
+        if not all(f in other.fmapper for f in self.fmapper):
+            return False
+
+        for f, hse0 in self.fmapper.items():
+            hse1 = other.fmapper[f]
+
+            # Clearly, `hse0`'s halos must be a subset of `hse1`'s halos...
+            if not hse0.halos.issubset(hse1.halos) or \
+               hse0.bundle is not hse1.bundle:
+                return False
+
+            # But now, to be a subset, `hse0`'s must be expecting such halos
+            # at a time index that is less than or equal to that of `hse1`
+            if hse0.loc_dirs != hse1.loc_dirs:
+                return False
+
+            loc_dirs = hse0.loc_dirs
+            loc_indices = {**hse0.loc_indices, **hse1.loc_indices}
+            projected_loc_indices, _ = process_loc_indices(loc_indices, loc_dirs)
+            if projected_loc_indices != hse1.loc_indices:
+                return False
+
+        return True
+
     def project(self, functions):
         """
         Create a new HaloScheme that only retains the HaloSchemeEntries corresponding

devito/passes/iet/mpi.py

@@ -56,73 +56,81 @@ def _drop_reduction_halospots(iet):
 
 def _hoist_invariant(iet):
     """
-    Hoist HaloSpots from inner to outer Iterations where all data dependencies
-    would be honored. This pass avoids redundant halo exchanges when the same
-    data is redundantly exchanged within the same Iteration tree level.
+    Hoist redundant HaloSpots out of Iterations.
 
     Example:
-                                   haloupd v[t0]
-    for time                       for time
-      W v[t1]- R v[t0]               W v[t1]- R v[t0]
-      haloupd v[t1]                  haloupd v[t1]
-      R v[t1]                        R v[t1]
-      haloupd v[t0]                  R v[t0]
-      R v[t0]
 
+                                    haloupd v[t0]
+                                    haloupd v[t0]'
+    for time                        for time
+      haloupd v[t0]
+      W v[t1]- R v[t0]   --->         W v[t1]- R v[t0]
+      haloupd v[t1]                   haloupd v[t1]
+      R v[t1]                         R v[t1]
+      haloupd v[t0]'
+      R v[t0]                         R v[t0]
+
+    Where `haloupd v[t0]` and `haloupd v[t0]'` are subsets of `haloupd v[t1]`.
     """
-
-    # Precompute scopes to save time
-    scopes = {i: Scope([e.expr for e in v]) for i, v in MapNodes().visit(iet).items()}
-
-    # Analysis
-    hsmapper = {}
-    imapper = defaultdict(list)
-
     cond_mapper = _make_cond_mapper(iet)
     iter_mapper = _filter_iter_mapper(iet)
 
+    hsmapper = defaultdict(list)
+    imapper = defaultdict(list)
     for it, halo_spots in iter_mapper.items():
+        scope = Scope([e.expr for e in FindNodes(Expression).visit(it)])
+
         for hs0, hs1 in combinations(halo_spots, r=2):
             if _check_control_flow(hs0, hs1, cond_mapper):
                 continue
 
-            # If there are overlapping loc_indices, skip
-            hs0_mdims = hs0.halo_scheme.loc_values
-            hs1_mdims = hs1.halo_scheme.loc_values
-            if hs0_mdims.intersection(hs1_mdims):
+            functions = set(hs0.fmapper) & set(hs1.fmapper)
+            if not functions:
+                continue
+            functions = sorted(functions, key=lambda f: f.name)  # Determinism
+
+            hsf0 = hs0.halo_scheme.project(functions)
+            hsf1 = hs1.halo_scheme.project(functions)
+
+            # NOTE: lexicographic order is important here, as if both `hsf0` and
+            # `hsf1` are iteration-carried, we must hoist the later one in the
+            # control flow (that is `hsf1`)
+            flag0 = hsf1.issubset(hsf0) and _is_iter_carried(hsf1, scope)
+            flag1 = hsf0.issubset(hsf1) and _is_iter_carried(hsf0, scope)
+            if flag0 and flag1:
+                # Special case: `_merge_halospots` will merge `hsf0` and `hsf1`
+                # anyway, so we can avoid a useless hoisting
+                continue
+            elif flag0:
+                hs, hsf = hs1, hsf1
+            elif flag1:
+                hs, hsf = hs0, hsf0
+            else:
+                # No hoisting possible, skip
                 continue
 
-            for f, v in hs1.fmapper.items():
-                if f not in hs0.functions:
-                    continue
+            # At this point, we must infer valid loc_indices for `found`
+            hhs = hsf.drop(functions)  # hhs -> hoisted HaloScheme
+            for f, hse in hsf.fmapper.items():
+                loc_indices = {}
+                for d, v in hse.loc_indices.items():
+                    if v in it.uindices:
+                        loc_indices[d] = v.symbolic_min.subs(it.dim, it.start)
+                    else:
+                        loc_indices[d] = v
+
+                hhs = hhs.add(f, hse._rebuild(loc_indices=loc_indices))
+
+            hsmapper[hs].extend(functions)
+            imapper[it].append(hhs)
 
-                for dep in scopes[it].d_flow.project(f):
-                    if not any(r(dep, hs1, v.loc_indices) for r in rules):
-                        break
-                else:
-                    # `hs1`` can be hoisted out of `it`, but we need to infer valid
-                    # loc_indices
-                    hse = hs1.halo_scheme.fmapper[f]
-                    loc_indices = {}
-
-                    for d, v in hse.loc_indices.items():
-                        if v in it.uindices:
-                            loc_indices[d] = v.symbolic_min.subs(it.dim, it.start)
-                        else:
-                            loc_indices[d] = v
-
-                    hse = hse._rebuild(loc_indices=loc_indices)
-                    hs1.halo_scheme.fmapper[f] = hse
-
-                    hsmapper[hs1] = hsmapper.get(hs1, hs1.halo_scheme).drop(f)
-                    imapper[it].append(hs1.halo_scheme.project(f))
-
-    mapper = {i: HaloSpot(i._rebuild(), HaloScheme.union(hss))
-              for i, hss in imapper.items()}
-    mapper.update({i: i.body if hs.is_void else i._rebuild(halo_scheme=hs)
-                   for i, hs in hsmapper.items()})
     # Transform the IET hoisting/dropping HaloSpots as according to the analysis
+    mapper = {hs: hs._rebuild(halo_scheme=hs.halo_scheme.drop(functions))
+              for hs, functions in hsmapper.items()}
+    mapper.update({i: HaloSpot(i._rebuild(), HaloScheme.union(hss))
+                   for i, hss in imapper.items()})
     iet = Transformer(mapper, nested=True).visit(iet)
+
     return iet
 
 
@@ -134,18 +142,16 @@ def _merge_halospots(iet):
 
     Example:
 
-    for time                       for time
-      haloupd v[t0]                  haloupd v[t0], h
-      W v[t1]- R v[t0]               W v[t1]- R v[t0]
+    for time                           for time
+      haloupd v[t0]                      haloupd v[t0], h
+      W v[t1]- R v[t0]      --->         W v[t1]- R v[t0]
       haloupd v[t0], h
-      W g[t1]- R v[t0], h            W g[t1]- R v[t0], h
+      W g[t1]- R v[t0], h                W g[t1]- R v[t0], h
     """
-
-    # Analysis
-    mapper = {}
     cond_mapper = _make_cond_mapper(iet)
     iter_mapper = _filter_iter_mapper(iet)
 
+    mapper = {}
     for it, halo_spots in iter_mapper.items():
         scope = Scope([e.expr for e in FindNodes(Expression).visit(it)])
 
@@ -157,19 +163,19 @@ def _merge_halospots(iet):
 
             for f, v in hs1.fmapper.items():
                 for dep in scope.d_flow.project(f):
-                    if not any(rule(dep, hs1, v.loc_indices) for rule in rules):
+                    if not any(r(dep, hs1, v.loc_indices) for r in rules):
                         break
                 else:
-                    # All good -- `hs1` can be merged with `hs0`
-                    hs = hs1.halo_scheme.project(f)
-                    mapper[hs0] = HaloScheme.union([mapper.get(hs0, hs0.halo_scheme), hs])
+                    # All good -- `hsf1` can be merged within `hs0`
+                    hsf1 = hs1.halo_scheme.project(f)
+                    mapper[hs0] = HaloScheme.union(
+                        [mapper.get(hs0, hs0.halo_scheme), hsf1]
+                    )
                     mapper[hs1] = mapper.get(hs1, hs1.halo_scheme).drop(f)
 
-    # Post-process analysis
+    # Transform the IET merging/dropping HaloSpots as according to the analysis
     mapper = {i: i.body if hs.is_void else i._rebuild(halo_scheme=hs)
               for i, hs in mapper.items()}
-
-    # Transform the IET merging/dropping HaloSpots as according to the analysis
     iet = Transformer(mapper, nested=True).visit(iet)
 
     return iet
@@ -375,6 +381,23 @@ def _check_control_flow(hs0, hs1, cond_mapper):
     return cond0 != cond1
 
 
+def _is_iter_carried(hs, scope):
+    """
+    True if the HaloScheme `hs` is iteration-carried, i.e., if the exchanged halo
+    *only* uses values computed in a past iteration, False otherwise.
+    """
+    for f, hse in hs.fmapper.items():
+        try:
+            writes = scope.writes[f]
+        except KeyError:
+            continue
+
+        if any(hse.loc_values.intersection(w.aindices) for w in writes):
+            return False
+
+    return True
+
+
 # Code motion rules -- if the retval is True, then it means the input `dep` is not
 # a stopper to moving the HaloSpot `hs` around
 

tests/test_mpi.py

@@ -1104,6 +1104,8 @@ def test_avoid_redundant_haloupdate(self, mode):
 
         calls = FindNodes(Call).visit(op)
         assert len(calls) == 1
+        calls = FindNodes(Call).visit(get_time_loop(op))
+        assert len(calls) == 1
 
     @pytest.mark.parallel(mode=1)
     def test_avoid_redundant_haloupdate_cond(self, mode):
@@ -1123,7 +1125,7 @@ def test_avoid_redundant_haloupdate_cond(self, mode):
                        # access `f` at `t`, not `t+1` through factor subdim!
                        Eq(g, f[t, j] + 1, implicit_dim=t_sub)])
 
-        calls = FindNodes(Call).visit(op)
+        calls = FindNodes(Call).visit(get_time_loop(op))
         assert len(calls) == 1
         assert calls[0].functions[0] is f
 
@@ -1243,6 +1245,8 @@ def test_hoist_haloupdate_if_no_flowdep(self, mode):
 
         calls = FindNodes(Call).visit(op)
         assert len(calls) == 1
+        calls = FindNodes(Call).visit(get_time_loop(op))
+        assert len(calls) == 1
 
         # Below, there is a flow-dependence along `x`, so a further halo update
         # before the Inc is required
@@ -1405,16 +1409,22 @@ def test_avoid_haloupdate_if_flowdep_along_other_dim(self, mode):
 
         calls = FindNodes(Call).visit(op)
         assert len(calls) == 1
+        calls = FindNodes(Call).visit(get_time_loop(op))
+        assert len(calls) == 1
 
         op.apply(time_M=1)
         glb_pos_map = f.grid.distributor.glb_pos_map
         R = 1e-07  # Can't use np.all due to rounding error at the tails
         if LEFT in glb_pos_map[x]:
-            assert np.allclose(f.data_ro_domain[0, :5], [5.6, 6.8, 5.8, 4.8, 4.8], rtol=R)
-            assert np.allclose(g.data_ro_domain[0, :5], [2., 5.8, 5.8, 4.8, 4.8], rtol=R)
+            assert np.allclose(f.data_ro_domain[0, :5], [5.6, 6.8, 5.8, 4.8, 4.8],
+                               rtol=R)
+            assert np.allclose(g.data_ro_domain[0, :5], [2., 5.8, 5.8, 4.8, 4.8],
+                               rtol=R)
         else:
-            assert np.allclose(f.data_ro_domain[0, 5:], [4.8, 4.8, 4.8, 4.8, 3.6], rtol=R)
-            assert np.allclose(g.data_ro_domain[0, 5:], [4.8, 4.8, 4.8, 4.8, 2.], rtol=R)
+            assert np.allclose(f.data_ro_domain[0, 5:], [4.8, 4.8, 4.8, 4.8, 3.6],
+                               rtol=R)
+            assert np.allclose(g.data_ro_domain[0, 5:], [4.8, 4.8, 4.8, 4.8, 2.],
+                               rtol=R)
 
     @pytest.mark.parallel(mode=2)
     def test_unmerge_haloupdate_if_no_locindices(self, mode):
@@ -1436,7 +1446,7 @@ def test_unmerge_haloupdate_if_no_locindices(self, mode):
         calls = FindNodes(Call).visit(op)
         assert len(calls) == 2
 
-        titer = op.body.body[-1].body[0]
+        titer = get_time_loop(op)
         assert titer.dim is grid.time_dim
         assert len(titer.nodes[0].body[0].body[0].body[0].body) == 1
         assert titer.nodes[0].body[0].body[0].body[0].body[0].is_Call
@@ -1483,8 +1493,8 @@ def test_merge_and_hoist_haloupdate_if_diff_locindices(self, mode):
         """
         This test is a revisited, more complex version of
         `test_merge_haloupdate_if_diff_locindices`, also checking hoisting.
-        And in addition to checking the generated code,
-        it also checks the numerical output.
+        And in addition to checking the generated code, it also checks the
+        numerical output.
 
         In the Operator there are three Eqs:
 
@@ -1851,11 +1861,12 @@ def test_haloupdate_buffer1_v2(self, mode):
         op.cfunction
 
         # Ensure there's a halo exchange over v2 before the rec interpolation
-        section1 = op.body.body[-1].body[1].nodes[1]
-        assert section1.is_Section
-        calls = FindNodes(HaloUpdateCall).visit(section1)
-        assert len(calls) == 1
-        assert calls[0].arguments[0] is v2
+        calls = FindNodes(HaloUpdateCall).visit(op)
+        assert len(calls) == 3
+        calls = FindNodes(HaloUpdateCall).visit(get_time_loop(op))
+        assert len(calls) == 2
+        assert calls[0].arguments[0] is v1
+        assert calls[1].arguments[0] is v2
 
 
 class TestOperatorAdvanced:
@@ -3162,6 +3173,14 @@ def test_halo_structure(self, mode):
         assert calls[0].functions[1].name == 'v'
 
 
+def get_time_loop(op):
+    iters = FindNodes(Iteration).visit(op)
+    for i in iters:
+        if i.dim.is_Time:
+            return i
+    assert False
+
+
 if __name__ == "__main__":
     # configuration['mpi'] = 'overlap'
     # TestDecomposition().test_reshape_left_right()