Sum factorisation on simplices

[pbrubeck]

Depends on firedrakeproject/fiat#262.

Summary

• gem.JaggedIndex represents simplex lattice indices with bounds depending on parent indices. TSFC lowers these to parameterized Loopy domains.
• gem.FlattenedTensor exposes a jagged lattice through a flat DoF index while retaining its tensor-product structure.
• gem.sparse_matrix represents sparse basis recombination.
• DuffyElement.duffy_evaluation() constructs sum-factorized Legendre and IntegratedLegendre tabulations; get_sparse_coeffs() provides DG ordering or CG recombination data.
• TSFC uses Duffy tabulations for coefficient and argument translation and selects collapsed quadrature for Duffy elements.

gem.optimise now:

• rewrites flat contractions as jagged lattice contractions with unflatten() and unflatten_returns();
• exposes both argument lattices before expanding separable derivative sums in bilinear forms;
• refactorizes eligible sparse-transform expressions as monomial sums and applies the same staged COFFEE sum factorization used by TSFC;
• retains the established local-distribution path for dense or residual argument transforms;
• cancels sparse row/column deltas during contraction; and
• lowers remaining FlattenedTensors to gather tables with replace_flattened().

For a three-dimensional collapsed basis, the reference derivative has 3 + 2 + 1 one-dimensional derivative paths. A bilinear Laplacian combines two such arguments and three physical gradient components. The previous unflatten_returns() expanded the first argument before exposing the second, cloning the second jagged lattice once per derivative term. The new ordering exposes both lattices jointly, then groups equal monomials and factorizes one contraction index at a time.

The final output loops remain genuinely jagged. In particular, the innermost tetrahedral domain has a bound of the form 0 <= k <= p - i - j; it is not replaced by a rectangular loop with masking.

Bilinear Laplacian code-generation benchmarks

These are tetrahedral Bernstein kernels generated in spectral mode. Structural counters and FLOPs are exact for the generated Loopy domains; times are representative local instrumented compiler runs.

Degree / scheme	Compile	unflatten_returns	FLOPs	Instructions	Temporaries	C for loops
3 collapsed, before	1.181 s	0.388 s	3,715,953	1,103	737	1,071
3 collapsed, optimized	0.254 s	0.044 s	439,113	220	189	92
3 Laplacian action reference	0.137 s	0.012 s	10,305	110	96	30
3 mass bilinear reference	0.076 s	0.006 s	43,768	23	26	39
10 canonical	0.489 s	0.001 s	495,076,099	52	54	4
10 collapsed, optimized	0.399 s	0.046 s	433,160,451	220	189	92

Degree 10 is the tested crossover where the sum-factorized bilinear Laplacian beats canonical quadrature in both generated FLOPs and compiler time. A regression test also checks the compact instruction/temporary/loop counts and requires a two-parent parameterized domain, which preserves the tetrahedral JaggedIndex loop.

One-element execution benchmarks

FLOPs use the exact generated jagged loop domains. Times are median assembly times after warm-up.

Case	Degree	Canonical FLOPs	Collapsed FLOPs	FLOP reduction	Canonical time	Collapsed time	Speedup
CG Laplacian action (2D)	16	317197	215681	32.0%	369.0 us	230.0 us	1.60x
Bernstein Laplacian action (2D)	16	317197	144331	54.5%	486.2 us	118.5 us	4.10x
CG Laplacian action (3D)	15	33162804	28707004	13.4%	77.65 ms	4.955 ms	15.67x
Bernstein Laplacian action (3D)	15	33162804	4057153	87.8%	31.58 ms	418.7 us	75.43x
DG mass matrix (2D)	12	2814534	247178	91.2%	501.1 us	292.9 us	1.71x
DG mass matrix (3D)	7	14807576	302041	98.0%	9.885 ms	0.440 ms	22.47x

AI assistance: OpenAI Codex was used for code analysis, implementation support, testing, and benchmarking.

[pbrubeck]

TSFC codegen diff should be as tight as possible. Add the new code in a separate finat submodule.

[rckirby]

This just checks that the code gives a correct answer. Do we have a way of checking whether the algorithm has the right complexity?

[pbrubeck]

This just checks that the code gives a correct answer. Do we have a way of checking whether the algorithm has the right complexity?

This PR adds both correctness and complexity tests. Complexity in flops is not enough, I also had to enforce tests on temporaries

[rckirby]

This PR adds both correctness and complexity tests. Complexity in flops is not enough, I also had to enforce tests on temporaries

I see those tests, was thinking about FIAT. We should also test out Bernstein in the one-element benchmarks -- it doesn't have the indirection internally that modified C^0 expansions have but can be directly (after Duffy) sum-factored.