compiler: fix various corner case of multi buffering

Author: mloubout

devito/ir/equations/equation.py

@@ -237,6 +237,14 @@ def __new__(cls, *args, **kwargs):
                     cond = d.relation(cond, GuardFactor(d))
                 conditionals[d] = cond
 
+        # 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})
+
         # Merge conditionals when possible. E.g if we have an implicit_dim
         # and there is a dimension with the same parent, we ca merged
         # its condition
@@ -247,19 +255,13 @@ def __new__(cls, *args, **kwargs):
                 if cd.parent == d.parent and cd != d:
                     cond = conditionals.pop(d)
                     mode = cd.relation and d.relation
-                    conditionals[cd] = mode(cond, conditionals[cd])
+                    if issubclass(mode, sympy.Or):
+                        conditionals[d] = cond
+                        conditionals.pop(cd)
+                    else:
+                        conditionals[cd] = mode(cond, conditionals[cd])
                     break
 
-        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/passes/clusters/buffering.py

@@ -3,7 +3,7 @@
 from itertools import chain
 
 import numpy as np
-from sympy import S, simplify
+from sympy import Mod, S, simplify
 
 from devito.exceptions import CompilationError
 from devito.ir import (
@@ -203,7 +203,7 @@ def callback(self, clusters, prefix):
                     guards = c.guards
 
                 properties = c.properties.sequentialize(d)
-                if not isinstance(d, BufferDimension):
+                if not isinstance(d, BufferDimension) and c.guards[d].has(Mod):
                     properties = properties.prefetchable(d)
                 # `c` may be a HaloTouch Cluster, so with no vision of the `bdims`
                 properties = properties.parallelize(v.bdims).affine(v.bdims)
@@ -377,7 +377,12 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs):
             buffer, = buffers
             xd = buffer.indices[dim]
         else:
-            size = infer_buffer_size(f, dim, clusters)
+            if len({c.guards[dim.root] for c in clusters}) > 1:
+                # Multiple clusters with different guards,
+                # will lead to conflicts in asynchrony with multiple (modulo) slots
+                size = 1
+            else:
+                size = infer_buffer_size(f, dim, clusters)
 
             if async_degree is not None:
                 if async_degree < size:

devito/symbolics/extended_sympy.py

@@ -18,7 +18,6 @@
 )
 from devito.types import Symbol
 from devito.types.basic import Basic
-from devito.types.relational import Ge
 
 __all__ = ['CondEq', 'CondNe', 'BitwiseNot', 'BitwiseXor', 'BitwiseAnd',  # noqa
            'LeftShift', 'RightShift', 'IntDiv', 'Terminal', 'CallFromPointer',
@@ -48,11 +47,6 @@ def canonical(self):
     def negated(self):
         return CondNe(*self.args, evaluate=False)
 
-    @property
-    def _as_min(self):
-        from devito.symbolics.extended_dtypes import INT
-        return INT(Ge(*self.args))
-
 
 class CondNe(sympy.Ne):
 

devito/types/relational.py

@@ -320,4 +320,5 @@ def _(expr, s):
 def _(expr, s):
     if isinstance(expr.lhs, sympy.Mod):
         return 0
-    return expr._as_min
+    from devito.symbolics.extended_dtypes import INT
+    return INT(Ge(*expr.args))

tests/test_buffering.py

@@ -818,3 +818,45 @@ def test_multi_cond_v1():
     for i in range(nt-1):
         assert np.allclose(f.data[i], i*2)
     assert np.allclose(f.data[nt-1], ntmod - 2)
+
+
+@pytest.mark.parametrize("factor", [1, 2, 3])
+def test_buffering_multi_cond(factor):
+    grid = Grid((16, 16))
+
+    nt = 5
+    ntmod = (nt - 1) * factor + 1
+
+    ct0 = ConditionalDimension(name="ct0", parent=grid.time_dim, factor=factor,
+                               relation=Or)
+    f = TimeFunction(grid=grid, name='f', time_order=0, space_order=0,
+                     time_dim=ct0, save=nt)
+    T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0)
+
+    eqs = []
+    eqs.append(Eq(T, grid.time_dim))
+
+    # conditional dimension for the last sample in the operator
+    ctend = ConditionalDimension(name="ctend", parent=grid.time_dim,
+                                 condition=CondEq(grid.time_dim, ntmod - 2),
+                                 relation=Or)
+
+    eqs.append(Eq(f, T))  # this to save times from 0 to nt - 2
+    # this to save the last time sample nt - 1
+    eqs.append(Eq(f.forward, T+1, implicit_dims=ctend))
+
+    # run operator with serialization
+    op = Operator(eqs, opt='buffering')
+    op.apply(time_m=0, time_M=ntmod-2)
+
+    # Now run backward as well with buffering
+
+    f_all = TimeFunction(grid=grid, name='f_all', time_order=0,
+                         space_order=0, time_dim=ct0, save=nt)
+
+    eq_all = [Eq(f_all, f)]
+    eq_all.append(Eq(f_all.forward, f.forward, implicit_dims=ctend))
+    op_all = Operator(eq_all, opt='buffering')
+    op_all.apply(time_m=0, time_M=ntmod-2)
+
+    assert np.allclose(f_all.data[:, 11, 11], factor * np.arange(nt))