compiler: Generalize ideriv lowering

Author: FabioLuporini

devito/passes/clusters/derivatives.py

@@ -3,7 +3,7 @@
 import numpy as np
 from sympy import S
 
-from devito.finite_differences import IndexDerivative
+from devito.finite_differences import IndexDerivative, Weights
 from devito.ir import Backward, Forward, Interval, IterationSpace, Queue
 from devito.passes.clusters.misc import fuse
 from devito.symbolics import BasicWrapperMixin, reuse_if_untouched, uxreplace
@@ -91,17 +91,39 @@ def _core(expr, c, ispace, weights, reusables, mapper, **kwargs):
 
 
 @_core.register(Symbol)
-@_core.register(Indexed)
 @_core.register(BasicWrapperMixin)
 def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
     return expr, []
 
 
+@_core.register(Indexed)
+def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
+    if not isinstance(expr.function, Weights):
+        return expr, []
+
+    # Lower or reuse a previously lowered Weights array
+    sregistry = kwargs['sregistry']
+    subs_user = kwargs['subs']
+
+    w0 = expr.function
+    k = tuple(w0.weights)
+    try:
+        w = weights[k]
+    except KeyError:
+        name = sregistry.make_name(prefix='w')
+        dtype = infer_dtype([w0.dtype, c.dtype])  # At least np.float32
+        initvalue = tuple(i.subs(subs_user) for i in k)
+        w = weights[k] = w0._rebuild(name=name, dtype=dtype, initvalue=initvalue)
+
+    rebuilt = expr._subs(w0.indexed, w.indexed)
+
+    return rebuilt, []
+
+
 @_core.register(IndexDerivative)
 def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
     sregistry = kwargs['sregistry']
     options = kwargs['options']
-    subs_user = kwargs['subs']
 
     try:
         cbk0 = deriv_schedule_registry[options['deriv-schedule']]
@@ -114,18 +136,10 @@ def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
 
     # Create the concrete Weights array, or reuse an already existing one
     # if possible
-    name = sregistry.make_name(prefix='w')
-    w0 = ideriv.weights.function
-    dtype = infer_dtype([w0.dtype, c.dtype])  # At least np.float32
-    k = tuple(w0.weights)
-    try:
-        w = weights[k]
-    except KeyError:
-        initvalue = tuple(i.subs(subs_user) for i in k)
-        w = weights[k] = w0._rebuild(name=name, dtype=dtype, initvalue=initvalue)
+    w, _ = _core(ideriv.weights, c, ispace, weights, reusables, mapper, **kwargs)
 
     # Replace the abstract Weights array with the concrete one
-    subs = {w0.indexed: w.indexed}
+    subs = {ideriv.weights.base: w.base}
     init = uxreplace(init, subs)
     ideriv = uxreplace(ideriv, subs)
 
@@ -155,10 +169,10 @@ def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
     # NOTE: created before recurring so that we ultimately get a sound ordering
     try:
         s = reusables.pop()
-        assert np.can_cast(s.dtype, dtype)
+        assert np.can_cast(s.dtype, w.dtype)
     except KeyError:
         name = sregistry.make_name(prefix='r')
-        s = Symbol(name=name, dtype=dtype)
+        s = Symbol(name=name, dtype=w.dtype)
 
     # Go inside `expr` and recursively lower any nested IndexDerivatives
     expr, processed = _core(expr, c, ispace1, weights, reusables, mapper, **kwargs)