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

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

Real-kernel benchmark (tsfc-compiled, DG triangle mass action)

degree	ndof	flops old→new	time old→new (µs)
3	10	878 → 618	15.3 → 15.3
7	36	10516 → 4490	18.9 → 15.5
11	78	48698 → 14698	39.7 → 17.7
16	153	186583 → 41048	90.0 → 24.6
20	231	424943 → 76744	202.4 → 32.6

Jagged index support (gem + tsfc)

gem.JaggedIndex is a subclass of Index carrying a parents tuple; its .extent stays the static rectangular bound, so every existing consumer that ignores jaggedness remains correct via the zero-padding invariant.

In tsfc/loopy.py, a For loop over a jagged index nested inside its parents now gets a parametrized ISL domain ([p] -> {[q]: 0 <= q < E - p}); otherwise it silently falls back to the rectangular bound. The proof is an executed-kernel test: a 2D Morton-gather contraction compiled through compile_gem → tsfc.loopy.generate produces for (i_2 = 0; i_2 <= 4 + -1*i_1; ...) in the C and matches numpy to machine precision. duffy_evaluation now returns jagged lattice indices, so its consumers get the triangular loops for free.

Matrix-free sum-factorized residual (tsfc/fem.py)

translate_coefficient/translate_argument take the sum-factorized (Duffy/lattice) tabulation path for elements mixing in finat.duffy.DuffyElement (DG Legendre, and CG IntegratedLegendre once the paired FIAT PR added C0 recombination), targeting O(p^{d+1}) flops for the coefficient-contraction path instead of dense O(p^{2d}). element.index_shape/argument_multiindices stay flat; the lattice ⟷ flat-dof reindexing happens entirely inside finat/duffy.py (paired PR), so no driver.py or kernel_interface changes were needed for this part.

Sum-factorise by default

tsfc/kernel_interface/common.py::set_quad_rule now defaults to scheme="collapsed" whenever any element in the form is a DuffyElement, instead of requiring dx(scheme="collapsed") spelled out explicitly.

Tests

tests/tsfc/test_codegen.py: duffy_evaluation/scatter/contract correctness against a dense reference, for both Legendre and IntegratedLegendre, triangle and tetrahedron.
tests/tsfc/test_sum_factorisation.py: flop-scaling for mass/Laplacian action (DG+CG, triangle/tetrahedron), and separately for the bilinear form with no action — both test and trial bases scattered simultaneously, the configuration that exposed the loopy scheduling bug fixed in the paired FIAT PR.
tests/firedrake/regression/test_quadrature.py: end-to-end assemble() comparison of dx vs dx(scheme="collapsed") for DG/CG mass matrices and a gradient residual.

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

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

Sum factorisation on simplices

9e10473