Preserve finite element linear maps during factorisation

[pbrubeck]

Goal: to inject finite element structure down to the C level.

Generate code that separates the pullbacks from the quadrature loop.

Stacked on #282. Needs firedrakeproject/firedrake#5362 to pass the argument
indices; without it everything here is inert, which the baseline below relies on.

What changes

• A pullback reaches TSFC as a sum over one argument axis: the physical
  gradient, divergence or curl of one basis function. Monomial collection
  distributed over that sum, so the geometry was pushed through the middle of
  the element tensor contraction and each argument axis mapped its own
  tabulation. collect_monomials now takes the argument axes, distributes only
  the sums spanning several of them, which separate form monomials, and keeps
  the one-axis sums atomic.

• COFFEE renames each argument axis to a canonical index. That exposes that
  test and trial apply the same map without inspecting the element family, and
  materialises the map once. This is generalised code motion: the basis
  transformation is evaluated once and both axes index its result.

• Removing ComponentTensors happens in collect_monomials, not inside the
  collector. The collector rebuilds nodes, and the maps identified here have to
  be the very nodes it then classifies.

• estimate_cost gives a caller a cost for a whole GEM DAG, so preserving a
  map is compared against expanding it rather than assumed better.
  has_linear_maps answers whether the comparison is worth making, for one
  traversal instead of a whole collection pass.

• flop_count counts a preserved map. It reaches scheduling as a
  ComponentTensor assignment rather than a loop nest, so the multiindex
  extents that no enclosing For supplies are counted here. It previously
  raised ValueError, and count_flops turned that into a silent zero.

• Literal compared and hashed without its dtype, so tables holding the same
  numbers at different dtypes were interchangeable wherever GEM memoizes on node
  identity, which the map sharing above relies on.

Benchmarks

Re-run against main for every case below. The bilinear form is
inner(u, v)*dx + inner(d(u), d(v))*dx1, with d = grad for CG and Q,
div for RT, and curl for NCE, dx1 forcing a separate quadrature degree
for the derivative term, on simplices and on an extruded quadrilateral mesh.
tsfc (s) is the TSFC compile time; build (s) is the isolated cold-cache C
build (compiler and linker only, no TSFC, no PyOP2); kernel (s) is the
isolated per-call kernel time: the compiled kernel is called directly,
bypassing PyOP2's Python wrapper and the local-to-global indirection, and
averaged over calls made in one second. array temps/entries/largest
count mutable loopy temporaries with a shape, their total entries, and the
largest single temporary. main is measured against the real, unpatched
firedrake, since it cannot import a tsfc that expects the gem API this PR
adds; the PR side is measured against the firedrake commit that adds exactly
that API (0bf4a4007).

Bilinear form

element	degree	dim	flops	array temps	entries	largest	AST lines	tsfc (s)	build (s)	kernel (s)
CG	1	2	159 -> 149	3	15	9	95 -> 96	0.033 -> 0.034	0.57 -> 0.56	0.000002 -> 0.000002
CG	3	2	5,547	4	220	100	116	0.034 -> 0.038	0.62 -> 0.66	0.000012 -> 0.000012
CG	5	2	53,827	4	924	441	116	0.037 -> 0.038	0.67 -> 0.68	0.000062 -> 0.000062
CG	1	3	429 -> 387	4	28	16	138 -> 135	0.057 -> 0.058	0.71 -> 0.68	0.000013 -> 0.000013
CG	3	3	49,592	5	860	400	169	0.061 -> 0.067	0.85 -> 0.83	0.000515 -> 0.000513
CG	5	3	1,340,021	5	6,440	3,136	169	0.064 -> 0.074	1.35 -> 1.33	0.005147 -> 0.005143
RT	1	2	259 -> 232	4	18	9	107 -> 103	0.041 -> 0.044	0.60 -> 0.58	0.000003 -> 0.000003
RT	3	2	18,098 -> 15,608	5	495	225	122 -> 118	0.045 -> 0.045	0.68 -> 0.65	0.000023 -> 0.000023
RT	5	2	210,463 -> 172,341	5	2,555	1,225	122 -> 118	0.045 -> 0.046	0.77 -> 0.76	0.000149 -> 0.000147
RT	1	3	852 -> 773	6 -> 5	36 -> 32	16	158 -> 134	0.060 -> 0.061	0.75 -> 0.68	0.000015 -> 0.000016
RT	3	3	281,708 -> 226,124	6	2,736	1,296	168 -> 150	0.068 -> 0.068	0.97 -> 0.92	0.000873 -> 0.000906
RT	5	3	10,384,508 -> 7,865,477	6	29,280	14,400	168 -> 150	0.097 -> 0.096	2.72 -> 2.65	0.020563 -> 0.020164
Q	1	3	7,774	55 -> 61	690 -> 708	64	684 -> 687	0.159 -> 0.165	2.19 -> 2.23	0.000015 -> 0.000014
Q	5	3	2,679,657	31 -> 37	23,149 -> 23,191	7,776	394	0.128 -> 0.133	1.62 -> 1.61	0.005725 -> 0.005670
Q	7	3	16,034,499	31 -> 37	87,199 -> 87,253	32,768	394	0.127 -> 0.131	1.90 -> 1.84	0.044677 -> 0.044074
NCE	1	3	66,447 -> 66,480	546 -> 552	4,375 -> 4,393	27	3,235 -> 3,253	1.077 -> 1.128	13.39 -> 13.15	0.000085 -> 0.000086
NCE	5	3	18,459,460 -> 18,459,493	239 -> 245	146,181 -> 146,223	6,480	1,962 -> 1,980	0.885 -> 0.892	8.16 -> 8.03	0.048325 -> 0.047772
NCE	7	3	115,963,956 -> 115,963,989	239 -> 245	604,385 -> 604,439	28,672	1,962 -> 1,980	0.833 -> 0.891	10.25 -> 10.23	0.452663 -> 0.444416

Matrix-free action

element	degree	dim	flops	array temps	entries	largest	AST lines	tsfc (s)	build (s)	kernel (s)
CG	1	2	84	1	3	3	88	0.033 -> 0.035	0.55 -> 0.57	0.000001 -> 0.000001
CG	3	2	1,107	2	20	10	107	0.035 -> 0.036	0.61 -> 0.61	0.000002 -> 0.000002
CG	5	2	5,133	2	42	21	107	0.036 -> 0.038	0.66 -> 0.67	0.000008 -> 0.000008
CG	1	3	191	1	4	4	129	0.056 -> 0.055	0.70 -> 0.69	0.000003 -> 0.000003
CG	3	3	4,994	2	40	20	145	0.055 -> 0.059	0.76 -> 0.78	0.000027 -> 0.000027
CG	5	3	47,954	2	112	56	145	0.057 -> 0.062	1.23 -> 1.23	0.000284 -> 0.000280
RT	1	2	151	1	3	3	92	0.042 -> 0.045	0.57 -> 0.59	0.000001 -> 0.000001
RT	3	2	2,413	2	30	15	111	0.042 -> 0.044	0.65 -> 0.64	0.000004 -> 0.000004
RT	5	2	12,018	2	70	35	111	0.043 -> 0.044	0.77 -> 0.75	0.000021 -> 0.000021
RT	1	3	405	1	4	4	132	0.055 -> 0.060	0.68 -> 0.67	0.000004 -> 0.000004
RT	3	3	15,642	2	72	36	145	0.061 -> 0.066	0.89 -> 0.92	0.000085 -> 0.000085
RT	5	3	172,974	2	240	120	145	0.088 -> 0.090	2.60 -> 2.60	0.001074 -> 0.001067
Q	1	3	3,500	30 -> 40	365 -> 395	27	447 -> 451	0.139 -> 0.142	1.51 -> 1.55	0.000004 -> 0.000004
Q	5	3	95,761	17 -> 23	355 -> 505	49	334 -> 342	0.117 -> 0.120	1.32 -> 1.30	0.000088 -> 0.000085
Q	7	3	254,223	17 -> 23	583 -> 829	81	334 -> 342	0.115 -> 0.119	1.35 -> 1.35	0.000184 -> 0.000169
NCE	1	3	18,135 -> 18,168	98 -> 104	1,200 -> 1,218	27	868 -> 886	0.447 -> 0.455	3.20 -> 3.18	0.000016 -> 0.000015
NCE	5	3	386,797 -> 386,830	57 -> 63	1,887 -> 1,929	245	800 -> 832	0.443 -> 0.445	3.03 -> 3.08	0.000258 -> 0.000256
NCE	7	3	972,999 -> 973,032	57 -> 63	3,631 -> 3,685	567	800 -> 832	0.430 -> 0.447	3.09 -> 3.08	0.000682 -> 0.000682

Raviart--Thomas is again where the pullback is worth preserving, and now
correctly measured against div, not curl: arithmetic falls 10.4%, 13.8%
and 18.1% at degree 1, 3 and 5 in 2D, and 9.3%, 19.7% and 24.3% in 3D. CG only
moves at degree 1, where preserving the map is cheaper than expanding it
(-6.3% in 2D, -9.8% in 3D); at degree 3 and 5 the cost comparison below picks
expansion instead, so flops are unchanged.

The tensor-product rows move only by what #282 already changed. Q is
flop-for-flop identical between main and this PR; NCE carries a small,
degree-independent +33 flops that is unrelated to this PR's mechanism (which
needs an argument axis NCE's tensor-product tabulation does not expose) --
it is an artifact of comparing against two different firedrake commits, one
of which main's FIAT cannot avoid (see above).

Zany elements, bilinear form

element	dim	flops	array temps	entries	largest	AST lines	tsfc (s)	build (s)	kernel (s)
Argyris	2	38,663	6	126	21	431	0.206 -> 0.227	2.63 -> 2.66	0.000041 -> 0.000041
Guzman--Neilan	2	10,041	8	72	9	372	0.157 -> 0.182	1.52 -> 1.51	0.000011 -> 0.000011
Guzman--Neilan	3	384,767 -> 371,974	18 -> 28	288 -> 448	16	1,803 -> 1,258	0.810 -> 0.841	6.69 -> 4.95	0.001689 -> 0.001634
Johnson--Mercier	2	21,660	14	630	225	454	0.233 -> 0.269	1.80 -> 1.81	0.000023 -> 0.000024
Johnson--Mercier	3	533,683	26	4,536	1,764	1,993	1.000 -> 1.269	9.83 -> 9.98	0.001538 -> 0.001540

Zany elements, matrix-free action

element	dim	flops	array temps	entries	largest	AST lines	tsfc (s)	build (s)	kernel (s)
Argyris	2	6,632	3	63	21	427	0.233 -> 0.249	2.61 -> 2.64	0.000006 -> 0.000006
Guzman--Neilan	2	3,549	4	36	9	343	0.159 -> 0.182	1.42 -> 1.41	0.000004 -> 0.000005
Guzman--Neilan	3	106,236	9	144	16	1,559	0.734 -> 0.864	5.52 -> 5.56	0.000976 -> 0.000971
Johnson--Mercier	2	3,753	9	135	15	444	0.267 -> 0.289	1.75 -> 1.72	0.000005 -> 0.000005
Johnson--Mercier	3	35,992	17	714	42	1,911	1.241 -> 1.409	10.02 -> 10.27	0.000205 -> 0.000201

Argyris, Guzman--Neilan in 2D, and Johnson--Mercier are unchanged by this PR
alone: none of their sums span exactly one argument axis yet, so
has_linear_maps finds nothing to preserve. Guzman--Neilan in 3D is the one
zany case this PR alone moves: flops fall 3.3%, AST lines fall 30.2% (1,803
-> 1,258), and the isolated build time falls 26.1%. Johnson--Mercier and
Argyris need #281's padded transformation before there is a map worth this
PR's cost comparison preserving; see #281.

The cost comparison is doing real work

estimate_cost is not a formality: it accepts preservation where it pays and
rejects it where it does not. Costing both plans for the same form gives

case	expanded	preserved	chosen
CG degree 1, 2D	152	142	preserved
CG degree 1, 3D	416	374	preserved
CG degree 3, 2D	8,125	8,305	expanded
RT degree 3, 3D	374,589	255,933	preserved
RT degree 5, 3D	12,965,694	8,711,463	preserved

CG degree 3 is why the earlier len(free_indices) > 1 heuristic was not enough:
preserving its map costs more arithmetic than distributing it, and only a cost
model sees that. On the RT rows the estimate matches the kernel's own
flop_count exactly.

Tests

test_sum_factorise.py builds a Laplacian element tensor from a mapped gradient
table, where test and trial apply the same map over their own argument index,
and covers that preserving the map leaves strictly fewer monomials than
distributing it; that COFFEE materialises one tensor for both axes; that the
preserved and expanded forms evaluate to the same matrix; that
has_linear_maps needs argument axes to find anything; and that estimate_cost
counts the contraction.

Validation

• pytest test/ --ignore=test/FIAT/regression (2378 passed, 26 skipped, 31 xfailed, 1 xpassed)
• tests/tsfc and the three zany regression suites in
  Preserve finite element pullbacks when collecting monomials firedrake#5362 (437 passed)

AI assistance

Claude Code was used for implementation, benchmarking, and drafting this
section. The human contributor remains responsible for understanding,
validating, and maintaining the changes.