Padded basis transformation for better codegen

[pbrubeck]

This PR vectorizes the sparse basis transformation.

Stacked on #284, which keeps the transformation whole through monomial
collection. Needs firedrakeproject/firedrake#5362.

The original sparse approach generated individual instructions for each row of the sparse-dense matrix product, which were poorly handled in loopy. Each row became its own temporary: Johnson--Mercier on tetrahedra declared 655 arrays of length num_quadrature_points, one per basis function and tabulation, where a loop over the basis index needs 24 arrays of shape (42, nqp).

In order to generate good loopy kernels, rows are padded so that they have same number of entries. A ragged index expresses this exactly, but loopy forbids one: it builds every iteration domain as an isl set whose bounds must be quasi-affine in the enclosing indices, and a row length read from a table is not.

Padding alone is not enough. collect_monomials distributes sums to reach a sum-of-products normal form, and the padded transformation reaches it as an IndexSum over the nonzeros of one basis row. Distributing over that sum expands the map into the Cartesian product of its entries before any loop is placed, and Johnson--Mercier on tetrahedra collects 3240 monomials for what is one linear operand. #284 keeps such a sum atomic and shares it between the argument axes, and costs that against expanding it; this PR is what makes the sum worth preserving.

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 gem API this PR's stack adds; the PR side
is measured against the firedrake commit that adds exactly the gem API this
PR's stack needs.

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.035	0.57 -> 0.58	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.039	0.67 -> 0.70	0.000062 -> 0.000061
CG	1	3	429 -> 387	4	28	16	138 -> 135	0.057 -> 0.056	0.71 -> 0.71	0.000013 -> 0.000013
CG	3	3	49,592	5	860	400	169	0.061 -> 0.068	0.85 -> 0.83	0.000515 -> 0.000514
CG	5	3	1,340,021	5	6,440	3,136	169	0.064 -> 0.071	1.35 -> 1.38	0.005147 -> 0.005188
RT	1	2	259 -> 232	4	18	9	107 -> 103	0.041 -> 0.042	0.60 -> 0.61	0.000003 -> 0.000003
RT	3	2	18,098 -> 15,608	5	495	225	122 -> 118	0.045 -> 0.044	0.68 -> 0.66	0.000023 -> 0.000023
RT	5	2	210,463 -> 172,341	5	2,555	1,225	122 -> 118	0.045 -> 0.047	0.77 -> 0.80	0.000149 -> 0.000148
RT	1	3	852 -> 773	6 -> 5	36 -> 32	16	158 -> 134	0.060 -> 0.062	0.75 -> 0.69	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.91	0.000873 -> 0.000906
RT	5	3	10,384,508 -> 7,865,477	6	29,280	14,400	168 -> 150	0.097 -> 0.097	2.72 -> 2.75	0.020563 -> 0.020631
Q	1	3	7,774	55 -> 61	690 -> 708	64	684 -> 687	0.159 -> 0.161	2.19 -> 2.30	0.000015 -> 0.000014
Q	5	3	2,679,657	31 -> 37	23,149 -> 23,191	7,776	394	0.128 -> 0.135	1.62 -> 1.68	0.005725 -> 0.005661
Q	7	3	16,034,499	31 -> 37	87,199 -> 87,253	32,768	394	0.127 -> 0.133	1.90 -> 1.79	0.044677 -> 0.044878
NCE	1	3	66,447 -> 66,480	546 -> 552	4,375 -> 4,393	27	3,235 -> 3,253	1.077 -> 1.130	13.39 -> 13.09	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.908	8.16 -> 7.98	0.048325 -> 0.047416
NCE	7	3	115,963,956 -> 115,963,989	239 -> 245	604,385 -> 604,439	28,672	1,962 -> 1,980	0.833 -> 0.877	10.25 -> 10.05	0.452663 -> 0.440325

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.037	0.55 -> 0.59	0.000001 -> 0.000001
CG	3	2	1,107	2	20	10	107	0.035 -> 0.036	0.61 -> 0.63	0.000002 -> 0.000002
CG	5	2	5,133	2	42	21	107	0.036 -> 0.037	0.66 -> 0.69	0.000008 -> 0.000008
CG	1	3	191	1	4	4	129	0.056 -> 0.057	0.70 -> 0.69	0.000003 -> 0.000003
CG	3	3	4,994	2	40	20	145	0.055 -> 0.058	0.76 -> 0.79	0.000027 -> 0.000027
CG	5	3	47,954	2	112	56	145	0.057 -> 0.062	1.23 -> 1.29	0.000284 -> 0.000290
RT	1	2	151	1	3	3	92	0.042 -> 0.043	0.57 -> 0.62	0.000001 -> 0.000001
RT	3	2	2,413	2	30	15	111	0.042 -> 0.045	0.65 -> 0.66	0.000004 -> 0.000004
RT	5	2	12,018	2	70	35	111	0.043 -> 0.046	0.77 -> 0.79	0.000021 -> 0.000021
RT	1	3	405	1	4	4	132	0.055 -> 0.059	0.68 -> 0.70	0.000004 -> 0.000004
RT	3	3	15,642	2	72	36	145	0.061 -> 0.064	0.89 -> 0.94	0.000085 -> 0.000085
RT	5	3	172,974	2	240	120	145	0.088 -> 0.091	2.60 -> 2.73	0.001074 -> 0.001092
Q	1	3	3,500	30 -> 40	365 -> 395	27	447 -> 451	0.139 -> 0.142	1.51 -> 1.58	0.000004 -> 0.000004
Q	5	3	95,761	17 -> 23	355 -> 505	49	334 -> 342	0.117 -> 0.121	1.32 -> 1.39	0.000088 -> 0.000083
Q	7	3	254,223	17 -> 23	583 -> 829	81	334 -> 342	0.115 -> 0.119	1.35 -> 1.32	0.000184 -> 0.000169
NCE	1	3	18,135 -> 18,168	98 -> 104	1,200 -> 1,218	27	868 -> 886	0.447 -> 0.462	3.20 -> 3.26	0.000016 -> 0.000015
NCE	5	3	386,797 -> 386,830	57 -> 63	1,887 -> 1,929	245	800 -> 832	0.443 -> 0.461	3.03 -> 3.02	0.000258 -> 0.000257
NCE	7	3	972,999 -> 973,032	57 -> 63	3,631 -> 3,685	567	800 -> 832	0.430 -> 0.449	3.09 -> 3.08	0.000682 -> 0.000689

CG, RT, Q and NCE are unaffected by this PR beyond what #284 already
contributes, and match #284's own table exactly: RT falls 10.4%/13.8%/18.1%
at degree 1/3/5 in 2D and 9.3%/19.7%/24.3% in 3D (now correctly measured
against div, not curl); CG only moves at degree 1 (-6.3% in 2D, -9.8% in
3D); Q and NCE are unchanged, since this PR's padded transformation only
applies to physically-mapped (zany) elements.

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. Arithmetic rises on every zany case (Argyris
+40.5%; Guzman--Neilan +26.8%/+8.6% in 2D/3D; Johnson--Mercier
+21.0%/+24.3%), and every action rises too (up to +128% for Guzman--Neilan in
3D), while AST lines fall 48-79% for the matrix and 18-65% for the action.
The largest temporary grows wherever the padded (ndof, nqp) table replaces
the per-row vectors main declares -- up to 2.7x for Guzman--Neilan in 3D --
and is unchanged for Johnson--Mercier, whose largest temporary already has
that shape. Isolated build time falls everywhere: 24-72% for the matrix,
18-74% for the action.

Validation

• focused flake8 over the touched files
• python -m pytest -q test/ --ignore=test/FIAT/regression (2379 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.