Sum factorisation on simplices

benchmarks/johnson_mercier.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+"""Measure simplex Johnson--Mercier code generation."""
+
+import argparse
+import hashlib
+
+import numpy
+
+from finat.ufl import FiniteElement, VectorElement
+from tsfc import compile_form
+from ufl import FunctionSpace, Mesh, TestFunction, TrialFunction, div, dx, inner
+from ufl.cell import Cell
+
+
+def compile_target(dim: int):
+    """Compile the JM mass-plus-divergence target on a simplex."""
+    cell = Cell(("triangle", "tetrahedron")[dim - 2])
+    mesh = Mesh(VectorElement("CG", cell, 1))
+    element = FiniteElement("Johnson-Mercier", cell, 1)
+    space = FunctionSpace(mesh, element)
+    u = TrialFunction(space)
+    v = TestFunction(space)
+    form = (inner(u, v) + inner(div(u), div(v))) * dx
+    return compile_form(form, parameters={"mode": "spectral"})[0]
+
+
+def temporary_metrics(kernel):
+    """Return counts for the statically allocated Loopy temporaries."""
+    temporaries = kernel.ast.default_entrypoint.temporary_variables.values()
+    sizes = []
+    for temporary in temporaries:
+        if temporary.shape:
+            sizes.append(numpy.prod(temporary.shape, dtype=int))
+    return len(sizes), sum(sizes), max(sizes, default=0)
+
+
+def main():
+    """Print compiler metrics as copyable Markdown."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dims", nargs="+", type=int, default=(2, 3))
+    args = parser.parse_args()
+    print("<!-- generated by benchmarks/johnson_mercier.py -->")
+    print("| dim | flops | temporaries | elements | bytes | largest | AST lines | hash |")
+    print("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |")
+    for dim in args.dims:
+        kernel = compile_target(dim)
+        source = str(kernel.ast)
+        digest = hashlib.sha256(source.encode()).hexdigest()[:12]
+        ntemp, nelem, largest = temporary_metrics(kernel)
+        print(f"| {dim} | {kernel.flop_count:.0f} | "
+              f"{ntemp} | {nelem} | {8 * nelem} | {largest} | "
+              f"{len(source.splitlines())} | {digest} |")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -23,7 +23,8 @@ dependencies = [
   # each Firedrake release to a specific UFL minor version (e.g. 2025.3.x)
   "fenics-ufl @ git+https://github.com/FEniCS/ufl.git@main",
   # TODO RELEASE
-  "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@main",
+  # DROP BEFORE MERGE: pinned to the paired FIAT branch for CI; revert to @main
+  "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/simplex-sum-factor",
   "h5py>3.12.1",
   "firedrake-rtree>=2026.2.0",
   "immutabledict",

tests/firedrake/regression/test_quadrature.py

@@ -52,3 +52,42 @@ def test_quadrature_element(mesh, family, mat_type, diagonal):
         a = inner(u, v) * dx
 
     assemble(a, mat_type=mat_type, diagonal=diagonal)
+
+
+@pytest.mark.parametrize("family", ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize("cell", ["triangle", "tetrahedron"])
+@pytest.mark.parametrize("degree", [1, 3])
+def test_collapsed_quadrature_sum_factorisation(cell, degree, family):
+    """Check sum-factorized residuals and matrices against dense tabulation."""
+    mesh = {"triangle": UnitSquareMesh(2, 2),
+            "tetrahedron": UnitCubeMesh(1, 1, 1)}[cell]
+    variant = None if family == "Bernstein" else "integral"
+    V = FunctionSpace(mesh, family, degree, variant=variant)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    rg = RandomGenerator(PCG64(seed=0))
+    w = rg.uniform(V, 0, 1)
+
+    # translate_coefficient path (forward transform): residual with a
+    # derivative, mixing both the coefficient and argument sum-factorized
+    # tabulations.
+    L = inner(grad(w), grad(v)) * dx(scheme="canonical")
+    L_collapsed = inner(grad(w), grad(v)) * dx(scheme="collapsed")
+    b = assemble(L)
+    b_collapsed = assemble(L_collapsed)
+    assert np.allclose(b.dat.data, b_collapsed.dat.data, rtol=1e-10, atol=1e-10)
+
+    # translate_argument path (backward transform): mass matrix.
+    a = inner(u, v) * dx(scheme="canonical")
+    a_collapsed = inner(u, v) * dx(scheme="collapsed")
+    M = assemble(a).M.values
+    M_collapsed = assemble(a_collapsed).M.values
+    assert np.allclose(M, M_collapsed, rtol=1e-10, atol=1e-10)
+
+    # Bilinear derivatives exercise two independently transformed argument
+    # lattices.
+    a = inner(grad(u), grad(v)) * dx(scheme="canonical")
+    a_collapsed = inner(grad(u), grad(v)) * dx(scheme="collapsed")
+    K = assemble(a).M.values
+    K_collapsed = assemble(a_collapsed).M.values
+    assert np.allclose(K, K_collapsed, rtol=1e-10, atol=1e-10)

tests/tsfc/test_codegen.py

@@ -1,6 +1,7 @@
+import numpy
 import pytest
 
-from gem import impero_utils
+from gem import gem, impero_utils
 from gem.gem import Index, Indexed, IndexSum, Product, Variable
 
 
@@ -24,6 +25,56 @@ def gencode(expr):
     assert len(gencode(e1).children) == len(gencode(e2).children)
 
 
+def test_jagged_index_codegen(monkeypatch):
+    import islpy as isl
+    import loopy as lp
+    import tsfc.loopy
+
+    # Execute the generated code so we check the numbers, not just the loop bounds
+    monkeypatch.setattr(tsfc.loopy, "target", lp.ExecutableCTarget())
+
+    n = 4
+    extent = n + 1
+    npts = 3
+    ndof = (n + 1) * (n + 2) // 2
+
+    rng = numpy.random.default_rng(7)
+    # Table zero-padded outside the simplex lattice p + q > n, and a
+    # clamped Morton index table for the coefficient gather
+    B = rng.random((extent, extent, npts))
+    morton = numpy.zeros((extent, extent), dtype=gem.uint_type)
+    for p_, q_ in numpy.ndindex(morton.shape):
+        if p_ + q_ > n:
+            B[p_, q_] = 0.0
+        else:
+            morton[p_, q_] = (p_ + q_) * (p_ + q_ + 1) // 2 + q_
+    c = rng.random(ndof)
+
+    i = Index(name="i", extent=npts)
+    p = Index(name="p", extent=extent)
+    q = gem.JaggedIndex(name="q", extent=extent, parents=(p,))
+
+    dof = gem.VariableIndex(Indexed(gem.Literal(morton, dtype=gem.uint_type), (p, q)))
+    integrand = Product(Indexed(Variable("c", (ndof,)), (dof,)),
+                        Indexed(gem.Literal(B), (p, q, i)))
+    expr = IndexSum(integrand, (p, q))
+
+    u = Variable("u", (npts,))
+    impero_c = impero_utils.compile_gem([(Indexed(u, (i,)), expr)], (i, p, q))
+    args = [lp.GlobalArg("u", dtype=numpy.float64, shape=(npts,)),
+            lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))]
+    knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64)
+
+    # The jagged loop must have a domain parametrized by its parent iname
+    assert any(dom.get_var_names(isl.dim_type.param)
+               for dom in knl.default_entrypoint.domains)
+
+    u_out = numpy.zeros(npts)
+    knl(c=c, u=u_out)
+    u_ref = numpy.tensordot(c[morton], B, axes=((0, 1), (0, 1)))
+    assert numpy.allclose(u_out, u_ref, rtol=1e-14)
+
+
 if __name__ == "__main__":
     import os
     import sys

tests/tsfc/test_pickle_gem.py

@@ -17,6 +17,20 @@ def test_pickle_gem(protocol):
     assert repr(expr) == repr(unpickled)
 
 
+@pytest.mark.parametrize('protocol', range(3))
+def test_pickle_jagged_index(protocol):
+    p = gem.Index(name='p', extent=4)
+    q = gem.JaggedIndex(name='q', extent=4, parents=(p,))
+    expr = gem.IndexSum(gem.Indexed(gem.Variable('A', (4, 4)), (p, q)), (p, q))
+
+    unpickled = pickle.loads(pickle.dumps(expr, protocol))
+    assert repr(expr) == repr(unpickled)
+    up, uq = unpickled.multiindex
+    assert isinstance(uq, gem.JaggedIndex)
+    assert uq.extent == 4
+    assert uq.parents == (up,)
+
+
 @pytest.mark.parametrize('protocol', range(3))
 def test_listtensor(protocol):
     expr = gem.ListTensor([gem.Variable('x', ()), gem.Zero()])

tests/tsfc/test_sum_factorisation.py

@@ -3,7 +3,7 @@
 
 from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction,
                  TensorProductCell, dx, action, interval, triangle,
-                 quadrilateral, curl, dot, div, grad)
+                 tetrahedron, quadrilateral, curl, dot, div, grad, inner)
 from finat.ufl import (FiniteElement, VectorElement, EnrichedElement,
                        TensorProductElement, HCurlElement, HDivElement)
 
@@ -168,7 +168,110 @@ def test_vector_laplace_action(cell, order):
     assert (rates < order).all()
 
 
-if __name__ == "__main__":
-    import os
-    import sys
-    pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:])
+def simplex_mass(cell, family, degree, scheme='collapsed'):
+    m = Mesh(VectorElement('CG', cell, 1))
+    variant = None if family == "Bernstein" else "integral"
+    V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant))
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    return inner(u, v) * dx(scheme=scheme)
+
+
+def simplex_laplacian(cell, family, degree, scheme='collapsed'):
+    m = Mesh(VectorElement('CG', cell, 1))
+    variant = None if family == "Bernstein" else "integral"
+    V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant))
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    return inner(grad(u), grad(v)) * dx(scheme=scheme)
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4)])
+def test_simplex_mass_action(cell, family, order):
+    degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8))
+    flops = [count_flops(action(simplex_mass(cell, family, degree)))
+             for degree in degrees]
+    rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees))
+    assert (rates < order).all()
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4.4)])
+def test_simplex_laplacian_action(cell, family, order):
+    degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8))
+    flops = [count_flops(action(simplex_laplacian(cell, family, degree)))
+             for degree in degrees]
+    rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees))
+    assert (rates < order).all()
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+def test_simplex_laplacian_action_compact_codegen(family):
+    form = action(simplex_laplacian(triangle, family, 3))
+    kernel, = compile_form(form, parameters=dict(mode='spectral'))
+    temporaries = kernel.ast.default_entrypoint.temporary_variables
+    assert len(temporaries) < 100
+
+
+def test_bernstein_laplacian_action_compact_literals():
+    degree = 5
+    form = action(simplex_laplacian(tetrahedron, "Bernstein", degree))
+    kernel, = compile_form(form, parameters=dict(mode='spectral'))
+    temporaries = kernel.ast.default_entrypoint.temporary_variables
+    literals = [numpy.asarray(temporary.initializer)
+                for temporary in temporaries.values()
+                if temporary.initializer is not None]
+    lattice_size = (degree + 1) ** 3
+    assert max(literal.size for literal in literals) <= lattice_size
+    assert sum(literal.size for literal in literals) < 10 * lattice_size
+
+
+def test_bernstein_laplacian_bilinear_compact_codegen():
+    import islpy as isl
+    import loopy as lp
+
+    degree = 10
+    collapsed = simplex_laplacian(
+        tetrahedron, "Bernstein", degree, scheme="collapsed")
+    canonical = simplex_laplacian(
+        tetrahedron, "Bernstein", degree, scheme="canonical")
+    collapsed_kernel, = compile_form(
+        collapsed, parameters=dict(mode="spectral"))
+    canonical_kernel, = compile_form(
+        canonical, parameters=dict(mode="spectral"))
+
+    # At this degree the lower asymptotic complexity of sum factorisation
+    # outweighs its setup cost.
+    assert collapsed_kernel.flop_count < canonical_kernel.flop_count
+
+    entrypoint = collapsed_kernel.ast.default_entrypoint
+    assert len(entrypoint.instructions) < 250
+    assert len(entrypoint.temporary_variables) < 250
+    code = lp.generate_code_v2(collapsed_kernel.ast).device_code()
+    assert sum(line.lstrip().startswith("for (")
+               for line in code.splitlines()) < 150
+
+    # A tetrahedral lattice has a loop whose bound depends on two parents.
+    assert max(domain.dim(isl.dim_type.param)
+               for domain in entrypoint.domains) >= 2
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)])
+def test_simplex_mass_bilinear(cell, family, order):
+    degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8))
+    flops = [count_flops(simplex_mass(cell, family, degree))
+             for degree in degrees]
+    rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees))
+    assert (rates < order).all()
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)])
+def test_simplex_laplacian_bilinear(cell, family, order):
+    degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8))
+    flops = [count_flops(simplex_laplacian(cell, family, degree))
+             for degree in degrees]
+    rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees))
+    assert (rates < order).all()

tsfc/kernel_interface/common.py

@@ -14,9 +14,10 @@
 import numpy
 from FIAT.reference_element import TensorProductCell
 from finat.cell_tools import max_complex
+from finat.duffy import DuffyElement
 from finat.quadrature import AbstractQuadratureRule
 from gem.node import traversal
-from gem.optimise import constant_fold_zero
+from gem.optimise import constant_fold_zero, unflatten_returns
 from gem.optimise import remove_componenttensors as prune
 from numpy import asarray
 from tsfc import fem
@@ -210,6 +211,8 @@ def compile_gem(self, ctx):
             assignments.extend(mode.flatten(var_reps.items(), ctx['index_cache']))
 
         if assignments:
+            # Rewrite flat FlattenedTensor scatters as jagged lattice loops
+            assignments = unflatten_returns(assignments)
             return_variables, expressions = zip(*assignments)
         else:
             return_variables = []
@@ -344,6 +347,11 @@ def set_quad_rule(params, cell, integral_type, functions):
         scheme = quad_rule
         fiat_cell = as_fiat_cell(cell)
         finat_elements = set(create_element(e) for e in elements if e.family() != "Real")
+        if (scheme == "default" and integral_type == "cell"
+                and any(isinstance(finat_el, DuffyElement)
+                        for finat_el in finat_elements)):
+            # Duffy tabulation requires a collapsed-coordinate point set.
+            scheme = "collapsed"
         fiat_cells = [fiat_cell] + [finat_el.complex for finat_el in finat_elements]
         if any(c.is_macrocell() for c in fiat_cells):
             if len(set(c.get_spatial_dimension() for c in fiat_cells)) > 1: