Optimise generic GEM sum factorisation

[pbrubeck]

Motivation

Chapter 4 of Luporini's thesis treats finite element kernel optimization as the coordination of sharing elimination, reduction pre-evaluation, factorization, and generalized loop-invariant code motion. Section 4.11 leaves several relevant directions open: systematically optimizing expressions outside the linear basis loops, exploiting redundancies in basis functions, relaxing restrictions on code motion, replacing a fixed memory threshold with a better cost model, and extending the method to jagged loop nests.

This PR addresses those limitations at the GEM level, where finite element and contraction structure are still explicit. GEM chooses the loop domain in which each reduction is evaluated; COFFEE then eliminates scalar sharing inside the resulting loop nest. Keeping those decisions separate prevents scalar algebra from obscuring a legal contraction and prevents a contraction rewrite from losing scalar factorization opportunities.

Johnson--Mercier is a useful target even though it is not sum-factorizable. Its physical basis transformation followed by mass and divergence contractions stresses exactly the boundary this PR changes: reference tabulation, sparse basis recombination, quadrature contraction, and scalar sharing.

Mathematical model

A normalized contraction is represented by a hypergraph whose factors are vertices and whose contracted indices connect the factors in which they occur. Independent connected components are planned separately. Within a component, subset dynamic programming chooses a product tree. An index is reduced at the first subtree containing every factor incident on that index, which is its earliest legal code motion.

Plans are ordered lexicographically by exact operation count, peak live intermediate storage, and total materialized storage. Rectangular and jagged iteration domains are counted from their actual index domains. This replaces an architecture-specific temporary-size threshold with explicit mathematical costs.

Physical basis recombination is kept as a transformation of the reference tabulation. The transformation is stored in compressed sparse row form and applied as one rectangular contraction, so the basis axis reaches the quadrature contraction as a single loop rather than as one expression per basis function. Rows shorter than the longest are padded with zero coefficients, and the generated kernel selects no basis function through an if branch. Explicit ComponentTensor value loops allow several mapped outputs to share scalar work without materializing that work as arrays.

Changes

• add gem.contraction as the single planner for associative scalar trees and indexed tensor contractions;
• preserve finite element linear maps while collecting COFFEE monomials;
• apply COFFEE sharing elimination at the reduction levels selected by GEM;
• apply sparse physical basis maps as one rectangular contraction over compressed sparse rows;
• distinguish index literals from value literals when comparing and hashing Literal;
• schedule bound component-tensor indices as explicit value loops;
• interpret and count rectangular and jagged domains consistently;
• add focused tests for contraction legality, cost estimation, sparse maps, map sharing, and temporary placement.

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/firedrake#5335. 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 generated basis-map loop computes shared scalar expressions once and writes those six outputs directly; the element-tensor loop then performs the quadrature contraction.

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 lives in firedrakeproject/firedrake#5335 and prints copyable Markdown.

Validation

• python -m pydocstyle .
• focused flake8 over all touched FIAT files
• python -m pytest -q test/ (2457 passed, 26 skipped, 31 xfailed)
• python -m pytest -q test/gem/test_simplify.py test/gem/test_sum_factorise.py (30 passed)
• paired TSFC tests in Optimise generic sum-factorisation lowering firedrake#5335 (368 passed)
• the physically mapped regression suites in Optimise generic sum-factorisation lowering firedrake#5335 (217 passed)
• 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

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.