compiler: Simplify handling of scalar aliases within guards

Author: FabioLuporini

devito/passes/clusters/aliases.py

@@ -25,7 +25,6 @@
     CustomDimension, Eq, Hyperplane, IncrDimension, Indexed, ModuloDimension, Size,
     StencilDimension, Symbol, Temp, TempArray, TempFunction
 )
-from devito.types.dimension import ConditionalDimension, SubsamplingFactor
 from devito.types.grid import MultiSubDimension
 
 __all__ = ['cire']
@@ -149,7 +148,7 @@ def _aliases_from_clusters(self, cgroup, exclude, meta):
         # Schedule -> [Clusters]_k
         processed, subs = lower_schedule(schedule, meta, self.sregistry,
                                          self.opt_ftemps, self.opt_min_dtype,
-                                         self.opt_minmem, nclusters=len(cgroup))
+                                         self.opt_minmem)
 
         # [Clusters]_k -> [Clusters]_k (optimization)
         if self.opt_multisubdomain:
@@ -281,6 +280,8 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None):
 
 class CireInvariants(CireTransformerLegacy, Queue):
 
+    _q_guards_in_key = True
+
     def __init__(self, sregistry, options, platform):
         super().__init__(sregistry, options, platform)
 
@@ -854,7 +855,7 @@ def optimize_schedule_rotations(schedule, sregistry):
 
 
 def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
-                   opt_minmem, nclusters=1):
+                   opt_minmem):
     """
     Turn a Schedule into a sequence of Clusters.
     """
@@ -925,26 +926,15 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
             callback = lambda idx: obj[  # noqa: B023
                 [i + s for i, s in zip(idx, shift, strict=True)]  # noqa: B023
             ]
-            guards = meta.guards
         else:
             # Degenerate case: scalar expression
             assert writeto.size == 0
 
-            is_cond = any(isinstance(d, (SubsamplingFactor, ConditionalDimension))
-                          for d in pivot.free_symbols)
-            if meta.guards and is_cond and nclusters > 1:
-                # Scalar alias that depends on a guard, unsafe to lift out of the guard
-                # Do not alias
-                expression = None
-                callback = lambda idx: uxreplace(pivot, subs)  # noqa: B023
-            else:
-                dtype = sympy_dtype(pivot, base=meta.dtype, smin=opt_min_dtype)
-                obj = Temp(name=name, dtype=dtype, is_const=True)
-                expression = Eq(obj, uxreplace(pivot, subs))
+            dtype = sympy_dtype(pivot, base=meta.dtype, smin=opt_min_dtype)
+            obj = Temp(name=name, dtype=dtype, is_const=True)
+            expression = Eq(obj, uxreplace(pivot, subs))
 
-                callback = lambda idx: obj  # noqa: B023
-                # Only keep the guard if there is no cross-cluster reuse of the scalar
-                guards = meta.guards if nclusters == 1 else None
+            callback = lambda idx: obj  # noqa: B023
 
         # Create the substitution rules for the aliasing expressions
         subs.update({
@@ -970,8 +960,7 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
             properties[Hyperplane(writeto.itdims)] = {SEPARABLE}
 
         # Finally, build the alias Cluster
-        if expression is not None:
-            clusters.append(Cluster(expression, ispace, guards, properties))
+        clusters.append(Cluster(expression, ispace, meta.guards, properties))
 
     return clusters, subs
 

devito/passes/clusters/misc.py

@@ -97,19 +97,14 @@ def callback(self, clusters, prefix):
 
             properties = c.properties.filter(key)
 
-            # Lifted scalar clusters cannot be guarded
-            # as they would not be in the scope of the guarded clusters
-            # unless the guard is for an outer dimension
-            guards = {} if c.is_scalar and not (prefix[:-1] and c.guards) else c.guards
-
-            _lifted = c.rebuild(ispace=ispace, properties=properties, guards=guards)
-            if guards and clusters[max(n-1, 0)].guards != guards and _lifted.is_scalar:
-                # Heuristic: if the lifted Cluster has different guards than the
-                # previous one, then we are likely to end up with a separate
-                # Cluster, hence give up on lifting
-                processed.append(_lifted)
+            # If `c` is made of scalar expressions within guards, then we must keep
+            # it close to the adjacent Clusters for correctness
+            if c.is_scalar and c.guards:
+                items = processed
             else:
-                lifted.append(_lifted)
+                items = lifted
+
+            items.append(c.rebuild(ispace=ispace, properties=properties))
 
         return lifted + processed
 

tests/test_dse.py

@@ -2705,10 +2705,10 @@ def test_split_cond(self):
 
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
-        # The alias should have been lifted out of the condition
+        # Each guard should have its own alias for cos(time)
         assert 'float r0 = cos(time);' in str(body0(op))
         scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
-        assert len(scalars) == 1
+        assert len(scalars) == 2
 
     def test_split_cond_multi_alias(self):
         grid = Grid((11, 11))
@@ -2728,10 +2728,11 @@ def test_split_cond_multi_alias(self):
 
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
-        # The alias should have been lifted out of the condition
-        assert 'float r3 = cos(time);' in str(body0(op))
+        # Each guard should have its own aliases for cos(time) and sin(time)
+        assert 'const float r0 = sin(time) + cos(time)' in str(body0(op))
+        assert 'const float r1 = cos(time);' in str(body0(op))
         scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
-        assert len(scalars) == 5
+        assert len(scalars) == 3
 
     def test_multi_cond_no_split(self):
         grid = Grid((11, 11))
@@ -2742,7 +2743,7 @@ def test_multi_cond_no_split(self):
         ct2 = ConditionalDimension(name='ct2', parent=time, factor=4)
 
         eq0 = Eq(u.forward, u + cos(time), implicit_dims=ct)
-        # Not hoisting the inite would have this equation split the time
+        # Not hoisting the init would have this equation split the time
         # loop to initialize the alias for sin(time)
         eq1 = Eq(u.forward, u.forward + sin(time), implicit_dims=ct2)
         eq2 = Eq(u.forward, u.forward - sin(time), implicit_dims=ct)
@@ -2778,9 +2779,29 @@ def test_alias_with_conditional(self):
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
 
-        # The alias should have been lifted out of the condition
+        # # Each guard should have its own alias for cos(time/ctf)
         scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
-        assert not scalars
+        assert len(scalars) == 2
+
+    def test_scalar_alias_interp(self):
+        grid = Grid(shape=(11, 11))
+        time = grid.time_dim
+
+        t_sub = ConditionalDimension(name='t_sub', parent=time, factor=3)
+
+        f = TimeFunction(name='f', grid=grid, space_order=4)
+        s = SparseTimeFunction(name='s', grid=grid, npoint=1, nt=100)
+
+        eq = Eq(f.forward, f.laplace + .002)
+
+        rec = s.interpolate(expr=f, implicit_dims=t_sub)
+
+        op = Operator(rec + [eq])
+
+        op.apply(time_M=3)
+
+        assert np.isclose(norm(f), 254292.75, atol=1e-1, rtol=0)
+        assert np.isclose(norm(s), 191.44644, atol=1e-1, rtol=0)
 
 
 class TestIsoAcoustic: