compiler: fix multi-cond for multi-layer

Author: mloubout

devito/ir/equations/equation.py

@@ -13,7 +13,9 @@
 from devito.tools import (
     Pickable, Tag, as_hashable, filter_sorted, frozendict, reuse_if_unchanged
 )
-from devito.types import Eq, Inc, ReduceMax, ReduceMin, ReduceMinMax, relational_min
+from devito.symbolics import IntDiv, limits_mapper, uxreplace
+from devito.tools import Pickable, Tag, frozendict
+from devito.types import Eq, Inc, ReduceMax, ReduceMin, ReduceMinMax, relational_min, relational_shift
 
 __all__ = [
     'ClusterizedEq',
@@ -308,13 +310,10 @@ def __new__(cls, *args, **kwargs):
             index = d.index
             if d.condition is not None and d in expr.free_symbols:
                 index = index - relational_min(d.condition, d.parent)
-                # If there is a condition we might access on a non-factor
-                # index and need to make sure we don't overwrite the previous
-                # index
-                num = index + d.symbolic_factor - 1
+                shift = relational_shift(d.condition, d.parent)
             else:
-                num = index
-            expr = uxreplace(expr, {d: IntDiv(num, d.symbolic_factor)})
+                shift = 0
+            expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift})
 
         conditionals = frozendict(conditionals)
 

devito/passes/clusters/asynchrony.py

@@ -240,7 +240,7 @@ def _actions_from_update_memcpy(c, d, clusters, actions, sregistry):
 
     fetch = e.rhs.indices[d]
     fshift = {Forward: 1, Backward: -1}.get(direction, 0)
-    findex = fetch._subs(pd, pd + fshift)
+    findex = fetch + fshift if fetch.find(IntDiv) else fetch._subs(pd, pd + fshift)
 
     # If fetching into e.g. `ub[t1]` we might need to prefetch into e.g. `ub[t0]`
     tindex0 = e.lhs.indices[d]

devito/passes/clusters/buffering.py

@@ -3,7 +3,7 @@
 from itertools import chain
 
 import numpy as np
-from sympy import S
+from sympy import S, simplify
 
 from devito.exceptions import CompilationError
 from devito.ir import (
@@ -774,7 +774,7 @@ def infer_buffer_size(f, dim, clusters):
         slots = [Vector(i) for i in slots]
         size = int((vmax(*slots) - vmin(*slots) + 1)[0])
 
-    return size
+    return simplify(size)
 
 
 def offset_from_centre(d, indices):

devito/symbolics/extended_sympy.py

@@ -19,6 +19,7 @@
 )
 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,6 +49,11 @@ 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

@@ -3,7 +3,8 @@
 
 import sympy
 
-__all__ = ['Ge', 'Gt', 'Le', 'Lt', 'Ne', 'relational_max', 'relational_min']
+__all__ = ['Ge', 'Gt', 'Le', 'Lt', 'Ne', 'relational_max', 'relational_min',
+           'relational_shift']
 
 
 class AbstractRel:
@@ -291,3 +292,17 @@ def _(expr, s):
         return expr.gts
     else:
         return sympy.S.Infinity
+
+
+def relational_shift(expr, s):
+    """
+    Infer shift incurred by the expression. Generally only
+    applies when a CondEq is used as it adds a single value.
+    """
+    if not expr.has(s):
+        return 0
+
+    try:
+        return expr._as_min
+    except (TypeError, AttributeError):
+        return 0