Sum factorisation on simplices by pbrubeck · Pull Request #5263 · firedrakeproject/firedrake

pbrubeck · 2026-07-18T23:36:11Z

Summary

gem.JaggedIndex represents simplex lattice indices with bounds depending on parent indices. TSFC lowers these to parameterized Loopy domains.
gem.FlattenedTensor exposes a jagged lattice through a flat DoF index while retaining its tensor-product structure.
gem.SparseMatrix represents the sparse C0_basis recombination.
DuffyElement.duffy_evaluation() constructs sum-factorized Legendre and IntegratedLegendre tabulations; get_sparse_coeffs() provides DG ordering or CG recombination data.
TSFC uses Duffy tabulations for coefficient and argument translation and selects collapsed quadrature for DuffyElements.

gem.optimise now:

rewrites flat contractions as jagged lattice contractions with unflatten() and unflatten_returns();
selectively distributes separable Duffy sums and sparse-matrix deltas;
cancels sparse row/column deltas during contraction;
sum-factorizes forward and transpose contractions; and
lowers remaining FlattenedTensors to gather tables with replace_flattened().

One-element benchmarks

FLOPs use the exact generated jagged loop domains. Times are median assembly times after warm-up.

Case	Degree	Canonical FLOPs	Collapsed FLOPs	FLOP reduction	Canonical time	Collapsed time	Speedup
CG Laplacian action (2D)	16	317197	215681	32.0%	369.0 us	230.0 us	1.60x
Bernstein Laplacian action (2D)	16	317197	144331	54.5%	486.2 us	118.5 us	4.10x
CG Laplacian action (3D)	15	33162804	28707004	13.4%	77.65 ms	4.955 ms	15.67x
Bernstein Laplacian action (3D)	15	33162804	4057153	87.8%	31.58 ms	418.7 us	75.43x
DG mass matrix (2D)	12	2814534	247178	91.2%	501.1 us	292.9 us	1.71x
DG mass matrix (3D)	7	14807576	302041	98.0%	9.885 ms	0.440 ms	22.47x

AI assistance: OpenAI Codex was used for code analysis, implementation support, testing, and benchmarking.

Extends tsfc/fem.py's translate_coefficient and translate_argument to take a sum-factorized (Duffy/lattice) tabulation path for simplicial DG Legendre elements under dx(scheme="collapsed"), targeting O(p^{d+1}) flops for the matrix-free residual instead of the dense O(p^{2d}). The lattice multiindex from Legendre.duffy_evaluation is gathered against coefficients and scattered back to the flat dof index through FIAT's existing Morton dof numbering (new morton_forward_table / morton_inverse_table in FIAT.expansions), entirely inside fem.py: element.index_shape and argument_multiindices stay flat, so no driver.py or kernel_interface changes are needed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

pbrubeck

TSFC codegen diff should be as tight as possible. Add the new code in a separate finat submodule.

pbrubeck · 2026-07-20T15:25:25Z

+    """Whether the sum-factorized (Duffy/lattice) coefficient contraction
+    applies.
+
+    This holds exactly when `element` is a simplicial DG element whose
+    nodal basis coincides with the Dubiner expansion set (any
+    `finat.duffy.DuffyElement`), evaluation points come from a collapsed
+    tensor-product quadrature rule (requested via
+    ``dx(scheme="collapsed")``), and the integral is over the cell
+    interior.  In that case `DuffyElement.duffy_contraction` contracts a
+    `Coefficient` against the element in O(p^d) space/time using the
+    lattice multi-index, whereas the standard dense contraction
+    materializes all O(p^d) basis functions at all O(p^d) points before
+    contracting.
+
+    The argument (basis evaluation) side needs no such dispatch: it is
+    handled transparently by `DuffyElement.basis_evaluation`, reached
+    through the usual `~.PointSetContext.basis_evaluation` call.
+
+    Parameters
+    ----------
+    element : finat.finiteelementbase.FiniteElementBase
+        The element being tabulated.
+    ctx : ContextBase
+        The translation context.
+
+    Returns
+    -------
+    bool
+        Whether to use `DuffyElement.duffy_contraction` in place of the
+        standard dense contraction.
+    """


One-line docstring

…rred Legendre's dof order is now lattice-lexicographic (permuted in FIAT), which simplified finat/duffy.py's index arithmetic but did not eliminate the VariableIndex gather/scatter: get_indices() must stay a flat index because it can't distinguish a cell-interior kernel from a facet-integral kernel, and facet tabulation always uses the dense, flat-(ndof,) FIAT path. Fully eliminating the gather/scatter would need a bespoke jagged gem.FlexiblyIndexed view in kernel_interface/ common.py's prepare_arguments/prepare_coefficient, shared code every Firedrake kernel depends on -- deliberately deferred as a separate, higher-risk follow-up. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

@main

Temporarily pins firedrake-fiat to pbrubeck/simplex-sum-factor so CI exercises this branch's paired FIAT changes (dof-order permutation, gem.Delta fix). Revert to @main once the FIAT PR merges. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

pbrubeck · 2026-07-21T09:26:50Z

+    u = TrialFunction(V)
+    v = TestFunction(V)
+    w = Function(V)
+    w.dat.data[:] = np.random.default_rng(0).random(w.dat.data.shape)


use rg = RandomGenerator(PCG64(seed=0))

translate_coefficient no longer dispatches on isinstance(element, DuffyElement): finat.duffy.DuffyElement.basis_evaluation now returns an already flat-dof-indexed tabulation, so the generic dense contraction path applies uniformly and recovers the same sum-factorized complexity without any special-casing. This also makes the Duffy fast path compose with Vector/TensorElement, which previously fell through to the slow path since TensorFiniteElement is never itself a DuffyElement instance. test_duffy_scatter_and_contract updated to contract generically instead of calling the now-removed duffy_contraction. test_collapsed_quadrature_sum_factorisation now compares against dx(scheme="canonical") instead of the default scheme: same collapsed Gauss-Jacobi points/weights as dx(scheme="collapsed"), but tabulated via the dense FIAT path, isolating the comparison to the sum-factorized tabulation itself rather than to a difference in quadrature rule. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

pbrubeck · 2026-07-21T13:24:38Z

+
+@pytest.mark.parametrize("element_name", ["Legendre", "IntegratedLegendre"])
+@pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)])
+def test_duffy_scatter_and_contract(cellname, degree, element_name):


Move to FIAT

rckirby

This just checks that the code gives a correct answer. Do we have a way of checking whether the algorithm has the right complexity?

pbrubeck · 2026-07-23T15:39:39Z

This just checks that the code gives a correct answer. Do we have a way of checking whether the algorithm has the right complexity?

This PR adds both correctness and complexity tests. Complexity in flops is not enough, I also had to enforce tests on temporaries

rckirby · 2026-07-23T16:09:59Z

This PR adds both correctness and complexity tests. Complexity in flops is not enough, I also had to enforce tests on temporaries

I see those tests, was thinking about FIAT. We should also test out Bernstein in the one-element benchmarks -- it doesn't have the indirection internally that modified C^0 expansions have but can be directly (after Duffy) sum-factored.

pbrubeck and others added 2 commits July 19, 2026 00:35

Sum factorisation on simplices

9e10473