Sum factorisation on simplices

.github/actions/install/action.yml

@@ -165,6 +165,13 @@ runs:
           --extra-index-url https://download.pytorch.org/whl/cpu \
           "./firedrake-repo[${{ inputs.deps }}]"
 
+        : # DROP BEFORE MERGE: the changes here need the FInAT and GEM changes
+        : # in the FIAT stack firedrakeproject/fiat#282 -> #284 -> #281 ->
+        : # #286 -> #262, whose head carries all five.  This has to land
+        : # before anything imports Firedrake, firedrake-clean below included.
+        pip install --no-deps --force-reinstall --ignore-installed \
+          git+https://github.com/firedrakeproject/fiat.git@pbrubeck/simplex-sum-factor
+
         firedrake-clean
         pip list
 

.github/workflows/core.yml

@@ -567,6 +567,14 @@ jobs:
           pip install --verbose -r ./firedrake-repo/requirements-build.txt
           CC=mpicc CXX=mpicxx \
             pip install --verbose --no-build-isolation './firedrake-repo[docs]'
+
+          : # DROP BEFORE MERGE: this job installs Firedrake itself rather than
+          : # going through .github/actions/install, so it needs its own copy of
+          : # the FIAT stack, and for the same reason: firedrake-clean below
+          : # imports Firedrake, and so the tsfc that needs those GEM changes.
+          pip install --no-deps --force-reinstall --ignore-installed \
+            git+https://github.com/firedrakeproject/fiat.git@pbrubeck/coffee-scalar-factor
+
           firedrake-clean
           pip list
 

benchmarks/bernstein_laplacian.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python
+"""Measure Bernstein Laplacian code generation on simplices."""
+
+import argparse
+import ctypes
+from math import comb
+import os
+from pathlib import Path
+import shlex
+import statistics
+import subprocess
+import tempfile
+import time
+
+import loopy as lp
+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 = tuple(
+        kernel.ast.default_entrypoint.temporary_variables.values())
+    shapes = [temporary.shape for temporary in temporaries]
+    storage_shapes = [temporary.shape for temporary in temporaries
+                      if temporary.base_storage is None]
+    array_sizes = [int(numpy.prod(shape)) for shape in shapes if shape]
+    storage_sizes = [int(numpy.prod(shape)) if shape else 1
+                     for shape in storage_shapes]
+    return (
+        sum(not shape for shape in shapes),
+        len(array_sizes),
+        sum(storage_sizes),
+        max(array_sizes, default=0),
+        max(map(len, shapes), default=0),
+    )
+
+
+def direct_runtime(
+        kernel: object, cell: Cell, degree: int, calls: int,
+        repeats: int) -> tuple[numpy.ndarray, float]:
+    """Compile and time direct calls to a generated C kernel.
+
+    Parameters
+    ----------
+    kernel
+        Compiled TSFC kernel.
+    cell
+        Reference simplex.
+    degree
+        Bernstein polynomial degree.
+    calls
+        Kernel calls in each timed sample.
+    repeats
+        Number of timed samples.
+
+    Returns
+    -------
+    output
+        Tensor produced by one kernel call.
+    elapsed
+        Median seconds per direct kernel call.
+    """
+    dimension = cell.topological_dimension
+    ndofs = comb(degree + dimension, dimension)
+    coordinates = numpy.vstack((
+        numpy.zeros((1, dimension)), numpy.eye(dimension)
+    )).astype(numpy.float64)
+    output = numpy.zeros((ndofs, ndofs), dtype=numpy.float64)
+    source = lp.generate_code_v2(kernel.ast).device_code()
+
+    with tempfile.TemporaryDirectory(prefix="bernstein-kernel-") as directory:
+        directory = Path(directory)
+        source_path = directory / "kernel.c"
+        library_path = directory / "kernel.so"
+        source_path.write_text(source)
+        compiler = shlex.split(os.environ.get("CC", "cc"))
+        subprocess.run(
+            [*compiler, "-O3", "-march=native", "-shared", "-fPIC",
+             str(source_path), "-lm", "-o", str(library_path)],
+            check=True,
+        )
+        library = ctypes.CDLL(str(library_path))
+        function = getattr(library, kernel.ast.default_entrypoint.name)
+        pointer = ctypes.POINTER(ctypes.c_double)
+        function.argtypes = (pointer, pointer)
+        output_pointer = output.ctypes.data_as(pointer)
+        coordinate_pointer = coordinates.ctypes.data_as(pointer)
+
+        output.fill(0)
+        function(output_pointer, coordinate_pointer)
+        reference = output.copy()
+        samples = []
+        for _ in range(repeats):
+            start = time.perf_counter()
+            for _ in range(calls):
+                function(output_pointer, coordinate_pointer)
+            samples.append((time.perf_counter() - start) / calls)
+    return reference, statistics.median(samples)
+
+
+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"))
+    parser.add_argument("--runtime-calls", type=int, default=5)
+    parser.add_argument("--runtime-repeats", type=int, default=3)
+    args = parser.parse_args()
+    cell = Cell(args.cell)
+
+    print("<!-- generated by benchmarks/bernstein_laplacian.py -->")
+    print("| cell | degree | scheme | compile (s) | runtime (ms/call) | "
+          "max error | flops | scalar temps | "
+          "array temps | allocated values | bytes | largest | max rank | "
+          "AST lines |")
+    print("| :--- | ---: | :--- | ---: | ---: | ---: | ---: | ---: | "
+          "---: | ---: | ---: | ---: | ---: | ---: |")
+    for degree in args.degrees:
+        rows = []
+        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)
+            output, runtime = direct_runtime(
+                kernel, cell, degree, args.runtime_calls,
+                args.runtime_repeats)
+            rows.append((scheme, output, runtime, elapsed, kernel.flop_count,
+                         nscalar, narray, nstored, largest, max_rank,
+                         len(source.splitlines())))
+        reference = next((output for scheme, output, *_ in rows
+                          if scheme == "canonical"), None)
+        for (scheme, output, runtime, elapsed, flops, nscalar, narray,
+             nstored, largest, max_rank, ast_lines) in rows:
+            error = (numpy.max(numpy.abs(output - reference))
+                     if reference is not None else numpy.nan)
+            print(
+                f"| {args.cell} | {degree} | {scheme} | {elapsed:.6f} | "
+                f"{1000 * runtime:.6f} | {error:.3e} | {flops:.0f} | "
+                f"{nscalar} | {narray} | "
+                f"{nstored} | {8 * nstored} | {largest} | {max_rank} | "
+                f"{ast_lines} |"
+            )
+
+
+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_impero_loopy_flop_counts.py

@@ -6,9 +6,9 @@
 import loopy
 from tsfc import compile_form
 from ufl import (FunctionSpace, Mesh, TestFunction,
-                 TrialFunction, dx, grad, inner,
+                 TrialFunction, div, dx, grad, inner,
                  interval, triangle, quadrilateral,
-                 TensorProductCell)
+                 tetrahedron, TensorProductCell)
 from finat.ufl import FiniteElement, VectorElement
 from tsfc.parameters import target
 
@@ -64,3 +64,20 @@ def test_flop_count(cell, parameters):
     loopy_flops = numpy.asarray(loopy_flops)
 
     assert all(new_flops == loopy_flops)
+
+
+@pytest.mark.parametrize("cell", [triangle, tetrahedron],
+                         ids=lambda cell: cell.cellname)
+def test_flop_count_mapped_tabulation(cell):
+    # Preserving a Piola map materialises it as a ComponentTensor.
+    # Scheduling emits that as an assignment, not a loop nest, so counting
+    # it needs its own extents.
+    mesh = Mesh(VectorElement("P", cell, 1))
+    for k in range(1, 4):
+        V = FunctionSpace(mesh, FiniteElement("RT", cell, k))
+        u = TrialFunction(V)
+        v = TestFunction(V)
+        a = inner(u, v)*dx + inner(div(u), div(v))*dx
+        kernel, = compile_form(a, prefix="form",
+                               parameters={"mode": "spectral"})
+        assert kernel.flop_count == count_loopy_flops(kernel)