Sum-factorisation on simplices by pbrubeck · Pull Request #262 · firedrakeproject/fiat

pbrubeck · 2026-07-18T23:27:49Z

TLDR

This PR adds sum-factorised kernels for Bernstein elements on triangles and
tetrahedra. You select them with dx(scheme="collapsed").

A collapsed kernel does far fewer operations. A degree-10 tetrahedral mass
matrix does 100 times fewer. It uses more temporary memory to do so.

Base: #286. Paired Firedrake PR: firedrakeproject/firedrake#5263.

How the loops work

A Bernstein basis on a simplex has a jagged lattice of functions. This PR keeps
that lattice in the compiler, as JaggedIndex and FlattenedTensor. The
compiler then contracts one axis at a time.

The compiler must also 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.

The compiler builds both loop orders. It measures the memory of each one. Then
it keeps the smaller one. Both orders do the same number of operations, so only
memory decides.

A separate bug in GEM

GEM sorted free indices by memory address. Addresses change when the program
allocates other objects. So the compiler made a different kernel each time you
compiled a different form first. The number of operations was not stable.

GEM now sorts the indices by the order in which it made them.

This bug is from 2016 and it is also on main. It affects all of TSFC. It is a
separate commit with its own test, so you can review it on its own.

Module layout

gem/optimise.py had grown to 1870 lines and held four jobs. This PR gives
three of them their own module.

gem/jagged.py holds everything about a jagged lattice: how a product of
simplex lattices is compacted into storage, and the rewrites that trade a flat
axis for the lattice it enumerates.
gem/driver.py holds the pipelines that compose several passes.
gem/cost.py holds the cost model.

gem/gem.py keeps the node classes and the lattice enumeration, which is what a
JaggedIndex means.

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 action	78,927	1,525,350	850	12
mass matrix	2,169,554	218,122,973	8,009	12
laplacian action	480,220	3,466,054	3,106	38
laplacian matrix	40,399,943	495,076,099	138,351	905

The mass matrix does 100 times fewer operations. The laplacian matrix does 12
times fewer.

The canonical kernel writes into the result directly, so it holds almost
nothing. The collapsed kernel keeps the result of each contraction step. Its
largest array holds 66 x 66 numbers, which is a pair of triangular lattices.

Speed

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 action	1.4x	2.1x	—	14.8x
laplacian matrix	0.7x	0.9x	2.0x	2.4x

Read this table with care:

The collapsed laplacian matrix is slower below degree 7. The extra loops cost
more than they save at a low degree.
The laplacian matrix does 12 times fewer operations but runs only 2.4 times
faster. Memory limits it, not arithmetic.
The table shows no number for the laplacian action at p8. The canonical kernel
has a defect at that degree, so the ratio has no meaning.

Tests

FIAT: 2505 pass, 0 fail.
make lint is clean.
Firedrake tests/tsfc: 400 pass, 0 fail.
firedrake-check passes.
The collapsed kernel and the canonical kernel agree to 6.2e-16.

AI assistance

OpenAI Codex and Claude Code helped to write, test, and measure this work. The
human contributor is responsible for it.

pbrubeck · 2026-07-21T12:24:53Z

+        tensor-product points.
+
+        The evaluation points are the image on the reference cell of the
+        tensor-product grid ``eta_pts`` (one ``(-1, 1)`` array per spatial


This eta transform is singular

That's not a problem, because we're only using the code path for collapsed-coordinate quadrature rules that avoid the singularity

pbrubeck · 2026-07-21T12:28:19Z

+class DuffyElement:
+    """Mixin for simplicial elements whose nodal basis coincides with the
+    Dubiner expansion set, enabling O(p^d) sum-factorized tabulation on
+    collapsed (Duffy) tensor-product point sets.


We need the Bernstein extension

In bernstein.py, we should implement BernsteinExpansionSet (which should implement tabulate_duffy)

pbrubeck · 2026-07-21T13:00:51Z

+from finat.point_set import CollapsedTensorProductPointSet
+
+
+def _element_scale(element):


Rename to get_sparse_coeffs

pbrubeck · 2026-07-21T13:02:21Z

+        for r, cols in enumerate(rows):
+            pad = cols[-1]
+            for t in range(k):
+                m = cols[t] if t < len(cols) else pad
+                row_multiindex[r, t] = raw_multiindices[m]
+                row_coeff[r, t] = R[r, m] if t < len(cols) else 0.0


use numpy.nonzeros

pbrubeck · 2026-07-22T08:05:01Z

+                          dtype=int, count=ndof * dim).reshape(ndof, dim)
+
+
+def lexicographic_offsets(dim, n):


rckirby · 2026-07-22T15:05:23Z

-    Dubiner expansion set, enabling O(p^d) sum-factorized tabulation on
-    collapsed (Duffy) tensor-product point sets.
+    """Mixin enabling O(p^d) sum-factorized tabulation on collapsed (Duffy)
+    coordinates, for simplicial elements whose nodal basis coincides with


I don't know what "nodal basis coinciding with Dubiner expansion set" means. Does this mean the nodal basis is the Dubiner basis (i.e. the nodes are moments against Dubiner polynomials)?

In this case, the element is unique (up to degree/dimension). But there are other things (e.g. modified C^0 expansion set, Bernstein) where the Duffy transform does a pull-back revealing interesting structure hence sum-factorizations.

There's are two assumptions that need to be satisfied:

The ExpansionSet implements tabulate_duffy and can be sum factorized, true for the hierarchical ExpansionSet and the new BersteinExpansionSet.

The matrix of coefficients is sparse. Only true for Legendre/IntegratedLegendre and Bernstein. Legendre and Bernstein should have the identity as the matrix of coefficients, but currently they have a permutation matrix. There are some numbering disagreements between the DualSet and the ExpansionSet that get resolved by entity_ids/coeffs.

Ideally we should just be working in the ordering that favors sum-factorization (lexicographical in the 1D jagged axes) but many parts of FIAT rely on the hierarchical ordering by total polynomial degree. The least invasive and most generic fix was to introduce a gather/scatter permutation for the tabulation and the coefficient that we contract them with.

rckirby · 2026-07-22T15:07:20Z

-        where ``phi`` is the raw per-lattice-multi-index tabulation
-        `duffy_evaluation` computes.
+        """Sparse recombination weights from the raw (continuity=None)
+        lattice tabulation into this element's nodal basis:


But we said above the nodal basis is Dubiner? Does this mean that what we actually have is the nodal basis is some nice linear combination of the Dubiner basis?

This is outdated. It should say Duffy-compatible ExpansionSet

rckirby · 2026-07-23T15:31:00Z

-class Bernstein(ScalarFiatElement):
-    # TODO: Replace this with a smarter implementation
-    def __init__(self, cell, degree):
+class Bernstein(DuffyElement, ScalarFiatElement):


This is a remarkably short change to get sum-factorization. I'm assuming (trying to parse) how the functions above get called through the Duffy mixing?

rckirby · 2026-07-23T15:32:23Z

+        tensor-product points.
+
+        The evaluation points are the image on the reference cell of the
+        tensor-product grid ``eta_pts`` (one ``(-1, 1)`` array per spatial


That's not a problem, because we're only using the code path for collapsed-coordinate quadrature rules that avoid the singularity

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

A simplex lattice is a triangular index set, but GEM had only rectangular indices. Lowering one either failed to build a convex iteration domain or ran every row out to the longest, so a simplex contraction paid for storage and arithmetic that the mathematics does not need. `gem.JaggedIndex` carries the parent indices that bound it. The node traversers, the interpreter, the flop counter and the Impero scheduler all read the exact domain from those parents, and component tensor value loops are scheduled explicitly rather than left to the consumer. `Literal` compared and hashed without its dtype, so an unsigned index literal and a floating point value holding the same number were interchangeable wherever GEM memoizes on node identity. It now separates them, which Mardal--Tai--Winther needs in order to compile. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

gem.unique built every free_indices tuple with sorted(..., key=id). Index defines neither __eq__ nor __hash__, so this ordered free indices by memory address. Addresses depend on everything the process has allocated, so compiling one form changed the free_indices order seen by the next one, and any pass that branches on that order made a different choice in a warm process than in a cold one. Sorting by Index.count orders by creation instead, which a uniform shift in allocation history preserves. The id sort dates to a5d4050 and affects all of TSFC; it only became visible where sum factorisation picks an index to unflatten by position in free_indices, which made the flop count of a compiled kernel depend on which forms had been compiled before it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

simplex_lattice_rank built a degree-d polynomial in the lattice indices, which tsfc inlined into every subscript, so the lattice layout did not survive into the generated code. simplex_lattice_ranks tabulates the rank of each lattice point instead, built from _lattice_points, the same enumeration FlattenedTensor uses to flatten a jagged tensor. The temporary layout and FlattenedTensor now share one definition of the bijection, and the layout stays visible in the AST. compact_index_layout no longer declines to compact a lone lattice. It compacts any simplex component strictly smaller than the box enclosing it, which drops the degenerate one-dimensional case that used to emit t6[11 - ((11 - i) // 1)] for what is just i. compile_gem loses assignment_group_size. It existed to cap the size of an output-shaped temporary that the scheduling change removes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

gem/optimise.py had grown to 1870 lines and held four different jobs. Give three of them their own module. gem/jagged.py holds everything about a JaggedIndex lattice: how a product of simplex lattices is compacted into storage, and the rewrites that trade a flat axis for the lattice it enumerates. The lattice combinatorics move out of gem/gem.py, which defines node classes; what stays there is the enumeration itself, which is what a JaggedIndex means, under public names. gem/driver.py holds the pipelines that compose several passes: contraction and unflatten_returns. Those reached gem.coffee and gem.refactorise through imports inside the function body, because a module-level import would have made a cycle. A pipeline sits above the primitives it uses, so from this module the imports are ordinary. contraction ran the same six lines twice, once before unflattening and once after. Run one flatten step twice instead. Every index now answers to `parents`, so the places that asked for it with getattr read the attribute. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NoTXWyj2fVdJTHnMDFB4k

…m-factor

pbrubeck mentioned this pull request Jul 18, 2026

Sum factorisation on simplices firedrakeproject/firedrake#5263

Open

pbrubeck force-pushed the pbrubeck/simplex-sum-factor branch 2 times, most recently from 0a6c400 to 103092e Compare July 20, 2026 09:09

pbrubeck commented Jul 20, 2026

View reviewed changes

Comment thread FIAT/expansions.py

pbrubeck commented Jul 21, 2026

View reviewed changes

Comment thread FIAT/expansions.py Outdated

pbrubeck commented Jul 21, 2026

View reviewed changes

Comment thread finat/duffy.py Outdated

pbrubeck commented Jul 21, 2026

View reviewed changes

pbrubeck force-pushed the pbrubeck/simplex-sum-factor branch from 8e644a8 to f872db4 Compare July 21, 2026 13:19

pbrubeck commented Jul 22, 2026

View reviewed changes

rckirby reviewed Jul 23, 2026

View reviewed changes

pbrubeck commented Jul 23, 2026

View reviewed changes

Comment thread gem/optimise.py Outdated