API: Subgrid (integer) support

devito/ir/support/basic.py

@@ -869,8 +869,15 @@ def is_storage_related(self, dims=None):
         """
         for d in self.findices:
             if d._defines & set(as_tuple(dims)):  # noqa: SIM102
-                if any(i.is_NonlinearDerived for i in d._defines) or \
-                   self.is_const(d):
+                # An injective Dimension maps distinct iterations to distinct
+                # locations by construction -- e.g. a factor-based
+                # ConditionalDimension, whose guard admits only
+                # `parent % factor == 0` and whose `i -> i//factor` is
+                # one-to-one on those. Unlike, say, a SteppingDimension, where
+                # `time % 2` deliberately revisits the same slot
+                nonlinear = [i for i in d._defines
+                             if i.is_NonlinearDerived and not i.is_Injective]
+                if nonlinear or self.is_const(d):
                     return True
         return False
 

devito/operator/operator.py

@@ -34,7 +34,7 @@
 from devito.passes import (
     Graph, error_mapper, finalize_args, generate_implicit, generate_macros, is_on_device,
     lower_dtypes, lower_index_derivatives, lower_local_sums, minimize_symbols,
-    optimize_pows, unevaluate
+    optimize_pows, subsampling, unevaluate
 )
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (
@@ -423,6 +423,11 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
         # Generate implicit Clusters from higher level abstractions
         clusters = generate_implicit(clusters)
 
+        # Strengthen factor-based ConditionalDimension guards into strided
+        # iteration. Before `lower_index_derivatives` so that the Dimension
+        # swap is seen by everything downstream
+        clusters = subsampling(clusters)
+
         # Lower all remaining high order symbolic objects
         clusters = lower_local_sums(clusters, **kwargs)
         clusters = lower_index_derivatives(clusters, **kwargs)

devito/passes/clusters/__init__.py

@@ -5,6 +5,7 @@
 from .factorization import *  # noqa
 from .fusion import *  # noqa
 from .blocking import *  # noqa
+from .subsampling import *  # noqa
 from .asynchrony import *  # noqa
 from .implicit import *  # noqa
 from .misc import *  # noqa

devito/passes/clusters/subsampling.py

@@ -0,0 +1,122 @@
+from devito.ir.clusters import Queue
+from devito.ir.support import GuardFactor
+from devito.symbolics import IntDiv, uxreplace
+from devito.tools import timed_pass
+
+__all__ = ['subsampling']
+
+
+@timed_pass()
+def subsampling(clusters):
+    """
+    Strengthen factor-based ConditionalDimension guards into strided iteration.
+
+    A Function subsampled by `factor` along `d` is written under a
+    `d % factor == 0` guard, so the enclosing loop visits every point of `d`
+    and discards all but one in `factor` of them. Where legal, drop the guard
+    and step the loop `factor` points at a time instead
+
+        for (x = x_m; x <= x_M; x += 1)      for (x = ...; x <= x_M; x += xf)
+          if (x%xf == 0)                 ->     g[x/xf] = ...
+            g[x/xf] = ...
+
+The loop is taken over the subsampled Dimension `cd` itself, with a *unit*
+    stride, and every access to the parent `d` is rescaled to `cd*factor`:
+
+        for (x = x_m; x <= x_M; x += 1)      for (xs = xs_m; xs <= xs_M; xs++)
+          if (x%xf == 0)                 ->     g[xs] = u[xs*xf]
+            g[x/xf] = u[x]
+
+    A non-unitary increment is deliberately avoided: it would escape the
+    Lamport-style dependence analysis the compiler relies on. The loop bounds
+    are `cd`'s own runtime values, which `ConditionalDimension._arg_values`
+    derives as `ceil(x_m/factor) .. floor(x_M/factor)`, so no anchoring of the
+    loop minimum is needed either.
+
+    Parameters
+    ----------
+    clusters : tuple of Clusters
+        Input Clusters, subject of the optimization pass.
+    """
+    return Subsampling().process(clusters)
+
+
+class Subsampling(Queue):
+
+    """
+    Replace uniform factor-based guards with strided iteration.
+
+    Clusters are grouped by iteration space prefix alone, deliberately *not*
+    by guard: a guard may only become a stride if every Cluster sharing the
+    prefix carries the very same one. Striding a loop nest that also hosts an
+    unguarded expression would run that expression on a subset of its
+    iterations -- and dropping the guard is what makes the nest shared in the
+    first place, since `stree_build` reuses a subtree across Clusters with
+    matching guards (see `reuse_whole_subtree`).
+    """
+
+    def callback(self, clusters, prefix):
+        if not prefix:
+            return clusters
+
+        d = prefix[-1].dim
+
+        guards = {c.guards.get(d) for c in clusters}
+        if len(guards) > 1:
+            # Non-uniformly guarded, hence `d`'s loop nest is shared with
+            # expressions that must keep running at every point
+            return clusters
+
+        guard = guards.pop()
+        if not all(_is_strideable(c, d, guard) for c in clusters):
+            return clusters
+
+        cd = guard.d
+        factor = cd.symbolic_factor
+
+        processed = []
+        for c in clusters:
+            # Iterate the subsampled Dimension itself, with unit stride, and
+            # rescale every parent access `d` to `cd*factor`
+            ispace = c.ispace.switch(d, cd)
+
+            # `IntDiv(d, factor)` is how the subsampled access was lowered;
+            # with `d -> cd*factor` it collapses to `cd`, so rewrite it
+            # directly rather than leaving `(cd*factor)/factor` behind
+            subs = {IntDiv(d, factor): cd, d: cd*factor}
+            exprs = [uxreplace(e, subs) for e in c.exprs]
+
+            # `cd` inherits `d`'s computational properties: it covers the very
+            # same points, just addressed in subsampled coordinates. Without
+            # this the loop would lose PARALLEL and so all of its OpenMP
+            properties = c.properties.drop(d)
+            properties = properties.add(cd, c.properties[d])
+
+            processed.append(c.rebuild(exprs=exprs, ispace=ispace,
+                                       guards=c.guards.popany(d),
+                                       properties=properties))
+
+        return processed
+
+
+def _is_strideable(cluster, d, guard):
+    """
+    True if iteration over `d` may step `guard`'s subsampling factor points at
+    a time in place of evaluating `guard` at every point.
+    """
+    # The guard must be exactly a subsampling guard. Anything `And`/`Or`-ed in
+    # (e.g. a user-supplied `condition`) must still be evaluated pointwise
+    if not isinstance(guard, GuardFactor):
+        return False
+
+    if d not in cluster.ispace.itdims:
+        return False
+
+    # Skipping points is only sound if no dependence is carried along `d`;
+    # otherwise the discarded iterations may be doing work that `d`'s
+    # remaining iterations rely upon
+    if not cluster.properties.is_parallel_relaxed(d):
+        return False
+
+    # Non-affine accesses may read the points that would be skipped
+    return cluster.properties.is_affine(d)

devito/types/dimension.py

@@ -116,6 +116,7 @@ class Dimension(ArgProvider):
     is_Custom = False
     is_Derived = False
     is_NonlinearDerived = False
+    is_Injective = False
     is_AbstractSub = False
     is_Sub = False
     is_MultiSub = False
@@ -1090,6 +1091,19 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None,
         self._indirect = indirect
         self._relation = relation
 
+    @property
+    def is_Injective(self):
+        """
+        True if distinct iterations of the parent that pass the guard map to
+        distinct indices of `self`, so that no two iterations can access the
+        same memory location through it.
+
+        A pure factor subsampling is injective: the guard admits only
+        `parent % factor == 0`, and `i -> i//factor` is one-to-one on those.
+        An arbitrary `condition` carries no such promise.
+        """
+        return self._factor is not None and self._condition is None
+
     @property
     def uses_symbolic_factor(self):
         return self._factor is not None
@@ -1153,13 +1167,18 @@ def _arg_values(self, interval, grid=None, args=None, **kwargs):
         if isinstance(fact, Constant):
             fact = fact.data
 
-        toint = lambda x: math.ceil(x / fact)
         vals = {}
+
+        # The first subsampled point at or after the parent's min...
         with suppress(KeyError, TypeError):
-            vals[self.min_name] = toint(kwargs.get(self.parent.min_name))
+            vals[self.min_name] = math.ceil(kwargs.get(self.parent.min_name) / fact)
 
+        # ...and the last one at or before its max. Note `floor` rather than
+        # `ceil`: with e.g. `factor=3` and `x_M=11` the last subsampled point
+        # is `x=9`, hence index 3, while `ceil(11/3) == 4` would overshoot the
+        # Function by one
         with suppress(KeyError, TypeError):
-            vals[self.max_name] = toint(kwargs.get(self.parent.max_name))
+            vals[self.max_name] = kwargs.get(self.parent.max_name) // fact
 
         vals[self.symbolic_factor.name] = fact