compiler: Match product keys in subs_if_composite via a factor index

devito/symbolics/manipulation.py

@@ -1,9 +1,11 @@
-from collections import OrderedDict
+from collections import Counter, OrderedDict
 from collections.abc import Iterable
-from functools import singledispatch
+from functools import partial, singledispatch
 
 import numpy as np
-from sympy import Add, Max, Min, Mul, Pow, S, SympifyError, Tuple, sympify
+from sympy import (
+    Add, Max, Min, Mul, Pow, S, SympifyError, Tuple, preorder_traversal, sympify
+)
 from sympy.core import Basic as SympyBasic
 from sympy.core.add import _addsort
 from sympy.core.mul import _mulsort
@@ -15,7 +17,7 @@
 from devito.symbolics.extended_sympy import DefFunction, rfunc
 from devito.symbolics.queries import q_leaf
 from devito.symbolics.search import (
-    retrieve_functions, retrieve_indexed, retrieve_terminals
+    retrieve_functions, retrieve_indexed, retrieve_terminals, search
 )
 from devito.symbolics.unevaluation import Add as UnevalAdd
 from devito.symbolics.unevaluation import Mul as UnevalMul
@@ -309,19 +311,146 @@ def subs_if_composite(expr, subs):
     `a*b` could be part of `a*b*c`, while `a[1]` cannot be part of a "larger
     Indexed"). Instead, if `subs` consists of just "primitive" expressions, then
     resort to the much faster `uxreplace`.
+
+    Product keys are matched in a single pass by `MulSubsIndex`.
     """
     if not subs:
         return expr
 
-    if type(expr) is tuple:
-        return reuse_if_untouched(expr, (subs_if_composite(e, subs) for e in expr))
-    elif type(expr) is list:
-        return reuse_if_untouched(expr, [subs_if_composite(e, subs) for e in expr])
-    elif all(isinstance(i, (Indexed, IndexDerivative)) for i in subs):
-        return uxreplace(expr, subs)
+    if all(isinstance(i, (Indexed, IndexDerivative)) for i in subs):
+        func = partial(uxreplace, rule=subs)
     else:
-        processed = expr.subs(subs)
-        return expr if processed == expr else processed
+        index = MulSubsIndex.build(subs)
+        func = index.apply if index is not None else partial(_subs_sympy, subs=subs)
+
+    return _subs_if_composite(expr, func)
+
+
+def _subs_if_composite(expr, func):
+    if type(expr) in (tuple, list):
+        return reuse_if_untouched(expr, (_subs_if_composite(e, func) for e in expr))
+    else:
+        return func(expr)
+
+
+def _subs_sympy(expr, subs):
+    processed = expr.subs(subs)
+    return expr if processed == expr else processed
+
+
+def _factor_base(i):
+    return i.base if i.is_Pow else i
+
+
+class Defer(Exception):
+    pass
+
+
+class MulSubsIndex:
+
+    """
+    Apply a substitution whose composite keys are all products in a single pass.
+
+    Each key is filed under the base of one of its factors (preferably an
+    Indexed), and matches a product if its factors are a sub-multiset of the
+    product's. Matches become a compound `uxreplace` rule. Wherever SymPy's
+    `subs` could produce a different outcome, `Defer` is raised and the
+    expression is handed to `expr.subs(subs)` instead.
+    """
+
+    def __init__(self, subs, atomic, index):
+        self.subs = subs
+        self.atomic = atomic
+        self.index = index
+
+    @classmethod
+    def build(cls, subs):
+        """
+        Return the index for `subs`, or None if `subs` cannot be indexed.
+        """
+        atomic = {}
+        index = {}
+        bases = set()
+        for k, v in subs.items():
+            if isinstance(k, (Indexed, IndexDerivative)):
+                atomic[k] = v
+                continue
+            if not k.is_Mul or not k.is_commutative:
+                return None
+
+            factors = Counter(k.args)
+            ordered = [_factor_base(i) for i in factors if not i.is_Number]
+            if not ordered:
+                return None
+            kbases = frozenset(ordered)
+            anchor = next((i for i in ordered if i.is_Indexed), ordered[0])
+
+            index.setdefault(anchor, []).append((k, factors, kbases, v))
+            bases.update(kbases)
+
+        # SymPy applies the keys one after the other, so an atomic key may consume
+        # a factor of a product key, and a replacement may be matched by another key
+        if not bases.isdisjoint(atomic):
+            return None
+        if not all(bases.isdisjoint(preorder_traversal(v)) for v in subs.values()):
+            return None
+
+        return cls(subs, atomic, index)
+
+    def apply(self, expr, strict=False):
+        """
+        Apply the substitution to `expr`. If `strict`, raise `Defer` rather than
+        falling back to SymPy.
+        """
+        rule = dict(self.atomic)
+        try:
+            for m in search(expr, lambda i: i.is_Mul, 'unique', 'bfs'):
+                if m in self.subs:
+                    rule[m] = self.subs[m]
+                elif submap := self._match(m):
+                    rule[m] = submap
+        except Defer:
+            if strict:
+                raise
+            return _subs_sympy(expr, self.subs)
+
+        return uxreplace(expr, rule)
+
+    def _match(self, m):
+        # Keyed by the key itself, as a base may occur in several factors
+        candidates = {}
+        for i in m.args:
+            for c in self.index.get(_factor_base(i), ()):
+                candidates[c[0]] = c
+        if not candidates:
+            return None
+
+        args = Counter(m.args)
+        mbases = {_factor_base(i) for i in args}
+        used = Counter()
+        submap = {}
+        for _, factors, kbases, v in candidates.values():
+            if not kbases <= mbases:
+                continue
+            if not m.is_commutative:
+                raise Defer
+            if not factors <= args:
+                # SymPy may still rewrite it, e.g. `a**2*b**2` or `4*a*b` with
+                # `a*b` or `2*a*b` respectively
+                raise Defer
+            if any(args[i] > 1 for i in factors):
+                # `uxreplace` would drop every copy of a matched factor
+                raise Defer
+            if any(used[i] for i in factors):
+                # Overlapping keys, so SymPy's outcome depends on its key order
+                raise Defer
+            used.update(factors)
+
+            first, *others = factors
+            submap[first] = v
+            submap.update(dict.fromkeys(others))
+
+        return submap
 
 
 def xreplace_indices(exprs, mapper, key=None):

tests/test_symbolics.py

@@ -10,7 +10,9 @@
     Operator, Real, SubDimension, SubDomain, TimeDimension, TimeFunction, configuration,
     cos, norm, sin, solve, switchconfig
 )
-from devito.finite_differences.differentiable import Mul, SafeInv, Weights
+from devito.finite_differences.differentiable import (
+    IndexDerivative, Mul, SafeInv, Weights
+)
 from devito.ir import Expression, FindNodes, ccode
 from devito.ir.support.guards import GuardExpr, pairwise_or, simplify_and
 from devito.mpi.halo_scheme import HaloTouch
@@ -21,6 +23,7 @@
     retrieve_derivatives, retrieve_functions, retrieve_indexed, subs_if_composite,
     uxreplace, xreplace_indices
 )
+from devito.symbolics.manipulation import Defer, MulSubsIndex
 from devito.tools import CustomDtype, as_tuple, dtypes_vector_mapper
 from devito.types import (
     Array, Bundle, ComponentAccess, FIndexed, LocalObject, Object, StencilDimension
@@ -1012,6 +1015,72 @@ def test_subs_if_composite_reuses_untouched_sequence(self):
         assert processed is not exprs
         assert processed[0] is not exprs[0]
 
+    @pytest.mark.parametrize('expr,subs,defer', [
+        ('a*b', '{a*b: r0}', False),
+        ('a*b*c + 1', '{a*b: r0}', False),
+        ('-a*b*c', '{a*b: r0}', False),
+        ('a*b*c + c*e*q', '{a*b: r0, e*q: r1}', False),
+        ('a*b*c + e', '{a*b: r0, e: r1}', False),
+        ('sin(a*b*c) + 1', '{a*b: r0}', False),
+        # SymPy may rewrite these, or its outcome depends on its key order
+        ('a**2*b**2*c', '{a*b: r0}', True),
+        ('c/(a*b)', '{a*b: r0}', True),
+        ('4*a*b*c', '{2*a*b: r0}', True),
+        ('a*b*c', '{a*b: r0, b*c: r1}', True),
+    ])
+    def test_subs_if_composite_mul_keys(self, expr, subs, defer):
+        grid = Grid(shape=(4, 4))
+        a, b, c, e, q, r0, r1 = [Function(name=n, grid=grid).indexify()  # noqa
+                                 for n in ('a', 'b', 'c', 'e', 'q', 'r0', 'r1')]
+
+        expr = Eq(q, eval(expr))
+        subs = eval(subs)
+        expected = expr.subs(subs)
+        index = MulSubsIndex.build(subs)
+
+        assert index is not None
+        if defer:
+            with pytest.raises(Defer):
+                index.apply(expr, strict=True)
+        else:
+            assert index.apply(expr, strict=True) == expected
+        assert subs_if_composite((expr,), subs)[0] == expected
+        assert expected != expr
+
+    def test_subs_if_composite_mul_keys_index_derivative(self):
+        grid = Grid(shape=(10,))
+        x, = grid.dimensions
+        i = StencilDimension('i', 0, 2)
+        f = Function(name='f', grid=grid, space_order=2)
+        g = Function(name='g', grid=grid, space_order=2)
+        r = Function(name='r', grid=grid)
+
+        fi = f.indexify().subs(x, x + i)
+        gi = g.indexify().subs(x, x + i)
+        w = Weights(name='w0', dimensions=i, initvalue=[-0.5, 0, 0.5])
+        ideriv = IndexDerivative(fi*gi*w, {x: i})
+
+        expr = Eq(r.indexify(), ideriv + 1)
+        subs = {ideriv.base: r[x + i]}
+        expected = expr.subs(subs)
+
+        assert MulSubsIndex.build(subs).apply(expr, strict=True) == expected
+        assert expected != expr
+
+    def test_subs_if_composite_mul_keys_untouched(self):
+        grid = Grid(shape=(4, 4))
+        a, b, c, e, r0, r1 = [Function(name=n, grid=grid).indexify()
+                              for n in ('a', 'b', 'c', 'e', 'r0', 'r1')]
+
+        exprs = (Eq(r0, a*c + b*e), Eq(r1, a + b))
+        items = list(exprs)
+        subs = {a*b: r0, c*e: r1}
+        index = MulSubsIndex.build(subs)
+
+        assert subs_if_composite(exprs, subs) is exprs
+        assert subs_if_composite(items, subs) is items
+        assert all(index.apply(e, strict=True) is e for e in exprs)
+
     def test_pow_to_mul_reuses_untouched_sequence(self):
         grid = Grid(shape=(4, 4))
         x, y = grid.dimensions