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.

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

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

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