api: fix handling of multiple conditions for buffering

devito/arch/archinfo.py

@@ -495,13 +495,15 @@ def parse_product_arch():
     return None
 
 
+device_vars = (
+    'CUDA_VISIBLE_DEVICES',
+    'NVIDIA_VISIBLE_DEVICES',
+    'ROCR_VISIBLE_DEVICES',
+    'HIP_VISIBLE_DEVICES'
+)
+
+
 def get_visible_devices():
-    device_vars = (
-        'CUDA_VISIBLE_DEVICES',
-        'NVIDIA_VISIBLE_DEVICES',
-        'ROCR_VISIBLE_DEVICES',
-        'HIP_VISIBLE_DEVICES'
-    )
     for v in device_vars:
         try:
             return v, tuple(int(i) for i in os.environ[v].split(','))

devito/ir/clusters/algorithms.py

@@ -20,7 +20,7 @@
 from devito.tools import (
     DefaultOrderedDict, Stamp, as_mapper, flatten, is_integer, split, timed_pass, toposort
 )
-from devito.types import Array, Eq, Symbol, Temp
+from devito.types import Array, ConditionalDimension, Eq, Symbol, Temp
 from devito.types.dimension import BOTTOM, ModuloDimension
 
 __all__ = ['clusterize']
@@ -259,10 +259,18 @@ def guard(clusters):
             # Separate out the indirect ConditionalDimensions, which only serve
             # the purpose of protecting from OOB accesses
             cds = [d for d in cds if not d.indirect]
+            modes = [cd.relation for cd in cds]
+            strict = ConditionalDimension.STRICT
+            if modes.count(strict) > 1:
+                raise CompilationError("Only one `strict` condition"
+                                       "can be used in an equation")
+            elif strict in modes:
+                mode = strict
+            else:
+                mode = sympy.And if sympy.And in modes else sympy.Or
 
             # Chain together all `cds` conditions from all expressions in `c`
             guards = {}
-            mode = sympy.Or
             for cd in cds:
                 # `BOTTOM` parent implies a guard that lives outside of
                 # any iteration space, which corresponds to the placeholder None
@@ -279,7 +287,6 @@ def guard(clusters):
 
                 # Pull `cd` from any expr
                 condition = guards.setdefault(k, [])
-                mode = mode and cd.relation
                 for e in exprs:
                     try:
                         condition.append(e.conditionals[cd])
@@ -296,7 +303,10 @@ def guard(clusters):
 
             # Combination `mode` is And by default.
             # If all conditions are Or then Or combination `mode` is used.
-            guards = {d: mode(*v, evaluate=False) for d, v in guards.items()}
+            if mode == strict:
+                guards = {d: v[0] for d, v in guards.items()}
+            else:
+                guards = {d: mode(*v, evaluate=False) for d, v in guards.items()}
 
             # Construct a guarded Cluster
             processed.append(c.rebuild(exprs=exprs, guards=guards))

devito/ir/equations/algorithms.py

@@ -2,14 +2,19 @@
 from functools import singledispatch
 
 from devito.data.allocators import DataReference
+from devito.finite_differences.differentiable import diff2sympy
+from devito.ir.support import GuardFactor
 from devito.logger import warning
 from devito.symbolics import (
-    retrieve_dimensions, retrieve_functions, retrieve_indexed, uxreplace
+    IntDiv, retrieve_dimensions, retrieve_functions, retrieve_indexed, uxreplace
 )
 from devito.tools import (
     Ordering, as_tuple, filter_ordered, filter_sorted, flatten, frozendict
 )
-from devito.types import ConditionalDimension, Dimension, Eq, IgnoreDimSort, SubDimension
+from devito.types import (
+    ConditionalDimension, Dimension, Eq, IgnoreDimSort, SubDimension, relational_min,
+    relational_shift
+)
 from devito.types.array import Array
 from devito.types.basic import AbstractFunction
 from devito.types.dimension import MultiSubDimension, Thickness
@@ -339,3 +344,51 @@ def _(d, mapper, rebuilt, sregistry):
     kwargs['functions'] = functions
 
     mapper[d] = d._rebuild(**kwargs)
+
+
+def generate_conditionals(expr, input_expr, ordering):
+    """
+    Generate the conditionals for the given expression,
+    based on the input expression and the ordering of dimensions.
+    """
+    # Construct the conditionals
+    conditionals = {}
+    for d in ordering:
+        if not d.is_Conditional:
+            continue
+
+        if d.condition is None:
+            conditionals[d] = GuardFactor(d)
+        else:
+            cond = diff2sympy(lower_exprs(d.condition))
+            if d._factor is not None:
+                cond = d.relation(cond, GuardFactor(d))
+            conditionals[d] = cond
+
+    if not conditionals and not input_expr.implicit_dims:
+        return expr, conditionals
+
+    # Merge conditionals when possible. E.g., if an implicit_dim shares
+    # its parent Dimension with another ConditionalDimension, the two
+    # conditions can be merged into a single guard.
+    for d in input_expr.implicit_dims:
+        if d not in conditionals:
+            continue
+        for cd in list(conditionals):
+            if cd.parent is d.parent and cd is not d:
+                if d.relation == ConditionalDimension.STRICT:
+                    conditionals[cd] = conditionals[d]
+                else:
+                    mode = cd.relation and d.relation
+                    conditionals[cd] = mode(conditionals.pop(d), conditionals[cd])
+                break
+
+    # Replace the ConditionalDimensions in `expr`
+    for d, cond in conditionals.items():
+        index = d.index
+        if d.condition is not None and expr.has(d):
+            index = index - relational_min(cond, d.parent)
+        shift = relational_shift(cond, d.parent)
+        expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift})
+
+    return expr, conditionals

devito/ir/equations/equation.py

@@ -5,15 +5,15 @@
 import sympy
 
 from devito.finite_differences.differentiable import diff2sympy
-from devito.ir.equations.algorithms import dimension_sort, lower_exprs
+from devito.ir.equations.algorithms import dimension_sort, generate_conditionals
 from devito.ir.support import (
-    GuardFactor, Interval, IntervalGroup, IterationSpace, Stencil, detect_accesses
+    Interval, IntervalGroup, IterationSpace, Stencil, detect_accesses
 )
-from devito.symbolics import IntDiv, limits_mapper, retrieve_accesses, uxreplace
+from devito.symbolics import limits_mapper, retrieve_accesses
 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.types import Eq, Inc, ReduceMax, ReduceMin, ReduceMinMax
 
 __all__ = [
     'ClusterizedEq',
@@ -292,26 +292,7 @@ def __new__(cls, *args, **kwargs):
                                   relations=ordering.relations, mode='partial')
         ispace = IterationSpace(intervals, iterators)
 
-        # Construct the conditionals and replace the ConditionalDimensions in `expr`
-        conditionals = {}
-        for d in ordering:
-            if not d.is_Conditional:
-                continue
-            if d.condition is None:
-                conditionals[d] = GuardFactor(d)
-            else:
-                cond = diff2sympy(lower_exprs(d.condition))
-                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)
-            expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor)})
-
-        conditionals = frozendict(conditionals)
-
+        expr, conditionals = generate_conditionals(expr, input_expr, ordering)
         # Lower all Differentiable operations into SymPy operations
         rhs = diff2sympy(expr.rhs)
 

devito/ir/stree/algorithms.py

@@ -168,8 +168,17 @@ def preprocess(clusters, options=None, **kwargs):
     processed = []
     for c in clusters:
         if c.is_halo_touch:
+            # Peel syncs (e.g. a WaitLock) off, they must survive an unbound halo.
+            # Restrict the IterationSpace to the syncs' prefix so we don't carry
+            # inner Dimensions (e.g. a layer's VirtualDimension) that have no
+            # content left, which would yield empty guarded/iteration nodes
+            if c.syncs:
+                key = lambda d: any(d in s._defines for s in c.syncs)  # noqa: B023
+                ispace = c.ispace.prefix(key)
+                processed.append(c.rebuild(exprs=[], ispace=ispace))
+
             hs = HaloScheme.union(e.rhs.halo_scheme for e in c.exprs)
-            queue.append(c.rebuild(exprs=[], halo_scheme=hs))
+            queue.append(c.rebuild(exprs=[], syncs={}, halo_scheme=hs))
 
         elif c.is_dist_reduce:
             processed.append(c)

devito/ir/support/guards.py

@@ -13,7 +13,8 @@
 from sympy.logic.boolalg import BooleanFunction
 
 from devito.ir.support.space import Forward, IterationDirection
-from devito.symbolics import CondEq, CondNe, search
+from devito.symbolics import CondEq, CondNe, IntDiv, search
+from devito.symbolics.manipulation import _uxreplace_handle, _uxreplace_registry
 from devito.tools import Pickable, as_tuple, frozendict, split
 from devito.types import Dimension, LocalObject
 
@@ -64,7 +65,14 @@ class GuardFactor(Guard, CondEq, Pickable):
 
     __rargs__ = ('d',)
 
-    def __new__(cls, d, **kwargs):
+    def __new__(cls, *args, **kwargs):
+        if len(args) != 1:
+            # Reconstruction with relational args (e.g. via sympy `_subs`): the
+            # factor semantics no longer hold, so degrade to a plain relational
+            base = CondNe if issubclass(cls, CondNe) else CondEq
+            return base(*args, **kwargs)
+
+        d, = args
         assert d.is_Conditional
 
         obj = super().__new__(cls, d.parent % d.symbolic_factor, 0)
@@ -138,54 +146,54 @@ class BaseGuardBoundNext(Guard, Pickable):
     given `direction`.
     """
 
-    __rargs__ = ('d', 'direction')
+    __rargs__ = ('d', 'index', 'direction')
+    __rkwargs__ = ('d_min', 'd_max')
 
-    def __new__(cls, d, direction, **kwargs):
+    def __new__(cls, d, index, direction,
+                d_min=None, d_max=None, **kwargs):
         assert isinstance(d, Dimension)
         assert isinstance(direction, IterationDirection)
 
-        if direction == Forward:
-            p0 = d.root
-            p1 = d.root.symbolic_max
+        # Always take the next index in the iteration direction
+        next_index = eval_next_index(index, d, direction)
 
-            if d.is_Conditional:
-                v = d.symbolic_factor
-                # Round `p0 + 1` up to the nearest multiple of `v`
-                p0 = Mul((((p0 + 1) + v - 1) / v), v, evaluate=False)
-            else:
-                p0 = p0 + 1
+        # The direction might be forward but accessing c - d
+        # making the access backward w.r.t
+        # Update direction according to access direction for valid guard
+        if index.has(-d):
+            direction = -direction
 
+        if direction == Forward:
+            p0 = next_index
+            p1 = d_max or d.root.symbolic_max
         else:
-            p0 = d.root.symbolic_min
-            p1 = d.root
-
-            if d.is_Conditional:
-                v = d.symbolic_factor
-                # Round `p1 - 1` down to the nearest sub-multiple of `v`
-                # NOTE: we use ABS to make sure we handle negative values properly.
-                # Once `p1 - 1` is negative (e.g. `iteration=time - 1` and `time=0`),
-                # as long as we get a negative number, rather than 0 and even if it's
-                # not `-v`, we're good
-                p1 = (p1 - 1) - abs(p1 - 1) % v
-            else:
-                p1 = p1 - 1
+            p0 = d_min if d_min is not None else d.root.symbolic_min
+            p1 = next_index
 
         try:
             if cls.__base__._eval_relation(p0, p1) is true:
                 return None
         except TypeError:
             pass
 
+        return cls._new(p0, p1, d, index, direction, d_min=d_min, d_max=d_max, **kwargs)
+
+    @classmethod
+    def _new(cls, p0, p1, d, index, direction, d_min=None, d_max=None):
+
         obj = super().__new__(cls, p0, p1, evaluate=False)
 
         obj.d = d
         obj.direction = direction
+        obj.index = index
+        obj.d_min = d_min
+        obj.d_max = d_max
 
         return obj
 
     @property
     def _args_rebuild(self):
-        return (self.d, self.direction)
+        return (self.d, self.index, self.direction)
 
 
 class GuardBoundNextLe(BaseGuardBoundNext, Le):
@@ -272,6 +280,15 @@ class Guards(frozendict):
     def get(self, d, v=true):
         return super().get(d, v)
 
+    def has(self, d, cls):
+        """
+        True if the guard registered for `d` contains an instance of `cls`.
+        """
+        g = super().get(d)
+        if g is None:
+            return False
+        return g.has(cls)
+
     def _reuse_if_untouched(self, mapper):
         return self if mapper == self else Guards(mapper)
 
@@ -569,3 +586,49 @@ def pairwise_or(*guards):
             pass
 
     return guard
+
+
+_uxreplace_registry.register(BaseGuardBoundNext)
+
+
+@_uxreplace_handle.register(BaseGuardBoundNext)
+def _(expr, args, kwargs):
+    p0, p1 = args
+    return expr._new(p0, p1, expr.d, expr.index, expr.direction,
+                     **kwargs)
+
+
+@singledispatch
+def eval_next_index(expr, dim, dir):
+    """
+    Evaluate `expr` at the next iteration point along `dim` in the given
+    `dir`-ection. The "next" point is obtained by substituting `dim` with
+    `dim + 1` for `Forward` and `dim - 1` for `Backward`.
+
+    For `IntDiv` expressions encoding subsampling (`dim.root // factor`),
+    the result is rounded to the next valid coarse-grained slot.
+    """
+    if dir == Forward:
+        return expr._subs(dim, dim + 1)
+    else:
+        return expr._subs(dim, dim - 1)
+
+
+@eval_next_index.register(Expr)
+def _(expr, dim, dir):
+    if not expr.args:
+        if dir == Forward:
+            return expr._subs(dim, dim + 1)
+        else:
+            return expr._subs(dim, dim - 1)
+    return expr.func(*[eval_next_index(a, dim, dir) for a in expr.args])
+
+
+@eval_next_index.register(IntDiv)
+def _(expr, dim, dir):
+    v = dim.symbolic_factor
+    p0 = dim.root
+    if dir == Forward:
+        return Mul((((p0 + 1) + v - 1) / v), v, evaluate=False)
+    else:
+        return (p0 - 1) - abs(p0 - 1) % v

devito/ir/support/space.py

@@ -633,6 +633,14 @@ def __repr__(self):
     def __hash__(self):
         return hash(self._name)
 
+    def __neg__(self):
+        if self._name == '++':
+            return Backward
+        elif self._name == '--':
+            return Forward
+        else:
+            return Any
+
 
 Forward = IterationDirection('++')
 """Forward iteration direction ('++')."""