Performance improvements for transformations

[wence-]

This PR does two things:

1. DAGifies eagerly when comparing UFL expressions for equality. If two expressions compare equal, then we can replace the operands of one with the operands of the other.
2. Provides a facility for map_expr_dags to reuse the result cache across calls. This makes a big difference in applying restrictions and derivatives, which have a dispatch rule that calls map_expr_dag that previously would not cache across calls into subtrees.

I am pretty sure that they are correct, but will run the Firedrake test suite to check for regressions.

Representative timing data for the below script

Before:

$ python slow-ufl-example.py
cfd took 9.76s

After eager DAGifying:

$ python slow-ufl-example.py
cfd took 1.25s

Reusing caches in apply restrictions

$ python slow-ufl-example.py
cfd took 0.17s

With all changes:

$ python slow-ufl-example.py
cfd took 0.15s

import time

from ufl import (Coefficient, Constant, FacetNormal, FiniteElement,
                 FunctionSpace, Identity, Mesh, VectorElement, avg, derivative,
                 det, diff, dS, grad, hexahedron, inner, ln, outer, variable)
from ufl.algorithms import compute_form_data

cell = hexahedron
mesh = Mesh(VectorElement("Q", cell, 1))
V = FunctionSpace(mesh, FiniteElement("NCF", cell, 2))
u = Coefficient(V)
mu = Constant(mesh)

psi = lambda F: (mu/2) * (inner(F, F) - u.ufl_shape[0]) - mu*ln(det(F))
flux = lambda F: diff(psi(F), F)

n = FacetNormal(mesh)

uxn = outer(u('+'), n('+')) + outer(u('-'), n('-'))
eye = Identity(u.ufl_shape[0])
grad_u = grad(u)

F_ = variable(grad_u + eye)

Fm = variable(grad_u('-') + eye)
Fp = variable(grad_u('+') + eye)

avg_flux = avg(flux(F_))

U = inner(avg_flux, uxn)*dS
a = derivative(derivative(U, u), u)

start = time.time()
fd = compute_form_data(a,
                       do_apply_function_pullbacks=True,
                       do_apply_default_restrictions=True,
                       do_apply_geometry_lowering=True,
                       do_apply_restrictions=True,
                       complex_mode=False)
end = time.time()

print(f"cfd took {end - start:.2f}s")

[wence-]

Firedrake test suit is happy.

@michalhabera, @IgorBaratta Since this makes some fairly core changes to the visitor mapping, can you check that fenicsx is happy?

[IgorBaratta]

Firedrake test suit is happy.

@michalhabera, @IgorBaratta Since this makes some fairly core changes to the visitor mapping, can you check that fenicsx is happy?

Dolfinx and ffcx seem happy. All tests are green on CI, using this branch.

[wence-]

Firedrake test suit is happy.
@michalhabera, @IgorBaratta Since this makes some fairly core changes to the visitor mapping, can you check that fenicsx is happy?

Dolfinx and ffcx seem happy. All tests are green on CI, using this branch.

Thanks. I'm reasonably happy to merge, but you may wish to look over the changes.