Padded basis transformation for better codegen

[pbrubeck]

TLDR

A physically mapped element applies a sparse matrix to its reference tabulation.
This PR makes that matrix a dense padded table with a regular loop, instead
of one instruction per matrix row.

The generated C shrinks a lot. Johnson--Mercier on tetrahedra drops from 1,993
lines to 424, and its cold C build time falls 72%. Arithmetic rises, because
padding adds work.

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

What this does

Every row of the transformation is padded to the same number of entries. A loop
over the basis index then has a plain rectangular domain, which Loopy can build.
Two small read-only tables say which reference column each padded entry selects.

Without padding, each row of the sparse matrix product became its own
instruction and its own temporary. Johnson--Mercier on tetrahedra declared 655
arrays of length nqp, one for each basis function and tabulation. A loop over
the basis index needs 24 arrays of shape (42, nqp).

The padded map reaches monomial collection as a sum over the nonzeros of one
basis row. #284 keeps that sum whole and shares it between the argument axes.
This PR is what makes such a sum exist.

Tidy-up. The delta rewrites are named for what each one does.
pull_back_indirect_delta finds and applies the pull-back;
cancel_nested_deltas walks a whole DAG. The guard on that walk is now
memoised, so it is linear in the size of the DAG.

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. The zany forms are (hess u, hess v) for Argyris,
(eps u, eps v) + (div u, div v) for Guzman--Neilan, and
(sym(u), sym(v)) + (div u, div v)*dx1 for Johnson--Mercier.

tsfc (s) is the TSFC compile time. build (s) is the cold-cache C build
alone. kernel (s) is the compiled kernel called directly, averaged over one
second of calls.

Zany elements, bilinear form

element	dim	flops	array temps	entries	largest	AST lines	tsfc (s)	build (s)	kernel (s)
Argyris	2	38,663 -> 54,308	6 -> 9	126 -> 218	21 -> 50	431 -> 223	0.206 -> 0.205	2.63 -> 2.01	0.000041 -> 0.000051
Guzman--Neilan	2	10,041 -> 12,727	8 -> 13	72 -> 123	9 -> 15	372 -> 196	0.157 -> 0.118	1.52 -> 0.91	0.000011 -> 0.000016
Guzman--Neilan	3	384,767 -> 417,910	18 -> 23	288 -> 395	16 -> 43	1,803 -> 502	0.810 -> 0.500	6.69 -> 2.35	0.001689 -> 0.001668
Johnson--Mercier	2	21,660 -> 26,217	14 -> 17	630 -> 668	225	454 -> 206	0.233 -> 0.220	1.80 -> 0.99	0.000023 -> 0.000028
Johnson--Mercier	3	533,683 -> 663,533	26 -> 33	4,536 -> 4,802	1,764	1,993 -> 424	1.000 -> 1.015	9.83 -> 2.71	0.001538 -> 0.002018

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,097	3 -> 7	63 -> 176	21 -> 50	427 -> 290	0.233 -> 0.171	2.61 -> 2.15	0.000006 -> 0.000014
Guzman--Neilan	2	3,549 -> 6,159	4 -> 5	36 -> 159	9 -> 36	343 -> 281	0.159 -> 0.130	1.42 -> 1.11	0.000004 -> 0.000011
Guzman--Neilan	3	106,236 -> 242,220	9 -> 7	144 -> 907	16 -> 144	1,559 -> 857	0.734 -> 0.453	5.52 -> 2.99	0.000976 -> 0.001672
Johnson--Mercier	2	3,753 -> 5,481	9 -> 7	135 -> 98	15	444 -> 250	0.267 -> 0.185	1.75 -> 1.11	0.000005 -> 0.000008
Johnson--Mercier	3	35,992 -> 71,260	17 -> 10	714 -> 392	42	1,911 -> 670	1.241 -> 0.546	10.02 -> 2.61	0.000205 -> 0.000384

The gain is in code generation, and the cost is arithmetic.

• AST lines fall 48-79% for the matrix and 18-65% for the action.
• The cold C build time falls 24-72% for the matrix and 18-74% for the action.
• Flops rise on every zany case: Argyris +40.5%, Guzman--Neilan +26.8% in 2D and
  +8.6% in 3D, Johnson--Mercier +21.0% in 2D and +24.3% in 3D. The action rises
  further, up to +128% for Guzman--Neilan in 3D.
• The biggest temporary grows where a padded (ndof, nqp) table replaces the
  per-row vectors, up to 2.7x for Guzman--Neilan in 3D. Johnson--Mercier already
  had a temporary of that shape, so it does not grow.

Elements this PR does not touch

CG, RT, Q and NCE match #284's table exactly, because the padded transformation
applies only to physically mapped elements.

element	degree	dim	flops	AST lines	tsfc (s)	build (s)
CG	3	3	49,592	169	0.061 -> 0.068	0.85 -> 0.83
RT	3	3	281,708 -> 226,124	168 -> 150	0.068 -> 0.068	0.97 -> 0.91
RT	5	3	10,384,508 -> 7,865,477	168 -> 150	0.097 -> 0.097	2.72 -> 2.75
Q	7	3	16,034,499	394	0.127 -> 0.133	1.90 -> 1.79
NCE	7	3	115,963,956 -> 115,963,989	1,962 -> 1,980	0.833 -> 0.877	10.25 -> 10.05

Validation

• FIAT suite, without test/FIAT/regression: 2379 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.

---

Update 2026-09-08: re-measured, and a correction

The tables above predate ufl#506, which changed the quadrature degree estimate
on tensor-product and zany cells, so their main column no longer matches
today's main. They also count read-only tables as temporaries. Superseding
them:

The gain is in code generation, and the cost is arithmetic that #286 gets
back.

• AST lines fall 48-79% for the matrix and 18-65% for the action.
• The cold C build time falls 24-72% for the matrix and 18-74% for the action.
• Flops rise on every zany case here, because a padded row contracts against
  zeros. The action rises furthest. Factor scalar maps through basis transformations #286 removes most of that: it reads the
  reference tabulation once instead of once per padded column, which takes the
  Johnson--Mercier 3D action to 32,806 flops against 44,248 on main. Judge the
  arithmetic at the top of the stack, not here.
• Working temporaries fall; read-only tables grow. The padded (ndof, nqp)
  tables are read-only arrays with initialisers, so they are shared across cells
  and are not working memory. Counting them together with the working
  temporaries reverses the conclusion: on the Johnson--Mercier 3D action the
  working temporaries fall from 798 elements to 224, while the tables grow from
  6,908 to 7,160.

End-to-end numbers for the whole stack are on #286 and
firedrakeproject/firedrake#5438.