diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index 40bcfee39b..489e1c28a0 100644 --- a/.github/actions/install/action.yml +++ b/.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 diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml index 2e1ea2735f..2accda6dcd 100644 --- a/.github/workflows/core.yml +++ b/.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 diff --git a/benchmarks/bernstein_laplacian.py b/benchmarks/bernstein_laplacian.py new file mode 100755 index 0000000000..d1820269c3 --- /dev/null +++ b/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("") + 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() diff --git a/pyproject.toml b/pyproject.toml index 0c942f8a82..a22354d8aa 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/firedrake/regression/test_quadrature.py b/tests/firedrake/regression/test_quadrature.py index 225a4b244d..bb7189ddc4 100644 --- a/tests/firedrake/regression/test_quadrature.py +++ b/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) diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py index 8d0bc79655..7dd97e2ba1 100644 --- a/tests/tsfc/test_codegen.py +++ b/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 diff --git a/tests/tsfc/test_impero_loopy_flop_counts.py b/tests/tsfc/test_impero_loopy_flop_counts.py index 240067bd6e..4bbdef0dc6 100644 --- a/tests/tsfc/test_impero_loopy_flop_counts.py +++ b/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) 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 85d9729e81..cff6e3b849 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -1,12 +1,16 @@ +from math import comb + import numpy import pytest from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, action, interval, triangle, - quadrilateral, hexahedron, curl, dot, div, grad) + quadrilateral, hexahedron, tetrahedron, curl, dot, div, + grad, inner) from finat.ufl import (FiniteElement, VectorElement, EnrichedElement, TensorProductElement, HCurlElement, HDivElement) +import tsfc.spectral from tsfc import compile_form @@ -79,6 +83,11 @@ def count_storage(form): and temporary.initializer is not None)) +def count_loops(form): + kernel, = compile_form(form, parameters=dict(mode='spectral')) + return len(kernel.ast.default_entrypoint.all_inames()) + + @pytest.mark.parametrize(('cell', 'order'), [(quadrilateral, 5), (TensorProductCell(interval, interval), 5), @@ -203,6 +212,250 @@ def test_equivalent_cells(cell, equivalent_cell, degree): assert count_flops(action(a)) == count_flops(action(b)) +@pytest.fixture +def expanded(monkeypatch): + """Force the expanded representation, for comparison.""" + def force(monkeypatch=monkeypatch): + collect_monomials = tsfc.spectral.collect_monomials + monkeypatch.setattr( + tsfc.spectral, "collect_monomials", + lambda expressions, classifier, _: collect_monomials( + expressions, classifier)) + return force + + +def piola_helmholtz(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement('RT', cell, degree)) + u = TrialFunction(V) + v = TestFunction(V) + return (inner(u, v) + inner(div(u), div(v)))*dx + + +@pytest.mark.parametrize('cell', [triangle, tetrahedron], + ids=lambda cell: cell.cellname) +@pytest.mark.parametrize('degree', [1, 2, 3]) +def test_piola_map_is_preserved(cell, degree, expanded): + # Test and trial apply the same Piola map, so preserving it evaluates + # the physical basis once instead of pushing the geometry through both + # argument axes. + form = piola_helmholtz(cell, degree) + selected = count_flops(form) + expanded() + assert selected < count_flops(form) + + +@pytest.mark.parametrize('cell', [triangle, tetrahedron], + ids=lambda cell: cell.cellname) +@pytest.mark.parametrize('degree', [1, 3]) +def test_preserving_a_map_is_never_worse(cell, degree, expanded): + # Expanding a map exposes scalar factorisation of its entries, which at + # some degrees beats sharing it. Selection costs both, so neither + # representation may regress the other. + form = helmholtz(cell, degree) + selected = count_flops(form) + expanded() + assert selected <= count_flops(form) + + +@pytest.mark.parametrize('cell', [triangle, tetrahedron], + ids=lambda cell: cell.cellname) +@pytest.mark.parametrize('degree', [1, 3]) +def test_shared_map_is_tabulated_in_one_loop(cell, degree, expanded): + # A map that both argument axes share is tabulated once, so it must be + # tabulated in one loop. An index per axis fissions the loop nest that + # the expanded representation keeps whole, which costs more than the + # flops it saves. + form = piola_helmholtz(cell, degree) + selected = count_loops(form) + expanded() + assert selected <= count_loops(form) + + +def bernstein_mass( + cell: object, degree: int, scheme: str = "collapsed") -> object: + """Construct a Bernstein mass form. + + Parameters + ---------- + cell + Reference simplex. + degree + Polynomial degree. + scheme + Quadrature scheme. + + Returns + ------- + object + UFL bilinear form. + """ + mesh = Mesh(VectorElement("CG", cell, 1)) + space = FunctionSpace(mesh, FiniteElement("Bernstein", cell, degree)) + u = TrialFunction(space) + v = TestFunction(space) + return inner(u, v) * dx(scheme=scheme) + + +def bernstein_laplacian( + cell: object, degree: int, scheme: str = "collapsed") -> object: + """Construct a Bernstein Laplacian form. + + Parameters + ---------- + cell + Reference simplex. + degree + Polynomial degree. + scheme + Quadrature scheme. + + Returns + ------- + object + UFL bilinear form. + """ + mesh = Mesh(VectorElement("CG", cell, 1)) + space = FunctionSpace(mesh, FiniteElement("Bernstein", cell, degree)) + u = TrialFunction(space) + v = TestFunction(space) + return inner(grad(u), grad(v)) * dx(scheme=scheme) + + +@pytest.mark.parametrize(("cell", "order"), [(triangle, 3), (tetrahedron, 4)]) +def test_bernstein_mass_action(cell: object, order: float) -> None: + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [ + count_flops(action(bernstein_mass(cell, degree))) + for degree in degrees + ] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize( + ("cell", "order"), [(triangle, 3), (tetrahedron, 4.4)]) +def test_bernstein_laplacian_action(cell: object, order: float) -> None: + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [ + count_flops(action(bernstein_laplacian(cell, degree))) + for degree in degrees + ] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +def test_bernstein_laplacian_action_compact_literals() -> None: + degree = 5 + form = action(bernstein_laplacian(tetrahedron, degree)) + kernel, = compile_form(form, parameters={"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() -> None: + import islpy as isl + import loopy as lp + + degree = 10 + collapsed = bernstein_laplacian( + tetrahedron, degree, scheme="collapsed") + canonical = bernstein_laplacian( + tetrahedron, degree, scheme="canonical") + collapsed_kernel, = compile_form( + collapsed, parameters={"mode": "spectral"}) + canonical_kernel, = compile_form( + canonical, parameters={"mode": "spectral"}) + + assert collapsed_kernel.flop_count < canonical_kernel.flop_count + + entrypoint = collapsed_kernel.ast.default_entrypoint + collapsed_temporaries = tuple(entrypoint.temporary_variables.values()) + canonical_temporaries = tuple( + canonical_kernel.ast.default_entrypoint.temporary_variables.values()) + collapsed_shapes = [ + temporary.shape + for temporary in collapsed_temporaries + ] + assert max(map(len, collapsed_shapes)) <= 5 + assert sum(map(numpy.prod, collapsed_shapes)) \ + < sum(numpy.prod(temporary.shape) + for temporary in canonical_temporaries) + + 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 + + +def test_bernstein_bilinear_contracts_into_the_output() -> None: + # The last one-dimensional contraction runs inside both argument + # lattices and accumulates into the scattered output. Nothing between + # the contraction and the output holds a whole element matrix, so the + # dominant loop nest is two simplex lattices plus one quadrature loop. + degree = 10 + dimension = 3 + kernel, = compile_form( + bernstein_mass(tetrahedron, degree, scheme="collapsed"), + parameters={"mode": "spectral"}) + entrypoint = kernel.ast.default_entrypoint + + nodes = comb(degree + dimension, dimension) + writable = [temporary + for temporary in entrypoint.temporary_variables.values() + if temporary.initializer is None] + assert max(numpy.prod(temporary.shape, dtype=int) + for temporary in writable) < nodes ** 2 + + scatters = [instruction for instruction in entrypoint.instructions + if "A" in instruction.write_dependency_names()] + assert scatters + scatter_inames, = {frozenset(instruction.within_inames) + for instruction in scatters} + arguments = 2 + assert len(scatter_inames) == arguments * dimension + # The quadrature reduction is nested inside that scatter. + assert any(instruction.within_inames > scatter_inames + for instruction in entrypoint.instructions) + + +@pytest.mark.parametrize(("cell", "order"), [(triangle, 5), (tetrahedron, 7)]) +def test_bernstein_mass_bilinear(cell: object, order: float) -> None: + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [ + count_flops(bernstein_mass(cell, degree)) + for degree in degrees + ] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize(("cell", "order"), [(triangle, 5), (tetrahedron, 7)]) +def test_bernstein_laplacian_bilinear( + cell: object, order: float) -> None: + degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8)) + flops = [ + count_flops(bernstein_laplacian(cell, degree)) + for degree in degrees + ] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + if __name__ == "__main__": import os import sys diff --git a/tsfc/coffee_mode.py b/tsfc/coffee_mode.py index 632b915b41..845b52ed1c 100644 --- a/tsfc/coffee_mode.py +++ b/tsfc/coffee_mode.py @@ -2,7 +2,8 @@ from gem.node import traversal, Memoizer from gem.gem import Failure, Sum, index_sum -from gem.optimise import replace_division, unroll_indexsum +from gem.optimise import (tabulate_indirect_contractions, replace_division, + unroll_indexsum) from gem.refactorise import collect_monomials from gem.unconcatenate import unconcatenate from gem.coffee import optimise_monomial_sum @@ -78,4 +79,5 @@ def optimise_expressions(expressions, argument_indices): classifier = partial(spectral.classify, set(argument_indices), delta_inside=Memoizer(spectral._delta_inside)) monomial_sums = collect_monomials(expressions, classifier) - return [optimise_monomial_sum(ms, argument_indices) for ms in monomial_sums] + return [tabulate_indirect_contractions( + optimise_monomial_sum(ms, argument_indices)) for ms in monomial_sums] diff --git a/tsfc/fem.py b/tsfc/fem.py index 943089052e..d1beaec041 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -16,6 +16,7 @@ from finat.point_set import PointSet, PointSingleton from finat.quadrature import make_quadrature from finat.element_factory import as_fiat_cell, create_element +from gem.driver import contraction from gem.node import traversal from gem.optimise import constant_fold_zero, ffc_rounding from gem.unconcatenate import unconcatenate @@ -781,7 +782,7 @@ def take_singleton(xs): for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) value = gem.IndexSum(gem.Product(expr, var), indices) - summands.append(gem.optimise.contraction(value)) + summands.append(contraction(value)) optimised_value = gem.optimise.make_sum(summands) value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 5d61a916aa..0d95a8e075 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -9,13 +9,16 @@ from ufl.domain import extract_unique_domain import gem +import gem.gem import gem.impero_utils as impero_utils import petsctools 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.driver import unflatten_returns from gem.optimise import constant_fold_zero from gem.optimise import remove_componenttensors as prune from numpy import asarray @@ -210,6 +213,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 = [] @@ -231,11 +236,18 @@ def compile_gem(self, ctx): active_variables = gem.extract_type(expressions, gem.Variable) # Construct ImperoC assignments = list(zip(return_variables, expressions)) - index_ordering = get_index_ordering(ctx['quadrature_indices'], return_variables) - try: - impero_c = impero_utils.compile_gem(assignments, index_ordering, remove_zeros=True) - except impero_utils.NoopError: + candidates = [] + for index_ordering in index_orderings( + ctx['quadrature_indices'], return_variables, assignments): + try: + candidates.append(impero_utils.compile_gem( + assignments, index_ordering, remove_zeros=True)) + except impero_utils.NoopError: + candidates.append(None) + if any(candidate is None for candidate in candidates): impero_c = None + else: + impero_c = min(candidates, key=_storage_cost) return impero_c, oriented, needs_cell_sizes, tabulations, active_variables def fem_config(self): @@ -344,6 +356,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: @@ -360,10 +377,73 @@ def set_quad_rule(params, cell, integral_type, functions): type(quad_rule)) -def get_index_ordering(quadrature_indices, return_variables): +def _spans_output(node, output_indices): + """Check whether a node reduces a summand that spans the whole output.""" + return (isinstance(node, gem.IndexSum) + and set(node.free_indices) == output_indices) + + +def _terminal_reductions(assignments): + """Quadrature indices whose reduction cannot shrink what it accumulates. + + Such a reduction has no tensor smaller than the output to accumulate + into, so running its loop outside the argument loops costs an + output-shaped temporary. Placing it innermost instead makes the + accumulator a scalar, at the cost of giving every stage that feeds it + a quadrature axis, so neither placement dominates and both are costed. + """ + terminal = set() + for variable, expression in assignments: + output_indices = set(variable.free_indices) + if _spans_output(expression, output_indices): + # A root reduction accumulates straight into the output, so + # the outer placement costs nothing extra. + continue + for node in traversal((expression,)): + if _spans_output(node, output_indices): + terminal.update(node.multiindex) + return terminal + + +def get_index_ordering(quadrature_indices, return_variables, assignments=()): + """Return the single best-guess outermost loop ordering. + + For callers that compile one ordering rather than costing several. + """ + return index_orderings(quadrature_indices, return_variables, + assignments)[0] + + +def index_orderings(quadrature_indices, return_variables, assignments=()): + """Return the candidate outermost loop orderings, best guess first. + + Quadrature loops run outside the argument loops, which keeps every + intermediate contraction stage free of a quadrature axis. When an + output-shaped reduction cannot accumulate into the output itself, + that placement also forces an output-shaped temporary, so the + ordering that runs it innermost is offered as an alternative. + """ split_argument_indices = tuple(chain(*(var.index_ordering() for var in return_variables))) - return tuple(quadrature_indices) + split_argument_indices + quadrature_indices = tuple(quadrature_indices) + default = quadrature_indices + split_argument_indices + terminal = _terminal_reductions(assignments) + if not terminal: + return (default,) + head = tuple(i for i in quadrature_indices if i not in terminal) + tail = tuple(i for i in quadrature_indices if i in terminal) + return (default, head + split_argument_indices + tail) + + +def _storage_cost(impero_c): + """Total declared size of an Impero program's temporaries.""" + total = 0 + for temporary in impero_c.temporaries: + if isinstance(temporary, gem.gem.Constant): + continue + shape, _ = gem.compact_index_layout(tuple(impero_c.indices[temporary])) + total += numpy.prod(shape + temporary.shape, dtype=int) + return total def get_index_names(quadrature_indices, argument_multiindices, index_cache): diff --git a/tsfc/loopy.py b/tsfc/loopy.py index d4a31a36cb..abba465ff6 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -8,6 +8,7 @@ from collections import defaultdict, OrderedDict from gem import gem, impero as imp +from gem.jagged import compact_index_layout, simplex_lattice_ranks from gem.node import Memoizer import islpy as isl @@ -123,7 +124,11 @@ 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.tabulated = (None, ()) # axes and inames of the preceding tabulation + self.index_parents = {} # iname -> parent inames bounding a jagged index self.gem_to_pymbolic = {} # gem node -> pymbolic variable + self.compact_indices = {} # temporary -> compact index layout + self.lattice_ranks = {} # lattice shape -> rank table and its entries self.name_gen = UniqueNameGenerator() self.target = target self.loop_priorities = set() # used to avoid disadvantageous loop interchanges @@ -173,6 +178,12 @@ def pymbolic_variable_and_destruct(self, node): # Generate pym variable or subscript def pymbolic_variable(self, node): pym = self._gem_to_pym_var(node) + if node in self.compact_indices: + indices = tuple( + self.lattice_rank(item) if isinstance(item, tuple) + else self.active_indices[item] + for item in self.compact_indices[node]) + return p.Subscript(pym, indices) if indices else pym if node in self.indices: indices = self.fetch_multiindex(self.indices[node]) if indices: @@ -188,6 +199,38 @@ def _gem_to_pym_var(self, node): self.gem_to_pymbolic[node] = pym return pym + def lattice_rank(self, component): + """Look up the compact rank of the active simplex lattice point. + + A simplex lattice is stored along one compact dimension, so a + temporary indexed by the lattice is subscripted by the rank of + the point rather than by the lattice indices themselves. The + ranks are tabulated once per lattice shape, which keeps the + lattice visible in the AST and keeps the polynomial that defines + the rank out of the loop body. + """ + # Equal degree and dimension describe a simplex lattice fully, + # so every lattice of one shape shares a single table. + shape = (component[0].extent, len(component)) + try: + table, _ = self.lattice_ranks[shape] + except KeyError: + table = p.Variable(self.name_gen("lattice_rank")) + self.lattice_ranks[shape] = ( + table, simplex_lattice_ranks(component)) + return p.Subscript(table, tuple(self.active_indices[index] + for index in component)) + + def lattice_rank_tables(self, address_space): + """Declare a read-only table for each tabulated simplex lattice.""" + return [ + lp.TemporaryVariable( + table.name, shape=ranks.shape, dtype=ranks.dtype, + initializer=ranks, address_space=address_space, + read_only=True) + for table, ranks in self.lattice_ranks.values() + ] + def active_inames(self): # Return all active indices return frozenset([i.name for i in self.active_indices.values()]) @@ -234,18 +277,31 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name # Create arguments data = list(args) - for i, (temp, dtype) in enumerate(assign_dtypes(impero_c.temporaries, scalar_type)): + descriptors = [] + for temp, dtype in assign_dtypes(impero_c.temporaries, scalar_type): + if isinstance(temp, gem.Constant): + shape, layout = temp.shape, None + else: + shape, layout = compact_index_layout( + tuple(ctx.indices[temp])) + shape += temp.shape + descriptors.append((temp, dtype, shape, layout)) + for i, (temp, dtype, shape, layout) in enumerate(descriptors): name = "t%d" % i if isinstance(temp, gem.Constant): data.append(lp.TemporaryVariable(name, shape=temp.shape, dtype=dtype, initializer=temp.array, address_space=lp.AddressSpace.LOCAL, read_only=True)) else: - shape = tuple([i.extent for i in ctx.indices[temp]]) + temp.shape - data.append(lp.TemporaryVariable(name, shape=shape, dtype=dtype, initializer=None, address_space=lp.AddressSpace.LOCAL, read_only=False)) + data.append(lp.TemporaryVariable( + name, shape=shape, dtype=dtype, initializer=None, + address_space=lp.AddressSpace.LOCAL, read_only=False)) + ctx.compact_indices[temp] = layout ctx.gem_to_pymbolic[temp] = p.Variable(name) # Create instructions instructions = statement(impero_c.tree, ctx) + data.extend(ctx.lattice_rank_tables(lp.AddressSpace.LOCAL)) + # add a no-op touching all kernel arguments to make sure they # are not silently dropped noop = lp.CInstruction( @@ -257,7 +313,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 +332,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("[] -> {[]}")] @@ -304,10 +374,27 @@ def statement(tree, ctx): raise AssertionError("cannot generate loopy from %s" % type(tree)) +def tabulated_axes(tree): + """The axes a statement binds, if it tabulates a tensor over its own.""" + if isinstance(tree, imp.Evaluate) \ + and isinstance(tree.expression, gem.ComponentTensor): + return tree.expression.multiindex + return None + + @statement.register(imp.Block) def statement_block(tree, ctx): - from itertools import chain - return list(chain(*(statement(child, ctx) for child in tree.children))) + # Tabulations of the same axes share a loop while they stay adjacent. + # Anything between them is a statement the schedule placed outside that + # loop, so the loop has to close before it and reopen after. + instructions = [] + ctx.tabulated = (None, ()) + for child in tree.children: + instructions.extend(statement(child, ctx)) + if tabulated_axes(child) is None: + ctx.tabulated = (None, ()) + ctx.tabulated = (None, ()) + return instructions @statement.register(imp.For) @@ -316,6 +403,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 +454,24 @@ 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_axes = tuple( + index for index in expr.multiindex + if index not in ctx.active_indices) + axes, implicit_values = ctx.tabulated + if axes != implicit_axes: + implicit_values = ctx.gem_to_pym_multiindex(implicit_axes) + ctx.tabulated = (implicit_axes, implicit_values) + implicit_indices = dict(zip(implicit_axes, implicit_values)) + value_indices = [] + for index in expr.multiindex: + if index in ctx.active_indices: + value_indices.append(ctx.active_indices[index]) + else: + value_indices.append(implicit_indices[index]) + 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) @@ -547,7 +655,18 @@ def _expression_variable(expr, ctx): @_expression.register(gem.Indexed) def _expression_indexed(expr, ctx): rank = ctx.fetch_multiindex(expr.multiindex) - var = expression(expr.children[0], ctx) + aggregate, = expr.children + if (isinstance(aggregate, gem.ComponentTensor) + and aggregate not in ctx.gem_to_pymbolic): + body, = aggregate.children + if body in ctx.gem_to_pymbolic: + replacements = dict(zip(aggregate.multiindex, expr.multiindex)) + multiindex = tuple(replacements.get(index, index) + for index in ctx.indices[body]) + rank = ctx.fetch_multiindex(multiindex) + return p.Subscript(ctx._gem_to_pym_var(body), rank) + + var = expression(aggregate, ctx) if isinstance(var, p.Subscript): rank = var.index + rank var = var.aggregate diff --git a/tsfc/spectral.py b/tsfc/spectral.py index a521fdb2fd..628cbcff95 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -4,7 +4,10 @@ from gem.gem import Delta, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg -from gem.optimise import filtered_replace_indices +from gem.cost import estimate_cost +from gem.optimise import (cancel_nested_deltas, + tabulate_indirect_contractions, + filtered_replace_indices, has_linear_maps) 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 @@ -34,6 +37,10 @@ def Integrals(expressions, quadrature_multiindex, argument_multiindices, paramet # Rewrite: a / b => a * (1 / b) expressions = replace_division(expressions) + # Cancel the Deltas that select a basis transformation's columns, so that + # monomial collection sees the resulting gather rather than the Delta. + expressions = [cancel_nested_deltas(e) for e in expressions] + # Unroll max_extent = parameters["unroll_indexsum"] if max_extent: @@ -52,34 +59,50 @@ def _delta_inside(node, self): for child in node.children) -def flatten(var_reps, index_cache): - quadrature_indices = OrderedDict() - - pairs = [] # assignment pairs - for variable, reps in var_reps: - # Extract argument indices - argument_indices, = set(r.argument_indices for r in reps) - assert set(variable.free_indices) == set(argument_indices) +def _group_key(pair): + variable, expression = pair + return frozenset(variable.free_indices) - # Extract and verify expressions - expressions = [r.expression for r in reps] - assert all(set(e.free_indices) <= set(argument_indices) - for e in expressions) - # Save assignment pair - pairs.append((variable, Sum(*expressions))) +def _preservable(pairs): + """Is there a linear map whose preservation could change a plan? - # Collect quadrature_indices - for r in reps: - quadrature_indices.update(zip_longest(r.quadrature_multiindex, ())) + Parameters + ---------- + pairs : tuple of tuple + Output variables and the integrands assigned to them. - # Split Concatenate nodes - pairs = unconcatenate(pairs, cache=index_cache) + Returns + ------- + bool + Whether any assignment contains a sum over one argument axis. - def group_key(pair): - variable, expression = pair - return frozenset(variable.free_indices) + """ + return any( + has_linear_maps([expression], set(free_indices)) + for free_indices, pair_group in groupby(pairs, _group_key) + for _, expression in pair_group) + + +def _factorise(pairs, quadrature_indices, preserve_maps): + """Factorise the arguments of each assignment and place its reductions. + + Parameters + ---------- + pairs : tuple of tuple + Output variables and the integrands assigned to them. + quadrature_indices : tuple of Index + Every quadrature index of the integral, in source order. + preserve_maps : bool + Keep a sum over one argument axis whole, as a finite element linear + map, rather than distributing it into scalar monomials. + + Returns + ------- + tuple of tuple + Output variables and their factorised GEM expressions. + """ # Common memoizer to remove ComponentTensors index_replacer = MemoizerArg(filtered_replace_indices) # Common memoizer to test for Deltas inside expressions @@ -89,11 +112,14 @@ def group_key(pair): # 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): + for free_indices, pair_group in groupby(pairs, _group_key): variables, expressions = zip(*pair_group) + argument_indices = set(free_indices) + classifier = partial(classify, argument_indices, delta_inside=delta_inside) # Argument factorise expressions - classifier = partial(classify, set(free_indices), delta_inside=delta_inside) - monomial_sums = collect_monomials(expressions, classifier) + monomial_sums = collect_monomials( + expressions, classifier, + argument_indices if preserve_maps else ()) # For each monomial, apply delta cancellation and insert # result into delta_simplified. for variable, monomial_sum in zip(variables, monomial_sums): @@ -103,18 +129,56 @@ def group_key(pair): delta_simplified[var].add(s, a, r) # Final factorisation + plan = [] 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] - # Apply sum factorisation combined with COFFEE technology + # Apply sum factorisation combined with COFFEE technology, then + # tabulate indirect contractions over the whole factorised assignment. expression = sum_factorise(variable, sum_indices, monomial_sum) - yield (variable, expression) + plan.append((variable, tabulate_indirect_contractions(expression))) + return tuple(plan) + + +def flatten(var_reps, index_cache): + quadrature_indices = OrderedDict() + + pairs = [] # assignment pairs + for variable, reps in var_reps: + # Extract argument indices + argument_indices, = set(r.argument_indices for r in reps) + assert set(variable.free_indices) == set(argument_indices) + + # Extract and verify expressions + expressions = [r.expression for r in reps] + assert all(set(e.free_indices) <= set(argument_indices) + for e in expressions) + + # Save assignment pair + pairs.append((variable, Sum(*expressions))) + + # Collect quadrature_indices + for r in reps: + quadrature_indices.update(zip_longest(r.quadrature_multiindex, ())) + + # Split Concatenate nodes + pairs = unconcatenate(pairs, cache=index_cache) + + # Expanding a linear map exposes scalar factorisation across its + # entries, preserving one exposes a tabulation that several argument + # axes share, and neither dominates. Cost both and keep the cheaper, + # skipping the second factorisation when there is no map to preserve. + plans = [_factorise(pairs, quadrature_indices, False)] + if _preservable(pairs): + plans.append(_factorise(pairs, quadrature_indices, True)) + return min(plans, key=lambda plan: estimate_cost( + expression for _, expression in plan)) -finalise_options = dict(replace_delta=False) +finalise_options = dict(replace_delta=False, remove_componenttensors=False) def classify(argument_indices, expression, delta_inside):