compiler: Fix HaloSpot merging

Author: FabioLuporini

devito/mpi/halo_scheme.py

@@ -47,17 +47,18 @@ def __hash__(self):
     def loc_values(self):
         return frozenset(self.loc_indices.values())
 
-    def union(self, other):
+    def merge(self, other):
         """
-        Return a new HaloSchemeEntry that is the union of this and `other`.
-        The `loc_indices` and `loc_dirs` must be the same, otherwise an
-        exception is raised.
+        Return a new HaloSchemeEntry that is the result of merging `other`
+        into `self`, irrespective of whether the `loc_indices` and `loc_dirs`
+        are the same or not. Thus, the returned HaloSchemeEntry will have
+        in general more halo exchanges than `self`, or exactly the same halo
+        exchanges in the worst case).
         """
-        if self.loc_indices != other.loc_indices or \
-           self.loc_dirs != other.loc_dirs or \
+        if self.loc_dirs != other.loc_dirs or \
            self.bundle is not other.bundle:
             raise HaloSchemeException(
-                "Inconsistency found while building a HaloScheme"
+                "Inconsistency found while merging HaloSchemeEntries"
             )
 
         halos = self.halos | other.halos
@@ -66,6 +67,22 @@ def union(self, other):
         return HaloSchemeEntry(self.loc_indices, self.loc_dirs, halos, dims,
                                bundle=self.bundle, getters=self.getters)
 
+    def union(self, other):
+        """
+        Return a new HaloSchemeEntry that is the union of this and `other`.
+        The `loc_indices` and `loc_dirs` must be the same, otherwise an
+        exception is raised. This is a more restrictive version of `merge`,
+        which is used when we want to ensure that the halo exchanges are
+        performed at the same time index, i.e., the same `loc_indices` are
+        are expected.
+        """
+        if self.loc_indices != other.loc_indices:
+            raise HaloSchemeException(
+                "Inconsistency found while taking the union of HaloSchemeEntries"
+            )
+
+        return self.merge(other)
+
 
 Halo = namedtuple('Halo', 'dim side')
 
@@ -481,6 +498,15 @@ def add(self, f, hse):
         fmapper[f] = hse
         return HaloScheme.build(fmapper, self.honored)
 
+    def merge(self, hs):
+        """
+        Create a new HaloScheme that is the result of merging `hs` into `self`.
+        """
+        fmapper = dict(self.fmapper)
+        for f, hse in hs.fmapper.items():
+            fmapper[f] = fmapper.get(f, hse).merge(hse)
+        return HaloScheme.build(fmapper, self.honored)
+
 
 def classify(exprs, ispace):
     """

devito/passes/iet/mpi.py

@@ -7,7 +7,6 @@
                            FindWithin, MapNodes, MapHaloSpots, Transformer,
                            retrieve_iteration_tree)
 from devito.ir.support import PARALLEL, Scope
-from devito.mpi.halo_scheme import HaloScheme
 from devito.mpi.reduction_scheme import DistReduce
 from devito.mpi.routines import HaloExchangeBuilder, ReductionBuilder
 from devito.passes.iet.engine import iet_pass
@@ -154,7 +153,7 @@ def _merge_halospots(iet):
 
                 # If the `loc_indices` differ, we rely on hoisting to optimize
                 # `hsf1` out of `it`, otherwise we just drop it
-                if hsf0.loc_values != hsf1.loc_values:
+                if not _semantical_eq_loc_indices(hsf0, hsf1):
                     continue
 
                 mapper.drop(hs1, f)
@@ -416,7 +415,7 @@ def add(self, node, hss):
         """
         v = self.get(node)
         if isinstance(v, HaloSpot):
-            hss = HaloScheme.union([v.halo_scheme, hss])
+            hss = v.halo_scheme.merge(hss)
             hs = v._rebuild(halo_scheme=hss)
         else:
             hs = HaloSpot(v._rebuild(), hss)
@@ -515,3 +514,21 @@ def _is_mergeable(hsf0, hsf1, scope):
 
     # Finally, check the data dependences would be satisfied
     return _is_iter_carried(hsf1, scope)
+
+
+def _semantical_eq_loc_indices(hsf0, hsf1):
+    if hsf0.loc_indices != hsf1.loc_indices:
+        return False
+
+    for v0, v1 in zip(hsf0.loc_values, hsf1.loc_values):
+        if v0 is v1:
+            continue
+
+        # Special case: they might be syntactically different, but semantically
+        # equivalent, e.g., `t0` and `t1` with same modulus
+        if v0.modulo == v1.modulo == 1:
+            continue
+
+        return False
+
+    return True

tests/test_mpi.py

@@ -1580,7 +1580,7 @@ def test_merge_and_hoist_haloupdate_if_diff_locindices(self, mode):
             assert np.allclose(h.data_ro_domain[0, 5:], [4.4, 4.4, 4.4, 3.4, 3.1], rtol=R)
 
     @pytest.mark.parallel(mode=1)
-    def test_merge_and_hoist_haloupdate_if_diff_locindices_v2(self, mode):
+    def test_merge_haloupdate_if_diff_but_equivalent_locindices(self, mode):
         grid = Grid(shape=(65, 65, 65))
 
         v1 = TimeFunction(name='v1', grid=grid, space_order=2, time_order=1,
@@ -1597,7 +1597,7 @@ def test_merge_and_hoist_haloupdate_if_diff_locindices_v2(self, mode):
         op = Operator(eqns)
         op.cfunction
 
-        calls, _ = check_halo_exchanges(op, 3, 2)
+        calls, _ = check_halo_exchanges(op, 2, 2)
         for i, v in enumerate([v2, v1]):
             assert calls[i].arguments[0] is v
 
@@ -2962,6 +2962,27 @@ def init(f, v=1):
         assert np.isclose(norm(u1), 12445251.87, rtol=1e-7)
         assert np.isclose(norm(v1), 147063.38, rtol=1e-7)
 
+    @pytest.mark.parallel(mode=1)
+    def test_interpolation_at_uforward(self, mode):
+        grid = Grid(shape=(10, 10, 10))
+        t = grid.stepping_dim
+
+        u = TimeFunction(name='u', grid=grid, space_order=2, time_order=2)
+
+        rec = SparseTimeFunction(name="rec", grid=grid, npoint=1, nt=10)
+
+        eqns = [Eq(u.forward, u.laplace + u.backward + 1),
+                rec.interpolate(expr=u.forward)]
+
+        op = Operator(eqns)
+
+        op.cfunction
+
+        calls, _ = check_halo_exchanges(op, 2, 1)
+        args = calls[0].arguments
+        assert args[-2].name == 't2'
+        assert args[-2].origin == t + 1
+
 
 def gen_serial_norms(shape, so):
     """