Preserve finite element linear maps during factorisation

[pbrubeck]

TLDR

A pullback reaches TSFC as a sum over one argument axis. That sum is the
finite element map. This PR keeps such a sum whole instead of multiplying it
out, and shares one map between the test and trial axes.

The compiler costs both plans and keeps the cheaper one. Raviart--Thomas gets up
to 24% fewer flops in 3D. Nothing else regresses.

Base: #282. Needs firedrakeproject/firedrake#5362 to pass the argument indices.

What this does

Keep a one-axis sum whole. Monomial collection used to distribute over every
sum. That pushed the geometry into the middle of the element tensor contraction,
and each argument axis then mapped its own tabulation. collect_monomials now
takes the argument axes. It distributes only the sums that span several of them,
because those separate the terms of the form. It leaves a one-axis sum whole.

Share one map between argument axes. COFFEE renames each argument axis to a
canonical index. Test and trial then look identical if they apply the same map,
whatever their element family. The compiler evaluates that map once, and both
axes read its result.

Choose by cost, not by rule. estimate_cost scores a whole GEM DAG, so the
compiler compares keeping a map against expanding it. has_linear_maps answers
in one walk whether the comparison is worth making at all.

Count a kept map. A kept map reaches scheduling as a ComponentTensor
assignment, not as a loop nest. flop_count now counts the index extents that
no enclosing loop supplies. It used to raise ValueError, which count_flops
turned into a silent zero.

Tell literal tables apart by dtype. Literal compared and hashed without its
dtype. Two tables holding the same numbers at different dtypes were therefore
interchangeable. Map sharing relies on that comparison being right.

Tidy-up. The per-node cost constants now live in one module, gem.cost.
gem.flop_count and estimate_cost read the same table. They weigh it
differently, and that is the only difference between them.

The cost comparison earns its place

Both plans, costed on the same form:

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

CG degree 3 is the case a fixed rule cannot get right. Keeping its map costs
more arithmetic than expanding it. Only a cost model sees that. On the RT rows
the estimate matches the kernel's own flop_count exactly.

Benchmarks

Measured on this branch against main. The 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 forces a separate quadrature degree for the
derivative term.

tsfc (s) is the TSFC compile time. build (s) is the cold-cache C build
alone: compiler and linker, no TSFC and no PyOP2. kernel (s) is the compiled
kernel called directly, without the PyOP2 wrapper, averaged over one second of
calls. array temps, entries and largest count mutable Loopy temporaries
that have a shape, their total entries, and the biggest single one.

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 where keeping the map pays. 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 moves
only at degree 1, where keeping the map is cheaper than expanding it: -6.3% in
2D and -9.8% in 3D. At degree 3 and 5 the cost comparison picks expansion, so
the flops do not change.

Q is identical to main, flop for flop. NCE carries 33 more flops at every
degree. That number does not grow with the degree. NCE's tensor-product
tabulation exposes no single argument axis, so this mechanism does not act on
it.

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

Guzman--Neilan in 3D is the one zany case this branch moves. Flops fall 3.3%,
AST lines fall 30.2%, and the isolated build time falls 26.1%. In the other zany
cases no sum spans exactly one argument axis, so has_linear_maps finds nothing
to keep.

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. It
covers that keeping the map leaves strictly fewer monomials than distributing
it, that COFFEE materialises one tensor for both axes, that the two forms
evaluate to the same matrix, that has_linear_maps needs argument axes to find
anything, and that estimate_cost counts the contraction.

Validation

• FIAT suite, without test/FIAT/regression: 2378 passed, 26 skipped, 31 xfailed, 1 xpassed
• make lint
• tests/tsfc and the three zany regression suites in Cost and fuse preserved pullbacks, and cancel indirect Deltas 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.