misc: cleanup and add comments

Author: mloubout

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']
@@ -260,11 +260,12 @@ def guard(clusters):
             # 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:
+            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'
+            elif strict in modes:
+                mode = strict
             else:
                 mode = sympy.And if sympy.And in modes else sympy.Or
 
@@ -302,7 +303,7 @@ def guard(clusters):
 
             # Combination `mode` is And by default.
             # If all conditions are Or then Or combination `mode` is used.
-            if mode == 'strict':
+            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()}

devito/ir/equations/algorithms.py

@@ -356,6 +356,7 @@ def generate_conditionals(expr, input_expr, ordering):
     for d in ordering:
         if not d.is_Conditional:
             continue
+
         if d.condition is None:
             conditionals[d] = GuardFactor(d)
         else:
@@ -364,16 +365,16 @@ def generate_conditionals(expr, input_expr, ordering):
                 cond = d.relation(cond, GuardFactor(d))
             conditionals[d] = cond
 
-    # Merge conditionals when possible. E.g if we have an implicit_dim
-    # and there is a dimension with the same parent, we can merge
-    # their conditions
+    # 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 == d.parent and cd != d:
+            if cd.parent == d.parent and cd is not d:
                 cond = conditionals.pop(d)
-                if d.relation == 'strict':
+                if d.relation == ConditionalDimension._STRICT:
                     conditionals[cd] = conditionals[d] = cond
                 else:
                     mode = cd.relation and d.relation
@@ -384,7 +385,7 @@ def generate_conditionals(expr, input_expr, ordering):
     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:
+        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})

devito/ir/support/guards.py

@@ -32,34 +32,6 @@
 ]
 
 
-@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
 
@@ -176,7 +148,7 @@ def __new__(cls, d, index, direction,
         assert isinstance(direction, IterationDirection)
 
         # Always take the next index in the iteration direction
-        next_index = bound_index(index, d, direction)
+        next_index = eval_next_index(index, d, direction)
 
         # The direction might be forward but accessing c - d
         # making the access backward w.r.t
@@ -601,3 +573,39 @@ def pairwise_or(*guards):
 @_uxreplace_handle.register(BaseGuardBoundNext)
 def _(expr, args, kwargs):
     return expr.func(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/passes/clusters/asynchrony.py

@@ -15,60 +15,6 @@
 __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
@@ -134,11 +80,15 @@ def callback(self, clusters, prefix):
             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)) \
-                    and not wraps_memcpy(c0):
-                # Explicit compute guards need no pipeline; memcpy clusters
-                # still need WithLock for the copy-back sync
-                continue
+            # Explicit compute guards need no pipeline; memcpy clusters
+            # still need WithLock for the copy-back sync
+            if g is not None and not wraps_memcpy(c0):
+                # An "explicit" guard is a plain reference to `d` (e.g., `d == K`)
+                # without subsampling -- subsampling guards (containing `Mod`)
+                # still require the standard async pipeline
+                explicit = not g.has(Mod) and d in retrieve_terminals(g)
+                if explicit:
+                    continue
             protected = self._schedule_waitlocks(c0, d, clusters, locks, syncs)
             self._schedule_withlocks(c0, d, protected, locks, syncs)
 
@@ -252,7 +202,7 @@ def memcpy_prefetch(clusters, key0, sregistry):
             continue
 
         if c.properties.is_prefetchable(d._defines):
-            _actions_from_update_memcpy(c, d, clusters, actions, sregistry, bounds)
+            _actions_from_update_memcpy(c, d, bounds, clusters, actions, sregistry)
         elif d.is_Custom and is_integer(c.ispace[d].size):
             _actions_from_init(c, d, actions)
         else:
@@ -308,25 +258,25 @@ def _actions_from_sync(c, d, actions):
     )
 
 
-def _actions_from_update_memcpy(c, d, clusters, actions, sregistry, bounds):
+def _actions_from_update_memcpy(c, d, bounds, clusters, actions, sregistry):
     pd = d.root  # E.g., `vd -> time`
     direction = c.ispace[pd].direction
 
     e = c.exprs[0]
-    function = e.rhs.function
+    f = e.rhs.function
     target = e.lhs.function
 
     fetch = e.rhs.indices[d]
     fshift = {Forward: 1, Backward: -1}.get(direction, 0)
-    findex = next_index(fetch, pd, fshift)
+    findex = eval_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
+    if f.dimensions[d].is_Conditional:
+        nslot = f.dimension_shape[d]
+        v = f.dimensions[d].symbolic_factor
         fd_max = bounds.setdefault(d, v * (nslot - 1))
     else:
-        fd_max = bounds.setdefault(d, function.dimension_shape[d] - 1)
+        fd_max = bounds.setdefault(d, f.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]
@@ -355,23 +305,26 @@ def _actions_from_update_memcpy(c, d, clusters, actions, sregistry, bounds):
     expr = uxreplace(e, {tindex0: tindex, fetch: findex})
 
     ispace = c.ispace.augment({pd: tindex}) if tindex is not tindex0 else c.ispace
+    # Insert a VirtualDimension nested under `pd` so that the access guard
+    # (`guard0`) below can be evaluated only after the tindex bound guard
+    # (`guard1`) has already been validated
+    vdnext = VirtualDimension(name=f'vdnext_{d.name}', parent=pd)
+    ispace = ispace.insert(pd, vdnext)
 
     guard0 = c.guards.get(d, true)._subs(fetch, findex)
-    guard1 = GuardBoundNext(function.indices[d], e.rhs.indices[d], direction,
+    guard1 = GuardBoundNext(f.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)
+    # First guard1; then, if guard1 is valid, we can safely evaluate guard0,
+    # which will then have valid indices into f
     # Check valid tindex first
     guards = c.guards.impose(d, guard1)
-    # THen check valid access
+    # Then check valid access
     guards = guards.impose(vdnext, guard0)
 
     syncs = {d: [
         ReleaseLock(handle, target),
-        PrefetchUpdate(handle, target, tindex, function, findex, d, 1, e.rhs)
+        PrefetchUpdate(handle, target, tindex, f, findex, d, 1, e.rhs)
     ]}
     syncs = {**c.syncs, **syncs}
 
@@ -409,3 +362,65 @@ def __init__(self, drop=False, syncs=None, insert=None):
 
 def wraps_memcpy(cluster):
     return len(cluster.exprs) == 1 and is_memcpy(cluster.exprs[0])
+
+
+@singledispatch
+def eval_next_index(expr, dim, dir):
+    """
+    Evaluate `expr` at the next iteration point along `dim` in the given
+    `dir`-ection, where `dir` is `+1` for forward and `-1` for backward.
+
+    For `IntDiv` subexpressions encoding subsampling, forward and backward
+    fetches are treated asymmetrically since piecewise correction terms may
+    already be applied at the current point.
+    """
+    return expr._subs(dim, dim + dir)
+
+
+@eval_next_index.register(Expr)
+def _(expr, dim, dir):
+    if not expr.args:
+        return expr._subs(dim, dim + dir)
+    return expr.func(*[eval_next_index(a, dim, dir) for a in expr.args])
+
+
+@eval_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