Sum factorisation on simplices

benchmarks/bernstein_laplacian.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python
+"""Measure Bernstein Laplacian code generation on simplices."""
+
+import argparse
+import time
+
+import numpy
+
+from finat.ufl import FiniteElement, VectorElement
+from tsfc import compile_form
+from ufl import FunctionSpace, Mesh, TestFunction, TrialFunction, dx, grad, inner
+from ufl.cell import Cell
+
+
+def compile_target(
+        cell: Cell, degree: int, scheme: str) -> tuple[object, float]:
+    """Compile a Bernstein Laplacian bilinear form.
+
+    Parameters
+    ----------
+    cell
+        Reference simplex.
+    degree
+        Polynomial degree.
+    scheme
+        Quadrature scheme.
+
+    Returns
+    -------
+    kernel
+        Compiled TSFC kernel.
+    elapsed
+        Compilation time in seconds.
+    """
+    mesh = Mesh(VectorElement("CG", cell, 1))
+    space = FunctionSpace(mesh, FiniteElement("Bernstein", cell, degree))
+    u = TrialFunction(space)
+    v = TestFunction(space)
+    form = inner(grad(u), grad(v)) * dx(scheme=scheme)
+    start = time.perf_counter()
+    kernel, = compile_form(form, parameters={"mode": "spectral"})
+    return kernel, time.perf_counter() - start
+
+
+def temporary_metrics(kernel: object) -> tuple[int, int, int, int, int]:
+    """Measure statically allocated Loopy temporaries.
+
+    Parameters
+    ----------
+    kernel
+        Compiled TSFC kernel.
+
+    Returns
+    -------
+    scalar_count
+        Number of scalar temporaries.
+    array_count
+        Number of array temporaries.
+    stored_values
+        Total scalar and array entries.
+    largest_array
+        Entries in the largest array temporary.
+    maximum_rank
+        Largest temporary tensor rank.
+    """
+    temporaries = kernel.ast.default_entrypoint.temporary_variables.values()
+    shapes = [temporary.shape for temporary in temporaries]
+    array_sizes = [int(numpy.prod(shape)) for shape in shapes if shape]
+    return (
+        sum(not shape for shape in shapes),
+        len(array_sizes),
+        sum(array_sizes) + sum(not shape for shape in shapes),
+        max(array_sizes, default=0),
+        max(map(len, shapes), default=0),
+    )
+
+
+def main() -> None:
+    """Print compiler metrics as copyable Markdown."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cell", choices=("triangle", "tetrahedron"),
+        default="tetrahedron")
+    parser.add_argument("--degrees", nargs="+", type=int, default=(10,))
+    parser.add_argument(
+        "--schemes", nargs="+", choices=("collapsed", "canonical"),
+        default=("collapsed", "canonical"))
+    args = parser.parse_args()
+    cell = Cell(args.cell)
+
+    print("<!-- generated by benchmarks/bernstein_laplacian.py -->")
+    print("| cell | degree | scheme | compile (s) | flops | scalar temps | "
+          "array temps | stored values | bytes | largest | max rank | "
+          "AST lines |")
+    print("| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | "
+          "---: | ---: | ---: | ---: |")
+    for degree in args.degrees:
+        for scheme in args.schemes:
+            kernel, elapsed = compile_target(cell, degree, scheme)
+            source = str(kernel.ast)
+            nscalar, narray, nstored, largest, max_rank = \
+                temporary_metrics(kernel)
+            print(
+                f"| {args.cell} | {degree} | {scheme} | {elapsed:.6f} | "
+                f"{kernel.flop_count:.0f} | {nscalar} | {narray} | "
+                f"{nstored} | {8 * nstored} | {largest} | {max_rank} | "
+                f"{len(source.splitlines())} |"
+            )
+
+
+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