dsl: Replace SparseEq kind string with Interpolation/Injection subclasses

The 'kind' attribute and the 'kind == ...' / 'isinstance(eq, Inc)' branches
that gated interpolation vs. injection behaviour are replaced by leaf
classes (Interpolation, IncrInterpolation, Injection) sharing a small
InterpolationMixin / InjectionMixin. The polymorphic surface (efunc_prefix,
field, is_head_eq, sparse_temps) lives on the leaves, so the IET pass and
the lowering pipeline drop their kind-discriminating conditionals.

Author: mloubout

devito/ir/equations/equation.py

@@ -12,14 +12,22 @@
 from devito.symbolics import IntDiv, limits_mapper, uxreplace
 from devito.tools import Pickable, Tag, frozendict
 from devito.types import (
-    Eq, Inc, ReduceMax, ReduceMin, ReduceMinMax, SparseEq, SparseOpMixin, relational_min
+    Eq, Inc, IncrInterpolation, Injection, InjectionMixin, Interpolation,
+    InterpolationMixin, ReduceMax, ReduceMin, ReduceMinMax, SparseEq, SparseOpMixin,
+    relational_min
 )
 
 __all__ = [
     'ClusterizedEq',
+    'ClusterizedIncrInterpolation',
+    'ClusterizedInjection',
+    'ClusterizedInterpolation',
     'ClusterizedSparseEq',
     'DummyEq',
     'LoweredEq',
+    'LoweredIncrInterpolation',
+    'LoweredInjection',
+    'LoweredInterpolation',
     'LoweredSparseEq',
     'OpInc',
     'OpMax',
@@ -283,12 +291,37 @@ class LoweredSparseEq(SparseOpMixin, LoweredEq):
 
     """
     The IR counterpart of ``SparseEq``: a regular ``LoweredEq`` that
-    also carries the ``interpolator``/``kind`` metadata used by the IET
-    pass ``lower_sparse_ops`` to wrap the resulting ``p_*, rp_*``
-    iteration nest in an ElementalFunction.
+    also carries the ``interpolator`` metadata used by the IET pass
+    ``lower_sparse_ops`` to wrap the resulting ``p_*, rp_*`` iteration
+    nest in an ElementalFunction. Subclassed by
+    ``LoweredInterpolation`` / ``LoweredIncrInterpolation`` /
+    ``LoweredInjection`` for the per-operation polymorphic behaviour.
     """
 
-    __rkwargs__ = LoweredEq.__rkwargs__ + ('interpolator', 'kind')
+    __rkwargs__ = LoweredEq.__rkwargs__ + ('interpolator',)
+
+
+class LoweredInterpolation(InterpolationMixin, LoweredSparseEq):
+    """IR counterpart of ``Interpolation``."""
+    pass
+
+
+class LoweredIncrInterpolation(InterpolationMixin, LoweredSparseEq):
+    """IR counterpart of ``IncrInterpolation``."""
+    pass
+
+
+class LoweredInjection(InjectionMixin, LoweredSparseEq):
+    """IR counterpart of ``Injection``."""
+    pass
+
+
+# Map user-level sparse-op classes to their IR-level counterparts.
+_lowered_sparse_cls = {
+    Interpolation: LoweredInterpolation,
+    IncrInterpolation: LoweredIncrInterpolation,
+    Injection: LoweredInjection,
+}
 
 
 class ClusterizedEq(IREq):
@@ -350,12 +383,35 @@ class ClusterizedSparseEq(SparseOpMixin, ClusterizedEq):
 
     """
     Frozen counterpart of ``LoweredSparseEq``: the same regular
-    ``ClusterizedEq`` augmented with ``interpolator``/``kind`` so the
-    IET pass ``lower_sparse_ops`` can identify and rewrite the sparse
-    op's iteration nest.
+    ``ClusterizedEq`` augmented with ``interpolator`` so the IET pass
+    ``lower_sparse_ops`` can identify and rewrite the sparse op's
+    iteration nest. Subclassed by ``ClusterizedInterpolation`` /
+    ``ClusterizedIncrInterpolation`` / ``ClusterizedInjection``.
     """
 
-    __rkwargs__ = ClusterizedEq.__rkwargs__ + ('interpolator', 'kind')
+    __rkwargs__ = ClusterizedEq.__rkwargs__ + ('interpolator',)
+
+
+class ClusterizedInterpolation(InterpolationMixin, ClusterizedSparseEq):
+    """Frozen counterpart of ``LoweredInterpolation``."""
+    pass
+
+
+class ClusterizedIncrInterpolation(InterpolationMixin, ClusterizedSparseEq):
+    """Frozen counterpart of ``LoweredIncrInterpolation``."""
+    pass
+
+
+class ClusterizedInjection(InjectionMixin, ClusterizedSparseEq):
+    """Frozen counterpart of ``LoweredInjection``."""
+    pass
+
+
+_clusterized_sparse_cls = {
+    LoweredInterpolation: ClusterizedInterpolation,
+    LoweredIncrInterpolation: ClusterizedIncrInterpolation,
+    LoweredInjection: ClusterizedInjection,
+}
 
 
 class DummyEq(ClusterizedEq):
@@ -408,11 +464,11 @@ def _(eq):
 
     if augmented != tuple(eq.implicit_dims or ()):
         eq = eq.func(eq.lhs, eq.rhs, interpolator=interp,
-                     kind=eq.kind, implicit_dims=augmented)
+                     implicit_dims=augmented)
 
-    obj = LoweredSparseEq(eq)
+    lowered_cls = _lowered_sparse_cls[type(eq)]
+    obj = lowered_cls(eq)
     obj._interpolator = interp
-    obj._kind = eq.kind
     return obj
 
 
@@ -427,6 +483,12 @@ def clusterize_eq(eq, **kwargs):
 
 
 @clusterize_eq.register(LoweredSparseEq)
+def _(eq, **kwargs):
+    return _clusterized_sparse_cls[type(eq)](eq, **kwargs)
+
+
 @clusterize_eq.register(ClusterizedSparseEq)
 def _(eq, **kwargs):
-    return ClusterizedSparseEq(eq, **kwargs)
+    # ``eq`` is already clusterized; rebuild via its own class to preserve
+    # the per-operation polymorphic behaviour.
+    return type(eq)(eq, **kwargs)

devito/operations/interpolators.py

@@ -15,7 +15,9 @@
 from devito.logger import warning
 from devito.symbolics import INT, retrieve_function_carriers, retrieve_functions
 from devito.tools import as_tuple, filter_ordered, memoized_meth
-from devito.types import Eq, Inc, SparseEq, SparseInc, SubFunction, Symbol
+from devito.types import (
+    Eq, Inc, IncrInterpolation, Injection, Interpolation, SubFunction, Symbol
+)
 from devito.types.utils import DimensionTuple
 
 __all__ = ['LinearInterpolator', 'PrecomputedInterpolator', 'SincInterpolator']
@@ -82,40 +84,38 @@ def _extract_subdomain(variables):
 
 def _build_interpolation(expr, increment, implicit_dims, self_subs, interpolator):
     """
-    Construct the SparseEq for an interpolation: the synthetic Eq is
-    ``Eq(sf[..., p_*], expr[..., rp_*])``; with ``increment`` it is
+    Construct the sparse-op Eq for an interpolation: the synthetic Eq
+    is ``Eq(sf[..., p_*], expr[..., rp_*])``; with ``increment`` it is
     an ``Inc``. User-supplied ``implicit_dims`` are carried as-is; the
     SparseFunction's iteration Dimensions are augmented in by
-    ``SparseEq.lower`` so the cluster pipeline sees them.
+    ``lower_eq`` so the cluster pipeline sees them.
     """
     eq = interpolator._interpolate(expr=expr, increment=increment,
                                    self_subs=self_subs,
                                    implicit_dims=None)
-    cls = SparseInc if isinstance(eq, Inc) else SparseEq
+    cls = IncrInterpolation if isinstance(eq, Inc) else Interpolation
     return cls(eq.lhs, eq.rhs, interpolator=interpolator,
-               kind='interpolate',
                implicit_dims=implicit_dims)
 
 
 def _build_injection(field, expr, implicit_dims, interpolator):
     """
-    Construct the SparseEq(s) for an injection: each synthetic Eq is
-    ``Inc(field[..., x, y, ...], weights * expr[..., rp_*])`` produced
-    by ``interpolator._inject``. A multi-field injection expands into
-    one ``SparseEq`` per ``(field, expr)`` pair so each target field is
-    individually visible to the cluster pipeline. User-supplied
-    ``implicit_dims`` are carried as-is; sparse-function iteration
-    Dimensions are augmented in by ``SparseEq.lower``.
+    Construct the ``Injection``(s) for an injection: each synthetic Eq
+    is ``Inc(field[..., x, y, ...], weights * expr[..., rp_*])``
+    produced by ``interpolator._inject``. A multi-field injection
+    expands into one ``Injection`` per ``(field, expr)`` pair so each
+    target field is individually visible to the cluster pipeline.
+    User-supplied ``implicit_dims`` are carried as-is; sparse-function
+    iteration Dimensions are augmented in by ``lower_eq``.
     """
     fields, exprs = as_tuple(field), as_tuple(expr)
     if len(exprs) == 1:
         exprs = tuple(exprs[0] for _ in fields)
     eqs = []
     for (f, e) in zip(fields, exprs, strict=True):
         inc = interpolator._inject(field=f, expr=e, implicit_dims=None)
-        eqs.append(SparseEq(inc.lhs, inc.rhs, interpolator=interpolator,
-                            kind='inject',
-                            implicit_dims=implicit_dims))
+        eqs.append(Injection(inc.lhs, inc.rhs, interpolator=interpolator,
+                             implicit_dims=implicit_dims))
     return eqs[0] if len(eqs) == 1 else eqs
 
 
@@ -261,6 +261,25 @@ def _positions(self, implicit_dims, shifts=None):
         return [Eq(v, INT(floor(k)), implicit_dims=implicit_dims)
                 for k, v in self.sfunction._position_map(shifts=shifts).items()]
 
+    def sparse_temps(self, rhs, implicit_dims, field=None):
+        """
+        Position/coefficient temps for a sparse op with right-hand side
+        ``rhs``. For an injection, ``field`` drives the per-Dimension
+        shifts so the temps' lhs (``pos*`` symbols) match the rhs of a
+        staggered injection; for an interpolation, ``field`` is None
+        and no shifts are applied.
+        """
+        if field is not None:
+            extras = [field] + list(retrieve_function_carriers(rhs))
+            shifts = self._field_shifts(field)
+        else:
+            extras = list(retrieve_function_carriers(rhs)) or None
+            shifts = None
+
+        implicit_dims = self._augment_implicit_dims(implicit_dims, extras=extras)
+        return list(self._positions(implicit_dims, shifts=shifts)) + \
+            list(self._coeff_temps(implicit_dims, shifts=shifts))
+
     def _interp_idx(self, variables, subdomain=None, shifts=None):
         """
         Generate the indirect-access index substitutions for the
@@ -289,30 +308,6 @@ def _interp_idx(self, variables, subdomain=None, shifts=None):
 
         return {v: v.subs(subs) for v in variables}
 
-    def _sparse_temps(self, kind, expr, field=None, implicit_dims=None):
-        """
-        Position/coefficient temps emitted alongside the radius
-        expansion. ``implicit_dims`` is augmented with the
-        SparseFunction's iteration dimensions (and any dim carried by
-        the operation inputs) so the temps share the radius-nest's
-        iteration space. For injection, ``field``'s staggering drives
-        the position-symbol shifts so the rhs `pos*` symbols match
-        the temps' lhs.
-        """
-        if kind == 'inject':
-            extras = list(as_tuple(field))
-            if expr is not None:
-                extras.extend(retrieve_function_carriers(expr))
-            shifts = self._field_shifts(field) if field is not None else None
-        else:
-            extras = list(retrieve_function_carriers(expr)) \
-                if expr is not None else None
-            shifts = None
-
-        implicit_dims = self._augment_implicit_dims(implicit_dims, extras=extras)
-        return list(self._positions(implicit_dims, shifts=shifts)) + \
-            list(self._coeff_temps(implicit_dims, shifts=shifts))
-
     @check_radius
     @check_coords
     def interpolate(self, expr, increment=False, self_subs=None, implicit_dims=None):

devito/passes/iet/sparse.py

@@ -32,7 +32,7 @@
     Transformer, make_callable
 )
 from devito.passes.iet.engine import iet_pass
-from devito.types import Eq, Symbol
+from devito.types import Eq, InjectionMixin, InterpolationMixin, Symbol
 
 __all__ = ['lower_sparse_ops']
 
@@ -53,7 +53,9 @@ def lower_sparse_ops(iet, sregistry=None, **kwargs):
     # from the original SparseEq (e.g. by ``factorize``/``cse``) are
     # left where they are inside the radius nest.
     sparse_exprs = [e for e in FindNodes(Expression).visit(iet)
-                    if e.expr.is_SparseOperation and _is_head(e.expr)]
+                    if e.expr.is_SparseOperation
+                    and type(e.expr).is_head_eq(e.expr,
+                                                e.expr.interpolator.sfunction)]
     if not sparse_exprs:
         return iet, {}
 
@@ -81,7 +83,7 @@ def lower_sparse_ops(iet, sregistry=None, **kwargs):
         new_nest = _materialise_nest(nest, exprs)
 
         lse = exprs[0].expr
-        prefix = f'{lse.kind}_{lse.interpolator.sfunction.name}'
+        prefix = f'{lse.efunc_prefix}_{lse.interpolator.sfunction.name}'
         efunc = make_callable(sregistry.make_name(prefix=prefix), new_nest)
         efuncs.append(efunc)
 
@@ -93,21 +95,6 @@ def lower_sparse_ops(iet, sregistry=None, **kwargs):
     return Transformer(mapper).visit(iet), {'efuncs': efuncs}
 
 
-def _is_head(eq):
-    """
-    True if ``eq`` is the "head" of its sparse op: the Expression
-    whose lhs is the SparseFunction (interpolation) or a
-    DiscreteFunction grid field (injection), as opposed to an
-    auxiliary scalar temporary extracted from the original SparseEq by
-    a cluster pass.
-    """
-    sf = eq.interpolator.sfunction
-    f = eq.lhs.function
-    if eq.kind == 'interpolate':
-        return f is sf
-    return f.is_DiscreteFunction and f is not sf
-
-
 def _find_outer_iteration(iet, expr):
     """
     Walk up the IET from ``expr`` and return the outermost Iteration
@@ -140,25 +127,22 @@ def _materialise_nest(nest, exprs):
     pattern. Multiple sparse-op Expressions sharing the same outer
     Iteration are materialised in one pass and reuse the same temps.
     """
-    sample = exprs[0].expr
-    interp = sample.interpolator
-
     # Position + coefficient temporaries as IET Expressions. These are
     # the same for every Expression in the group, so we emit them once.
-    field = sample.lhs.function if sample.kind == 'inject' else None
-    temps = interp._sparse_temps(sample.kind, sample.rhs, field=field,
-                                 implicit_dims=sample.implicit_dims)
+    # The sample's leaf class (Interpolation/Injection) drives whether
+    # the temps carry staggering shifts.
+    sample = exprs[0].expr
     temp_exprs = tuple(Expression(DummyEq(e.lhs, e.rhs))
-                       for e in lower_exprs(temps))
+                       for e in lower_exprs(sample.sparse_temps()))
 
     # The radius nest is what runs once per sparse point. For each
     # interpolation Expression in the group, build its
     # accumulator-wrapped copy of the radius nest. Injection Exprs
     # share a single copy of the radius nest (their ``Inc`` already
     # carries the right ``weights * rhs`` form).
     inner = nest.nodes[0] if len(nest.nodes) == 1 else List(body=nest.nodes)
-    interp_exprs = [e for e in exprs if e.expr.kind == 'interpolate']
-    inject_exprs = [e for e in exprs if e.expr.kind == 'inject']
+    interp_exprs = [e for e in exprs if isinstance(e.expr, InterpolationMixin)]
+    inject_exprs = [e for e in exprs if isinstance(e.expr, InjectionMixin)]
 
     body = []
     for expr in interp_exprs: