compiler: Patch CSE in presence of conditionals

devito/passes/clusters/cse.py

@@ -98,24 +98,26 @@ def _cse(maybe_exprs, make, min_cost=1, mode='default'):
 
         # Create temporaries
         hit = max(targets.values())
-        temps = [Eq(make(), k) for k, v in targets.items() if v == hit]
+        chosen = [(k, make()) for k, v in targets.items() if v == hit]
 
         # Apply replacements
         # The extracted temporaries are inserted before the first expression
         # that contains it
+        scheduled = []
         updated = []
         for e in processed:
             pe = e
-            for t in temps:
-                pe, changed = _uxreplace(pe, {t.rhs: t.lhs})
-                if changed and t not in updated:
-                    updated.append(t)
+            for k, v in chosen:
+                pe, changed = _uxreplace(pe, {k: v})
+                if changed and v not in scheduled:
+                    updated.append(pe.func(v, k, operation=None))
+                    scheduled.append(v)
             updated.append(pe)
         processed = updated
 
         # Update `exclude` for the same reasons as above -- to rule out CSE across
         # Dimension-independent data dependences
-        exclude.update({t.lhs for t in temps})
+        exclude.update(scheduled)
 
     # At this point we may have useless temporaries (e.g., r0=r1). Let's drop them
     processed = _compact_temporaries(processed, exclude)

tests/test_dse.py

@@ -13,7 +13,7 @@
                     ConditionalDimension, DefaultDimension, Grid, Operator,
                     norm, grad, div, dimensions, switchconfig, configuration,
                     centered, first_derivative, solve, transpose, Abs, cos,
-                    sin, sqrt)
+                    sin, sqrt, Ge)
 from devito.exceptions import InvalidArgument, InvalidOperator
 from devito.finite_differences.differentiable import diffify
 from devito.ir import (Conditional, DummyEq, Expression, Iteration, FindNodes,
@@ -168,6 +168,29 @@ def test_cse_temp_order():
     assert type(args[2]) is CTemp
 
 
+def test_cse_w_conditionals():
+    grid = Grid(shape=(10, 10, 10))
+    x, _, _ = grid.dimensions
+
+    cd = ConditionalDimension(name='cd', parent=x, condition=Ge(x, 4),
+                              indirect=True)
+
+    f = Function(name='f', grid=grid)
+    g = Function(name='g', grid=grid)
+    h = Function(name='h', grid=grid)
+    a0 = Function(name='a0', grid=grid)
+    a1 = Function(name='a1', grid=grid)
+
+    eqns = [Eq(h, a0, implicit_dims=cd),
+            Eq(a0, a0 + f*g, implicit_dims=cd),
+            Eq(a1, a1 + f*g, implicit_dims=cd)]
+
+    op = Operator(eqns)
+
+    assert_structure(op, ['x,y,z'], 'xyz')
+    assert len(FindNodes(Conditional).visit(op)) == 1
+
+
 @pytest.mark.parametrize('expr,expected', [
     ('2*fa[x] + fb[x]', '2*fa[x] + fb[x]'),
     ('fa[x]**2', 'fa[x]*fa[x]'),