Cost and fuse preserved pullbacks, and cancel indirect Deltas

[pbrubeck]

TLDR

This is the TSFC half of firedrakeproject/fiat#284 and firedrakeproject/fiat#281.

It does three things. TSFC costs two factorisation plans and keeps the cheaper
one, instead of choosing by a rule. Loopy puts adjacent shared tabulations in
one loop, instead of one loop each. And TSFC cancels the indirect Deltas that
select a padded basis transformation's columns, which is what makes fiat#281's
transformation viable at all.

Raviart--Thomas in 3D loses a third of its operations. Q and NCE do not change.

Base: main. Needs firedrakeproject/fiat#281. The TSFC half of
firedrakeproject/fiat#286 is stacked on this branch in #5438.

What this does

Cost the two plans (tsfc/spectral.py).

A pullback reaches TSFC as a sum over one argument axis. There are two ways to
handle it.

• Expand it. This exposes scalar factorisation across its entries.
• Keep it. This exposes one tabulation that several argument axes can share.

Neither wins everywhere. So flatten builds both plans and keeps the cheaper
one by estimate_cost. It skips the second plan when has_linear_maps says
there is no map to keep, so the extra work is only paid where it can help.

On a CG degree 3 Laplacian in 2D, keeping the map costs 8,305 operations against
8,125 for expanding it. Only a cost model rejects it there.

Fuse adjacent tabulations (tsfc/loopy.py).

One kept map per ComponentTensor gave one loop each, so a single fused nest
became three loops over the same extent. TSFC now reuses one iname per bound
index and fuses them back.

The iname is reused only between adjacent tabulations. Impero can place
other statements between two tabulations that depend on each other, and one
iname cannot be both inside and outside such a statement. statement_block
therefore resets the memo on any statement that is not a ComponentTensor
Evaluate.

Cancel the indirect Deltas (tsfc/spectral.py).

fiat#281 selects a padded transformation's columns with a Delta over a variable
index. Monomial collection only cancels a Delta that surfaces as a factor of a
monomial, and one buried in a kept linear map never does, so it would reach code
generation as a full-width contraction against an identity. Integrals cancels
the indirect Deltas up front and leaves the plain-index ones to monomial
collection.

This is not an optimisation, it is what makes the padded transformation usable.
Paired with fiat#281 and without it, the Johnson--Mercier action in 3D costs
319,963 operations against 44,248 on main.

Effect

Raviart--Thomas is where keeping the pullback pays. On
inner(u, v)*dx + inner(curl(u), curl(v))*dx, main against this PR:

family	dim	degree	flops main	flops PR	AST lines main	AST lines PR
RT	2	3	23,318	18,038	90	81
RT	2	5	253,051	192,564	90	81
RT	3	3	374,589	255,933	136	109
RT	3	5	12,965,694	8,711,463	136	109

Arithmetic falls 32% at degree 3 and 33% at degree 5 in 3D. There are 31% fewer
scalar temporaries and 20% fewer lines of C.

Q and NCE do not change. On a tensor-product cell no sum depends on exactly one
argument index, so the second plan is never built.

The zany elements need the whole stack

This PR paired with fiat#281 regresses the Johnson--Mercier and Argyris actions,
because a padded transformation reads its reference tabulation once per padded
column. fiat#286 and #5438 are what tabulate that contraction one time. Numbers
are in #5438; do not merge this pair on its own.

Validation

tests/tsfc paired with fiat#281: 386 passed. That pairing did not build
before this split, because the branch also carried fiat#286's call site.

AI assistance

Claude Code was used for the split, the benchmarking and this description. The
human contributor remains responsible for understanding, validating and
maintaining the changes.

🤖 Generated with Claude Code

https://claude.ai/code/session_01A1NNAv95LANgpBX61ozPE2