api: fix buffering with multiple conditions

Author: mloubout

devito/ir/equations/equation.py

@@ -213,7 +213,7 @@ def __new__(cls, *args, **kwargs):
                                   relations=ordering.relations, mode='partial')
         ispace = IterationSpace(intervals, iterators)
 
-        # Construct the conditionals and replace the ConditionalDimensions in `expr`
+        # Construct the conditionals
         conditionals = {}
         for d in ordering:
             if not d.is_Conditional:
@@ -225,14 +225,6 @@ def __new__(cls, *args, **kwargs):
                 if d._factor is not None:
                     cond = d.relation(cond, GuardFactor(d))
                 conditionals[d] = cond
-            # Replace dimension with index
-            index = d.index
-            if d.condition is not None and d in expr.free_symbols:
-                index = index - relational_min(d.condition, d.parent)
-                shift = relational_shift(d.condition, d.parent)
-            else:
-                shift = 0
-            expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift})
 
         # Merge conditionals when possible. E.g if we have an implicit_dim
         # and there is a dimension with the same parent, we ca merged
@@ -249,6 +241,14 @@ def __new__(cls, *args, **kwargs):
 
         conditionals = frozendict(conditionals)
 
+        # Replace the ConditionalDimensions in `expr`
+        for d, cond in conditionals.items():
+            # Replace dimension with index
+            index = d.index
+            index = index - relational_min(cond, d.parent)
+            shift = relational_shift(cond, d.parent)
+            expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift})
+
         # Lower all Differentiable operations into SymPy operations
         rhs = diff2sympy(expr.rhs)
 

devito/ir/support/vector.py

@@ -128,6 +128,7 @@ def __lt__(self, other):
                         return True
                     elif q_positive(i):
                         return False
+
                 raise TypeError("Non-comparable index functions") from e
 
         return False
@@ -164,6 +165,7 @@ def __gt__(self, other):
                         return True
                     elif q_negative(i):
                         return False
+
                 raise TypeError("Non-comparable index functions") from e
 
         return False
@@ -203,6 +205,7 @@ def __le__(self, other):
                         return True
                     elif q_positive(i):
                         return False
+
                 raise TypeError("Non-comparable index functions") from e
 
         # Note: unlike `__lt__`, if we end up here, then *it is* <=. For example,

devito/types/relational.py

@@ -302,7 +302,22 @@ def relational_shift(expr, s):
     if not expr.has(s):
         return 0
 
-    try:
-        return expr._as_min
-    except (TypeError, AttributeError):
+    return _relational_shift(expr, s)
+
+
+@singledispatch
+def _relational_shift(s, expr):
+    return 0
+
+
+@_relational_shift.register(sympy.Or)
+@_relational_shift.register(sympy.And)
+def _(expr, s):
+    return sum([_relational_shift(e, s) for e in expr.args])
+
+
+@_relational_shift.register(sympy.Eq)
+def _(expr, s):
+    if isinstance(expr.lhs, sympy.Mod):
         return 0
+    return expr._as_min

tests/test_buffering.py

@@ -754,7 +754,7 @@ def test_buffer_reuse():
     assert all(np.all(vsave.data[i-1] == i + 1) for i in range(1, nt + 1))
 
 
-def test_multi_cond():
+def test_multi_cond_v0():
     grid = Grid((3, 3))
     nt = 5
 
@@ -774,14 +774,47 @@ def test_multi_cond():
     T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0)
 
     eqs = [Eq(T, grid.time_dim)]
-    # this to save times from 0 to nt - 2
+    # This saves
+    # - All subsampled times since ct1 is the dimension of f
+    # - The last time step (ntmod - 2) through ctend (since it's set as ct1 or ctend)
+    eqs.append(Eq(f, T, implicit_dims=ctend))
+
+    # run operator with buffering
+    op = Operator(eqs, opt='buffering')
+    op.apply(time_m=0, time_M=ntmod-2)
+
+    for i in range(nt-1):
+        assert np.allclose(f.data[i], i*2)
+    assert np.allclose(f.data[nt-1], ntmod - 2)
+
+
+def test_multi_cond_v1():
+    grid = Grid((3, 3))
+    nt = 5
+
+    x, y = grid.dimensions
+
+    factor = 2
+    ntmod = (nt - 1) * factor + 1
+
+    ct1 = ConditionalDimension(name="ct1", parent=grid.time_dim,
+                               factor=factor, relation=Or,
+                               condition=CondEq(grid.time_dim, ntmod - 2))
+
+    f = TimeFunction(grid=grid, name='f', time_order=0,
+                     space_order=0, save=nt, time_dim=ct1)
+    T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0)
+
+    eqs = [Eq(T, grid.time_dim)]
+    # This saves
+    # - All subsampled times since ct1 is the dimension of f with factor 2
+    # - The last time step (ntmod - 2) since ct1 also has the condition for ntmod - 2
     eqs.append(Eq(f, T))
-    # this to save the last time sample nt - 1
-    eqs.append(Eq(f.forward, T+1, implicit_dims=ctend))
 
     # run operator with buffering
     op = Operator(eqs, opt='buffering')
     op.apply(time_m=0, time_M=ntmod-2)
 
-    for i in range(nt):
+    for i in range(nt-1):
         assert np.allclose(f.data[i], i*2)
+    assert np.allclose(f.data[nt-1], ntmod - 2)