Factor scalar maps through basis transformations

[pbrubeck]

TLDR

A padded basis transformation stores the same coefficient many times, and writes
each copy into the kernel as its own instruction. This PR stores each distinct
coefficient once and looks the rest up through a small table.

It also factors a reduction through a repeated indirect lookup, so the
kernel sums over the reference table once and then gathers.

Johnson--Mercier on tetrahedra drops from 1,993 lines of C to 358, and its cold
C build time falls 68%. Flops fall on the elements this PR touches.

Base: #281. Needs firedrakeproject/firedrake#5362.

What this does

Intern equal coefficients. MappedTabulation collects the distinct symbolic
coefficients of the transformation. Two read-only tables then pick the reference
column and the coefficient for each padded entry. The 3D Guzman--Neilan matrix
no longer holds a mutable 16-by-9 coefficient matrix, and its distinct
coefficient workspace has length 43.

Give an additive map an order-free key. COFFEE keys an additive linear map
on its set of summands, not on the order they were built in. It then moves a
uniform multiplicity out into the scalar part of the monomial. A diagonal strain
component simplifies from 0.5 (Gii + Gii) to Gii. An off-diagonal component
keeps its multiplicity.

Factor a reduction through an indirect lookup. For an indirect column c(i),
linearity gives

sum_q B[c(i), q] f[q] = (sum_q B[r, q] f[q])[c(i)]

The kernel computes a dense reference residual once per table row, then gathers
it through the padded map. A cost model decides when this is worth doing.

Factor common scalars, but only inside coefficients. Sums that carry free
indices stay with the contraction planner.

Count an empty contraction. A contraction of an empty product contributes the
size of its index space. This is the identity sum_i 1 = dim(i). It comes up
when a selector makes one part of a contraction independent of an index.

Nothing here branches on element family, cell dimension, degree or form operator.

Benchmarks

Measured on this branch against main. The forms and columns are the same as
in #281.

Zany elements, bilinear form

element	dim	flops	array temps	entries	largest	AST lines	tsfc (s)	build (s)	kernel (s)
Argyris	2	38,663 -> 42,738	6 -> 8	126 -> 197	21 -> 50	431 -> 203	0.206 -> 0.279	2.63 -> 2.08	0.000041 -> 0.000048
Guzman--Neilan	2	10,041 -> 10,963	8 -> 11	72 -> 105	9 -> 15	372 -> 162	0.157 -> 0.143	1.52 -> 0.94	0.000011 -> 0.000014
Guzman--Neilan	3	384,767 -> 344,078	18 -> 20	288 -> 347	16 -> 43	1,803 -> 440	0.810 -> 0.696	6.69 -> 2.58	0.001689 -> 0.001458
Johnson--Mercier	2	21,660 -> 21,645	14 -> 16	630 -> 653	225	454 -> 182	0.233 -> 0.269	1.80 -> 1.03	0.000023 -> 0.000026
Johnson--Mercier	3	533,683 -> 488,173	26 -> 30	4,536 -> 4,676	1,764	1,993 -> 358	1.000 -> 1.538	9.83 -> 3.16	0.001538 -> 0.001638

Zany elements, matrix-free action

element	dim	flops	array temps	entries	largest	AST lines	tsfc (s)	build (s)	kernel (s)
Argyris	2	6,632 -> 10,095	3 -> 7	63 -> 176	21 -> 50	427 -> 290	0.233 -> 0.255	2.61 -> 2.30	0.000006 -> 0.000014
Guzman--Neilan	2	3,549 -> 4,431	4 -> 5	36 -> 111	9 -> 36	343 -> 273	0.159 -> 0.178	1.42 -> 1.11	0.000004 -> 0.000008
Guzman--Neilan	3	106,236 -> 147,180	9 -> 7	144 -> 547	16 -> 144	1,559 -> 853	0.734 -> 0.834	5.52 -> 3.43	0.000976 -> 0.001006
Johnson--Mercier	2	3,753 -> 5,481	9 -> 7	135 -> 98	15	444 -> 250	0.267 -> 0.291	1.75 -> 1.18	0.000005 -> 0.000007
Johnson--Mercier	3	35,992 -> 71,260	17 -> 10	714 -> 392	42	1,911 -> 670	1.241 -> 1.051	10.02 -> 3.22	0.000205 -> 0.000386

Interning removes real arithmetic where the map has repeated coefficients:
Guzman--Neilan in 3D falls 10.6%, Johnson--Mercier falls 0.1% in 2D and 8.5% in
3D. AST lines fall further than in #281: Argyris to 203 lines, 52.9% below
main; Guzman--Neilan in 3D to 440, down 75.6%; Johnson--Mercier in 3D to 358,
down 82.0%. The isolated build time improves in every matrix and action case,
by 12-68%.

The action still costs more arithmetic than main, up to +98.0% for
Johnson--Mercier in 3D. An action has no second argument axis to share
coefficients across, so interning has nothing to remove there.

Elements this PR does not touch

element	degree	dim	flops	array temps	entries	AST lines	tsfc (s)
CG	3	3	49,592	5	860	169	0.061 -> 0.069
RT	5	3	10,384,508 -> 7,865,477	6	29,280	168 -> 142	0.097 -> 0.097
Q	7	3	16,034,499	31 -> 37	87,199 -> 87,253	394	0.127 -> 0.141
NCE	7	3	115,963,956 -> 115,963,989	239 -> 245	604,385 -> 604,439	1,962 -> 1,980	0.833 -> 1.180

Q and NCE flops do not change. Interning only touches the scalar coefficients
inside a basis-transformation matrix, and Q and NCE have none.

NCE compile time rises 32-66%. This is the one measured cost on a case the PR
does not otherwise touch. The new cost comparison walks NCE's much larger
expression DAG and finds nothing to rewrite.

Bernstein control

Degree-4 Bernstein Laplacians are structurally identical to main.

cell	flops main	flops PR	scalar temps	array temps	stored values	AST lines
triangle	11,965	11,965	22	5	424	82
tetrahedron	181,454	181,454	47	7	2,590	135

Validation

• focused FIAT and GEM tests: 135 passed
• make lint
• resource-bounded, single-threaded code generation and cold-cache Firedrake
  benchmarks

AI assistance

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