compiler: support mutli-buffering

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -259,10 +259,17 @@ 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]
+            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 +286,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 +302,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/equation.py

@@ -237,14 +237,6 @@ 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
@@ -254,14 +246,22 @@ def __new__(cls, *args, **kwargs):
             for cd in dict(conditionals):
                 if cd.parent == d.parent and cd != d:
                     cond = conditionals.pop(d)
-                    mode = cd.relation and d.relation
-                    if issubclass(mode, sympy.Or):
-                        conditionals[d] = cond
-                        conditionals.pop(cd)
+                    if d.relation == 'strict':
+                        conditionals[cd] = conditionals[d] = cond
                     else:
+                        mode = cd.relation and d.relation
                         conditionals[cd] = mode(cond, conditionals[cd])
                     break
 
+        # Replace the ConditionalDimensions in `expr`
+        for d, cond in conditionals.items():
+            # Replace dimension with index
+            index = d.index
+            if d.condition is not None and d in expr.free_symbols:
+                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/guards.py

@@ -13,7 +13,7 @@
 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.tools import Pickable, as_tuple, frozendict, split
 from devito.types import Dimension, LocalObject
 
@@ -31,6 +31,34 @@
 ]
 
 
+@singledispatch
+def bound_index(expr, dim, dir):
+    if dir == Forward:
+        return expr._subs(dim, dim + 1)
+    else:
+        return expr._subs(dim, dim - 1)
+
+
+@bound_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(*[bound_index(a, dim, dir) for a in expr.args])
+
+
+@bound_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
+
+
 class AbstractGuard:
     pass
 
@@ -138,37 +166,29 @@ 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 = bound_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:
@@ -180,6 +200,9 @@ def __new__(cls, d, direction, **kwargs):
 
         obj.d = d
         obj.direction = direction
+        obj.index = index
+        obj.d_min = d_min
+        obj.d_max = d_max
 
         return obj
 

devito/ir/support/space.py

@@ -601,6 +601,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 ('++')."""

devito/passes/clusters/asynchrony.py

@@ -1,19 +1,74 @@
 from collections import defaultdict
+from functools import singledispatch
 
-from sympy import true
+from sympy import Expr, Mod, true
 
 from devito.ir import (
     Backward, Forward, GuardBoundNext, PrefetchUpdate, Queue, ReleaseLock, SyncArray,
     WaitLock, WithLock, normalize_syncs
 )
 from devito.passes.clusters.utils import in_critical_region, is_memcpy
-from devito.symbolics import IntDiv, uxreplace
+from devito.symbolics import IntDiv, retrieve_terminals, uxreplace
 from devito.tools import OrderedSet, is_integer, timed_pass
-from devito.types import CustomDimension, Lock
+from devito.types import CustomDimension, Lock, VirtualDimension
 
 __all__ = ['memcpy_prefetch', 'tasking']
 
 
+@singledispatch
+def next_index(expr, dim, dir):
+    return expr._subs(dim, dim + dir)
+
+
+@next_index.register(Expr)
+def _(expr, dim, dir):
+    if not expr.args:
+        return expr._subs(dim, dim + dir)
+    return expr.func(*[next_index(a, dim, dir) for a in expr.args])
+
+
+@next_index.register(IntDiv)
+def _(expr, dim, dir):
+    """
+    Handle forward and backward fetches separately to handle non-canonical index
+    expressions of the form:
+
+        t//factor + cond(t)
+
+    where ``cond(t)`` is a piecewise correction term.
+
+    The forward fetch advances to the next coarse-grained slot while evaluating
+    the correction at the next time point:
+
+        t//factor + cond(t)
+            -> (t//factor + 1) + cond(t + 1)
+
+    The backward fetch is not, in general, the inverse transformation obtained by
+    replacing ``+1`` with ``-1``. The correction may already be applied at the
+    current time point, causing the forward and backward fetches to be asymmetric.
+
+    For example, with ``factor=2`` and ``cond(t) := (t == a)``, the index at
+    ``t=a=3`` is:
+
+        3//2 + 1 = 2
+
+    while the previous index is:
+
+        2//2 + 0 = 1
+
+    A symmetric backward transformation would instead yield:
+
+        3//2 - 1 + 0 = 0
+    """
+    if expr.lhs._defines & dim._defines:
+        if dir == 1:
+            return expr + dir
+        else:
+            return expr._subs(dim, dim + dir)
+    else:
+        return expr
+
+
 def async_trigger(c, dims):
     """
     Return the Dimension in `c`'s IterationSpace that triggers the
@@ -78,7 +133,9 @@ def callback(self, clusters, prefix):
             d = self.key0(c0)
             if d is not dim:
                 continue
-
+            g = c0.guards.get(d)
+            if g is not None and (not g.has(Mod) and d in retrieve_terminals(g)):
+                continue
             protected = self._schedule_waitlocks(c0, d, clusters, locks, syncs)
             self._schedule_withlocks(c0, d, protected, locks, syncs)
 
@@ -181,6 +238,7 @@ def memcpy_prefetch(clusters, key0, sregistry):
     """
     _, key = keys(key0)
     actions = defaultdict(Actions)
+    bounds = {}
 
     for c in clusters:
         d = key(c)
@@ -191,9 +249,9 @@ def memcpy_prefetch(clusters, key0, sregistry):
             continue
 
         if c.properties.is_prefetchable(d._defines):
-            _actions_from_update_memcpy(c, d, clusters, actions, sregistry)
+            _actions_from_update_memcpy(c, d, clusters, actions, sregistry, bounds)
         elif d.is_Custom and is_integer(c.ispace[d].size):
-            _actions_from_init(c, d, clusters, actions)
+            _actions_from_init(c, d, actions)
 
     # Attach the computed Actions
     processed = []
@@ -214,7 +272,7 @@ def memcpy_prefetch(clusters, key0, sregistry):
     return processed
 
 
-def _actions_from_init(c, d, clusters, actions):
+def _actions_from_init(c, d, actions):
     e = c.exprs[0]
     function = e.rhs.function
     target = e.lhs.function
@@ -230,7 +288,7 @@ def _actions_from_init(c, d, clusters, actions):
     )
 
 
-def _actions_from_update_memcpy(c, d, clusters, actions, sregistry):
+def _actions_from_update_memcpy(c, d, clusters, actions, sregistry, bounds):
     pd = d.root  # E.g., `vd -> time`
     direction = c.ispace[pd].direction
 
@@ -240,7 +298,15 @@ 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 + fshift if fetch.find(IntDiv) else fetch._subs(pd, pd + fshift)
+    findex = next_index(fetch, pd, fshift)
+
+    # Maximum allowed access along d
+    if function.dimensions[d].is_Conditional:
+        nslot = function.dimension_shape[d]
+        v = function.dimensions[d].symbolic_factor
+        fd_max = bounds.setdefault(d, v * (nslot - 1))
+    else:
+        fd_max = bounds.setdefault(d, function.dimension_shape[d] - 1)
 
     # If fetching into e.g. `ub[t1]` we might need to prefetch into e.g. `ub[t0]`
     tindex0 = e.lhs.indices[d]
@@ -271,8 +337,17 @@ def _actions_from_update_memcpy(c, d, clusters, actions, sregistry):
     ispace = c.ispace.augment({pd: tindex}) if tindex is not tindex0 else c.ispace
 
     guard0 = c.guards.get(d, true)._subs(fetch, findex)
-    guard1 = GuardBoundNext(function.indices[d], direction)
-    guards = c.guards.impose(d, guard0 & guard1)
+    guard1 = GuardBoundNext(function.indices[d], e.rhs.indices[d], direction,
+                            d_min=0, d_max=fd_max)
+
+    # First guard1 then if guard1 is valid we can safely evaluate guard0
+    # that will have valid indices into f
+    vdnext = VirtualDimension(name=f'vdnext_{d.name}', parent=pd)
+    ispace = ispace.insert(pd, vdnext)
+    # Check valid tindex first
+    guards = c.guards.impose(d, guard1)
+    # THen check valid access
+    guards = guards.impose(vdnext, guard0)
 
     syncs = {d: [
         ReleaseLock(handle, target),