diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index 54c1411d7f..065d2162ef 100644 --- a/.github/actions/install/action.yml +++ b/.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 diff --git a/benchmarks/johnson_mercier.py b/benchmarks/johnson_mercier.py new file mode 100644 index 0000000000..7d347616e5 --- /dev/null +++ b/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"") + 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("") + 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() diff --git a/benchmarks/metrics.py b/benchmarks/metrics.py new file mode 100644 index 0000000000..473ccfacfd --- /dev/null +++ b/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)) diff --git a/benchmarks/sum_factorisation.py b/benchmarks/sum_factorisation.py new file mode 100644 index 0000000000..5333250187 --- /dev/null +++ b/benchmarks/sum_factorisation.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python +"""Measure tensor-product sum-factorisation code generation and assembly. + +The cases follow ``docs/notebooks/10-sum-factorisation.py``: a Laplacian on +hexahedra, and a curl-curl form on the H(curl) conforming NCE element. Each +is compiled in three representations, which separate what sum factorisation +contributes from what collocated quadrature contributes: + +vanilla + No factorisation, the O(p^9) reference. +spectral + Sum factorisation on a canonical Gauss--Legendre rule, O(p^7). +gll + Sum factorisation on the collocated Gauss--Lobatto--Legendre rule, whose + identity tabulations reduce the same form to O(p^5). +""" + +import argparse +import cProfile +import pstats +import time +import types + +from metrics import (dump_kernel, isolate_caches, kernel_metrics, + time_kernel) + + +OPERATORS = {} + + +def build_operators() -> None: + """Populate the case table with UFL operators. + + Notes + ----- + UFL is imported through Firedrake, so the table cannot be built until + the caches have been redirected. + """ + from firedrake import curl, dot, grad + OPERATORS.update( + CG=lambda u, v: dot(grad(u), grad(v)), + NCE=lambda u, v: dot(curl(u), curl(v)), + ) + + +def gauss_lobatto_legendre_line_rule(degree: int) -> object: + """Build the GLL rule on the reference interval. + + Parameters + ---------- + degree + Polynomial degree the rule integrates. + + Returns + ------- + object + FInAT quadrature rule. + """ + import FIAT + import finat + fiat_rule = FIAT.quadrature.GaussLobattoLegendreQuadratureLineRule( + FIAT.ufc_simplex(1), degree + 1) + points = finat.point_set.GaussLobattoLegendrePointSet( + fiat_rule.get_points()) + return finat.quadrature.QuadratureRule(points, fiat_rule.get_weights()) + + +def gauss_lobatto_legendre_cube_rule(dimension: int, degree: int) -> object: + """Build the GLL rule on the reference hypercube. + + Parameters + ---------- + dimension + Topological dimension of the cube. + degree + Polynomial degree the rule integrates. + + Returns + ------- + object + FInAT tensor product quadrature rule. + """ + import finat + rule = gauss_lobatto_legendre_line_rule(degree) + for _ in range(1, dimension): + rule = finat.quadrature.TensorProductQuadratureRule( + [rule, gauss_lobatto_legendre_line_rule(degree)]) + return rule + + +MODES = { + "vanilla": dict(mode="vanilla", variant=None, collocated=False), + "spectral": dict(mode="spectral", variant=None, collocated=False), + "gll": dict(mode="spectral", variant="spectral", collocated=True), +} + + +def build_form(mesh: object, family: str, degree: int, mode: str, + operator: str) -> tuple[object, object]: + """Build one benchmark form on a hexahedral mesh. + + Parameters + ---------- + mesh + Extruded hexahedral mesh. + family + Element family, a key of ``OPERATORS``. + degree + Polynomial degree. + mode + Representation, a key of ``MODES``. + operator + ``"bilinear"`` for the operator itself, ``"action"`` for its action + on a coefficient. + + Returns + ------- + form + The form to compile. + space + The function space it is posed on. + + Notes + ----- + Timing the bilinear form measures the insertion of a dense element + matrix as much as the local assembly that produced it, and at high + degree the insertion dominates. Its action assembles a vector, so the + contraction being factorized is what the timing sees. + """ + from firedrake import (FiniteElement, Function, FunctionSpace, + TestFunction, TrialFunction, action, dx) + settings = MODES[mode] + element = FiniteElement(family, mesh.ufl_cell(), degree=degree, + variant=settings["variant"]) + space = FunctionSpace(mesh, element) + u = TrialFunction(space) + v = TestFunction(space) + measure = dx + if settings["collocated"]: + measure = dx(scheme=gauss_lobatto_legendre_cube_rule( + mesh.topological_dimension, degree)) + form = OPERATORS[family](u, v) * measure + if operator == "action": + form = action(form, Function(space)) + return form, space + + +def build_mesh(size: int) -> object: + """Extrude a quadrilateral mesh into hexahedra. + + Parameters + ---------- + size + Number of cells along each axis. + + Returns + ------- + object + The extruded mesh. + """ + from firedrake import ExtrudedMesh, UnitSquareMesh + return ExtrudedMesh(UnitSquareMesh(size, size, quadrilateral=True), size) + + +def measure(mesh: object, family: str, degree: int, mode: str, + operator: str, repeats: int, + dump: str = "") -> types.SimpleNamespace: + """Compile and time one case. + + Parameters + ---------- + mesh + Extruded hexahedral mesh. + family + Element family, a key of ``OPERATORS``. + degree + Polynomial degree. + mode + Representation, a key of ``MODES``. + operator + ``"bilinear"`` or ``"action"``. + repeats + Number of kernel calls to average over, zero to skip execution. + dump + Directory to write the loopy and C sources into, empty to skip. + + Returns + ------- + types.SimpleNamespace + Compiled kernel metrics, compile time, and mean kernel time. + """ + from tsfc import compile_form + form, space = build_form(mesh, family, degree, mode, operator) + parameters = {"mode": MODES[mode]["mode"]} + + start = time.perf_counter() + kernel, = compile_form(form, parameters=parameters) + compile_time = time.perf_counter() - start + + run = kernel_metrics(kernel) + run.family = family + run.degree = degree + run.mode = mode + run.operator = operator + run.dofs = space.dim() + run.cells = space.mesh().num_cells() + run.compile_time = compile_time + run.warm = (time_kernel(form, parameters, repeats) + if repeats else float("nan")) + if dump: + dump_kernel(kernel, form, parameters, dump, + f"{family}{degree}-{mode}-{operator}") + return run + + +def report(runs: list[types.SimpleNamespace]) -> None: + """Print measurements as copyable Markdown. + + Parameters + ---------- + runs + Measurements to tabulate. + """ + print("| family | degree | mode | form | dofs | compile (s) | kernel (s) | " + "Gflop/s | flops | scalar temps | mutable arrays | " + "mutable elements | largest mutable | tables | table elements | " + "AST lines |") + print("| ---: " * 16 + "|") + for run in runs: + print(f"| {run.family} | {run.degree} | {run.mode} | " + f"{run.operator} | {run.dofs} | " + f"{run.compile_time:.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} | {run.largest_mutable} | " + f"{run.ntables} | {run.ntable_elements} | {run.ast_lines} |") + + +def profile(mesh: object, family: str, degree: int, mode: str, + operator: str, count: int) -> None: + """Print the hottest calls in one compilation. + + Parameters + ---------- + mesh + Extruded hexahedral mesh. + family + Element family, a key of ``OPERATORS``. + degree + Polynomial degree. + mode + Representation, a key of ``MODES``. + operator + ``"bilinear"`` or ``"action"``. + count + Number of lines of profile output to print. + """ + from tsfc import compile_form + form, _ = build_form(mesh, family, degree, mode, operator) + profiler = cProfile.Profile() + profiler.enable() + compile_form(form, parameters={"mode": MODES[mode]["mode"]}) + profiler.disable() + print(f"") + pstats.Stats(profiler).sort_stats("tottime").print_stats(count) + + +def main() -> None: + """Print benchmark measurements as copyable Markdown.""" + parser = argparse.ArgumentParser() + parser.add_argument("--family", nargs="+", default=("CG",), + choices=("CG", "NCE")) + parser.add_argument("--degrees", nargs="+", type=int, + default=(1, 2, 3, 4, 5, 6)) + parser.add_argument("--modes", nargs="+", default=tuple(MODES), + choices=tuple(MODES)) + parser.add_argument("--forms", nargs="+", default=("action",), + choices=("bilinear", "action")) + parser.add_argument("--size", type=int, default=4) + parser.add_argument("--repeats", type=int, default=10, + help="kernel calls to average, 0 to compile only") + parser.add_argument("--warm-cache", action="store_true", + help="reuse the on-disk kernel caches") + parser.add_argument("--dump", default="", metavar="DIR", + help="write the loopy and C sources of each kernel") + parser.add_argument("--profile", type=int, default=0, metavar="LINES", + help="profile compilation instead of timing it") + args = parser.parse_args() + + if not args.warm_cache: + isolate_caches("sum-factorisation-") + build_operators() + mesh = build_mesh(args.size) + + if args.profile: + for family in args.family: + for degree in args.degrees: + for mode in args.modes: + for operator in args.forms: + profile(mesh, family, degree, mode, operator, + args.profile) + return + + print("") + report([ + measure(mesh, family, degree, mode, operator, args.repeats, + args.dump) + for family in args.family + for degree in args.degrees + for mode in args.modes + for operator in args.forms + ]) + + +if __name__ == "__main__": + main() diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index 8d0bc79655..d41e94af0d 100644 --- a/tests/tsfc/test_codegen.py +++ b/tests/tsfc/test_codegen.py @@ -1,6 +1,10 @@ +import numpy import pytest +import gem +from finat.physically_mapped import MappedTabulation from gem import impero_utils +from gem.flop_count import count_flops from gem.gem import Index, Indexed, IndexSum, Product, Variable @@ -24,6 +28,50 @@ def gencode(expr): assert len(gencode(e1).children) == len(gencode(e2).children) +def sparse_map_flops(lengths, ncolumns=4, npoints=2): + """Count the flops of applying a sparse basis map. + + Parameters + ---------- + lengths : list of int + Number of nonzeros in each row of the map. + ncolumns : int + Number of reference basis functions. + npoints : int + Number of points the reference tabulation holds. + + Returns + ------- + int + Flops the lowered tabulation costs. + """ + rows = [] + for row, length in enumerate(lengths): + entries = [gem.Zero()] * ncolumns + for column in range(length): + entries[column] = gem.Variable(f"c_{row}_{column}", ()) + rows.append(entries) + M = gem.ListTensor(numpy.asarray(rows, dtype=object)) + table = gem.Literal(numpy.ones((ncolumns, npoints))) + + mapped = MappedTabulation(M, {None: table})[None] + i, j = gem.indices(2) + expr, = impero_utils.preprocess_gem([Indexed(mapped, (i, j))]) + result = Indexed(Variable("A", (len(lengths), npoints)), (i, j)) + return count_flops(impero_utils.compile_gem([(result, expr)], (i, j))) + + +@pytest.mark.parametrize("lengths", [[3, 1, 1], [1, 3, 1], [1, 1, 3]]) +def test_sparse_basis_map_is_one_rectangle(lengths): + """Apply a sparse basis map as one rectangular contraction. + + Every row contracts over the same number of entries. The cost follows + the longest row, and not the arrangement of the nonzeros. + """ + assert sparse_map_flops(lengths) == sparse_map_flops([3, 3, 3]) + assert sparse_map_flops([2, 2, 2]) < sparse_map_flops([3, 3, 3]) + + if __name__ == "__main__": import os import sys diff --git a/tests/tsfc/test_coffee_optimise.py b/tests/tsfc/test_coffee_optimise.py index 065bb22ee4..81d0bd1b5f 100644 --- a/tests/tsfc/test_coffee_optimise.py +++ b/tests/tsfc/test_coffee_optimise.py @@ -45,30 +45,30 @@ def test_loop_optimise(): Z = Variable('z', ()) - # Bj*Ek + Bj*Fk => (Ek + Fk)*Bj + # Bj*Ek + Bj*Fk => Bj*(Ek + Fk) expr = Sum(Product(Bj, Ek), Product(Bj, Fk)) result, = optimise_expressions([expr], (j, k)) - expected = Product(Sum(Ek, Fk), Bj) + expected = Product(Bj, Sum(Ek, Fk)) assert result == expected # Bj*Ek + Bj*Fk + Bj*Gk + Cj*Ek + Cj*Fk => - # (Ek + Fk + Gk)*Bj + (Ek+Fk)*Cj + # Bj*(Ek + Fk + Gk) + Cj*(Ek + Fk) expr = Sum(Sum(Sum(Sum(Product(Bj, Ek), Product(Bj, Fk)), Product(Bj, Gk)), Product(Cj, Ek)), Product(Cj, Fk)) result, = optimise_expressions([expr], (j, k)) - expected = Sum(Product(Sum(Sum(Ek, Fk), Gk), Bj), Product(Sum(Ek, Fk), Cj)) + expected = Sum(Product(Bj, Sum(Sum(Ek, Fk), Gk)), Product(Cj, Sum(Ek, Fk))) assert result == expected # Z*A1i*Bj*Ek + Z*A2i*Bj*Ek + A3i*Bj*Ek + Z*A1i*Bj*Fk => - # Bj*(Ek*(Z*A1i + Z*A2i) + A3i) + Z*A1i*Fk) + # Bj*(Ek*(Z*A1i + Z*A2i + A3i) + Fk*(Z*A1i)) expr = Sum(Sum(Sum(Product(Z, Product(A1i, Product(Bj, Ek))), Product(Z, Product(A2i, Product(Bj, Ek)))), Product(A3i, Product(Bj, Ek))), Product(Z, Product(A1i, Product(Bj, Fk)))) result, = optimise_expressions([expr], (j, k)) - expected = Product(Sum(Product(Ek, Sum(Sum(Product(Z, A1i), Product(Z, A2i)), A3i)), - Product(Fk, Product(Z, A1i))), Bj) + expected = Product(Bj, Sum(Product(Ek, Sum(Sum(Product(Z, A1i), Product(Z, A2i)), A3i)), + Product(Fk, Product(Z, A1i)))) assert result == expected diff --git a/tests/tsfc/test_pickle_gem.py b/tests/tsfc/test_pickle_gem.py index beb101f912..b68905cb0d 100644 --- a/tests/tsfc/test_pickle_gem.py +++ b/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()]) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 891cf1c6cc..d7fc40165c 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -1,13 +1,20 @@ import numpy import pytest -from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, +import gem +import gem.coffee +import tsfc.spectral +from gem.gem import one +from gem.contraction import estimate_cost +from gem.refactorise import MonomialSum +from ufl import (Coefficient, Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, action, interval, triangle, - quadrilateral, curl, dot, div, grad) + quadrilateral, curl, dot, div, grad, inner) from finat.ufl import (FiniteElement, VectorElement, EnrichedElement, TensorProductElement, HCurlElement, HDivElement) from tsfc import compile_form +from tsfc.spectral import _plans def helmholtz(cell, degree): @@ -168,6 +175,142 @@ def test_vector_laplace_action(cell, order): assert (rates < order).all() +def test_shared_physically_mapped_tabulation( + monkeypatch: pytest.MonkeyPatch) -> None: + """Share a mapped tabulation between both argument axes. + + Parameters + ---------- + monkeypatch + Pytest fixture used to disable the sharing pass for comparison. + + Notes + ----- + Johnson--Mercier has six mapped basis outputs in two dimensions. The + seventh writable vector holds geometry data, and the eighth holds the + coefficients of the basis transformation. Algebra shared while + constructing those outputs belongs inside their common basis-row loop + and must therefore remain scalar. + """ + mesh = Mesh(VectorElement("CG", triangle, 1)) + element = FiniteElement("Johnson-Mercier", triangle, 1) + space = FunctionSpace(mesh, element) + u = TrialFunction(space) + v = TestFunction(space) + form = (inner(u, v) + inner(div(u), div(v))) * dx + + optimized, = compile_form(form, parameters={"mode": "spectral"}) + monkeypatch.setattr( + gem.coffee, "_share_linear_maps", + lambda monomial_sum, indices: monomial_sum) + baseline, = compile_form(form, parameters={"mode": "spectral"}) + + optimized_shapes = [ + temporary.shape for temporary in + optimized.ast.default_entrypoint.temporary_variables.values() + if not (temporary.read_only + and temporary.initializer is not None)] + baseline_shapes = [ + temporary.shape for temporary in + baseline.ast.default_entrypoint.temporary_variables.values() + if not (temporary.read_only + and temporary.initializer is not None)] + + assert optimized.flop_count < baseline.flop_count + assert sum(not shape for shape in optimized_shapes) \ + < sum(not shape for shape in baseline_shapes) + assert [shape for shape in optimized_shapes if shape] == [(15,)] * 8 + + +def test_linear_map_representation_is_costed( + monkeypatch: pytest.MonkeyPatch) -> None: + """Do not preserve a linear map when expansion costs fewer FLOPs. + + Parameters + ---------- + monkeypatch + Pytest fixture used to force the expanded representation for + comparison. + + Notes + ----- + A Bernstein tabulation is separable, so preserving every one-axis sum + produces a compact loop nest but can hide profitable scalar + factorization. Plan selection must compare that representation with the + expanded polynomial instead of committing to either transformation. + + """ + mesh = Mesh(VectorElement("CG", triangle, 1)) + element = FiniteElement("Bernstein", triangle, 2) + space = FunctionSpace(mesh, element) + u = TrialFunction(space) + v = TestFunction(space) + form = inner(grad(u), grad(v)) * dx(scheme="canonical") + + optimized, = compile_form(form, parameters={"mode": "spectral"}) + collect_monomials = tsfc.spectral.collect_monomials + monkeypatch.setattr( + tsfc.spectral, "collect_monomials", + lambda expressions, classifier, _: collect_monomials( + expressions, classifier)) + expanded, = compile_form(form, parameters={"mode": "spectral"}) + + assert optimized.flop_count <= expanded.flop_count + + +def test_sum_factorisation_order() -> None: + """Select the least-cost quadrature contraction ordering.""" + i, j, q0, q1 = (gem.Index(extent=4) for _ in range(4)) + inner = gem.Indexed(gem.Variable("inner", (4, 4)), (i, q0)) + outer = gem.Indexed( + gem.Variable("outer", (4, 4, 4)), (i, j, q1)) + variable = gem.Indexed( + gem.Variable("result", (4, 4)), (i, j)) + monomial_sum = MonomialSum() + monomial_sum.add((q0, q1), (inner * outer,), one) + + separate, _ = _plans(((variable, monomial_sum),), (q1, q0)) + (_, expression), = separate + candidates = [ + estimate_cost((tsfc.spectral.sum_factorise( + variable, ordering, monomial_sum),)) + for ordering in ((q1, q0), (q0, q1)) + ] + + assert estimate_cost((expression,)) == min(candidates) + + +def test_shared_contraction_ordering_bounds_storage(monkeypatch) -> None: + """Share one contraction ordering when separate ones cost storage.""" + cell = TensorProductCell(quadrilateral, interval) + mesh = Mesh(VectorElement('Q', cell, 1)) + space = FunctionSpace(mesh, FiniteElement('NCE', cell, 3)) + u = TrialFunction(space) + v = TestFunction(space) + form = action(dot(curl(u), curl(v)) * dx, Coefficient(space)) + + def stored(kernel): + temporaries = kernel.ast.default_entrypoint.temporary_variables + return sum( + numpy.prod(temporary.shape, dtype=int) + for temporary in temporaries.values() + if temporary.shape and not (temporary.read_only + and temporary.initializer is not None)) + + chosen, = compile_form(form, parameters={'mode': 'spectral'}) + + # Deny the selection its shared-ordering candidate, leaving the plan + # that lets every assignment minimise its own arithmetic. + plans = tsfc.spectral._plans + monkeypatch.setattr( + tsfc.spectral, '_plans', + lambda assignments, quadrature_indices: 2 * plans( + assignments, quadrature_indices)[:1]) + separate, = compile_form(form, parameters={'mode': 'spectral'}) + + assert stored(chosen) < stored(separate) + + if __name__ == "__main__": import os import sys diff --git a/tsfc/loopy.py b/tsfc/loopy.py index d4a31a36cb..4304e58422 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -123,6 +123,7 @@ def __init__(self, target=None): self.indices = {} # indices for declarations and referencing values, from ImperoC self.active_indices = {} # gem index -> pymbolic variable self.index_extent = OrderedDict() # pymbolic variable for indices -> extent + self.index_parents = {} # iname -> parent inames bounding a jagged index self.gem_to_pymbolic = {} # gem node -> pymbolic variable self.name_gen = UniqueNameGenerator() self.target = target @@ -257,7 +258,7 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name instructions, event_name, preamble = profile_insns(kernel_name, instructions, log) # Create domains - domains = create_domains(ctx.index_extent.items()) + domains = create_domains(ctx.index_extent.items(), ctx.index_parents) # Create loopy kernel knl = lp.make_kernel( @@ -276,16 +277,30 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name return knl, event_name -def create_domains(indices): - """ Create ISL domains from indices +def create_domains(indices, index_parents=None): + """Create ISL domains for independent and dependent indices. - :arg indices: iterable of (index_name, extent) pairs - :returns: A list of ISL sets representing the iteration domain of the indices.""" + Parameters + ---------- + indices : iterable of tuple + Index names and their static extents. + index_parents : mapping, optional + Parent inames for simplex-lattice bounds. + Returns + ------- + list of isl.Set + Iteration domains for Loopy. + """ domains = [] for idx, extent in indices: - inames = isl.make_zero_and_vars([idx]) - domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(inames[0] + extent)))) + parents = index_parents.get(idx, ()) if index_parents else () + inames = isl.make_zero_and_vars([idx], parents) + bound = inames[0] + extent + for parent in parents: + bound = bound - inames[parent] + domains.append(inames[0].le_set(inames[idx]) + & inames[idx].lt_set(bound)) if not domains: domains = [isl.BasicSet("[] -> {[]}")] @@ -316,6 +331,13 @@ def statement_for(tree, ctx): assert extent idx = ctx.name_gen(ctx.index_names[tree.index]) ctx.index_extent[idx] = extent + if isinstance(tree.index, gem.JaggedIndex) and \ + all(parent in ctx.active_indices for parent in tree.index.parents): + # Tighten the loop bound of a jagged index nested inside its parents. + # If a parent loop is not in scope, the rectangular bound `extent` + # remains correct: jagged expressions are zero-padded. + ctx.index_parents[idx] = tuple(ctx.active_indices[parent].name + for parent in tree.index.parents) with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active: return statement(tree.children[0], ctx_active) @@ -360,10 +382,19 @@ def statement_evaluate(leaf, ctx): elif isinstance(expr, gem.Constant): return [] elif isinstance(expr, gem.ComponentTensor): - idx = ctx.gem_to_pym_multiindex(expr.multiindex) + implicit_indices = {} + value_indices = [] + for index in expr.multiindex: + if index in ctx.active_indices: + value_indices.append(ctx.active_indices[index]) + else: + value, = ctx.gem_to_pym_multiindex((index,)) + implicit_indices[index] = value + value_indices.append(value) + value_indices = tuple(value_indices) var, sub_idx = ctx.pymbolic_variable_and_destruct(expr) - lhs = p.Subscript(var, sub_idx + idx) - with active_indices(dict(zip(expr.multiindex, idx)), ctx) as ctx_active: + lhs = p.Subscript(var, sub_idx + value_indices) + with active_indices(implicit_indices, ctx) as ctx_active: return [lp.Assignment(lhs, expression(expr.children[0], ctx_active), within_inames=ctx_active.active_inames())] elif isinstance(expr, gem.Inverse): idx = ctx.pymbolic_multiindex(expr.shape) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 69e471104e..d503b74873 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -1,10 +1,25 @@ +"""Apply structure-preserving optimization to finite element integrals. + +The order of transformations is part of the algorithm. Argument +factorization first identifies finite element linear maps without expanding +their basis transformations. Delta cancellation then exposes the legal +contractions. Sum factorization places quadrature reductions, and COFFEE +eliminates scalar sharing at every reduction level. In this form, +sum-factorization is generalized code motion on a spectral loop nest rather +than a separate algebraic optimization pipeline. +""" + from collections import OrderedDict, defaultdict, namedtuple from functools import partial -from itertools import chain, zip_longest +from itertools import chain, permutations, zip_longest + +import numpy -from gem.gem import Delta, Indexed, Sum, index_sum, one +from gem import impero_utils +from gem.gem import Conditional, Delta, Indexed, Node, Sum, index_sum, one +from gem.contraction import estimate_cost from gem.node import Memoizer, MemoizerArg -from gem.optimise import filtered_replace_indices +from gem.optimise import constant_fold_zero, filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import replace_division, unroll_indexsum from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials @@ -17,6 +32,11 @@ 'quadrature_multiindex', 'argument_indices']) +Plan = tuple[tuple[Node, Node], ...] + +# Cache-resident working set: 8192 doubles = 64 KiB of contraction temporaries. +storage_budget = 8192 + def Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): """Constructs an integral representation for each GEM integrand @@ -52,6 +72,188 @@ def _delta_inside(node, self): for child in node.children) +def _declared_storage(plan: Plan, quadrature_indices: tuple) -> int: + """Count the temporary entries this plan makes Impero declare. + + Only a schedule fixes how wide a temporary must be, because that width + follows from the loop nest a value outlives, which an expression DAG + does not record. + + Parameters + ---------- + plan + Output variables and their factorized GEM expressions. + quadrature_indices + Every quadrature index of the integral, in source order. + + Returns + ------- + int + Scalar entries in the temporaries, or zero if the plan schedules + to nothing. + + """ + variables = [variable for variable, _ in plan] + expressions = impero_utils.preprocess_gem( + constant_fold_zero([expression for _, expression in plan]), + **finalise_options) + ordering = quadrature_indices + tuple(chain.from_iterable( + variable.index_ordering() for variable in variables)) + try: + impero_c = impero_utils.compile_gem( + list(zip(variables, expressions)), ordering, remove_zeros=True) + except impero_utils.NoopError: + return 0 + return sum( + numpy.prod([index.extent for index in impero_c.indices[temporary]], + dtype=int) + for temporary in impero_c.temporaries) + + +def _factorise( + pairs: tuple[tuple[Node, Node], ...], + preserve_maps: bool) -> tuple[tuple[Node, MonomialSum], ...]: + """Argument factorize and delta cancel one representation of the maps. + + Parameters + ---------- + pairs + Output variables and their integral expressions. + preserve_maps + Keep one-axis sums as finite element linear operands when true; + expose their scalar polynomial structure when false. + + Returns + ------- + tuple of tuple + Output variables and their delta-cancelled monomial sums. + + """ + index_replacer = MemoizerArg(filtered_replace_indices) + delta_inside = Memoizer(_delta_inside) + narrow_variables = OrderedDict() + delta_simplified = defaultdict(MonomialSum) + + groups = groupby( + pairs, key=lambda pair: frozenset(pair[0].free_indices)) + for free_indices, pair_group in groups: + variables, expressions = zip(*pair_group) + argument_indices = set(free_indices) + classifier = partial( + classify, argument_indices, delta_inside=delta_inside) + monomial_sums = collect_monomials( + expressions, classifier, + argument_indices if preserve_maps else ()) + for variable, monomial_sum in zip(variables, monomial_sums): + for monomial in monomial_sum: + var, indices, atomics, rest = delta_elimination( + variable, *monomial, index_replacer) + narrow_variables.setdefault(var) + delta_simplified[var].add( + indices, atomics, rest) + + return tuple((variable, delta_simplified[variable]) + for variable in narrow_variables) + + +def _plans( + assignments: tuple[tuple[Node, MonomialSum], ...], + quadrature_indices: tuple) -> tuple[Plan, Plan]: + """Place quadrature reductions, one plan per contraction strategy. + + Separate orderings minimize arithmetic; one shared ordering spans the + assignments over a single loop nest, keeping the values they share + narrow. Both searches are exhaustive in the quadrature axes, of which + a reference cell supplies one per direction. + + Parameters + ---------- + assignments + Output variables and their delta-cancelled monomial sums. + quadrature_indices + Every quadrature index of the integral, in source order. + + Returns + ------- + separate + Plan giving each assignment its cheapest ordering. + shared + Plan contracting every assignment in one common ordering. + + """ + contracted = [] + for _, monomial_sum in assignments: + summed = set(chain.from_iterable( + monomial.sum_indices for monomial in monomial_sum)) + contracted.append( + tuple(index for index in quadrature_indices if index in summed)) + + factorised = { + (position, ordering): sum_factorise(variable, ordering, monomial_sum) + for position, (variable, monomial_sum) in enumerate(assignments) + for ordering in permutations(contracted[position])} + + def plan(orderings) -> Plan: + return tuple( + (variable, factorised[position, orderings[position]]) + for position, (variable, _) in enumerate(assignments)) + + separate = plan([ + min(permutations(axes), + key=lambda ordering, p=position: estimate_cost( + (factorised[p, ordering],))) + for position, axes in enumerate(contracted)]) + shared = min( + (plan([tuple(index for index in ordering if index in axes) + for axes in contracted]) + for ordering in permutations(quadrature_indices)), + key=lambda candidate: estimate_cost( + expression for _, expression in candidate)) + return separate, shared + + +def _select_plan( + pairs: tuple[tuple[Node, Node], ...], + quadrature_indices: tuple) -> Plan: + """Minimize arithmetic among plans whose temporaries fit the budget. + + Preserving a linear map exposes tabulation reuse, expanding it exposes + scalar factorization, and neither dominates. Overflowing cache is a + cliff rather than a gradient, so storage bounds the search instead of + trading against arithmetic. When no plan fits, take the narrowest. + + Parameters + ---------- + pairs + Output variables and their integral expressions. + quadrature_indices + Every quadrature index of the integral, in source order. + + Returns + ------- + Plan + Optimized output variables and GEM expressions. + + """ + candidates = list(dict.fromkeys( + plan + for preserve_maps in (False, True) + for plan in _plans(_factorise(pairs, preserve_maps), + quadrature_indices))) + if len(candidates) == 1: + return candidates[0] + + storage = {plan: _declared_storage(plan, quadrature_indices) + for plan in candidates} + feasible = [plan for plan in candidates + if storage[plan] <= storage_budget] + if not feasible: + return min(candidates, key=storage.get) + return min(feasible, + key=lambda plan: estimate_cost( + expression for _, expression in plan)) + + def flatten(var_reps, index_cache): quadrature_indices = OrderedDict() @@ -76,60 +278,36 @@ def flatten(var_reps, index_cache): # Split Concatenate nodes pairs = unconcatenate(pairs, cache=index_cache) - def group_key(pair): - variable, expression = pair - return frozenset(variable.free_indices) - - # Common memoizer to remove ComponentTensors - index_replacer = MemoizerArg(filtered_replace_indices) - # Common memoizer to test for Deltas inside expressions - delta_inside = Memoizer(_delta_inside) - # Variable ordering after delta cancellation - narrow_variables = OrderedDict() - # Assignments are variable -> MonomialSum map - delta_simplified = defaultdict(MonomialSum) - # Group assignment pairs by argument indices - for free_indices, pair_group in groupby(pairs, group_key): - variables, expressions = zip(*pair_group) - # Argument factorise expressions - classifier = partial(classify, set(free_indices), delta_inside=delta_inside) - monomial_sums = collect_monomials(expressions, classifier) - # For each monomial, apply delta cancellation and insert - # result into delta_simplified. - for variable, monomial_sum in zip(variables, monomial_sums): - for monomial in monomial_sum: - var, s, a, r = delta_elimination(variable, *monomial, index_replacer) - narrow_variables.setdefault(var) - delta_simplified[var].add(s, a, r) - - # Final factorisation - for variable in narrow_variables: - monomial_sum = delta_simplified[variable] - # Collect sum indices applicable to the current MonomialSum - sum_indices = set(chain.from_iterable(m.sum_indices for m in monomial_sum)) - # Put them in a deterministic order - sum_indices = [i for i in quadrature_indices if i in sum_indices] - # Sort for increasing index extent, this obtains the good - # factorisation for triangle x interval cells. Python sort is - # stable, so in the common case when index extents are equal, - # the previous deterministic ordering applies which is good - # for getting smaller temporaries. - sum_indices = sorted(sum_indices, key=lambda index: index.extent) - # Apply sum factorisation combined with COFFEE technology - expression = sum_factorise(variable, sum_indices, monomial_sum) - yield (variable, expression) + return _select_plan(tuple(pairs), tuple(quadrature_indices)) -finalise_options = dict(replace_delta=False) +finalise_options = dict(replace_delta=True, remove_componenttensors=False) def classify(argument_indices, expression, delta_inside): - """Classifier for argument factorisation""" + """Classify one expression for multilinear factorization. + + Parameters + ---------- + argument_indices : set of Index + Free argument indices. + expression : Node + Expression to classify. + delta_inside : callable + Predicate detecting delta nodes. + Returns + ------- + str + Refactorization label. + """ n = len(argument_indices.intersection(expression.free_indices)) if n == 0: return OTHER elif n == 1: - if isinstance(expression, (Delta, Indexed)) and not delta_inside(expression): + if isinstance(expression, Conditional): + return ATOMIC + if isinstance(expression, (Delta, Indexed)) \ + and not delta_inside(expression): return ATOMIC else: return COMPOUND @@ -162,36 +340,14 @@ def prune(factors): variable = factors.pop() args = [f for f in factors if f != one] - assert set(var_indices) == set(variable.free_indices) + assert set(var_indices) <= set(variable.free_indices) + # A delta may replace a variable index by a contraction index. That + # index now describes a scatter in the assignment, not a sum. + sum_indices = [i for i in sum_indices if i not in variable.free_indices] return variable, sum_indices, args, rest def sum_factorise(variable, tail_ordering, monomial_sum): - if tail_ordering: - key_ordering = OrderedDict() - sub_monosums = defaultdict(MonomialSum) - for sum_indices, atomics, rest in monomial_sum: - # Pull out those sum indices that are not contained in the - # tail ordering, together with those atomics which do not - # share free indices with the tail ordering. - # - # Based on this, split the monomial sum, then recursively - # optimise each sub monomial sum with the first tail index - # removed. - tail_indices = tuple(i for i in sum_indices if i in tail_ordering) - tail_atomics = tuple(a for a in atomics - if set(tail_indices) & set(a.free_indices)) - head_indices = tuple(i for i in sum_indices if i not in tail_ordering) - head_atomics = tuple(a for a in atomics if a not in tail_atomics) - key = (head_indices, head_atomics) - key_ordering.setdefault(key) - sub_monosums[key].add(tail_indices, tail_atomics, rest) - sub_monosums = [(k, sub_monosums[k]) for k in key_ordering] - - monomial_sum = MonomialSum() - for (sum_indices, atomics), monosum in sub_monosums: - new_rest = sum_factorise(variable, tail_ordering[1:], monosum) - monomial_sum.add(sum_indices, atomics, new_rest) - - # Use COFFEE algorithm to optimise the monomial sum - return optimise_monomial_sum(monomial_sum, variable.index_ordering()) + return optimise_monomial_sum( + monomial_sum, variable.index_ordering(), + tuple(tail_ordering))