Skip to content

Padded basis transformation for better codegen - #281

Open
pbrubeck wants to merge 7 commits into
pbrubeck/structured-codegenfrom
pbrubeck/zany-matvec
Open

Padded basis transformation for better codegen#281
pbrubeck wants to merge 7 commits into
pbrubeck/structured-codegenfrom
pbrubeck/zany-matvec

Conversation

@pbrubeck

@pbrubeck pbrubeck commented Aug 19, 2026

Copy link
Copy Markdown

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 over
the 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_delta finds and applies the pull-back;
cancel_nested_deltas walks a whole DAG. The guard on that walk is now
memoised, so it is linear in the size of the DAG.

Benchmarks

Measured on this branch against main. The form is
inner(u, v)*dx + inner(d(u), d(v))*dx1, with d = grad for CG and Q, div
for RT and curl for 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)*dx1 for Johnson--Mercier.

tsfc (s) is the TSFC compile time. build (s) is the cold-cache C build
alone. kernel (s) is the compiled kernel called directly, averaged over one
second of calls.

Zany elements, bilinear form

element dim flops array temps entries largest AST lines tsfc (s) build (s) kernel (s)
Argyris 2 38,663 -> 54,308 6 -> 9 126 -> 218 21 -> 50 431 -> 223 0.206 -> 0.205 2.63 -> 2.01 0.000041 -> 0.000051
Guzman--Neilan 2 10,041 -> 12,727 8 -> 13 72 -> 123 9 -> 15 372 -> 196 0.157 -> 0.118 1.52 -> 0.91 0.000011 -> 0.000016
Guzman--Neilan 3 384,767 -> 417,910 18 -> 23 288 -> 395 16 -> 43 1,803 -> 502 0.810 -> 0.500 6.69 -> 2.35 0.001689 -> 0.001668
Johnson--Mercier 2 21,660 -> 26,217 14 -> 17 630 -> 668 225 454 -> 206 0.233 -> 0.220 1.80 -> 0.99 0.000023 -> 0.000028
Johnson--Mercier 3 533,683 -> 663,533 26 -> 33 4,536 -> 4,802 1,764 1,993 -> 424 1.000 -> 1.015 9.83 -> 2.71 0.001538 -> 0.002018

Zany elements, matrix-free action

element dim flops array temps entries largest AST lines tsfc (s) build (s) kernel (s)
Argyris 2 6,632 -> 10,097 3 -> 7 63 -> 176 21 -> 50 427 -> 290 0.233 -> 0.171 2.61 -> 2.15 0.000006 -> 0.000014
Guzman--Neilan 2 3,549 -> 6,159 4 -> 5 36 -> 159 9 -> 36 343 -> 281 0.159 -> 0.130 1.42 -> 1.11 0.000004 -> 0.000011
Guzman--Neilan 3 106,236 -> 242,220 9 -> 7 144 -> 907 16 -> 144 1,559 -> 857 0.734 -> 0.453 5.52 -> 2.99 0.000976 -> 0.001672
Johnson--Mercier 2 3,753 -> 5,481 9 -> 7 135 -> 98 15 444 -> 250 0.267 -> 0.185 1.75 -> 1.11 0.000005 -> 0.000008
Johnson--Mercier 3 35,992 -> 71,260 17 -> 10 714 -> 392 42 1,911 -> 670 1.241 -> 0.546 10.02 -> 2.61 0.000205 -> 0.000384

The gain is in code generation, and the cost is arithmetic.

  • AST lines fall 48-79% for the matrix and 18-65% for the action.
  • The cold C build time falls 24-72% for the matrix and 18-74% for the action.
  • Flops rise on every zany case: Argyris +40.5%, Guzman--Neilan +26.8% in 2D and
    +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.
  • The biggest temporary grows where a padded (ndof, nqp) table replaces the
    per-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.

element degree dim flops AST lines tsfc (s) build (s)
CG 3 3 49,592 169 0.061 -> 0.068 0.85 -> 0.83
RT 3 3 281,708 -> 226,124 168 -> 150 0.068 -> 0.068 0.97 -> 0.91
RT 5 3 10,384,508 -> 7,865,477 168 -> 150 0.097 -> 0.097 2.72 -> 2.75
Q 7 3 16,034,499 394 0.127 -> 0.133 1.90 -> 1.79
NCE 7 3 115,963,956 -> 115,963,989 1,962 -> 1,980 0.833 -> 0.877 10.25 -> 10.05

Validation

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.

@pbrubeck
pbrubeck force-pushed the pbrubeck/zany-matvec branch 2 times, most recently from 0080000 to 4969bd6 Compare August 19, 2026 09:08
@pbrubeck
pbrubeck requested a review from rckirby August 19, 2026 09:42
@pbrubeck
pbrubeck changed the base branch from main to pbrubeck/structured-codegen August 20, 2026 15:01
@pbrubeck
pbrubeck force-pushed the pbrubeck/zany-matvec branch from 0b2cfe8 to 51ee6d4 Compare August 20, 2026 15:01
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
pbrubeck force-pushed the pbrubeck/zany-matvec branch from 51ee6d4 to c3eb76b Compare August 22, 2026 15:39
@pbrubeck
pbrubeck force-pushed the pbrubeck/zany-matvec branch from c6ccec2 to a780820 Compare August 26, 2026 11:42
pbrubeck and others added 7 commits August 29, 2026 15:43
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
pbrubeck force-pushed the pbrubeck/zany-matvec branch from b1ad6dc to 9b47d06 Compare August 29, 2026 15:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant