diff --git a/devito/finite_differences/differentiable.py b/devito/finite_differences/differentiable.py index 88e3cb214e..7bf7effbe8 100644 --- a/devito/finite_differences/differentiable.py +++ b/devito/finite_differences/differentiable.py @@ -22,7 +22,7 @@ from devito.logger import warning from devito.tools import ( as_tuple, extract_dtype, filter_ordered, flatten, frozendict, infer_dtype, is_integer, - is_number, split + is_number, memoized_func, split ) from devito.types import Array, DimensionTuple, Evaluable, StencilDimension from devito.types.basic import AbstractFunction, Indexed @@ -489,6 +489,19 @@ def has_free(self, *patterns): return all(i in self.free_symbols for i in patterns) +@memoized_func(scope='build') +def deep_priority(expr): + """ + The highest `_fd_priority` among the Functions inside `expr`. + + `expr._fd_priority` does not give this: an `Add` or `Mul` falls back on a + generic value, so `mu*tau_xx` reports .75 rather than `tau_xx`'s 2.1. + """ + prio = getattr(expr, '_fd_priority', 0) + return max([prio] + [deep_priority(i) + for i in getattr(expr, '_args_diff', ())]) + + def highest_priority(diff_op, candidates=None): """ The Function whose location a product should be evaluated at. @@ -504,7 +517,9 @@ def highest_priority(diff_op, candidates=None): # We also need to make sure that the object with the largest # set of dimensions is used when multiple ones with the same # priority appear - prio = lambda x: (getattr(x, '_fd_priority', 0), len(x.dimensions)) + # `deep_priority`, not `_fd_priority`: an `Add` or `Mul` reports a generic + # fallback, so operands would tie and `sorted` would pick on argument order + prio = lambda x: (deep_priority(x), len(x.dimensions)) prio_func = sorted(args_diff, key=prio, reverse=True)[0] # The highest priority must be a Function diff --git a/tests/test_differentiable.py b/tests/test_differentiable.py index 444c8456e4..efbcdd33ae 100644 --- a/tests/test_differentiable.py +++ b/tests/test_differentiable.py @@ -6,7 +6,8 @@ from devito import NODE, Differentiable, Eq, Function, Grid, Operator from devito.finite_differences.differentiable import ( - Add, EvalDerivative, Mul, Pow, SafeInv, diffify, interp_for_fd + Add, EvalDerivative, Mul, Pow, SafeInv, deep_priority, diffify, highest_priority, + interp_for_fd ) @@ -224,6 +225,33 @@ def test_mul_three_funcs(self, interp_mode, targets): assert b.name in evaluated_str assert c.name in evaluated_str + def test_mul_gather_priority(self): + """ + The gather lands on the highest-priority Function inside the operands. + + Without `deep_priority` both operands report the `Differentiable` + fallback, tie, and the winner comes down to SymPy's argument order. + """ + grid = Grid((11, 11)) + funcs = self._all_funcs(grid) + node, fx, fy, fxy = (funcs['node'], funcs['x'], + funcs['y'], funcs['xy']) + + # Two sums, only the first carrying a NODE Function + with_node = node * fxy + fxy + staggered = fx + fy + + assert deep_priority(with_node) == node._fd_priority + assert deep_priority(staggered) == fx._fd_priority + assert deep_priority(with_node) > deep_priority(staggered) + + # Whichever way the product sorts, NODE wins + for prod in (with_node * staggered, staggered * with_node): + assert highest_priority( + prod, candidates=[with_node, staggered] + ) is node + assert prod.indices_ref == node.indices_ref + @pytest.mark.parametrize('interp_mode', ['direct', 'symmetric']) @pytest.mark.parametrize('targets', [ ('node', 'x', 'xy'),