Keep factorised contractions atomic in dual evaluation

[pbrubeck]

Stacked on #269.

TLDR: interpolation was undoing sum-factorisation from the source expression and refactorising it after computing the degrees of freedom. This PR preserves the original sum-factorization by stopping early when traversing a contracted GEM expression.

TSFC hands FInAT's dual evaluation coefficient evaluations that are already sum factorised, and traverse_product flattened them back into the surrounding contraction, undoing sum-factorization. That discarded the factorisation along with the subexpressions the factors shared, and multiplied the indices to search over, so a product of a few evaluations exceeded the sum factorisation limit.

This PR makes contractions atomic: the new gem.optimise.is_contraction predicate is passed as stop_at, so an already factorised contraction is not flattened. Bounding the index count this way retires the ignore argument of contraction, along with the shape-index workaround that was its only user.

Kernel flop counts for interpolation on a hexahedron, where error is NotImplementedError: Too many indices for sum factorisation!:

interpolate	before	after
f*f*f, CG1 → CG1	96	96
f*f, CG4 → DG3	5008	2568
dot(u, u)*u, vector CG4 → vector CG4	error	23875
dot(A, A), tensor CG2 → tensor CG3	error	16776
inner(A, A), tensor CG2 → CG2	error	4131
grad(f)[0], CG4 → CG4	13981	14106
div(u), vector CG4 → DG3	23100	23548

Tests are added for the GEM and FInAT paths.

AI declaration: written with Claude Code (Claude Opus 5).

[connorjward]

Is it always right to do this? Sum factorisation might benefit from seeing all of the indices no?

[pbrubeck]

Is it always right to do this? Sum factorisation might benefit from seeing all of the indices no?

For this particular situation sum-factorization happens first in the source expression, and we run it again after we contract the quadrature weights to compute degrees of freedom. This is right because it is somehow preserving the mathematical structures, althoug it might not be optimal.

The optimal approach is to factor everything in one go, but that will give you much harder problem, and you might pay a price, which is a considerably longer compile time.

[pbrubeck]

The optimal approach is to factor everything in one go, but that will give you much harder problem, and you might pay a price, which is a considerably longer compile time.

Well, actually, I don't think that the splitting should make things (asymptotically) suboptimal, since sum-factorization should be working independently on the connected components of the contraction graph. It does make sense to split the task in two stages.

The only way I can explain the small regressions is that scalar operations (i.e. Sum as opposed to IndexSum) are being reorganized, and we don't apply many optimizations for those.

[connorjward]

The optimal approach is to factor everything in one go, but that will give you much harder problem, and you might pay a price, which is a considerably longer compile time.

Now that we split things into connected components are things actually that slow?

I am hesitant to approve this PR because it feels hacky and I want to be confident that we actually need the hack in the first place.

[pbrubeck]

Now that we split things into connected components are things actually that slow?

It shouldn't be that slow. But we can try to shut off the inner sum-factorization and see if we can do it in a single outer pass.

[pbrubeck]

And you might be right. This PR is hacky because the original code was very hacky.

[pbrubeck]

@connorjward I tried one single pass of gem.optimise.contraction (which itself calls sum_factorise). I basically got rid of the sum_factorise call we have above it. This hits the error NotImplementedError: Too many indices for sum factorisation!, so naturally I increased the limit from 6 to 8. I ran benchmarks, and everything indicates that we still need the split.

• There are a few cases where single pass is 8-60x slower than splitting.
• Only in some cases, the single pass yields marginal gains in flops (which is what it optimises). In other cases, single pass regresses by a factor of 2 in flops.
• The allocated memory in the single-pass kernel can be up to 10x larger than the split kernel.

The moral of the story is that gem.optimise.contraction needs to be used be care. We have much better optimisation machinery in Form assembly than we have for interpolation. And the plan is to move all codegen for same-mesh interpolate into assemble: firedrakeproject/firedrake#5258

Flops (ERR = NotImplementedError: Too many indices for sum factorisation!)

interpolate on hex	#269	#280	single pass
f*f*f, CG1 → CG1	96	96	96
f*f*f*f, CG1 → CG1	96	104	96
f*f, CG4 → DG3	5008	2568	5008
dot(u, u)*u, vector CG4 → vector CG4	ERR	23875	35125
dot(A, A), tensor CG2 → tensor CG3	ERR	16776	16776
inner(A, A), tensor CG2 → CG2	ERR	4131	4374
grad(f)[0], CG4 → CG4	13981	14106	13981
div(u), vector CG4 → DG3	23100	23548	23100

Temporaries, as count / total entries / largest

interpolate on hex	#269	#280	single pass
f*f*f, CG1 → CG1	9 / 11 / 2	9 / 11 / 2	9 / 21 / 4
f*f*f*f, CG1 → CG1	10 / 12 / 2	9 / 11 / 2	10 / 22 / 4
f*f, CG4 → DG3	7 / 82 / 25	4 / 51 / 25	7 / 382 / 100
dot(u, u)*u, vector CG4 → vector CG4	ERR	14 / 210 / 75	23 / 2283 / 375
dot(A, A), tensor CG2 → tensor CG3	ERR	81 / 240 / 4	81 / 1536 / 48
inner(A, A), tensor CG2 → CG2	ERR	120 / 126 / 3	120 / 360 / 9
grad(f)[0], CG4 → CG4	91 / 291 / 25	92 / 292 / 25	91 / 291 / 25
div(u), vector CG4 → DG3	109 / 390 / 25	110 / 391 / 25	109 / 390 / 25

Compile time (s, best of 3), and the largest connected contraction handed to
the exhaustive search:

interpolate on hex	#269	#280	single pass	indices #269	indices #280	indices single pass
f*f*f, CG1 → CG1	0.017	0.011	0.016	3	3	3
f*f*f*f, CG1 → CG1	0.020	0.011	0.018	3	3	3
f*f, CG4 → DG3	0.013	0.008	0.012	3	3	3
dot(u, u)*u, vector CG4 → vector CG4	ERR	0.016	0.988	—	3	7
dot(A, A), tensor CG2 → tensor CG3	ERR	0.090	1.112	—	3
inner(A, A), tensor CG2 → CG2	ERR	0.126	8.630	—	3	8
grad(f)[0], CG4 → CG4	0.061	0.066	0.063	3	3	3
div(u), vector CG4 → DG3	0.076	0.078	0.075	3	3	3

[connorjward]

In other cases, single pass regresses by a factor of 2 in flops.

I don't understand how this can happen. Isn't the search space for the single pass just bigger than the other? And hence will reach at least the same optimum just slower?

[pbrubeck]

In other cases, single pass regresses by a factor of 2 in flops.

I don't understand how this can happen. Isn't the search space for the single pass just bigger than the other? And hence will reach at least the same optimum just slower?

I don't have a good answer for this. I will get this properly fixed, but that needs its own PR.

[pbrubeck]

I think that there's much more going on behind the scenes and the transformations applied by sum-factorization do not consider the entire search space of possible transformations.

We do not only split sum-factorization in two calls, but we essentially split gem.optimise.contraction (which does Delta cancellation + sum-factorization)

[connorjward]

I think you've convinced me that there isn't an immediately better solution than this. If you create an issue and reference it in the code, with a big comment, then I am happy with this.

I just want to make sure we don't have what we usually have which is a block of magic code that makes no sense to anyone.

[pbrubeck]

I think you've convinced me that there isn't an immediately better solution than this. If you create an issue and reference it in the code, with a big comment, then I am happy with this.

What if I open the PR that actually fixes this?

[connorjward]

I think you've convinced me that there isn't an immediately better solution than this. If you create an issue and reference it in the code, with a big comment, then I am happy with this.

What if I open the PR that actually fixes this?

Even better!

I will get this properly fixed, but that needs its own PR.

I read this as meaning that you'd be unwilling to do that right now.

[connorjward]

Just the typo, happy otherwise