Padded basis transformation for better codegen - #281
Open
pbrubeck wants to merge 7 commits into
Open
Conversation
pbrubeck
force-pushed
the
pbrubeck/zany-matvec
branch
2 times, most recently
from
August 19, 2026 09:08
0080000 to
4969bd6
Compare
pbrubeck
force-pushed
the
pbrubeck/zany-matvec
branch
from
August 20, 2026 15:01
0b2cfe8 to
51ee6d4
Compare
pbrubeck
added a commit
to firedrakeproject/firedrake
that referenced
this pull request
Aug 20, 2026
This branch needs the linear-map preservation added in firedrakeproject/fiat#281. Revert this commit once that lands on FIAT main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
pbrubeck
force-pushed
the
pbrubeck/zany-matvec
branch
from
August 22, 2026 15:39
51ee6d4 to
c3eb76b
Compare
pbrubeck
force-pushed
the
pbrubeck/zany-matvec
branch
from
August 26, 2026 11:42
c6ccec2 to
a780820
Compare
A padded basis transformation tabulates each facet as an IndexSum, so selecting one by a variable facet index reached _select_expression with a type it could not factorise. Rewrite the summands over one shared multiindex and select inside the reduction, which the equal extents on every facet make well defined. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The padded transformation was applied by building its row-padded gather directly, which fixed the orientation at construction: contracting the coefficient against a mapped tabulation then costs one gather per quadrature point, six times the symmetric test-side scatter. Represent M instead as a rank-2 expression, an interned entry summed over the padded row against a Delta selecting its column. Cancelling that Delta reproduces the gather, so the mat-mat is unchanged, while contracting M's own axes first pulls a coefficient back to the reference basis once per cell. Guzman-Neilan 3D action: 147180 -> 51920 flops, largest working temporary 144 -> 24 entries. Delta now propagates the free indices of a VariableIndex operand, and substitution folds a variable index that has become constant. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
MappedTabulation minted a fresh index on every call, so the tabulations of different derivative orders contracted over distinct indices of equal extent. Expressions that are structurally equal then hash apart, and the scheduler gives each its own loop nest. Reuse one index per instance for the reference basis and for the padded row. Equal tabulations now share a subexpression, and their loops fuse without any change to the scheduler. The row index of a tabulation stays per call, since a ComponentTensor binds it and sharing it only forces redundant materialisation. Guzman-Neilan 3D action: 51920 -> 51296 flops, 17 -> 15 array temporaries; 2D action: 2676 -> 2586 flops. Four groups of sibling loops over equal extents collapse to one loop each, and the Argyris and Johnson-Mercier actions lose theirs likewise. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
MappedTabulation retains its selector Delta until the contraction optimizer runs, so its raw DAG legitimately contains two Product nodes. Exercise the production contraction path before asserting the sparse arithmetic structure.
Three functions carried names for delta cancellation and one of them was a wrapper. Fold the search for a pull-back candidate into the pull-back itself, and call it pull_back_indirect_delta. Its callers now run it and delta_elimination in sequence, which is what they always did. Rename the whole-DAG traversal to cancel_nested_deltas, so that it no longer reads as a synonym of delta_elimination. Its guard walked the subtree again at every enclosing contraction. A memoised map from a node to the Delta axes below it answers the same question once per node. Use one helper for the cardinality of an index space, and promote the traversal child rule in gem.node so that other modules can share it. Drop MappedTabulation.matrix(); nothing calls it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k
pbrubeck
force-pushed
the
pbrubeck/zany-matvec
branch
from
August 29, 2026 15:21
b1ad6dc to
9b47d06
Compare
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.
TLDR
A physically mapped element applies a sparse matrix to its reference tabulation.
This PR makes that matrix a dense padded table with a regular loop, instead
of one instruction per matrix row.
The generated C shrinks a lot. Johnson--Mercier on tetrahedra drops from 1,993
lines to 424, and its cold C build time falls 72%. Arithmetic rises, because
padding adds work.
Base: #284. Needs firedrakeproject/firedrake#5362.
What this does
Every row of the transformation is padded to the same number of entries. A loop
over the basis index then has a plain rectangular domain, which Loopy can build.
Two small read-only tables say which reference column each padded entry selects.
Without padding, each row of the sparse matrix product became its own
instruction and its own temporary. Johnson--Mercier on tetrahedra declared 655
arrays of length
nqp, one for each basis function and tabulation. A loop overthe basis index needs 24 arrays of shape
(42, nqp).The padded map reaches monomial collection as a sum over the nonzeros of one
basis row. #284 keeps that sum whole and shares it between the argument axes.
This PR is what makes such a sum exist.
Tidy-up. The delta rewrites are named for what each one does.
pull_back_indirect_deltafinds and applies the pull-back;cancel_nested_deltaswalks a whole DAG. The guard on that walk is nowmemoised, so it is linear in the size of the DAG.
Benchmarks
Measured on this branch against
main. The form isinner(u, v)*dx + inner(d(u), d(v))*dx1, withd=gradfor CG and Q,divfor RT and
curlfor NCE. The zany forms are(hess u, hess v)for Argyris,(eps u, eps v) + (div u, div v)for Guzman--Neilan, and(sym(u), sym(v)) + (div u, div v)*dx1for Johnson--Mercier.tsfc (s)is the TSFC compile time.build (s)is the cold-cache C buildalone.
kernel (s)is the compiled kernel called directly, averaged over onesecond of calls.
Zany elements, bilinear form
Zany elements, matrix-free action
The gain is in code generation, and the cost is arithmetic.
+8.6% in 3D, Johnson--Mercier +21.0% in 2D and +24.3% in 3D. The action rises
further, up to +128% for Guzman--Neilan in 3D.
(ndof, nqp)table replaces theper-row vectors, up to 2.7x for Guzman--Neilan in 3D. Johnson--Mercier already
had a temporary of that shape, so it does not grow.
Elements this PR does not touch
CG, RT, Q and NCE match #284's table exactly, because the padded transformation
applies only to physically mapped elements.
Validation
test/FIAT/regression: 2379 passed, 26 skipped, 31 xfailed, 1 xpassedmake linttests/tsfcand the three zany regression suites in Preserve finite element pullbacks when collecting monomials firedrake#5362: 437 passedAI assistance
Claude Code was used for implementation, benchmarking, and drafting this
section. The human contributor remains responsible for understanding,
validating, and maintaining the changes.