Merge pull request #3016 from devitocodes/fix-fd-gather-priority

Fix fd gather priority

Author: mloubout

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

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'),