Preserve finite element pullbacks when collecting monomials - #5362
Open
pbrubeck wants to merge 5 commits into
Open
Preserve finite element pullbacks when collecting monomials#5362pbrubeck wants to merge 5 commits into
pbrubeck wants to merge 5 commits into
Conversation
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>
pbrubeck
force-pushed
the
pbrubeck/zany-matvec
branch
from
August 20, 2026 15:01
39f1e41 to
eee9def
Compare
The TSFC changes here need the GEM changes in the FIAT stack firedrakeproject/fiat#282 -> #284 -> #281, whose head carries all three. 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>
pbrubeck
force-pushed
the
pbrubeck/zany-matvec
branch
from
August 20, 2026 15:10
eee9def to
a458257
Compare
pbrubeck
commented
Aug 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The TSFC side of the FIAT stack firedrakeproject/fiat#282 -> #284 -> #281.
One PR, because the two commits below are what make that stack observable at
all: the GEM changes are inert unless TSFC passes the argument indices.
What changes
Cost linear map preservation against expansion (
tsfc/spectral.py).A pullback reaches TSFC as a sum over one argument axis. Expanding it exposes
scalar factorisation across its entries; preserving it exposes a tabulation
that several argument axes share. Neither dominates, so
flattenbuilds bothplans and keeps the cheaper by
estimate_cost, skipping the secondfactorisation when
has_linear_mapssays there is no map to preserve.This replaces a
len(free_indices) > 1heuristic. The heuristic was notmerely inelegant: on a CG degree 3 Laplacian in 2D, preserving the map costs
8,305 operations against 8,125 for expanding it, so preservation has to be
rejected on cost, which counting argument axes cannot do.
Tabulate adjacent shared maps in one loop (
tsfc/loopy.py).One preserved map per
ComponentTensorgave one loop each, so a single fusednest became three loops over the same extent. Reusing one iname per bound
index fuses them back.
The iname is reused only between adjacent tabulations. Sharing it whenever
the index matches breaks 13 tests with
pytools.graph.CycleError: imperointerleaves other statements between tabulations that depend on them, and one
iname cannot be both inside and outside the statement between.
statement_blocktherefore resets the memo on any statement that is not a
ComponentTensorEvaluate. Scoping by loop depth instead is not enough, as the cycle recursone level down.
Effect
Raviart--Thomas is where the pullback is worth preserving. On
inner(u, v)*dx + inner(curl(u), curl(v))*dx,mainto this PR:Arithmetic falls 32% at degree 3 and 33% at degree 5 in 3D, with 31% fewer
scalar temporaries and 20% fewer lines of C. Q and NCE are untouched: on a
tensor-product cell no sum depends on exactly one argument index, so the second
plan is never built. The full 18-case table is in firedrakeproject/fiat#284.
On time this reaches parity, not a speedup. Best of four interleaved
repetitions:
-fno-fast-mathPyOP2 compiles with
-ffast-math, so gcc may reassociatea*x + a*yintoa*(x + y)and hoist the sum out of the loop, which is exactly the savingpreservation makes structurally. Forbid the reassociation and the two separate:
mainloses 7.0% and 9.2% while this branch loses 2.4% and 0.6%. The counted32-33% is real, and the C compiler was already collecting most of it. What
remains is a third fewer operations to schedule and a representation that does
not depend on gcc choosing to reassociate.
Before the loop fusion the branch was 7-11% slower on these cases, so that
commit is load-bearing rather than a tidy-up.
Validation
tests/tsfcand the zany helmholtz, interpolate and projection regressionsuites (437 passed)
test_impero_loopy_flop_counts.pycovers theComponentTensorcase thatcount_flopspreviously turned into a silent zero;test_sum_factorisation.pycovers the plan selection.Notes
A DROP BEFORE MERGE commit adds a step to
.github/actions/install/action.ymlthat installs the head of that FIAT stack over the one
pyproject.tomlresolvesfrom
main, so CI exercises both halves together. Revert it once they land.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.