api: add mul interp mode

Author: mloubout

devito/core/cpu.py

@@ -87,6 +87,7 @@ def _normalize_kwargs(cls, **kwargs):
 
         # Code generation options for derivatives
         o['expand'] = oo.pop('expand', cls.EXPAND)
+        o['eval-mul-first'] = oo.pop('eval-mul-first', cls.MUL_FIRST)
         o['deriv-collect'] = oo.pop('deriv-collect', cls.DERIV_COLLECT)
         o['deriv-schedule'] = oo.pop('deriv-schedule', cls.DERIV_SCHEDULE)
         o['deriv-unroll'] = oo.pop('deriv-unroll', False)

devito/core/gpu.py

@@ -102,6 +102,7 @@ def _normalize_kwargs(cls, **kwargs):
 
         # Code generation options for derivatives
         o['expand'] = oo.pop('expand', cls.EXPAND)
+        o['eval-mul-first'] = oo.pop('eval-mul-first', cls.MUL_FIRST)
         o['deriv-collect'] = oo.pop('deriv-collect', cls.DERIV_COLLECT)
         o['deriv-schedule'] = oo.pop('deriv-schedule', cls.DERIV_SCHEDULE)
         o['deriv-unroll'] = oo.pop('deriv-unroll', False)

devito/core/operator.py

@@ -125,6 +125,12 @@ class BasicOperator(Operator):
     finite-difference derivatives.
     """
 
+    MUL_FIRST = False
+    """
+    When evaluating expressions location, prioritize multiplication
+    operations.
+    """
+
     DERIV_COLLECT = True
     """
     Factorize finite-difference derivatives exploiting the linearity of the FD

devito/finite_differences/derivative.py

@@ -472,7 +472,7 @@ def T(self):
 
         return self._rebuild(transpose=adjoint)
 
-    def _eval_at(self, func):
+    def _eval_at(self, func, **kwargs):
         """
         Evaluates the derivative at the location of `func`. It is necessary for staggered
         setup where one could have Eq(u(x + h_x/2), v(x).dx)) in which case v(x).dx
@@ -522,7 +522,7 @@ def _eval_at(self, func):
                 return self._rebuild(self.expr, **rkw)
             args = [self.expr.func(*v) for v in mapper.values()]
             args.extend([a for a in self.expr.args if a not in self.expr._args_diff])
-            args = [self._rebuild(a)._eval_at(func) for a in args]
+            args = [self._rebuild(a)._eval_at(func, **kwargs) for a in args]
             return self.expr.func(*args)
         elif self.expr.is_Mul:
             # For Mul, We treat the basic case `u(x + h_x/2) * v(x) which is what appear
@@ -593,7 +593,7 @@ def _eval_fd(self, expr, **kwargs):
             res = generic_derivative(expr, self.dims[0], self.fd_order[0],
                                      self.deriv_order[0], weights=self.weights,
                                      side=self.side, matvec=self.transpose,
-                                     x0=self.x0, expand=expand)
+                                     x0=x0_deriv, expand=expand)
 
         # Step 4: Apply substitutions
         for e in self._ppsubs:

devito/finite_differences/differentiable.py

@@ -184,13 +184,18 @@ def coefficients(self):
         key = lambda x: coeff_priority.get(x, -1)
         return sorted(coefficients, key=key, reverse=True)[0]
 
-    def _eval_at(self, func):
+    def _eval_at(self, func, **kwargs):
         if not func.is_Staggered:
             # Cartesian grid, do no waste time
             return self
+<<<<<<< HEAD
         return self.func(*[
             getattr(a, '_eval_at', lambda x: a)(func) for a in self.args  # noqa: B023
         ])  # false positive
+=======
+        return self.func(*[getattr(a, '_eval_at', lambda x, **kw: a)(func, **kwargs)
+                           for a in self.args])
+>>>>>>> 198aedce4 (api: add mul interp mode)
 
     def _subs(self, old, new, **hints):
         if old == self:
@@ -501,12 +506,20 @@ def highest_priority(diff_op):
     # set of dimensions is used when multiple ones with the same
     # priority appear
     prio = lambda x: (getattr(x, '_fd_priority', 0), len(x.dimensions))
+<<<<<<< HEAD
     prio_func = sorted(diff_op._args_diff, key=prio, reverse=True)[0]
 
     # The highest priority must be a Function
     if not isinstance(prio_func, AbstractFunction):
         return highest_priority(prio_func)
     return prio_func
+=======
+    args = DiffOp._args_diff
+    if not args:
+        return DiffOp
+    else:
+        return sorted(DiffOp._args_diff, key=prio, reverse=True)[0]
+>>>>>>> 198aedce4 (api: add mul interp mode)
 
 
 class DifferentiableOp(Differentiable):
@@ -574,11 +587,16 @@ class DifferentiableFunction(DifferentiableOp):
     def __new__(cls, *args, **kwargs):
         return cls.__sympy_class__.__new__(cls, *args, **kwargs)
 
+<<<<<<< HEAD
     @property
     def _fd_priority(self):
         if highest_priority(self) is self:
             return super()._fd_priority
         return highest_priority(self)._fd_priority
+=======
+    def _eval_at(self, func, **kwargs):
+        return self
+>>>>>>> 198aedce4 (api: add mul interp mode)
 
 
 class Add(DifferentiableOp, sympy.Add):
@@ -669,6 +687,56 @@ def _gather_for_diff(self):
             other = self.func(*other)._eval_at(highest_priority(self))
             return self.func(other, *derivs)
 
+    def _eval_at(self, func, mul_first=False, **kwargs):
+        # Dont evaluate mul first
+        if not mul_first:
+            return super()._eval_at(func, mul_first=mul_first)
+
+        # Not a basic a*b*c... expression, just defer to superclass
+        if any(isinstance(f, DifferentiableOp) for f in self.args):
+            return super()._eval_at(func, mul_first=mul_first)
+
+        # Split Derivative and Differentiable args
+        derivs, other = split(self.args, lambda e: isinstance(e, sympy.Derivative))
+
+        if derivs:
+            derivs = Differentiable._eval_at(self.func(*derivs), func,
+                                             mul_first=mul_first)
+        else:
+            derivs = 1
+
+        if not other:
+            return derivs
+        elif len(other) > 1:
+            expr = self.func(*other)._gather_for_diff
+        else:
+            expr = other[0]
+
+        # Non differentiable expr (e.g., number)
+        if not isinstance(expr, Differentiable):
+            return self.func(derivs, expr)
+
+        # Build mapper for dimensions that need to be interpolated
+        mapper = {}
+        for d in self.dimensions:
+            try:
+                if self.indices_ref[d] is not func.indices_ref[d]:
+                    mapper[d] = func.indices_ref[d]
+            except KeyError:
+                pass
+
+        # Nothing to interpolate
+        if not mapper:
+            return super()._eval_at(func, mul_first=mul_first)
+
+        # Interpolate expr at the required indices
+        interp = expr.diff(*mapper.keys(), deriv_order=[0 for _ in mapper],
+                           fd_order=[self.interp_order for _ in mapper],
+                           x0=mapper)
+
+        # Return the full expression with Derivatives
+        return self.func(derivs, interp)
+
 
 class Pow(DifferentiableOp, sympy.Pow):
     _fd_priority = 0
@@ -1020,7 +1088,7 @@ def _subs(self, old, new, **hints):
 
 class DiffDerivative(IndexDerivative, DifferentiableOp):
 
-    def _eval_at(self, func):
+    def _eval_at(self, func, **kwargs):
         # Like EvalDerivative, a DiffDerivative must have already been evaluated
         # at a valid x0 and should not be re-evaluated at a different location
         return self
@@ -1074,7 +1142,7 @@ def _new_rawargs(self, *args, **kwargs):
         kwargs.pop('is_commutative', None)
         return self.func(*args, **kwargs)
 
-    def _eval_at(self, func):
+    def _eval_at(self, func, **kwargs):
         # An EvalDerivative must have already been evaluated at a valid x0
         # and should not be re-evaluated at a different location
         return self

devito/operator/operator.py

@@ -342,6 +342,7 @@ def _lower_exprs(cls, expressions, **kwargs):
             * Shift indices for domain alignment.
         """
         expand = kwargs['options'].get('expand', True)
+        mul_first = kwargs['options'].get('eval-mul-first', False)
 
         # Specialization is performed on unevaluated expressions
         expressions = cls._specialize_dsl(expressions, **kwargs)
@@ -352,7 +353,8 @@ def _lower_exprs(cls, expressions, **kwargs):
         # ModuloDimensions
         if not expand:
             expand = lambda d: d.is_Stepping
-        expressions = flatten([i._evaluate(expand=expand) for i in expressions])
+        expressions = flatten([i._evaluate(expand=expand, mul_first=mul_first)
+                               for i in expressions])
 
         # Scalarize the tensor equations, if any
         expressions = [j for i in expressions for j in i._flatten]

devito/types/dense.py

@@ -1116,7 +1116,7 @@ def __fd_setup__(self):
 
     @cached_property
     def _fd_priority(self):
-        return 1 if self.staggered.on_node else 2
+        return 1.2 if self.staggered.on_node else 1.1
 
     def _eval_at(self, func):
         if self.staggered == func.staggered or self.interp_order == 0:
@@ -1545,7 +1545,7 @@ def __shape_setup__(cls, **kwargs):
 
     @cached_property
     def _fd_priority(self):
-        return 2.1 if self.staggered.on_node else 2.2
+        return 2.1 if self.staggered.on_node else 2
 
     @property
     def time_order(self):

devito/types/equation.py

@@ -110,7 +110,7 @@ def _evaluate(self, **kwargs):
         """
         try:
             lhs = self.lhs._evaluate(**kwargs)
-            rhs = self.rhs._eval_at(self.lhs)._evaluate(**kwargs)
+            rhs = self.rhs._eval_at(self.lhs, **kwargs)._evaluate(**kwargs)
         except AttributeError:
             lhs, rhs = self._evaluate_args(**kwargs)
         eq = self.func(lhs, rhs, subdomain=self.subdomain,

devito/types/sparse.py

@@ -695,7 +695,7 @@ def _dist_scatter(self, alias=None, data=None):
                 mapper.update(self._dist_subfunc_scatter(sf))
         return mapper
 
-    def _eval_at(self, func):
+    def _eval_at(self, func, **kwargs):
         return self
 
     def _halo_exchange(self):

devito/types/tensor.py

@@ -170,12 +170,13 @@ def __getattr__(self, name):
                 f'{self.__class__!r} object has no attribute {name!r}'
             ) from e
 
-    def _eval_at(self, func):
+    def _eval_at(self, func, **kwargs):
         """
         Evaluate tensor at func location
         """
         def entries(i, j, func):
-            return getattr(self[i, j], '_eval_at', lambda x: self[i, j])(func[i, j])
+            return getattr(self[i, j], '_eval_at',
+                           lambda x: self[i, j])(func[i, j], **kwargs)
         entry = lambda i, j: entries(i, j, func)
         return self._new(self.rows, self.cols, entry)