compiler: Introduce SparseLocalSum

Author: FabioLuporini

devito/finite_differences/differentiable.py

@@ -21,8 +21,8 @@
 from devito.finite_differences.tools import coeff_priority, make_shift_x0
 from devito.logger import warning
 from devito.tools import (
-    Tag, as_tuple, extract_dtype, filter_ordered, flatten, frozendict, infer_dtype,
-    is_integer, is_number, memoized_func, split
+    Pickable, Tag, as_tuple, extract_dtype, filter_ordered, flatten, frozendict,
+    infer_dtype, is_integer, is_number, memoized_func, split
 )
 from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
 from devito.types.basic import AbstractFunction, Indexed
@@ -36,6 +36,7 @@
     'IndexDerivative',
     'IndexDerivativeProperty',
     'Real',
+    'SparseLocalSum',
     'Weights',
 ]
 
@@ -947,6 +948,82 @@ def free_symbols(self):
     func = DifferentiableOp._rebuild
 
 
+class SparseLocalSum(IndexSum, Pickable):
+
+    """
+    A zero-initialized sum over the guarded taps of a sparse interpolation.
+
+    `dimensions` are bound StencilDimensions; `predicates` contains their bounds
+    conditions in the same order. Masked taps contribute zero.
+    The sum remains symbolic until Cluster lowering chooses its implementation.
+
+    Examples
+    --------
+    For bilinear interpolation, `posx` and `posy` are the grid indices of sparse
+    point `p`, and `wx` and `wy` hold its interpolation weights::
+
+        i = StencilDimension('i', 0, 1)
+        j = StencilDimension('j', 0, 1)
+        value = SparseLocalSum(
+            wx[p, i]*wy[p, j]*f[posx + i, posy + j],
+            dimensions=(i, j), dtype=f.dtype,
+            predicates=(And(posx + i >= x_m, posx + i <= x_M),
+                        And(posy + j >= y_m, posy + j <= y_M))
+        )
+        Eq(rcv[p], value)
+
+    The scalar lowering has the following semantics (pseudocode)::
+
+        acc = 0
+        for i in range(2):
+            for j in range(2):
+                if x_m <= posx + i <= x_M and y_m <= posy + j <= y_M:
+                    acc += wx[p, i]*wy[p, j]*f[posx + i, posy + j]
+        rcv[p] = acc
+
+    `i` and `j` are local to the sum; `p` remains an outer iteration dimension.
+    If every tap is masked, `rcv[p]` receives zero.
+    """
+
+    __rargs__ = ('expr',)
+    __rkwargs__ = ('dimensions', 'dtype', 'predicates')
+
+    def __new__(cls, expr, dimensions, dtype, predicates, **kwargs):
+        # Unlike IndexSum, the summand can be independent of a bound index, e.g. zero
+        obj = sympy.Expr.__new__(cls, expr)
+        obj._expr = expr
+        obj._dimensions = as_tuple(dimensions)
+        obj._dtype = np.dtype(dtype).type
+        obj._predicates = as_tuple(predicates)
+        return obj
+
+    def _hashable_content(self):
+        return super()._hashable_content() + (self.dtype, self.predicates)
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @property
+    def predicates(self):
+        return self._predicates
+
+    @property
+    def free_symbols(self):
+        symbols = super().free_symbols.union(*[c.free_symbols for c in self.predicates])
+        return symbols.difference(self.dimensions)
+
+    def _evaluate(self, **kwargs):
+        return self._rebuild(Evaluable._evaluate_maybe_nested(self.expr, **kwargs))
+
+    def _xreplace(self, rule):
+        from devito.symbolics import uxreplace
+        rebuilt = uxreplace(self, {k: sympy.sympify(v) for k, v in rule.items()})
+        return rebuilt, rebuilt is not self
+
+    __reduce_ex__ = Pickable.__reduce_ex__
+
+
 class WeightsIndexed(Indexed):
 
     @property

devito/ir/clusters/cluster.py

@@ -152,6 +152,11 @@ def dist_dimensions(self):
     def scope(self):
         return Scope(self.exprs)
 
+    @cached_property
+    def sparse_sums(self):
+        """The sparse sums in equation order, retaining occurrences across equations."""
+        return tuple(s for e in self.exprs for s in e.sparse_sums)
+
     @cached_property
     def functions(self):
         return self.scope.functions
@@ -674,6 +679,10 @@ def rebuild(self, **kwargs):
     def exprs(self):
         return flatten(c.exprs for c in self)
 
+    @cached_property
+    def sparse_sums(self):
+        return tuple(s for c in self for s in c.sparse_sums)
+
     @cached_property
     def scope(self):
         return Scope(exprs=self.exprs)

devito/ir/equations/algorithms.py

@@ -33,7 +33,9 @@ def handle_indexed(indexed):
         for i in indexed.indices:
             try:
                 # Assume it's an AffineIndexAccessFunction...
-                relation.append(i.d)
+                # It may contain only a scalar offset and bound stencil indices
+                if i.d:
+                    relation.append(i.d)
             except AttributeError:
                 # It's not! Maybe there are some nested Indexeds (e.g., the
                 # situation is A[B[i]])

devito/ir/equations/equation.py

@@ -4,14 +4,15 @@
 import numpy as np
 import sympy
 
-from devito.finite_differences.differentiable import diff2sympy
+from devito.finite_differences.differentiable import SparseLocalSum, diff2sympy
 from devito.ir.equations.algorithms import dimension_sort, generate_conditionals
 from devito.ir.support import (
     Interval, IntervalGroup, IterationSpace, Stencil, detect_accesses
 )
-from devito.symbolics import limits_mapper, retrieve_accesses
+from devito.symbolics import limits_mapper, retrieve_accesses, search
 from devito.tools import (
-    Pickable, Tag, as_hashable, filter_sorted, frozendict, reuse_if_unchanged
+    Pickable, Tag, as_hashable, filter_ordered, filter_sorted, frozendict,
+    reuse_if_unchanged
 )
 from devito.types import Eq, Inc, ReduceMax, ReduceMin, ReduceMinMax
 
@@ -50,6 +51,11 @@ def ispace(self):
     def dimensions(self):
         return set(self.ispace.dimensions)
 
+    @cached_property
+    def sparse_sums(self):
+        """The sparse sums in dependency order, with nested sums first."""
+        return tuple(filter_ordered(search(self, SparseLocalSum, mode='all')))
+
     @property
     def implicit_dims(self):
         return self._implicit_dims

devito/operations/interpolators.py

@@ -11,14 +11,16 @@
 except ImportError:
     from numpy import i0
 
-from devito.finite_differences.differentiable import Mul
+from devito.finite_differences.differentiable import Mul, SparseLocalSum
 from devito.finite_differences.elementary import floor
 from devito.logger import warning
 from devito.symbolics import INT, retrieve_function_carriers, retrieve_functions
 from devito.tools import (
     Pickable, as_fp64_decimal, as_list, as_tuple, filter_ordered, flatten, memoized_meth
 )
-from devito.types import CustomDimension, Eq, Evaluable, Inc, SubFunction, Symbol
+from devito.types import (
+    CustomDimension, Eq, Evaluable, Inc, StencilDimension, SubFunction
+)
 from devito.types.utils import DimensionTuple
 
 __all__ = ['LinearInterpolator', 'NearestInterpolator',
@@ -453,19 +455,21 @@ def _interp_idx(self, variables, implicit_dims=None, subdomain=None,
 
         return idx_subs, temps
 
-    def _local_accumulator(self, expr, idx_subs, implicit_dims=None, subdomain=None):
+    def _local_accumulator(self, expr, idx_subs, subdomain=None):
         """
-        Generate a local accumulator for the interpolation/injection operation.
+        Represent the local sum of weighted interpolation contributions.
         """
-        # Accumulate point-wise contributions into a temporary
-        rhs = Symbol(name=f'sum{self.sfunction.name}', dtype=self.sfunction.dtype)
-        summands = [Eq(rhs, 0., implicit_dims=implicit_dims)]
-        # Substitute coordinate base symbols into the interpolation coefficients
         weights = self._weights(subdomain=subdomain)
-        summands.extend([Inc(rhs, (weights * expr).xreplace(idx_subs),
-                             implicit_dims=implicit_dims)])
-
-        return summands, rhs
+        rdims = self._rdim(subdomain=subdomain)
+        # Bound stencil indices are independent of the outer sparse-point dimension
+        dimensions = tuple(StencilDimension(d.name, d.parent.symbolic_min,
+                                            d.parent.symbolic_max) for d in rdims)
+        aliases = {a: i for d, i in zip(rdims, dimensions, strict=True)
+                   for a in (d, d.parent)}
+        summand = (weights * expr).xreplace(idx_subs).xreplace(aliases)
+        predicates = tuple(d.condition.xreplace(aliases) for d in rdims)
+        return SparseLocalSum(summand, dimensions, self.sfunction.dtype,
+                              predicates=predicates)
 
     @check_radius
     @check_coords
@@ -539,16 +543,13 @@ def _interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None
         idx_subs, temps = self._interp_idx(variables, implicit_dims=implicit_dims,
                                            subdomain=subdomain)
 
-        # Local scalar for accumulation over radius
-        summands, rhs = self._local_accumulator(expr, idx_subs,
-                                                implicit_dims=implicit_dims,
-                                                subdomain=subdomain)
+        rhs = self._local_accumulator(expr, idx_subs, subdomain=subdomain)
         # Write/Incr `self`
         lhs = self.sfunction.subs(self_subs)
         ecls = Inc if increment else Eq
         last = [ecls(lhs, rhs, implicit_dims=implicit_dims)]
 
-        return temps + summands + last
+        return temps + last
 
     def _inject(self, field, expr, increment=True, implicit_dims=None):
         """
@@ -785,8 +786,8 @@ class NearestInterpolator(LinearInterpolator):
 
     _name = 'nearest'
 
-    def _local_accumulator(self, expr, idx_subs, implicit_dims=None, subdomain=None):
-        return [], expr.xreplace(idx_subs)
+    def _local_accumulator(self, expr, idx_subs, subdomain=None):
+        return expr.xreplace(idx_subs)
 
     @memoized_meth
     def _rdim(self, subdomain=None, shifts=None):

devito/operator/operator.py

@@ -33,7 +33,8 @@
 from devito.parameters import configuration
 from devito.passes import (
     Graph, error_mapper, finalize_args, generate_implicit, generate_macros, is_on_device,
-    lower_dtypes, lower_index_derivatives, minimize_symbols, optimize_pows, unevaluate
+    lower_dtypes, lower_index_derivatives, lower_sparse_sums, minimize_symbols,
+    optimize_pows, unevaluate
 )
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (
@@ -423,6 +424,7 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs):
         clusters = generate_implicit(clusters)
 
         # Lower all remaining high order symbolic objects
+        clusters = lower_sparse_sums(clusters, **kwargs)
         clusters = lower_index_derivatives(clusters, **kwargs)
 
         # Turn pows into multiplications. This must happen as late as possible

devito/passes/clusters/__init__.py

@@ -9,4 +9,5 @@
 from .implicit import *  # noqa
 from .misc import *  # noqa
 from .derivatives import *  # noqa
+from .sparse import *  # noqa
 from .unevaluate import *  # noqa

devito/passes/clusters/cse.py

@@ -11,7 +11,7 @@
     # Moved in 1.13
     from sympy.core.basic import ordering_of_classes
 
-from devito.finite_differences.differentiable import IndexDerivative
+from devito.finite_differences.differentiable import IndexSum
 from devito.ir import Cluster, Scope, cluster_pass
 from devito.symbolics import (
     DefFunction, Reserved, estimate_cost, q_leaf, q_terminal, search
@@ -427,7 +427,7 @@ def _(expr):
     return {}
 
 
-@_catch.register(IndexDerivative)
+@_catch.register(IndexSum)
 def _(expr):
     """
     Handler for symbol-binding objects. There can be many of them and therefore

devito/passes/clusters/sparse.py

@@ -0,0 +1,82 @@
+from functools import partial
+
+from devito.ir import ClusterizedEq, Interval, IterationSpace
+from devito.symbolics import uxreplace
+from devito.tools import timed_pass
+from devito.types import ConditionalDimension, Eq, Inc, Temp
+
+__all__ = ['lower_sparse_sum', 'lower_sparse_sums']
+
+
+@timed_pass()
+def lower_sparse_sums(clusters, sregistry=None, **kwargs):
+    """
+    Lower SparseLocalSums into a private initializer and guarded accumulation.
+
+    For example, with `i = StencilDimension('i', 0, 1)`::
+
+        Eq(rcv[p], SparseLocalSum(
+            w[p, i]*f[pos + i], dimensions=(i,), dtype=f.dtype,
+            predicates=(And(pos + i >= x_m, pos + i <= x_M),)
+        ))
+
+    becomes the following computation at each sparse point `p` (pseudocode)::
+
+        sum0 = 0
+        for i in range(2):
+            if x_m <= pos + i <= x_M:
+                sum0 += w[p, i]*f[pos + i]
+        rcv[p] = sum0
+
+    The initializer and result stay outside the tap guard, so even a fully
+    masked stencil assigns zero to `rcv[p]`.
+    """
+    processed = []
+    for c in clusters:
+        if not c.sparse_sums:
+            processed.append(c)
+            continue
+
+        for e in c.exprs:
+            subs = {}
+            for reduction in e.sparse_sums:
+                init, update, value = lower_sparse_sum(
+                    c, uxreplace(reduction, subs), sregistry
+                )
+                processed.extend([init, update])
+                subs[reduction] = value
+
+            expr = e.apply(partial(uxreplace, rule=subs))
+            processed.append(c.rebuild(exprs=[expr]))
+
+    return processed
+
+
+def lower_sparse_sum(cluster, reduction, sregistry):
+    """
+    Construct the private initializer and guarded accumulation for one sum.
+    """
+    value = Temp(name=sregistry.make_name(prefix='sum'), dtype=reduction.dtype)
+
+    dims = reduction.dimensions
+    inner = IterationSpace([Interval(d) for d in dims])
+    ispace = IterationSpace.union(
+        cluster.ispace, inner, relations=(cluster.ispace.itdims + dims,)
+    )
+
+    init = cluster.rebuild(exprs=[Eq(value, 0)])
+    conditionals = {
+        ConditionalDimension(d.name, parent=d, condition=v, indirect=True): v
+        for d, v in zip(dims, reduction.predicates, strict=True)
+    }
+    # Cluster construction does not infer conditionals from a plain Eq/Inc,
+    # so attach the per-tap guards explicitly
+    expr = ClusterizedEq(
+        Inc(value, reduction.expr), ispace=ispace, conditionals=conditionals
+    )
+
+    properties = cluster.properties.sequentialize(dims)
+
+    update = cluster.rebuild(exprs=expr, ispace=ispace, properties=properties)
+
+    return init, update, value

devito/symbolics/inspection.py

@@ -7,7 +7,7 @@
 from sympy.core.numbers import ImaginaryUnit
 
 from devito.finite_differences import Derivative
-from devito.finite_differences.differentiable import IndexDerivative
+from devito.finite_differences.differentiable import IndexDerivative, IndexSum
 from devito.logger import warning
 from devito.symbolics.extended_dtypes import INT
 from devito.symbolics.extended_sympy import CallFromPointer, Cast, DefFunction, Reserved
@@ -265,7 +265,7 @@ def _(expr, estimate, seen):
     return _estimate_cost(expr._evaluate(expand=False), estimate, seen)
 
 
-@_estimate_cost.register(IndexDerivative)
+@_estimate_cost.register(IndexSum)
 @dont_count_if_seen
 def _(expr, estimate, seen):
     flops, _ = _estimate_cost(expr.expr, estimate, seen)
@@ -275,7 +275,7 @@ def _(expr, estimate, seen):
 
     # To be multiplied by the number of points this index sum implicitly
     # iterates over
-    flops *= prod(i._size for i in expr.dimensions)
+    flops *= prod(i.symbolic_size for i in expr.dimensions)
 
     return flops, False