Factor scalar maps through basis transformations

[pbrubeck]

Stacked on #281. The complete FIAT stack is #282 -> #284 -> #281 ->
this PR, with the TSFC and Loopy integration in
firedrakeproject/firedrake#5362.

Motivation

A physically mapped finite element kernel first applies a sparse basis
transformation to a reference tabulation and then contracts the transformed
tabulation with quadrature data. The compiler should retain those as distinct
linear operations: reference tabulation -> physical basis -> quadrature
contraction.

The padded representation from #281 gives the sparse map a regular iteration
domain, but equal symbolic coefficients were still materialised as every entry
of a mutable matrix. In the 3D Guzman--Neilan matrix this appeared as 144
individual assignments to a 16-by-9 temporary before the map loop.

What changes

• COFFEE gives additive linear maps a commutative key and extracts uniform
  multiplicity into the scalar monomial. Thus a diagonal strain component
  simplifies naturally from 0.5 (Gii + Gii) to Gii, while off-diagonal
  components retain their multiplicity.

• GEM factors a reduction through a repeated indirect lookup when its cost
  model finds the rewrite profitable. 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)].
  

  This computes a dense reference residual once per table row, then gathers
  through the padded map.

• MappedTabulation interns structurally equal symbolic coefficients. Two
  read-only selector tables choose the reference column and one of the distinct
  coefficients. The generated 3D Guzman--Neilan matrix now has no mutable
  16-by-9 coefficient matrix; the distinct coefficient workspace has length
  43.

• Common factors are extracted only inside scalar basis-transformation
  coefficients. Sums carrying free indices remain with the contraction
  planner, avoiding the action-kernel regression caused by globally
  redistributing scalar factors.

• A contraction of an empty product explicitly contributes the cardinality of
  its index space. This is the identity sum_i 1 = dim(i), which arises after a
  selector makes one contraction component independent of an index.

There are no branches on element family, cell dimension, polynomial degree, or
form operator.

The quadrature weight times abs(detJ) is computed once per quadrature point.
The component weights are applied to one argument vector before the final
outer-product sum. Moving that common scalar outside the final matrix update
would replace O(ndof) scaling with O(ndof^2) scaling, so the generated form
retains the cheaper placement.

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; the zany forms are (hess u, hess v) for Argyris
(hess = sym(grad(grad(.)))), (eps u, eps v) + (div u, div v) for
Guzman--Neilan 1st kind H1 (eps = sym(grad(.))), and
(sym(u), sym(v)) + (div u, div v)*dx1 for Johnson--Mercier. 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's stack adds; the PR
side is measured against firedrake pbrubeck/zany-matvec HEAD, which adds
exactly that API.

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 -> 92	0.033 -> 0.036	0.57 -> 0.56	0.000002 -> 0.000002
CG	3	2	5,547	4	220	100	116	0.034 -> 0.038	0.62 -> 0.65	0.000012 -> 0.000012
CG	5	2	53,827	4	924	441	116	0.037 -> 0.040	0.67 -> 0.69	0.000062 -> 0.000061
CG	1	3	429 -> 387	4	28	16	138 -> 127	0.057 -> 0.057	0.71 -> 0.68	0.000013 -> 0.000013
CG	3	3	49,592	5	860	400	169	0.061 -> 0.069	0.85 -> 0.84	0.000515 -> 0.000514
CG	5	3	1,340,021	5	6,440	3,136	169	0.064 -> 0.071	1.35 -> 1.37	0.005147 -> 0.005150
RT	1	2	259 -> 232	4	18	9	107 -> 99	0.041 -> 0.042	0.60 -> 0.59	0.000003 -> 0.000003
RT	3	2	18,098 -> 15,608	5	495	225	122 -> 114	0.045 -> 0.046	0.68 -> 0.65	0.000023 -> 0.000023
RT	5	2	210,463 -> 172,341	5	2,555	1,225	122 -> 114	0.045 -> 0.049	0.77 -> 0.79	0.000149 -> 0.000145
RT	1	3	852 -> 773	6 -> 5	36 -> 32	16	158 -> 126	0.060 -> 0.064	0.75 -> 0.68	0.000015 -> 0.000015
RT	3	3	281,708 -> 226,124	6	2,736	1,296	168 -> 142	0.068 -> 0.070	0.97 -> 0.89	0.000873 -> 0.000890
RT	5	3	10,384,508 -> 7,865,477	6	29,280	14,400	168 -> 142	0.097 -> 0.097	2.72 -> 2.76	0.020563 -> 0.019859
Q	1	3	7,774	55 -> 61	690 -> 708	64	684 -> 687	0.159 -> 0.166	2.19 -> 2.25	0.000015 -> 0.000014
Q	5	3	2,679,657	31 -> 37	23,149 -> 23,191	7,776	394	0.128 -> 0.143	1.62 -> 1.65	0.005725 -> 0.005609
Q	7	3	16,034,499	31 -> 37	87,199 -> 87,253	32,768	394	0.127 -> 0.141	1.90 -> 1.87	0.044677 -> 0.044620
NCE	1	3	66,447 -> 66,480	546 -> 552	4,375 -> 4,393	27	3,235 -> 3,253	1.077 -> 1.786	13.39 -> 14.05	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 -> 1.166	8.16 -> 8.49	0.048325 -> 0.047264
NCE	7	3	115,963,956 -> 115,963,989	239 -> 245	604,385 -> 604,439	28,672	1,962 -> 1,980	0.833 -> 1.180	10.25 -> 10.61	0.452663 -> 0.443888

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.56	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.036	0.66 -> 0.66	0.000008 -> 0.000008
CG	1	3	191	1	4	4	129	0.056 -> 0.057	0.70 -> 0.72	0.000003 -> 0.000003
CG	3	3	4,994	2	40	20	145	0.055 -> 0.060	0.76 -> 0.81	0.000027 -> 0.000027
CG	5	3	47,954	2	112	56	145	0.057 -> 0.063	1.23 -> 1.23	0.000284 -> 0.000289
RT	1	2	151	1	3	3	92	0.042 -> 0.043	0.57 -> 0.60	0.000001 -> 0.000001
RT	3	2	2,413	2	30	15	111	0.042 -> 0.045	0.65 -> 0.67	0.000004 -> 0.000004
RT	5	2	12,018	2	70	35	111	0.043 -> 0.046	0.77 -> 0.78	0.000021 -> 0.000021
RT	1	3	405	1	4	4	132	0.055 -> 0.061	0.68 -> 0.71	0.000004 -> 0.000004
RT	3	3	15,642	2	72	36	145	0.061 -> 0.066	0.89 -> 0.95	0.000085 -> 0.000088
RT	5	3	172,974	2	240	120	145	0.088 -> 0.089	2.60 -> 2.64	0.001074 -> 0.001048
Q	1	3	3,500	30 -> 40	365 -> 395	27	447 -> 451	0.139 -> 0.152	1.51 -> 1.55	0.000004 -> 0.000004
Q	5	3	95,761	17 -> 23	355 -> 505	49	334 -> 342	0.117 -> 0.125	1.32 -> 1.37	0.000088 -> 0.000084
Q	7	3	254,223	17 -> 23	583 -> 829	81	334 -> 342	0.115 -> 0.128	1.35 -> 1.40	0.000184 -> 0.000170
NCE	1	3	18,135 -> 18,168	98 -> 104	1,200 -> 1,218	27	868 -> 886	0.447 -> 0.580	3.20 -> 3.43	0.000016 -> 0.000015
NCE	5	3	386,797 -> 386,830	57 -> 63	1,887 -> 1,929	245	800 -> 832	0.443 -> 0.517	3.03 -> 3.20	0.000258 -> 0.000256
NCE	7	3	972,999 -> 973,032	57 -> 63	3,631 -> 3,685	567	800 -> 832	0.430 -> 0.532	3.09 -> 3.23	0.000682 -> 0.000683

CG, RT, Q and NCE are unaffected by this PR beyond what #284 and #281 already
contribute: RT falls 10.4%/13.8%/18.1% at degree 1/3/5 in 2D and
9.3%/19.7%/24.3% in 3D; CG only moves at degree 1 (-6.3%/-9.8%); Q and NCE
flops are unchanged, since this PR's coefficient-interning mechanism only
changes the scalar coefficients inside a basis-transformation matrix, and
Q/NCE have none. NCE's tsfc (s) column is the one place this PR's stack
costs something on a case it does not touch: compile time rises 32-66%,
consistent with the new indirect-reduction-factoring cost comparison walking
NCE's much larger expression DAG without finding anything to rewrite.

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

Flops keep rising over #281 in the matrix (Argyris +10.5%, Guzman--Neilan
+9.2% in 2D) except where interning removes real arithmetic: Guzman--Neilan
in 3D falls 10.6% and Johnson--Mercier falls 0.1%/8.5% in 2D/3D. AST lines
fall further still: Argyris to 203 lines (-52.9% from main), Guzman--Neilan
in 3D to 440 (-75.6%), Johnson--Mercier in 3D to 358 (-82.0%, the largest
drop of any case). Isolated build time improves in every matrix and action
case measured (12-68%). The action keeps costing arithmetic everywhere (up to
+98.0% for Johnson--Mercier in 3D), which tracks the padding-related
arithmetic from #281 rather than anything this PR adds: this PR's
interning only removes redundant coefficient work, and an action has no
argument axis left to share coefficients across.

Bernstein control

Degree-4 Bernstein Laplacians remain structurally identical to main.

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

Validation

• focused FIAT/GEM tests: 135 passed
• focused flake8 over every touched file
• pydocstyle over the touched implementation modules
• resource-bounded, single-threaded code-generation and cold-cache Firedrake
  benchmarks

AI assistance

OpenAI Codex was used for implementation, benchmarking, and drafting this
description. The human contributor remains responsible for understanding,
validating, and maintaining the changes.