Performance improvements for transformations

ufl/algorithms/apply_derivatives.py

@@ -7,6 +7,7 @@
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
+from collections import defaultdict
 
 from ufl.log import error, warning
 
@@ -1048,6 +1049,9 @@ def coordinate_derivative(self, o):
 class DerivativeRuleDispatcher(MultiFunction):
     def __init__(self):
         MultiFunction.__init__(self)
+        # caches for reuse in the dispatched transformers
+        self.vcaches = defaultdict(dict)
+        self.rcaches = defaultdict(dict)
 
     def terminal(self, o):
         return o
@@ -1059,24 +1063,41 @@ def derivative(self, o):
 
     def grad(self, o, f):
         rules = GradRuleset(o.ufl_shape[-1])
-        return map_expr_dag(rules, f)
+        key = (GradRuleset, o.ufl_shape[-1])
+        return map_expr_dag(rules, f,
+                            vcache=self.vcaches[key],
+                            rcache=self.rcaches[key])
 
     def reference_grad(self, o, f):
         rules = ReferenceGradRuleset(o.ufl_shape[-1])  # FIXME: Look over this and test better.
-        return map_expr_dag(rules, f)
+        key = (ReferenceGradRuleset, o.ufl_shape[-1])
+        return map_expr_dag(rules, f,
+                            vcache=self.vcaches[key],
+                            rcache=self.rcaches[key])
 
     def variable_derivative(self, o, f, dummy_v):
-        rules = VariableRuleset(o.ufl_operands[1])
-        return map_expr_dag(rules, f)
+        op = o.ufl_operands[1]
+        rules = VariableRuleset(op)
+        key = (VariableRuleset, op)
+        return map_expr_dag(rules, f,
+                            vcache=self.vcaches[key],
+                            rcache=self.rcaches[key])
 
     def coefficient_derivative(self, o, f, dummy_w, dummy_v, dummy_cd):
         dummy, w, v, cd = o.ufl_operands
         rules = GateauxDerivativeRuleset(w, v, cd)
-        return map_expr_dag(rules, f)
+        key = (GateauxDerivativeRuleset, w, v, cd)
+        return map_expr_dag(rules, f,
+                            vcache=self.vcaches[key],
+                            rcache=self.rcaches[key])
 
     def coordinate_derivative(self, o, f, dummy_w, dummy_v, dummy_cd):
         o_ = o.ufl_operands
-        return CoordinateDerivative(map_expr_dag(self, o_[0]), o_[1], o_[2], o_[3])
+        key = (CoordinateDerivative, o_[0])
+        return CoordinateDerivative(map_expr_dag(self, o_[0],
+                                                 vcache=self.vcaches[key],
+                                                 rcache=self.rcaches[key]),
+                                    o_[1], o_[2], o_[3])
 
     def indexed(self, o, Ap, ii):  # TODO: (Partially) duplicated in generic rules
         # Reuse if untouched
@@ -1209,6 +1230,8 @@ def jacobian(self, o):
 class CoordinateDerivativeRuleDispatcher(MultiFunction):
     def __init__(self):
         MultiFunction.__init__(self)
+        self.vcache = defaultdict(dict)
+        self.rcache = defaultdict(dict)
 
     def terminal(self, o):
         return o
@@ -1227,17 +1250,17 @@ def reference_grad(self, o):
     def coefficient_derivative(self, o):
         return o
 
-    def coordinate_derivative(self, o):
+    def coordinate_derivative(self, o, f, w, v, cd):
         from ufl.algorithms import extract_unique_elements
         spaces = set(c.family() for c in extract_unique_elements(o))
         unsupported_spaces = {"Argyris", "Bell", "Hermite", "Morley"}
         if spaces & unsupported_spaces:
             error("CoordinateDerivative is not supported for elements of type %s. "
                   "This is because their pullback is not implemented in UFL." % unsupported_spaces)
-        f, w, v, cd = o.ufl_operands
-        f = self(f)  # transform f
+        _, w, v, cd = o.ufl_operands
         rules = CoordinateDerivativeRuleset(w, v, cd)
-        return map_expr_dag(rules, f)
+        key = (CoordinateDerivativeRuleset, w, v, cd)
+        return map_expr_dag(rules, f, vcache=self.vcache[key], rcache=self.rcache[key])
 
 
 def apply_coordinate_derivatives(expression):

ufl/algorithms/apply_restrictions.py

@@ -21,6 +21,9 @@ def __init__(self, side=None):
         MultiFunction.__init__(self)
         self.current_restriction = side
         self.default_restriction = "+"
+        # Caches for propagating the restriction with map_expr_dag
+        self.vcaches = {"+": {}, "-": {}}
+        self.rcaches = {"+": {}, "-": {}}
         if self.current_restriction is None:
             self._rp = {"+": RestrictionPropagator("+"),
                         "-": RestrictionPropagator("-")}
@@ -32,8 +35,10 @@ def restricted(self, o):
         if self.current_restriction is not None:
             error("Cannot restrict an expression twice.")
         # Configure a propagator for this side and apply to subtree
-        # FIXME: Reuse cache between these calls!
-        return map_expr_dag(self._rp[o.side()], o.ufl_operands[0])
+        side = o.side()
+        return map_expr_dag(self._rp[side], o.ufl_operands[0],
+                            vcache=self.vcaches[side],
+                            rcache=self.rcaches[side])
 
     # --- Reusable rules
 

ufl/corealg/map_dag.py

@@ -14,32 +14,56 @@
 from ufl.corealg.multifunction import MultiFunction
 
 
-def map_expr_dag(function, expression, compress=True):
+def map_expr_dag(function, expression,
+                 compress=True,
+                 vcache=None,
+                 rcache=None):
     """Apply a function to each subexpression node in an expression DAG.
 
     If *compress* is ``True`` (default) the output object from
     the function is cached in a ``dict`` and reused such that the
     resulting expression DAG does not contain duplicate objects.
 
+    If the same funtion is called multiple times in a transformation
+    (as for example in apply_derivatives), then to reuse caches across
+    the call, provide these two arguments:
+
+    :arg vcache: Optional dict for caching results of intermediate transformations
+    :arg rcache: Optional dict for caching results for compression.
+
     Return the result of the final function call.
     """
-    result, = map_expr_dags(function, [expression], compress=compress)
+    result, = map_expr_dags(function, [expression], compress=compress,
+                            vcache=vcache,
+                            rcache=rcache)
     return result
 
 
-def map_expr_dags(function, expressions, compress=True):
+def map_expr_dags(function, expressions,
+                  compress=True,
+                  vcache=None,
+                  rcache=None):
     """Apply a function to each subexpression node in an expression DAG.
 
     If *compress* is ``True`` (default) the output object from
     the function is cached in a ``dict`` and reused such that the
     resulting expression DAG does not contain duplicate objects.
 
+    If the same funtion is called multiple times in a transformation
+    (as for example in apply_derivatives), then to reuse caches across
+    the call, provide these two arguments:
+
+    :arg vcache: Optional dict for caching results of intermediate transformations
+    :arg rcache: Optional dict for caching results for compression.
+
     Return a list with the result of the final function call for each expression.
     """
 
     # Temporary data structures
-    vcache = {}  # expr -> r = function(expr,...),  cache of intermediate results
-    rcache = {}  # r -> r,  cache of result objects for memory reuse
+    # expr -> r = function(expr,...),  cache of intermediate results
+    vcache = {} if vcache is None else vcache
+    # r -> r,  cache of result objects for memory reuse
+    rcache = {} if rcache is None else rcache
 
     # Build mapping typecode:bool, for which types to skip the subtree of
     if isinstance(function, MultiFunction):

ufl/exprequals.py

@@ -145,6 +145,8 @@ def nonrecursive_expr_equals(self, other):
                 left.append((s, o))
 
     # Equal if we get out of the above loop!
+    # Eagerly DAGify to reduce the size of the tree.
+    self.ufl_operands = other.ufl_operands
     return True