Group integrals with common integrand in IntegralData

[jorgensd]

MWE:

from ufl import FiniteElement, TrialFunction, TestFunction, dx, Mesh, VectorElement, triangle, FunctionSpace, inner, grad
import ufl.algorithms
import ufl.classes
mesh = Mesh(VectorElement("Lagrange", triangle, 1))
V = FunctionSpace(mesh, FiniteElement("Lagrange", triangle, 1))
u = TrialFunction(V)
v = TestFunction(V)
a = inner(u, v)*dx((1, 2, 3)) + inner(grad(u), grad(v))*dx(2) + inner(u, v)*dx

form_data = ufl.algorithms.compute_form_data(
    a,
    do_apply_function_pullbacks=True,
    do_apply_integral_scaling=True,
    do_apply_geometry_lowering=True,
    preserve_geometry_types=(ufl.classes.Jacobian,),
    do_apply_restrictions=True,
    do_append_everywhere_integrals=False)
for integral in form_data.integral_data:
    print(str(integral))

On this branch, we get two IntegralData, one for inner(u,v)*dx((1,3,"everywhere")) and one for (inner(u,v) + inner(grad(u), grad(v)))*dx((2,)).
On main one would get four integral datas; inner(u,v)*dx(i) for i=1,3,"everywhere" and (inner(u,v) + inner(grad(u), grad(v)))*dx((2,)).

This is relevant when loading complex meshes with hundreds of subdomain_ids, where a tuple of subdomain ids is used to indicate a union of domains to integrate over. This would cause the compilers in Firedrake and DOLFINx to generate hundreds of duplicate kernels, which makes the generated code unreadable and slow to generate.

[wence-]

@connorjward could you try if this works for us. I think it will need some minor adaptation in the tsfc pipeline to allow the subdomain_id to be a tuple of subdomain_ids (the construction of the measure_set to integrate over I think will just work then).

[connorjward]

@connorjward could you try if this works for us. I think it will need some minor adaptation in the tsfc pipeline to allow the subdomain_id to be a tuple of subdomain_ids (the construction of the measure_set to integrate over I think will just work then).

Sure will do.

[jorgensd]

@connorjward @wence- I’ve find a tiny bug when running it through some of my extensions. I’ll let you know once I’ve fixed it

[jorgensd]

@connorjward @wence- I've now corrected some mistakes I made the first time around, and added a test checking that merging integral data does what we expect.

[connorjward]

I've now got a branch of TSFC where this works. I haven't implemented the necessary bits to achieve any performance improvements since we touch code using subdomain_id in a bunch of places (e.g. Slate, multigrid) and I haven't had time to track them all down.

[jorgensd]

I've now got a branch of TSFC where this works. I haven't implemented the necessary bits to achieve any performance improvements since we touch code using subdomain_id in a bunch of places (e.g. Slate, multigrid) and I haven't had time to track them all down.

@connorjward any update on this?

[connorjward]

I've now got a branch of TSFC where this works. I haven't implemented the necessary bits to achieve any performance improvements since we touch code using subdomain_id in a bunch of places (e.g. Slate, multigrid) and I haven't had time to track them all down.

@connorjward any update on this?

I'm afraid I haven't done anything more on this. I was only intending to work on optimising TSFC once this PR was merged into UFL since they didn't need to land together. Sorry if there was a miscommunication here.

[jorgensd]

I've now got a branch of TSFC where this works. I haven't implemented the necessary bits to achieve any performance improvements since we touch code using subdomain_id in a bunch of places (e.g. Slate, multigrid) and I haven't had time to track them all down.

@connorjward any update on this?

I'm afraid I haven't done anything more on this. I was only intending to work on optimising TSFC once this PR was merged into UFL since they didn't need to land together. Sorry if there was a miscommunication here.

Ok, I'll get the PR into a mergable shape and then we can get it merged. I thought it was a blocker for getting it merged if Firedrake didn't get the speedup.

[connorjward]

I've now got a branch of TSFC where this works. I haven't implemented the necessary bits to achieve any performance improvements since we touch code using subdomain_id in a bunch of places (e.g. Slate, multigrid) and I haven't had time to track them all down.

@connorjward any update on this?

I'm afraid I haven't done anything more on this. I was only intending to work on optimising TSFC once this PR was merged into UFL since they didn't need to land together. Sorry if there was a miscommunication here.

Ok, I'll get the PR into a mergable shape and then we can get it merged. I thought it was a blocker for getting it merged if Firedrake didn't get the speedup.

Go right ahead. Firedrake follows a fork of UFL (that we regularly update) so this can even land without needing to coordinate anything with us. It won't break anything for our users.

[jorgensd]

@dham @connorjward I'm planning to merge this sometime this week, as the FEniCSx integration tests are passing, and this yields a significant speedup for form-compilation of forms with many subdomains.

Note that if TSFC decides not to adapt this in the recent future, changes made by @mscroggs to UFL might not be easily picked up as breaking changes for TSFC

[connorjward]

@dham @connorjward I'm planning to merge this sometime this week, as the FEniCSx integration tests are passing, and this yields a significant speedup for form-compilation of forms with many subdomains.

Note that if TSFC decides not to adapt this in the recent future, changes made by @mscroggs to UFL might not be easily picked up as breaking changes for TSFC

I am on it. I was able to get most things working (with the optimisation this time) pretty quickly.

[jorgensd]

As @chrisrichardson is doing a major restructuring of FFCx, I'll merge this very soon, as it keeps getting tedious to resolve merge conflicts in FFCx.

[mscroggs]

Does this fix #70?

[jorgensd]

I dont think so