compiler: fix terminal detection from dynamic dim bounds

Author: mloubout

devito/mpi/halo_scheme.py

@@ -5,7 +5,7 @@
 from operator import attrgetter
 
 import sympy
-from sympy import Max, Min, S
+from sympy import Max, Min
 
 from devito import configuration
 from devito.data import CENTER, CORE, LEFT, OWNED, RIGHT
@@ -514,31 +514,6 @@ def merge(self, hs):
             fmapper[f] = fmapper.get(f, hse).merge(hse)
         return HaloScheme.build(fmapper, self.honored)
 
-    def _is_iter_carried(self, scope):
-        """
-        True if the HaloScheme is iteration-carried, i.e., it induces
-        a halo exchange that requires values from the previous iteration(s); False
-        otherwise.
-        """
-
-        def rule0(dep):
-            # E.g., `dep=W<f,[t1, x]> -> R<f,[t0, x-1]>`, `d=t` => OK
-            return not any(dep.distance_mapper[d] is S.Infinity for d in dep.cause)
-
-        def rule1(dep, loc_indices):
-            # E.g., `dep=W<f,[t1, x+1]> -> R<f,[t1, xl+1]>`, `loc_indices={t: t0}` => OK
-            return any(dep.distance_mapper[d] == 0 and
-                       dep.source[d] is not v and
-                       dep.sink[d] is not v
-                       for d, v in loc_indices.items())
-
-        for f, v in self.fmapper.items():
-            for dep in scope.d_flow.project(f):
-                if not rule0(dep) and not rule1(dep, v.loc_indices):
-                    return False
-
-        return True
-
 
 def classify(exprs, ispace):
     """

devito/passes/iet/mpi.py

@@ -206,7 +206,7 @@ def _hoist_invariant(iet):
                 # Ensure there's another HaloScheme that could cover for
                 # us should we get hoisted while still satisfying the
                 # data dependences
-                if hsf1.issubset(hsf0) and hsf1._is_iter_carried(scope):
+                if hsf1.issubset(hsf0) and _is_iter_carried(hsf1, scope):
                     hs, hsf = hs1, hsf1
                 elif hsf0.issubset(hsf1) and hs0 is halo_spots[0]:
                     # Special case
@@ -474,6 +474,32 @@ def _check_control_flow(hs0, hs1, cond_mapper):
     return cond0 != cond1
 
 
+def _is_iter_carried(hsf, scope):
+    """
+    True if the provided HaloScheme `hsf` is iteration-carried, i.e., it induces
+    a halo exchange that requires values from the previous iteration(s); False
+    otherwise.
+    """
+
+    def rule0(dep):
+        # E.g., `dep=W<f,[t1, x]> -> R<f,[t0, x-1]>`, `d=t` => OK
+        return not any(dep.distance_mapper[d] is S.Infinity for d in dep.cause)
+
+    def rule1(dep, loc_indices):
+        # E.g., `dep=W<f,[t1, x+1]> -> R<f,[t1, xl+1]>`, `loc_indices={t: t0}` => OK
+        return any(dep.distance_mapper[d] == 0 and
+                   dep.source[d] is not v and
+                   dep.sink[d] is not v
+                   for d, v in loc_indices.items())
+
+    for f, v in hsf.fmapper.items():
+        for dep in scope.d_flow.project(f):
+            if not rule0(dep) and not rule1(dep, v.loc_indices):
+                return False
+
+    return True
+
+
 def _is_mergeable(hsf0, hsf1, scope):
     """
     True if `hsf1` can be merged into `hsf0`, i.e., if they are compatible
@@ -489,7 +515,7 @@ def _is_mergeable(hsf0, hsf1, scope):
         return False
 
     # Finally, check the data dependences would be satisfied
-    return hsf1._is_iter_carried(scope)
+    return _is_iter_carried(hsf1, scope)
 
 
 def _semantical_eq_loc_indices(hsf0, hsf1):

devito/symbolics/search.py

@@ -60,6 +60,8 @@ def __init__(self, query: Callable[[Expression], bool], deep: bool = False) -> N
     def _next(self, expr: Expression) -> Iterable[Expression]:
         if self.deep and expr.is_Indexed:
             return expr.indices
+        elif self.deep and q_dimension(expr):
+            return expr.bound_symbols
         elif q_leaf(expr):
             return ()
         return expr.args

tests/test_dle.py

@@ -755,7 +755,7 @@ def test_dynamic_nthreads(self):
         ('[Eq(f, 2*f)]', [2, 0, 0], False),
         ('[Eq(u, 2*u)]', [0, 2, 0, 0], False),
         ('[Eq(u, 2*u + f)]', [0, 3, 0, 0, 0, 0, 0], True),
-        ('[Eq(u, 2*u), Eq(f, u.dzr)]', [0, 2, 0, 0, 0], False)
+        ('[Eq(u, 2*u), Eq(f, u.dzr)]', [0, 2, 0, 0, 2, 0, 0], False)
     ])
     def test_collapsing(self, eqns, expected, blocking):
         grid = Grid(shape=(3, 3, 3))

tests/test_operator.py

@@ -1582,12 +1582,12 @@ def test_no_fission_as_illegal(self, exprs):
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z+2])',
           'Eq(tu[t,x,y,0], tu[t,x,y,0] + 1.)'),
-         '+++++', ['txyz', 'txyz', 'txy'], 'txyzz'),
+         '+++++++', ['txyz', 'txyz', 'txy'], 'txyzxyz'),
         # 7) WAR 1->2, 2->3
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z+2])',
           'Eq(tw[t,x,y,z], tv[t,x,y,z-1] + 1.)'),
-         '++++++', ['txyz', 'txyz', 'txyz'], 'txyzzz'),
+         '++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzz'),
         # 8) WAR 1->2; WAW 1->3
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x+2,y,z])',
@@ -1597,7 +1597,7 @@ def test_no_fission_as_illegal(self, exprs):
         (('Eq(tu[t,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z-2])',
           'Eq(tw[t,x,y,z], tv[t,x,y+1,z] + 1.)'),
-         '+++++++', ['txyz', 'txyz', 'txyz'], 'txyzzyz'),
+         '+++++++++', ['txyz', 'txyz', 'txyz'], 'txyzxyzyz'),
         # 10) WAR 1->2; WAW 1->3
         (('Eq(tu[t-1,x,y,z], tu[t,x,y,z] + tv[t,x,y,z])',
           'Eq(tv[t,x,y,z], tu[t,x,y,z+2])',