Sum factorisation on simplices by pbrubeck · Pull Request #5263 · firedrakeproject/firedrake

pbrubeck · 2026-07-18T23:36:11Z

TLDR

TSFC can now lower a jagged simplex lattice to Loopy. This makes the collapsed
Bernstein kernels work.

A degree-10 tetrahedral mass matrix does 100 times fewer operations and
assembles 9.9 times faster.

Base: #5362. Paired FIAT PR:
firedrakeproject/fiat#262, which holds the element and GEM half. Read that PR
first.

pyproject.toml pins FIAT to the paired branch for CI. Drop that commit before
you merge.

Where the quadrature loops go

TSFC must choose where to put the quadrature loops.

Outside the basis-function loops, each step of the contraction can drop its
quadrature axis. This is best for squares and cubes.
Inside them, a sum over the whole result adds into a scalar. This is best for
triangles and tetrahedra.

index_orderings builds both orders. _storage_cost measures the memory of
each one. TSFC keeps the smaller one. Both orders do the same number of
operations, so only memory decides.

_storage_cost adds up the declared arrays. It does not measure how many are
alive at one time. The two measures choose the same order in every case tested,
and the gap between the two orders is much larger than the gap between the
measures.

Lattice tables

A simplex lattice is stored along one compact axis. TSFC needs the position of
each lattice point to index it.

TSFC reads that position from a small table. It builds one table for each
lattice shape. The lattice therefore stays visible in the Loopy code, and no
polynomial appears in the loop body.

Results

Degree 10 on a tetrahedron. Memory is the total size of the temporary arrays, in
words.

kernel	operations collapsed	operations canonical	memory collapsed	memory canonical
mass matrix	2,169,554	218,122,973	8,009	12
laplacian matrix	40,399,943	495,076,099	138,351	905

The collapsed kernel does far fewer operations. It uses more temporary memory,
because it keeps the result of each contraction step. Its largest array holds
66 x 66 numbers, which is a pair of triangular lattices.

Assembly time, divided by the canonical time. A larger number is better.

kernel	p4	p6	p8	p10
mass action	1.1x	1.9x	6.8x	12.3x
mass matrix	1.3x	2.4x	6.2x	9.9x
laplacian matrix	0.7x	0.9x	2.0x	2.4x

The collapsed laplacian matrix is slower below degree 7. The extra loops cost
more than they save at a low degree.

Tests

tests/tsfc: 400 pass, 0 fail.
firedrake-check passes.
FIAT: 2504 pass, 1 fails. That failure also fails on the FIAT base branch, so
it does not come from this stack.

AI assistance

OpenAI Codex and Claude Code helped to write, test, and measure this work. The
human contributor is responsible for it.

pbrubeck

TSFC codegen diff should be as tight as possible. Add the new code in a separate finat submodule.

rckirby

This just checks that the code gives a correct answer. Do we have a way of checking whether the algorithm has the right complexity?

pbrubeck · 2026-07-23T15:39:39Z

This just checks that the code gives a correct answer. Do we have a way of checking whether the algorithm has the right complexity?

This PR adds both correctness and complexity tests. Complexity in flops is not enough, I also had to enforce tests on temporaries

rckirby · 2026-07-23T16:09:59Z

This PR adds both correctness and complexity tests. Complexity in flops is not enough, I also had to enforce tests on temporaries

I see those tests, was thinking about FIAT. We should also test out Bernstein in the one-element benchmarks -- it doesn't have the indirection internally that modified C^0 expansions have but can be directly (after Duffy) sum-factored.

Expanding a pullback exposes scalar factorisation across its entries; preserving it exposes a physical basis that both argument axes share. Neither dominates: preserving wins on the Piola mapped families, where the geometry otherwise crosses the element tensor contraction twice, and expanding wins on Lagrange at moderate degree, where the entries carry enough structure to fold. Factorisation is therefore parameterised on that choice and run twice, and the cheaper plan by estimate_cost is kept. The second run is skipped when no sum spans exactly one argument axis, which is every tensor product cell here: sum factorisation has already split the basis into one dimensional factors and contracted the Jacobian into the per point geometry, so no mapped tabulation exists to share and the search would cost compile time for an identical plan. Preserved maps must survive finalisation to be shared at all, so spectral mode keeps its ComponentTensors. Measured on inner(u, v)*dx + inner(d(u), d(v))*dx, d the family's derivative, against the same tree without this change: RT tetrahedra degree 5 12,965,694 -> 8,711,463 flops, 33% fewer RT tetrahedra degree 3 374,589 -> 255,933 flops, 32% fewer RT triangles degree 5 253,051 -> 192,564 flops, 24% fewer CG tetrahedra degree 1 432 -> 390 flops, 10% fewer Q, NCE hexahedra unchanged, no map to preserve Scalar temporaries fall 31% and AST lines 11% on RT tetrahedra, and no case regresses on flops. Compile time rises by up to 17% where a second plan is built, and is unchanged elsewhere. Fewer operations do not yet make a faster kernel. The RT bilinear kernel runs 11% slower on triangles and 4.5% slower on tetrahedra at degree 3, reproducibly: each shared map becomes its own ComponentTensor, and scheduling gives each one its own loop, so one fused loop over the basis becomes three over the same extent. Stacking the maps that share an extent into one tensor is what this needs next. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

A preserved linear map is materialised as a ComponentTensor, and the loopy backend minted a fresh iname for each one, so several maps over one extent became several loops where expansion emits a single fused nest. That fission cost 7-11% of the bilinear kernel. Reuse the iname between tabulations that the schedule places side by side. Only adjacent ones: impero interleaves statements that depend on a tabulation, and one iname can not sit both inside and outside such a statement, which loopy reports as a scheduling cycle. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The TSFC changes here need the GEM changes in the FIAT stack firedrakeproject/fiat#282 -> #284 -> #281 -> #286, whose head carries all four. Install it over the one pyproject.toml resolves from main, so that CI exercises both halves together. Revert this commit once the FIAT stack lands. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

GEM no longer factors reductions through indirect gathers inside optimise_monomial_sum, which the recursive sum_factorise calls at every level. Apply the traversal once to each finished assignment instead, so plan costing still sees its effect. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

A basis transformation is now a contraction against a Delta, and delta_elimination only inspects top-level product factors, so the Delta inside a preserved linear map never reaches it. Cancel those before monomial collection, which recovers the gather the transformation used to build directly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

A jagged index is bounded by its parents, so its ISL domain must be built against them rather than as an independent axis. Loopy generation now carries the parent inames alongside the extents and constrains each dependent index inside the loops that bound it, which is what lets a sparse basis map and a simplex lattice reach the generated kernel without a data-dependent loop bound. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

get_index_ordering put every quadrature index outermost. Impero derives every loop nest from that one global order, so a reduction ordered outside its own free indices has to accumulate into a temporary carrying all of them. Where the argument lattice is factorised those free indices are the whole output, so the degree-10 tetrahedral mass bilinear staged the entire (286, 286) element matrix and then traversed both lattices a second time only to combine and scatter it. Quadrature-outermost is right in general: it is what lets each stage of a tensor-product contraction shed its quadrature axis, and flipping it globally regresses test_contraction_storage_rate on quadrilaterals and tensor-product cells from O(p) to O(p^2) stages. So offer both orderings and keep the cheaper. _terminal_reductions finds quadrature indices whose reduction spans the whole output and is not the root of its assignment; a root reduction is already absorbed by ReturnAccumulate, a buried one is not. index_orderings returns the default plus one that moves exactly those indices innermost, and the candidate with less declared temporary storage wins. The two orderings have identical flop counts, so the choice is only about where the intermediates live. Peak live storage would be the better metric in principle, but the two metrics agree on every case measured -- 32 two-candidate choices across mass and Laplacian, bilinear and action, triangles and tetrahedra at degrees 3 to 10 -- and the margin between orderings is two orders of magnitude wider than the margin between metrics. With the output-shaped temporary gone, two workarounds are no longer needed: the batch size of eight in _has_product_lattice_scatter, and the lifetime aliasing in _temporary_base_storage, which was saving 583 words out of 157756. The loop-count assertion goes back to the original < 150; the degree-10 Laplacian emits 136 loops, against 254 with aliasing. Degree 10 on tetrahedra, declared temporary storage in words: mass bilinear 87241 to 7997, Laplacian bilinear 3089666 to 138204. The largest staged matrix is (66, 66), a product of two 2-simplex lattices. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

…fact # Conflicts: # tests/tsfc/test_sum_factorisation.py

compile_expression_dual_evaluation compiles one ordering rather than costing several, so it needs the scalar form rather than the candidate list. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

# Conflicts: # tests/tsfc/test_sum_factorisation.py

…edrake The install step ends with firedrake-clean, which imports Firedrake, and so imports the tsfc that needs the GEM changes in the FIAT stack. Swapping the stack in from a step after that one leaves firedrake-clean to run against the FIAT that pyproject.toml resolved from main: ImportError: cannot import name 'eliminate_deltas' from 'gem.optimise' The job never reached firedrake-check. Install the stack inline instead, after the Firedrake install and before firedrake-clean, which is what pbrubeck/form-interp-tsfc does with its own siblings. Revert this commit once the FIAT stack lands. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

The stack head installed here, pbrubeck/coffee-scalar-factor, carries the GEM changes the TSFC half needs but not the FInAT half of this PR, so the job died importing Firedrake: ImportError: cannot import name 'CollapsedTensorProductPointSet' from 'finat.point_set' firedrakeproject/fiat#262 sits on top of that head and carries both. Install it instead. Revert this commit once the FIAT stack lands. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

The docs job installs Firedrake itself rather than going through .github/actions/install, so the stack this branch needs never reached it, and it failed the same way the test jobs did before 34c5f58: ImportError: cannot import name 'eliminate_deltas' from 'gem.optimise' raised by firedrake-clean at the end of its own install step. Give it its own copy of the install, in the same place relative to firedrake-clean. Revert this along with 34c5f58 once the FIAT stack lands. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

FIAT renamed eliminate_deltas to cancel_nested_deltas, and moved the cost model to gem.cost. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

FIAT split gem.optimise. The pipelines that compose several passes are in gem.driver, and everything about a jagged lattice is in gem.jagged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k