Optimise generic sum-factorisation lowering

.github/actions/install/action.yml

@@ -160,6 +160,7 @@ runs:
           --extra-index-url https://download.pytorch.org/whl/cpu \
           "./firedrake-repo[${{ inputs.deps }}]"
 
+        pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/fiat.git@pbrubeck/optimise-sum-factor
         firedrake-clean
         pip list
 

benchmarks/johnson_mercier.py

@@ -0,0 +1,154 @@
+#!/usr/bin/env python
+"""Measure simplex Johnson--Mercier code generation and assembly."""
+
+import argparse
+import cProfile
+import pstats
+import time
+import types
+
+from metrics import isolate_caches, kernel_metrics, time_kernel
+
+
+def build_form(dim: int, size: int) -> tuple[object, object]:
+    """Build the JM mass-plus-divergence form on a simplex mesh.
+
+    Parameters
+    ----------
+    dim
+        Topological dimension.
+    size
+        Number of mesh cells along each axis.
+
+    Returns
+    -------
+    form
+        The bilinear form.
+    space
+        The Johnson--Mercier function space it is posed on.
+    """
+    from firedrake import (FunctionSpace, TestFunction, TrialFunction,
+                           UnitCubeMesh, UnitSquareMesh, div, dx, inner)
+    mesh = (UnitSquareMesh, UnitCubeMesh)[dim - 2](*(size,) * dim)
+    space = FunctionSpace(mesh, "Johnson-Mercier", 1)
+    u = TrialFunction(space)
+    v = TestFunction(space)
+    return (inner(u, v) + inner(div(u), div(v))) * dx, space
+
+
+def compile_target(form: object) -> object:
+    """Compile a form through TSFC.
+
+    Parameters
+    ----------
+    form
+        The bilinear form.
+
+    Returns
+    -------
+    object
+        Compiled TSFC kernel.
+    """
+    from tsfc import compile_form
+    return compile_form(form, parameters={"mode": "spectral"})[0]
+
+
+def measure(dim: int, size: int, repeats: int) -> types.SimpleNamespace:
+    """Time compilation, cold assembly and the generated cell kernel.
+
+    Parameters
+    ----------
+    dim
+        Topological dimension.
+    size
+        Number of mesh cells along each axis.
+    repeats
+        Number of kernel calls to average over.
+
+    Returns
+    -------
+    types.SimpleNamespace
+        Timings, problem size and compiled kernel metrics.
+    """
+    from firedrake import assemble
+    form, space = build_form(dim, size)
+
+    start = time.perf_counter()
+    kernel = compile_target(form)
+    compile_time = time.perf_counter() - start
+
+    start = time.perf_counter()
+    assemble(form)
+    cold = time.perf_counter() - start
+
+    run = kernel_metrics(kernel)
+    run.dim = dim
+    run.dofs = space.dim()
+    run.cells = space.mesh().num_cells()
+    run.compile_time = compile_time
+    run.cold = cold
+    run.warm = time_kernel(form, {"mode": "spectral"}, repeats)
+    return run
+
+
+def profile(dim: int, size: int, count: int) -> None:
+    """Print the hottest calls in a cold assembly.
+
+    Parameters
+    ----------
+    dim
+        Topological dimension.
+    size
+        Number of mesh cells along each axis.
+    count
+        Number of lines of profile output to print.
+    """
+    from firedrake import assemble
+    form, _ = build_form(dim, size)
+    profiler = cProfile.Profile()
+    profiler.enable()
+    assemble(form)
+    profiler.disable()
+    print(f"<!-- cold assemble profile, dim {dim} -->")
+    pstats.Stats(profiler).sort_stats("tottime").print_stats(count)
+
+
+def main() -> None:
+    """Print benchmark measurements as copyable Markdown."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dims", nargs="+", type=int, default=(2, 3))
+    parser.add_argument("--size", type=int, default=8)
+    parser.add_argument("--repeats", type=int, default=20)
+    parser.add_argument("--warm-cache", action="store_true",
+                        help="reuse the on-disk kernel caches")
+    parser.add_argument("--profile", type=int, default=0, metavar="LINES",
+                        help="profile a cold assemble instead of timing it")
+    args = parser.parse_args()
+
+    if not args.warm_cache:
+        isolate_caches("johnson-mercier-")
+
+    if args.profile:
+        for dim in args.dims:
+            profile(dim, args.size, args.profile)
+        return
+
+    print("<!-- generated by benchmarks/johnson_mercier.py -->")
+    print("| dim | cells | dofs | compile (s) | assemble cold (s) | "
+          "kernel (s) | Gflop/s | flops | scalar temps | mutable arrays | "
+          "mutable elements | mutable bytes | largest mutable | tables | "
+          "table elements | AST lines |")
+    print("| ---: " * 16 + "|")
+    for dim in args.dims:
+        run = measure(dim, args.size, args.repeats)
+        print(f"| {run.dim} | {run.cells} | {run.dofs} | "
+              f"{run.compile_time:.6f} | {run.cold:.6f} | {run.warm:.6f} | "
+              f"{run.flops * run.cells / run.warm / 1e9:.2f} | "
+              f"{run.flops:.0f} | {run.nscalar} | {run.nmutable} | "
+              f"{run.nmutable_elements} | {8 * run.nmutable_elements} | "
+              f"{run.largest_mutable} | {run.ntables} | "
+              f"{run.ntable_elements} | {run.ast_lines} |")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/metrics.py

@@ -0,0 +1,188 @@
+"""Instrument generated finite element kernels."""
+
+import os
+import tempfile
+import time
+import types
+
+import numpy
+
+
+def isolate_caches(prefix: str) -> None:
+    """Point the TSFC and PyOP2 caches at a fresh directory.
+
+    Parameters
+    ----------
+    prefix
+        Prefix for the temporary cache directory.
+
+    Notes
+    -----
+    Must run before ``import firedrake``, which fills these variables in if
+    they are unset.  A warm disk cache hides code generation, which is part
+    of the quantity being measured.
+    """
+    cache = tempfile.mkdtemp(prefix=prefix)
+    os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] = os.path.join(cache, "tsfc")
+    os.environ["PYOP2_CACHE_DIR"] = os.path.join(cache, "pyop2")
+
+
+def kernel_metrics(kernel: object) -> types.SimpleNamespace:
+    """Separate writable intermediates from immutable tables.
+
+    Parameters
+    ----------
+    kernel
+        Compiled TSFC kernel.
+
+    Returns
+    -------
+    types.SimpleNamespace
+        ``flops``, the scalar and array temporary counts, the entries they
+        hold, and the length of the generated AST.
+
+    Notes
+    -----
+    Loopy represents compile-time quadrature and tabulation data as
+    initialized temporary variables.  Those arrays are kernel inputs in the
+    finite element algorithm, not writable contraction intermediates, so
+    combining them would overstate the working set created by factorization.
+    """
+    temporaries = tuple(
+        kernel.ast.default_entrypoint.temporary_variables.values())
+    mutable = [
+        temporary for temporary in temporaries
+        if temporary.shape
+        and not (temporary.read_only and temporary.initializer is not None)
+    ]
+    tables = [
+        temporary for temporary in temporaries
+        if temporary.shape
+        and temporary.read_only and temporary.initializer is not None
+    ]
+    mutable_sizes = [
+        numpy.prod(temporary.shape, dtype=int) for temporary in mutable
+    ]
+    table_sizes = [
+        numpy.prod(temporary.shape, dtype=int) for temporary in tables
+    ]
+    return types.SimpleNamespace(
+        flops=kernel.flop_count,
+        nscalar=sum(not temporary.shape for temporary in temporaries),
+        nmutable=len(mutable_sizes),
+        nmutable_elements=sum(mutable_sizes),
+        largest_mutable=max(mutable_sizes, default=0),
+        ntables=len(table_sizes),
+        ntable_elements=sum(table_sizes),
+        ast_lines=len(str(kernel.ast).splitlines()),
+    )
+
+
+def hottest_global_kernel(form: object, parameters: dict) -> tuple:
+    """Assemble a form and return the global kernel that did most work.
+
+    Parameters
+    ----------
+    form
+        The form to assemble.
+    parameters
+        Form compiler parameters, which must match the ones the measured
+        kernel was compiled with; ``assemble`` otherwise silently uses the
+        defaults and every mode times the same generated code.
+
+    Returns
+    -------
+    kernel
+        The PyOP2 global kernel with the highest local flop count.
+    comm
+        Communicator it was called on.
+    arguments
+        Arguments it was called with.
+    """
+    from firedrake import assemble
+    from pyop2.global_kernel import GlobalKernel
+
+    calls = []
+    original = GlobalKernel.__call__
+
+    def record(self, comm, *arguments):
+        calls.append((self, comm, arguments))
+        return original(self, comm, *arguments)
+
+    GlobalKernel.__call__ = record
+    try:
+        assemble(form, form_compiler_parameters=parameters)
+    finally:
+        GlobalKernel.__call__ = original
+
+    return max(calls, key=lambda call: call[0].local_kernel.num_flops)
+
+
+def time_kernel(form: object, parameters: dict, repeats: int) -> float:
+    """Time repeated calls to the compiled cell kernel.
+
+    Parameters
+    ----------
+    form
+        The form to assemble.
+    parameters
+        Form compiler parameters.
+    repeats
+        Number of calls to average over.
+
+    Returns
+    -------
+    float
+        Mean seconds per call to the generated code.
+
+    Notes
+    -----
+    Calling the compiled function directly measures the cell loop without
+    the Python and PETSc work that surrounds a call to ``assemble``.
+    """
+    from pyop2.global_kernel import compile_global_kernel
+
+    kernel, comm, arguments = hottest_global_kernel(form, parameters)
+    execute = compile_global_kernel(kernel, comm)
+
+    execute(*arguments)
+    start = time.perf_counter()
+    for _ in range(repeats):
+        execute(*arguments)
+    return (time.perf_counter() - start) / repeats
+
+
+def dump_kernel(kernel: object, form: object, parameters: dict,
+                directory: str, name: str) -> None:
+    """Write the loopy and C forms of a kernel for inspection.
+
+    Parameters
+    ----------
+    kernel
+        Compiled TSFC kernel.
+    form
+        The form it came from.
+    parameters
+        Form compiler parameters.
+    directory
+        Directory to write into.
+    name
+        Basename identifying the case.
+
+    Notes
+    -----
+    The local kernel shows the loop nest sum factorisation produced; the
+    PyOP2 wrapper shows the C a compiler actually sees.
+    """
+    import loopy
+    from pyop2.global_kernel import _generate_code_from_global_kernel
+
+    os.makedirs(directory, exist_ok=True)
+    with open(os.path.join(directory, f"{name}.loopy"), "w") as handle:
+        handle.write(str(kernel.ast))
+    with open(os.path.join(directory, f"{name}.c"), "w") as handle:
+        handle.write(loopy.generate_code_v2(kernel.ast).device_code())
+
+    global_kernel, comm, _ = hottest_global_kernel(form, parameters)
+    with open(os.path.join(directory, f"{name}.wrapper.c"), "w") as handle:
+        handle.write(_generate_code_from_global_kernel(global_kernel, comm))