API: Subgrid interpolation (non-integer factors)

devito/finite_differences/interpolation.py

@@ -1,6 +1,104 @@
 from contextlib import suppress
 
-__all__ = ['interp_at', 'interp_mapper', 'post_x0_indices']
+import sympy
+
+from devito.types.lazy import Evaluable
+
+__all__ = ['SubGridInterp', 'interp_at', 'interp_mapper', 'is_subgrid_read',
+           'post_x0_indices']
+
+
+class SubGridInterp(sympy.Expr, Evaluable):
+
+    """
+    A SubGrid Function read at the location of a parent-Grid Function,
+    unevaluated.
+
+    A SubGrid point `i` sits at parent coordinate `i*factor`, so reading it at
+    the parent point `x` means evaluating at `x/factor` -- a fractional offset
+    that varies with the iteration index, which the finite-difference
+    machinery cannot turn into a compile-time stencil window.
+
+    What makes it a stencil after all is that the offset is *periodic*. With
+    `factor = p/q` the coordinate is `x*q/p`, so
+
+        base = (x*q) // p        frac = ((x*q) % p) / p
+
+    and `frac` depends on `x` only through `x % p`: there are `p` phases, and
+    within a phase the tap weights are constants. Deriving that needs the
+    iteration space, so it is done by a Cluster pass
+    (`passes/clusters/subgrid.py`); this object just records *what* to relocate
+    and survives FD lowering untouched, in the manner of `IndexDerivative`.
+
+    Parameters
+    ----------
+    func : Function
+        The SubGrid Function being read.
+    target : Function
+        The Function whose location to read it at.
+    """
+
+    __rargs__ = ('func', 'target')
+
+    is_commutative = True
+
+    def __new__(cls, func, target=None, **kwargs):
+        # `target` rides along as a sympy arg so that generic reconstruction
+        # (`expr.func(*expr.args)`, as `uxreplace` does) round-trips
+        obj = sympy.Expr.__new__(cls, func, target)
+        obj._func = func
+        obj._target = target
+
+        return obj
+
+    @property
+    def func_(self):
+        """The SubGrid Function being read."""
+        return self._func
+
+    @property
+    def target(self):
+        """The Function whose location it is read at."""
+        return self._target
+
+    @property
+    def subgrid(self):
+        return self._func.function.grid
+
+    @property
+    def dtype(self):
+        return self._func.function.dtype
+
+    def __repr__(self):
+        return f"SubGridInterp({self._func}, at={self._target})"
+
+    __str__ = __repr__
+
+    def _evaluate(self, **kwargs):
+        # Deliberately inert: the relocation restructures the iteration space,
+        # so it is lowered by a Cluster pass
+        return self
+
+
+def is_subgrid_read(func, target):
+    """
+    Whether reading `func` at `target`'s location is a SubGrid relocation.
+
+    True when `func` lives on a SubGrid and `target` does not live on that same
+    SubGrid. The exception matters: a Function read at a shift of its own
+    Dimensions, as a recurrence does, is not a relocation.
+
+    Only the Dimensions `target` actually iterates can be relocated, so a
+    SubGrid sharing none of them is not a relocation either.
+    """
+    grid = getattr(func, 'grid', None)
+    if grid is None or not getattr(grid, 'is_SubGrid', False):
+        return False
+
+    if getattr(target, 'grid', None) is grid:
+        return False
+
+    return any(d in target.dimensions for d in grid.parent.dimensions)
 
 
 def interp_mapper(source, target, dims):

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, subsampling, unevaluate
+    optimize_pows, subgrid_interp, subsampling, unevaluate
 )
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (
@@ -423,9 +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
+        # Lower SubGrid reads into stencils, and strengthen factor-based
+        # ConditionalDimension guards into strided iteration. Both rewrite the
+        # iteration space, so they run before `lower_index_derivatives` and
+        # everything else downstream sees the result
+        clusters = subgrid_interp(clusters, **kwargs)
         clusters = subsampling(clusters)
 
         # Lower all remaining high order symbolic objects

devito/passes/clusters/__init__.py

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

devito/passes/clusters/subgrid.py

@@ -0,0 +1,193 @@
+from fractions import Fraction
+
+import sympy
+
+from devito.finite_differences.interpolation import SubGridInterp
+from devito.ir.clusters import Queue
+from devito.symbolics import search, uxreplace
+from devito.tools import timed_pass
+from devito.types import Array, CustomDimension
+
+__all__ = ['subgrid_interp']
+
+
+@timed_pass()
+def subgrid_interp(clusters, sregistry=None, **kwargs):
+    """
+    Lower a Cluster reading an interpolated SubGrid Function into a plain
+    stencil.
+
+    A SubGrid point `i` sits at parent coordinate `i*factor`, so reading it at
+    the parent point `x` means evaluating at `x/factor`. With `factor = p/q`
+    that coordinate is `x*q/p`, hence
+
+        base = (x*q) // p        frac = ((x*q) % p) / p
+
+    Both are plain integer arithmetic, and `frac` depends on `x` only through
+    `x % p` -- so there are `p` phases, and within a phase the tap weights are
+    *constants*. Only the weights vary between phases, so they go in a small
+    compile-time table indexed by `Mod(x, p)`.
+
+    No phase loop is therefore needed, and nothing is added to the iteration
+    space: the SubGrid Dimensions are simply dropped and the inner loops stay
+    unit-stride and branch-free.
+    """
+    return SubGridInterpolation(sregistry).process(clusters)
+
+
+class SubGridInterpolation(Queue):
+
+    def __init__(self, sregistry=None, **kwargs):
+        super().__init__(**kwargs)
+        self.sregistry = sregistry
+
+    def callback(self, clusters, prefix):
+        if prefix:
+            return clusters
+
+        processed = []
+        for c in clusters:
+            processed.append(_lower(c, self.sregistry))
+
+        return processed
+
+
+def _lower(cluster, sregistry):
+    """Expand the SubGridInterp markers in one Cluster, or return it as is."""
+    markers = set()
+    for e in cluster.exprs:
+        markers.update(search(e, SubGridInterp))
+
+    if not markers:
+        return cluster
+
+    ispace = cluster.ispace
+    weights = {}       # (factor, dtype) -> (Array, ntaps, nphases)
+    subs = {}
+
+    # Blocking may already have run, in which case a parent Dimension appears
+    # in the IterationSpace only through its block Dimensions (`y -> y0_blk0`).
+    # Promote those back so the membership test below sees the roots
+    itdims = set(ispace.promote(lambda d: d.is_Block).itdims)
+
+    for m in markers:
+        grid = m.subgrid
+        func = m.func_.function
+
+        terms = []
+        for ds, d, fac in zip(grid.dimensions, grid.parent.dimensions,
+                              grid.factor, strict=True):
+            if d not in itdims:
+                continue
+
+            frac = Fraction(float(fac)).limit_denominator(10**6)
+            key = (float(fac), func.dtype)
+            if key not in weights:
+                name = sregistry.make_name(prefix='w') if sregistry else f'w{d.name}'
+                weights[key] = _phase_weights(fac, func.dtype, name)
+            w, ntaps, nphases = weights[key]
+
+            # `nphases == 1` means every parent point lands exactly on a
+            # SubGrid point: a single tap of unit weight, no interpolation
+            ph = sympy.Mod(d, nphases) if nphases > 1 else None
+            terms.append((ds, d, ph, w, ntaps, frac))
+
+        if not terms:
+            continue
+
+        # The relocated read: a weighted sum over the taps of every relocated
+        # Dimension, with the weights taken from each Dimension's phase row
+        subs[m] = _stencil(func, terms)
+
+    exprs = [uxreplace(e, subs) for e in cluster.exprs]
+
+    # The SubGrid Dimensions are no longer iterated: the relocated reads are
+    # indexed by the parent's Dimensions alone. An axis of factor 1 is the
+    # exception -- there the SubGrid shares the parent's own Dimension, which
+    # stays in the IterationSpace
+    gone = {ds for m in markers
+            for ds, d in zip(m.subgrid.dimensions, m.subgrid.parent.dimensions,
+                             strict=True)
+            if ds is not d}
+    ispace = ispace.project(lambda d: d not in gone)
+
+    return cluster.rebuild(exprs=exprs, ispace=ispace)
+
+
+def _stencil(func, terms):
+    """
+    `func` read as a weighted sum over the taps of every relocated Dimension,
+    with the weights taken from each Dimension's phase row.
+    """
+    from itertools import product
+
+    retval = []
+    ranges = [range(t[4]) if t[2] is not None else range(1) for t in terms]
+    for offsets in product(*ranges):
+        weight = sympy.S.One
+        indices = {}
+        for (ds, d, ph, w, ntaps, frac), o in zip(terms, offsets,
+                                                  strict=True):
+            if ph is not None:
+                weight = weight*w[ntaps*ph + o]
+            indices[ds] = _base(d, frac) + o
+
+        # `func` is indexed directly here, which bypasses the halo shift that
+        # index lowering applies to ordinary reads -- so add it explicitly
+        idxs = [indices.get(d, d) + h for d, h in
+                zip(func.dimensions, func._size_nodomain.left, strict=True)]
+        retval.append(weight*func.indexed[tuple(idxs)])
+
+    return sympy.Add(*retval)
+
+
+def _base(d, frac):
+    """
+    The SubGrid base index for parent index `d`: `(d*q)//p` for a
+    `factor = p/q`, in integer arithmetic.
+
+    `d` is the raw iteration Dimension: index lowering adds each Function's
+    halo afterwards, so neither halo takes part in the division.
+    """
+    from devito.symbolics import IntDiv
+    return IntDiv(frac.denominator*d, frac.numerator)
+
+
+def _phase_weights(factor, dtype, name):
+    """
+    The per-phase tap weights for *reading* a SubGrid of the given `factor` at
+    its parent's points, as a compile-time table.
+
+    A SubGrid point sits at parent coordinate `i*factor`, so conversely the
+    parent point `x` sits at SubGrid coordinate `x/factor`. With
+    `factor = p/q` that is `x*q/p`, hence
+
+        base = (x*q)//p        frac = ((x*q) % p)/p
+
+    `base` is the integer part, and `frac` depends on `x` only through
+    `x % p`. So there are `p` phases, and within a phase the two linear
+    weights `(1 - frac, frac)` are constants -- an ordinary constant-offset
+    stencil, read at `base` and `base + 1`.
+
+    Returns `(weights, ntaps, nphases)`: the `Array` holding `nphases*ntaps`
+    values, the taps per phase, and the number of phases.
+    """
+    frac = Fraction(float(factor)).limit_denominator(10**6)
+    p, q = frac.numerator, frac.denominator
+
+    # Reading the SubGrid at the parent's points inverts the factor, so the
+    # phase period is `p`
+    nphases, ntaps = p, 2
+
+    values = []
+    for ph in range(nphases):
+        f = float(Fraction((ph*q) % p, p))
+        values.extend(sympy.Float(v) for v in (1. - f, f))
+
+    d = CustomDimension(name=f'{name}d', symbolic_min=0,
+                        symbolic_max=nphases*ntaps - 1,
+                        symbolic_size=nphases*ntaps)
+    weights = Array(name=name, dimensions=(d,), shape=(nphases*ntaps,),
+                    dtype=dtype, scope='stack', initvalue=values)
+
+    return weights, ntaps, nphases

devito/types/dense.py

@@ -15,7 +15,9 @@
 from devito.deprecations import deprecations
 from devito.exceptions import InvalidArgument
 from devito.finite_differences import Differentiable, generate_fd_shortcuts
-from devito.finite_differences.interpolation import interp_mapper
+from devito.finite_differences.interpolation import (
+    SubGridInterp, interp_mapper, is_subgrid_read
+)
 from devito.finite_differences.tools import fd_weights_registry
 from devito.logger import debug, warning
 from devito.mpi import MPI
@@ -1130,6 +1132,15 @@ def _fd_priority(self):
         return 1.2 if self.staggered.on_node else 1.1
 
     def _eval_at(self, func, **kwargs):
+        # A Function on a SubGrid lives on different Dimensions altogether, so
+        # relocating it onto `func` means re-expressing its indices in the
+        # parent's terms. Same convention as staggering: the RHS is evaluated
+        # at the LHS's location
+        if is_subgrid_read(self, func):
+            # The relocation depends on the iteration index, so it cannot be
+            # turned into a stencil here; defer it to the Cluster pass
+            return SubGridInterp(self, func)
+
         if self.staggered == func.staggered or self.interp_order == 0:
             return self