Optimise generic sum-factorisation lowering

[pbrubeck]

Motivation

Chapter 4 of Luporini's thesis derives finite element kernel optimization from sharing elimination, reduction pre-evaluation, factorization, and generalized loop-invariant code motion. Section 4.11 calls out the need to optimize expressions outside the linear basis loops systematically, use domain knowledge about basis functions, improve the memory cost model, and generalize the method to spectral loop nests.

The paired FIAT PR, firedrakeproject/fiat#276, expresses that extension in GEM. This PR preserves the selected contraction and finite element map structure through TSFC and into Loopy. GEM chooses where indexed reductions execute; COFFEE factors scalar expressions within those loops; Loopy receives explicit value loops rather than large materialized intermediates.

Johnson--Mercier is not sum-factorizable. It is the generic benchmark because its physical basis map followed by mass and divergence contractions reveals whether the compiler preserves the intended finite element stages: transform reference tabulations first, then contract them with quadrature and geometry.

Changes

• retain shared physically mapped tabulations during argument factorization;
• compare legal quadrature contraction orders using GEM's exact operation and storage model;
• apply scalar sharing elimination at each selected reduction level;
• expose ComponentTensor value indices as explicit loops so mapped outputs share scalar temporaries;
• distinguish immutable tabulation tables from writable kernel intermediates in metrics;
• add structural tests for the mapped-tabulation loop and contraction ordering;
• add a benchmark that prints copyable Johnson--Mercier kernel metrics;

Johnson--Mercier code generation

The target is

(inner(u, v) + inner(div(u), div(v))) * dx

for degree-one Johnson--Mercier elements on a triangle and tetrahedron. The comparison is main against this PR together with firedrakeproject/fiat#276. Compile time, kernel run time, and cold-cache assemble time are pinned-core averages over repeated runs, with the kernel timed by calling the compiled cell loop directly rather than through assemble; the remaining values are deterministic properties of the generated Loopy kernel.

metric	2D: main → PR	change	3D: main → PR	change
compile time (s)	0.314 → 0.226	−28.0%	1.870 → 1.418	−24.2%
kernel run time (s)	0.000560 → 0.000565	+0.9%	0.3614 → 0.2180	−39.7%
cold-cache assemble (s)	2.623 → 1.063	2.5x faster	25.46 → 4.66	5.5x faster
FLOPs	49,571 → 40,834	−17.6%	1,859,003 → 1,032,470	−44.5%
scalar temporaries	162 → 51	3.2x fewer	1,015 → 157	6.5x fewer
writable arrays	18 → 8	2.2x fewer	48 → 15	3.2x fewer
writable entries	270 → 120	2.2x fewer	2,016 → 630	3.2x fewer
writable bytes	2,160 → 960	2.2x fewer	16,128 → 5,040	3.2x fewer
largest writable array	15 → 15	unchanged	42 → 42	unchanged
read-only table arrays	135 → 12	11.2x fewer	1,009 → 28	36x fewer
read-only table entries	1,215 → 1,254	+3.2%	16,144 → 16,270	+0.8%
AST lines	652 → 188	3.5x fewer	4,150 → 525	7.9x fewer

The eight writable 2D arrays are six mapped basis outputs, one geometry vector, and the coefficients of the basis transformation. The mapped-tabulation loop evaluates shared scalar expressions once and writes those six outputs directly. The following loop performs the element-tensor contraction. No basis function is selected through an if branch.

Reproduce each side by checking out both repositories at either main or their pbrubeck/optimise-sum-factor branches and running:

OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
taskset -c 0 python benchmarks/johnson_mercier.py --dims 2 3

The script prints the table as Markdown.

Tensor-product code generation

The targets are

dot(grad(u), grad(v)) * dx    # CG
dot(curl(u), curl(v)) * dx    # NCE

at degree seven on an extruded hexahedral mesh, following docs/notebooks/10-sum-factorisation.py. Both are measured as operator actions, so the timing sees the local contraction rather than the insertion of an element matrix. Each is compiled twice: once on the default Gauss--Legendre rule, and once on the collocated Gauss--Lobatto--Legendre rule, whose tabulations are Kronecker deltas that cancel. Timings follow the protocol above.

CG degree 7, Laplacian action:

metric	Gauss--Legendre: main → PR	change	collocated GLL: main → PR	change
compile time (s)	0.126 → 0.158	+25.5%	0.131 → 0.119	−8.5%
kernel run time (s)	0.000663 → 0.000638	−3.8%	0.000232 → 0.000107	2.2x faster
FLOPs	835,386 → 835,386	unchanged	210,564 → 93,908	2.2x fewer
scalar temporaries	96 → 90	−6.3%	116 → 77	−33.6%
writable arrays	17 → 23	+35.3%	46 → 32	−30.4%
writable entries	1,039 → 1,465	+41.0%	2,664 → 4,176	+56.8%
largest writable array	225 → 225	unchanged	512 → 512	unchanged
AST lines	329 → 337	+2.4%	482 → 345	−28.4%

NCE degree 7, curl-curl action:

metric	Gauss--Legendre: main → PR	change	collocated GLL: main → PR	change
compile time (s)	0.406 → 0.717	+76.6%	0.890 → 0.644	−27.7%
kernel run time (s)	0.002997 → 0.003219	+7.4%	0.003192 → 0.000565	5.6x faster
FLOPs	4,001,499 → 4,001,532	+0.001%	3,108,806 → 557,535	5.6x fewer
scalar temporaries	204 → 188	−7.8%	735 → 165	4.5x fewer
writable arrays	43 → 68	+58.1%	164 → 82	2.0x fewer
writable entries	6,643 → 7,018	+5.6%	13,130 → 7,958	−39.4%
largest writable array	1,575 → 1,575	unchanged	448 → 512	+14.3%
AST lines	723 → 766	+5.9%	2,492 → 736	3.4x fewer

Read-only table arrays and entries are unchanged in all four cases.

A collocated rule makes the value tabulation an identity. This PR turns that identity into a shorter contraction: on that rule the curl-curl kernel does 5.6x fewer FLOPs and runs 5.6x faster, from an AST 3.4x shorter that compiles 27.7% faster, and the Laplacian does 2.2x fewer FLOPs and runs 2.2x faster.

The Gauss--Legendre rule carries no such structure, and the FLOP counts are held there: exactly for CG, and within 0.001% for NCE. Three costs come with that. The NCE kernel is 7.4% slower, with its FLOP count and its largest temporary both matching main, so the remaining gap is not storage and is not yet explained. Compilation is 76.6% slower for NCE and 25.5% slower for CG, because selecting a contraction plan schedules each candidate to measure the storage it declares, which main does not do; compilation results are cached on disk. Writable entries rise, by 41.0% for CG and 5.6% for NCE.

Reproduce each side as above, with:

OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
taskset -c 0 python benchmarks/sum_factorisation.py \
    --family CG NCE --degrees 7 --modes spectral gll --repeats 200 --size 2

The script prints the table as Markdown, and --degrees sweeps polynomial degree.

Validation

• make srclint over the touched TSFC files, and FIAT prose checks
• tests/tsfc (368 passed)
• the physically mapped regression suites: test_interpolate_zany,
  test_projection_zany, test_helmholtz_zany, test_stress_elements,
  test_mtw, test_macro_interp_project, test_macro_multigrid,
  test_stokes_macroelements, test_embedded_transfer, test_submesh_solve
  (217 passed)
• paired FIAT suite in Optimise generic GEM sum factorisation fiat#276 (2457 passed, 26 skipped,
  31 xfailed)
• exact 2D and 3D JM kernel inspection and metric comparison
• Bernstein triangle degree 10 with canonical quadrature matches main at
  1,782,525 FLOPs, 22 scalar temporaries, 5 arrays, and 13,454 stored values

Stack

• paired generic FIAT PR: Optimise generic GEM sum factorisation fiat#276
• simplex/Bernstein Firedrake PR stacked on this branch: Sum factorisation on simplices #5263
• paired simplex FIAT PR: Sum-factorisation on simplices fiat#262

AI assistance

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