From 02f808222ffe0f87500df9a956bf18c6878c9a2b Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 2 Aug 2026 15:14:05 +0100 Subject: [PATCH 01/26] WIP: extract generic sum-factorisation lowering --- benchmarks/johnson_mercier.py | 87 ++++++++++ tests/tsfc/test_pickle_gem.py | 14 ++ tsfc/loopy.py | 23 ++- tsfc/spectral.py | 302 +++++++++++++++++++++++++++++----- 4 files changed, 380 insertions(+), 46 deletions(-) create mode 100644 benchmarks/johnson_mercier.py diff --git a/benchmarks/johnson_mercier.py b/benchmarks/johnson_mercier.py new file mode 100644 index 0000000000..033663fcbe --- /dev/null +++ b/benchmarks/johnson_mercier.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python +"""Measure simplex Johnson--Mercier code generation.""" + +import argparse +import hashlib + +import numpy + +from finat.ufl import FiniteElement, VectorElement +from tsfc import compile_form +from ufl import FunctionSpace, Mesh, TestFunction, TrialFunction, div, dx, inner +from ufl.cell import Cell + + +def compile_target(dim: int) -> object: + """Compile the JM mass-plus-divergence target on a simplex. + + Parameters + ---------- + dim + Topological dimension. + + Returns + ------- + object + Compiled TSFC kernel. + """ + cell = Cell(("triangle", "tetrahedron")[dim - 2]) + mesh = Mesh(VectorElement("CG", cell, 1)) + element = FiniteElement("Johnson-Mercier", cell, 1) + space = FunctionSpace(mesh, element) + u = TrialFunction(space) + v = TestFunction(space) + form = (inner(u, v) + inner(div(u), div(v))) * dx + return compile_form(form, parameters={"mode": "spectral"})[0] + + +def temporary_metrics(kernel: object) -> tuple[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. + array_elements + Total entries in array temporaries. + largest_array + Entries in the largest array temporary. + """ + temporaries = kernel.ast.default_entrypoint.temporary_variables.values() + shapes = [temporary.shape for temporary in temporaries] + sizes = [numpy.prod(shape, dtype=int) for shape in shapes if shape] + return ( + sum(not shape for shape in shapes), + len(sizes), + sum(sizes), + max(sizes, default=0), + ) + + +def main() -> None: + """Print compiler metrics as copyable Markdown.""" + parser = argparse.ArgumentParser() + parser.add_argument("--dims", nargs="+", type=int, default=(2, 3)) + args = parser.parse_args() + print("") + print("| dim | flops | scalar temps | array temps | elements | bytes | largest | AST lines | hash |") + print("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |") + for dim in args.dims: + kernel = compile_target(dim) + source = str(kernel.ast) + digest = hashlib.sha256(source.encode()).hexdigest()[:12] + nscalar, narray, nelem, largest = temporary_metrics(kernel) + print(f"| {dim} | {kernel.flop_count:.0f} | " + f"{nscalar} | {narray} | {nelem} | {8 * nelem} | {largest} | " + f"{len(source.splitlines())} | {digest} |") + + +if __name__ == "__main__": + main() 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/tsfc/loopy.py b/tsfc/loopy.py index d4a31a36cb..e41b8348fc 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,23 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name return knl, event_name -def create_domains(indices): +def create_domains(indices, index_parents=None): """ Create ISL domains from indices :arg indices: iterable of (index_name, extent) pairs + :arg index_parents: optional mapping from index_name to a tuple of parent + index names; the domain of a jagged index is parametrized by its + parents, with upper bound extent minus the sum of the parents. :returns: A list of ISL sets representing the iteration domain of the indices.""" 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 +324,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) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 69e471104e..57a54c5953 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -1,21 +1,27 @@ from collections import OrderedDict, defaultdict, namedtuple from functools import partial +import itertools from itertools import chain, zip_longest -from gem.gem import Delta, Indexed, Sum, index_sum, one +from gem.gem import Conditional, Delta, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination -from gem.optimise import replace_division, unroll_indexsum +from gem.optimise import ( + estimate_cost, factorisation_group_options, replace_division, + unroll_indexsum, +) from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials from gem.unconcatenate import unconcatenate -from gem.coffee import optimise_monomial_sum +from gem.coffee import sum_factorise_monomial_sum from gem.utils import groupby Integral = namedtuple('Integral', ['expression', 'quadrature_multiindex', 'argument_indices']) +FactorisationCandidates = namedtuple( + 'FactorisationCandidates', ['alternatives', 'baseline']) def Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): @@ -52,6 +58,217 @@ def _delta_inside(node, self): for child in node.children) +def _factorisation_candidates( + expression, argument_indices, + delta_inside) -> tuple[FactorisationCandidates, ...]: + """Build alternative pre-expansion grouping plans. + + Parameters + ---------- + expression : Node + Multilinear integrand to factorize. + argument_indices : set of Index + Free argument indices. + delta_inside : callable + Memoized predicate detecting delta nodes. + + Returns + ------- + tuple of FactorisationCandidates + Distinct monomial representations for each independent group of + choices, together with the all-or-nothing grouping plans. + """ + terms = factorisation_group_options( + expression, argument_indices) + layouts = OrderedDict() + for term, layout, options in terms: + option_sums = layouts.setdefault( + layout, [MonomialSum() for _ in options]) + assert len(option_sums) == len(options) + for position, groups in enumerate(options): + classifier = partial( + classify, argument_indices, + delta_inside=delta_inside, groups=groups) + monomial_sum, = collect_monomials([term], classifier) + option_sums[position] = MonomialSum.sum( + option_sums[position], monomial_sum) + + options = tuple(tuple(option_sums) for option_sums in layouts.values()) + supports = tuple( + frozenset( + atomic + for option in option_sums + for monomial in option + for atomic in monomial.atomics) + for option_sums in options) + + remaining = set(range(len(options))) + components = [] + while remaining: + component = {remaining.pop()} + support = set().union(*(supports[i] for i in component)) + while True: + neighbours = { + i for i in remaining if supports[i].intersection(support)} + if not neighbours: + break + component.update(neighbours) + remaining.difference_update(neighbours) + support.update(*(supports[i] for i in neighbours)) + components.append(tuple(sorted(component))) + + result = [] + for component in components: + candidates = OrderedDict() + baseline = OrderedDict() + choices = ( + tuple(enumerate(options[i])) for i in component) + for selected in itertools.product(*choices): + positions, monomial_sums = zip(*selected) + candidate = MonomialSum.sum(*monomial_sums) + key = tuple(candidate) + candidates.setdefault(key, candidate) + if all(position in {0, len(options[i]) - 1} + for i, position in zip(component, positions)): + baseline.setdefault(key, candidate) + result.append(FactorisationCandidates( + tuple(candidates.values()), tuple(baseline.values()))) + return tuple(result) + + +def _optimise_candidate( + variable, monomial_sum, quadrature_indices, + index_replacer) -> tuple[tuple, ...]: + """Apply delta elimination and contraction optimization to one plan. + + Parameters + ---------- + variable : Node + Assignment variable. + monomial_sum : MonomialSum + Candidate polynomial representation. + quadrature_indices : tuple of Index + Preferred quadrature contraction order. + index_replacer : MemoizerArg + Shared index-substitution mapper. + + Returns + ------- + tuple + Optimized assignment pairs for cost evaluation. + """ + narrow_variables = OrderedDict() + simplified = defaultdict(MonomialSum) + for monomial in monomial_sum: + var, indices, atomics, rest = delta_elimination( + variable, *monomial, index_replacer) + narrow_variables.setdefault(var) + simplified[var].add(indices, atomics, rest) + + pairs = [] + for var in narrow_variables: + candidate = simplified[var] + contracted = set(chain.from_iterable( + monomial.sum_indices for monomial in candidate)) + ordering = sorted( + (index for index in quadrature_indices + if index in contracted), + key=lambda index: index.extent) + pairs.append(( + var, sum_factorise(var, ordering, candidate))) + return tuple(pairs) + + +def _candidate_score(pairs: tuple[tuple, ...]) -> tuple[int, ...]: + """Estimate work and storage in a GEM contraction candidate. + + Parameters + ---------- + pairs + Assignment pairs to schedule. + + Returns + ------- + tuple of int + Operations, total contraction storage, largest contraction, and + expression-node count. Lexicographic ordering prioritizes arithmetic + work. + """ + return estimate_cost(expression for _, expression in pairs) + + +def _select_factorisation_plan( + variable, candidate_groups, quadrature_indices, + index_replacer) -> MonomialSum: + """Choose a minimum-work plan under the prior DAG-size budget. + + The all-or-nothing grouping choices define the complexity budget of the + previous factorisation algorithm. Independent groups are combined with a + dynamic program, retaining the least expensive plan for each total DAG + size. This exposes partial groupings without allowing expression growth. + + Parameters + ---------- + variable : Node + Assignment variable. + candidate_groups : tuple of FactorisationCandidates + Independent factorisation choices. + quadrature_indices : tuple of Index + Preferred quadrature contraction order. + index_replacer : MemoizerArg + Shared index-substitution mapper. + + Returns + ------- + MonomialSum + Selected factorisation plan. + """ + plans = [] + node_budget = 0 + for group in candidate_groups: + alternatives = [] + scores = {} + for candidate in group.alternatives: + score = _candidate_score(_optimise_candidate( + variable, candidate, quadrature_indices, index_replacer)) + alternatives.append((score, candidate)) + scores[tuple(candidate)] = score + plans.append(alternatives) + node_budget += min( + (scores[tuple(candidate)] for candidate in group.baseline), + key=lambda score: score)[3] + + # score -> operations, total storage, largest intermediate, DAG nodes + states = {(0, 0, 0, 0): ()} + for alternatives in plans: + updated = {} + for left, selected in states.items(): + for right, candidate in alternatives: + score = ( + left[0] + right[0], + left[1] + right[1], + max(left[2], right[2]), + left[3] + right[3], + ) + if score[3] <= node_budget: + previous = updated.get(score[3]) + if previous is None or score < previous[0]: + updated[score[3]] = (score, selected + (candidate,)) + + # A state with both a larger DAG and a worse lexicographic cost can + # never become optimal as subsequent component costs are additive. + states = {} + best = None + for size in sorted(updated): + score, selected = updated[size] + if best is None or score[:3] < best: + states[score] = selected + best = score[:3] + + score = min(states) + return MonomialSum.sum(*states[score]) + + def flatten(var_reps, index_cache): quadrature_indices = OrderedDict() @@ -88,15 +305,17 @@ def group_key(pair): narrow_variables = OrderedDict() # Assignments are variable -> MonomialSum map delta_simplified = defaultdict(MonomialSum) + quadrature_indices = tuple(quadrature_indices) # 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): + argument_indices = set(free_indices) + for variable, expression in zip(variables, expressions): + candidate_groups = _factorisation_candidates( + expression, argument_indices, delta_inside) + monomial_sum = _select_factorisation_plan( + variable, candidate_groups, quadrature_indices, + index_replacer) for monomial in monomial_sum: var, s, a, r = delta_elimination(variable, *monomial, index_replacer) narrow_variables.setdefault(var) @@ -120,16 +339,38 @@ def group_key(pair): yield (variable, expression) -finalise_options = dict(replace_delta=False) +finalise_options = dict(replace_delta=True) + + +def classify(argument_indices, expression, delta_inside, groups=frozenset()): + """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. + groups : frozenset of Node + Algebraic groups selected by contraction-plan optimization. -def classify(argument_indices, expression, delta_inside): - """Classifier for argument factorisation""" + Returns + ------- + str + Refactorization label. + """ + if expression in groups: + return ATOMIC 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 +403,13 @@ 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 sum_factorise_monomial_sum( + monomial_sum, tuple(tail_ordering), variable.index_ordering()) From bd7439c4eb347d31d1dfed18a8dd0b394fa8a598 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 2 Aug 2026 23:01:12 +0100 Subject: [PATCH 02/26] WIP --- tsfc/spectral.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 57a54c5953..655e306750 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -3,7 +3,7 @@ import itertools from itertools import chain, zip_longest -from gem.gem import Conditional, Delta, Indexed, Sum, index_sum, one +from gem.gem import Conditional, Delta, Indexed, IndexSum, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination @@ -59,7 +59,7 @@ def _delta_inside(node, self): def _factorisation_candidates( - expression, argument_indices, + expression, argument_indices, quadrature_indices, delta_inside) -> tuple[FactorisationCandidates, ...]: """Build alternative pre-expansion grouping plans. @@ -69,6 +69,8 @@ def _factorisation_candidates( Multilinear integrand to factorize. argument_indices : set of Index Free argument indices. + quadrature_indices : tuple of Index + Indices contracted by quadrature. delta_inside : callable Memoized predicate detecting delta nodes. @@ -87,7 +89,7 @@ def _factorisation_candidates( assert len(option_sums) == len(options) for position, groups in enumerate(options): classifier = partial( - classify, argument_indices, + classify, argument_indices, quadrature_indices, delta_inside=delta_inside, groups=groups) monomial_sum, = collect_monomials([term], classifier) option_sums[position] = MonomialSum.sum( @@ -312,7 +314,8 @@ def group_key(pair): argument_indices = set(free_indices) for variable, expression in zip(variables, expressions): candidate_groups = _factorisation_candidates( - expression, argument_indices, delta_inside) + expression, argument_indices, quadrature_indices, + delta_inside) monomial_sum = _select_factorisation_plan( variable, candidate_groups, quadrature_indices, index_replacer) @@ -339,16 +342,19 @@ def group_key(pair): yield (variable, expression) -finalise_options = dict(replace_delta=True) +finalise_options = dict(replace_delta=True, remove_componenttensors=False) -def classify(argument_indices, expression, delta_inside, groups=frozenset()): +def classify(argument_indices, quadrature_indices, expression, delta_inside, + groups=frozenset()): """Classify one expression for multilinear factorization. Parameters ---------- argument_indices : set of Index Free argument indices. + quadrature_indices : tuple of Index + Indices contracted by quadrature. expression : Node Expression to classify. delta_inside : callable @@ -367,6 +373,9 @@ def classify(argument_indices, expression, delta_inside, groups=frozenset()): if n == 0: return OTHER elif n == 1: + if isinstance(expression, IndexSum) and set( + expression.multiindex).isdisjoint(quadrature_indices): + return ATOMIC if isinstance(expression, Conditional): return ATOMIC if isinstance(expression, (Delta, Indexed)) \ From 36b605ed7976794c613a8ac83f920b7798572977 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 3 Aug 2026 13:31:21 +0100 Subject: [PATCH 03/26] hoist_linear_index --- tsfc/spectral.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 655e306750..36e7b34758 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -8,7 +8,8 @@ from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import ( - estimate_cost, factorisation_group_options, replace_division, + estimate_cost, factorisation_group_options, hoist_linear_index, + replace_division, unroll_indexsum, ) from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials @@ -339,6 +340,8 @@ def group_key(pair): 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) + expression = hoist_linear_index( + expression, variable.free_indices) yield (variable, expression) From 42f54c191aca3584a09f3c9c7a164a497f340d98 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 3 Aug 2026 14:54:57 +0100 Subject: [PATCH 04/26] Test shared physically mapped tabulations --- tests/tsfc/test_sum_factorisation.py | 42 +++++++++++++++++++++++++++- tsfc/spectral.py | 19 ++++--------- 2 files changed, 46 insertions(+), 15 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 891cf1c6cc..8147ac0a2c 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -1,9 +1,10 @@ import numpy import pytest +import tsfc.spectral from ufl import (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) @@ -168,6 +169,45 @@ def test_vector_laplace_action(cell, order): assert (rates < order).all() +def test_shared_physically_mapped_tabulation( + monkeypatch: pytest.MonkeyPatch) -> None: + """Check that a mapped tabulation is shared by both argument axes. + + Parameters + ---------- + monkeypatch + Pytest fixture used to disable the sharing pass for comparison. + """ + 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( + tsfc.spectral, "hoist_linear_index", + lambda expression, indices: expression) + baseline, = compile_form(form, parameters={"mode": "spectral"}) + + optimized_source = str(optimized.ast) + baseline_source = str(baseline.ast) + optimized_shapes = [ + temporary.shape for temporary in + optimized.ast.default_entrypoint.temporary_variables.values()] + baseline_shapes = [ + temporary.shape for temporary in + baseline.ast.default_entrypoint.temporary_variables.values()] + + 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 sum(shape == (15,) for shape in optimized_shapes) \ + > sum(shape == (15,) for shape in baseline_shapes) + assert optimized_source.count(" if ") < baseline_source.count(" if ") + + if __name__ == "__main__": import os import sys diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 36e7b34758..ecc0cb3ce7 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -3,7 +3,7 @@ import itertools from itertools import chain, zip_longest -from gem.gem import Conditional, Delta, Indexed, IndexSum, Sum, index_sum, one +from gem.gem import Conditional, Delta, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination @@ -60,7 +60,7 @@ def _delta_inside(node, self): def _factorisation_candidates( - expression, argument_indices, quadrature_indices, + expression, argument_indices, delta_inside) -> tuple[FactorisationCandidates, ...]: """Build alternative pre-expansion grouping plans. @@ -70,8 +70,6 @@ def _factorisation_candidates( Multilinear integrand to factorize. argument_indices : set of Index Free argument indices. - quadrature_indices : tuple of Index - Indices contracted by quadrature. delta_inside : callable Memoized predicate detecting delta nodes. @@ -90,7 +88,7 @@ def _factorisation_candidates( assert len(option_sums) == len(options) for position, groups in enumerate(options): classifier = partial( - classify, argument_indices, quadrature_indices, + classify, argument_indices, delta_inside=delta_inside, groups=groups) monomial_sum, = collect_monomials([term], classifier) option_sums[position] = MonomialSum.sum( @@ -315,8 +313,7 @@ def group_key(pair): argument_indices = set(free_indices) for variable, expression in zip(variables, expressions): candidate_groups = _factorisation_candidates( - expression, argument_indices, quadrature_indices, - delta_inside) + expression, argument_indices, delta_inside) monomial_sum = _select_factorisation_plan( variable, candidate_groups, quadrature_indices, index_replacer) @@ -348,16 +345,13 @@ def group_key(pair): finalise_options = dict(replace_delta=True, remove_componenttensors=False) -def classify(argument_indices, quadrature_indices, expression, delta_inside, - groups=frozenset()): +def classify(argument_indices, expression, delta_inside, groups=frozenset()): """Classify one expression for multilinear factorization. Parameters ---------- argument_indices : set of Index Free argument indices. - quadrature_indices : tuple of Index - Indices contracted by quadrature. expression : Node Expression to classify. delta_inside : callable @@ -376,9 +370,6 @@ def classify(argument_indices, quadrature_indices, expression, delta_inside, if n == 0: return OTHER elif n == 1: - if isinstance(expression, IndexSum) and set( - expression.multiindex).isdisjoint(quadrature_indices): - return ATOMIC if isinstance(expression, Conditional): return ATOMIC if isinstance(expression, (Delta, Indexed)) \ From 6018fb0e86accfbfb930b25b233c32ce99b61cec Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 5 Aug 2026 16:30:44 +0100 Subject: [PATCH 05/26] Order contractions by retained support --- tests/tsfc/test_sum_factorisation.py | 19 +++++++++ tsfc/spectral.py | 60 ++++++++++++++++++++++++---- 2 files changed, 71 insertions(+), 8 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 8147ac0a2c..1c77d8b095 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -1,7 +1,10 @@ import numpy import pytest +import gem import tsfc.spectral +from gem.gem import one +from gem.refactorise import MonomialSum from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, action, interval, triangle, quadrilateral, curl, dot, div, grad, inner) @@ -9,6 +12,7 @@ TensorProductElement, HCurlElement, HDivElement) from tsfc import compile_form +from tsfc.spectral import _sum_factorisation_order def helmholtz(cell, degree): @@ -208,6 +212,21 @@ def test_shared_physically_mapped_tabulation( assert optimized_source.count(" if ") < baseline_source.count(" if ") +def test_sum_factorisation_order() -> None: + """Contract the quadrature direction with least argument support first.""" + 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)) + monomial_sum = MonomialSum() + monomial_sum.add((q0, q1), (inner * outer,), one) + + ordering = _sum_factorisation_order( + (q1, q0), monomial_sum) + + assert ordering == (q0, q1) + + if __name__ == "__main__": import os import sys diff --git a/tsfc/spectral.py b/tsfc/spectral.py index ecc0cb3ce7..5229fe683b 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -1,10 +1,12 @@ from collections import OrderedDict, defaultdict, namedtuple +from collections.abc import Iterable from functools import partial import itertools from itertools import chain, zip_longest +import math -from gem.gem import Conditional, Delta, Indexed, Sum, index_sum, one -from gem.node import Memoizer, MemoizerArg +from gem.gem import Conditional, Delta, Index, Indexed, Sum, index_sum, one +from gem.node import Memoizer, MemoizerArg, traversal from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import ( @@ -137,6 +139,52 @@ def _factorisation_candidates( return tuple(result) +def _sum_factorisation_order( + indices: Iterable[Index], + monomial_sum: MonomialSum) -> tuple[Index, ...]: + """Order quadrature contractions by their retained index support. + + A quadrature direction belongs earlier when its minimal dependency + frontier retains fewer non-quadrature indices. COFFEE can then isolate + that contraction before introducing more strongly coupled factors. The + support sizes are equal for ordinary tensor products, retaining the + established extent-based stable ordering, while nested simplex factors + naturally order themselves from least to most coupled. + + Parameters + ---------- + indices : iterable of Index + Quadrature indices to contract. + monomial_sum : MonomialSum + Factorized integrand. + + Returns + ------- + tuple of Index + Contraction indices from outermost to innermost stage. + """ + indices = tuple(indices) + contraction_indices = frozenset(indices) + factors = tuple( + factor + for monomial in monomial_sum + for factor in (*monomial.atomics, monomial.rest)) + nodes = tuple(traversal(factors)) + + def support_size(index: Index) -> int: + support = set() + for node in nodes: + if (index in node.free_indices + and not any(index in child.free_indices + for child in node.children)): + support.update( + set(node.free_indices) - contraction_indices) + return math.prod(argument.extent for argument in support) + + return tuple(sorted( + indices, key=lambda index: (support_size(index), index.extent))) + + def _optimise_candidate( variable, monomial_sum, quadrature_indices, index_replacer) -> tuple[tuple, ...]: @@ -329,12 +377,8 @@ def group_key(pair): 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) + sum_indices = _sum_factorisation_order( + sum_indices, monomial_sum) # Apply sum factorisation combined with COFFEE technology expression = sum_factorise(variable, sum_indices, monomial_sum) expression = hoist_linear_index( From f7a928a3096cc213cc6e15584424aa0efb234336 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 6 Aug 2026 22:43:26 +0100 Subject: [PATCH 06/26] Report reproducible JM codegen metrics --- benchmarks/johnson_mercier.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/benchmarks/johnson_mercier.py b/benchmarks/johnson_mercier.py index 033663fcbe..689362f915 100644 --- a/benchmarks/johnson_mercier.py +++ b/benchmarks/johnson_mercier.py @@ -2,7 +2,7 @@ """Measure simplex Johnson--Mercier code generation.""" import argparse -import hashlib +import time import numpy @@ -71,16 +71,19 @@ def main() -> None: parser.add_argument("--dims", nargs="+", type=int, default=(2, 3)) args = parser.parse_args() print("") - print("| dim | flops | scalar temps | array temps | elements | bytes | largest | AST lines | hash |") - print("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | --- |") + print("| dim | compile (s) | flops | scalar temps | array temps | " + "elements | bytes | largest | AST lines |") + print("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | " + "---: |") for dim in args.dims: + start = time.perf_counter() kernel = compile_target(dim) + elapsed = time.perf_counter() - start source = str(kernel.ast) - digest = hashlib.sha256(source.encode()).hexdigest()[:12] nscalar, narray, nelem, largest = temporary_metrics(kernel) - print(f"| {dim} | {kernel.flop_count:.0f} | " + print(f"| {dim} | {elapsed:.6f} | {kernel.flop_count:.0f} | " f"{nscalar} | {narray} | {nelem} | {8 * nelem} | {largest} | " - f"{len(source.splitlines())} | {digest} |") + f"{len(source.splitlines())} |") if __name__ == "__main__": From 2a8ffd512c848291384ce42e7a70a4cc80cdd8c0 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 7 Aug 2026 16:38:37 +0100 Subject: [PATCH 07/26] Prune dominated factorisation candidates --- tsfc/spectral.py | 111 +++++++++++++++++++++++++++++++++++++---------- 1 file changed, 87 insertions(+), 24 deletions(-) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 5229fe683b..affa60511c 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -14,7 +14,8 @@ replace_division, unroll_indexsum, ) -from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials +from gem.refactorise import (ATOMIC, COMPOUND, OTHER, ExpansionLimitExceeded, + MonomialSum, collect_monomials) from gem.unconcatenate import unconcatenate from gem.coffee import sum_factorise_monomial_sum from gem.utils import groupby @@ -26,6 +27,25 @@ FactorisationCandidates = namedtuple( 'FactorisationCandidates', ['alternatives', 'baseline']) +def _factor_dag_size(monomial_sum): + """Count nodes in the irreducible factors of a polynomial. + + Parameters + ---------- + monomial_sum : MonomialSum + Polynomial whose atomic and scalar factors are counted. + + Returns + ------- + int + Number of distinct GEM nodes reachable from the factors. + """ + factors = tuple( + factor + for monomial in monomial_sum + for factor in (*monomial.atomics, monomial.rest)) + return len(tuple(traversal(factors))) + def Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): """Constructs an integral representation for each GEM integrand @@ -81,29 +101,60 @@ def _factorisation_candidates( Distinct monomial representations for each independent group of choices, together with the all-or-nothing grouping plans. """ - terms = factorisation_group_options( - expression, argument_indices) + terms = factorisation_group_options(expression, argument_indices) layouts = OrderedDict() for term, layout, options in terms: - option_sums = layouts.setdefault( - layout, [MonomialSum() for _ in options]) - assert len(option_sums) == len(options) - for position, groups in enumerate(options): + records = layouts.setdefault(layout, []) + records.append((term, options)) + + def collect(records, position, max_monomials=None): + result = MonomialSum() + for term, options in records: + groups = options[position] classifier = partial( classify, argument_indices, delta_inside=delta_inside, groups=groups) - monomial_sum, = collect_monomials([term], classifier) - option_sums[position] = MonomialSum.sum( - option_sums[position], monomial_sum) + monomial_sum, = collect_monomials( + [term], classifier, max_monomials=max_monomials) + result = MonomialSum.sum(result, monomial_sum) + if (max_monomials is not None + and len(result) > max_monomials): + raise ExpansionLimitExceeded + return result + + grouped = { + layout: collect(records, -1) + for layout, records in layouts.items() + } + expansion_budget = _factor_dag_size( + MonomialSum.sum(*grouped.values())) + + options = [] + endpoints = [] + for layout, records in layouts.items(): + noptions = len(records[0][1]) + assert all(len(term_options) == noptions + for _, term_options in records) + available = [] + for position in range(noptions): + if position == noptions - 1: + monomial_sum = grouped[layout] + else: + try: + monomial_sum = collect(records, position, expansion_budget) + except ExpansionLimitExceeded: + continue + available.append((position, monomial_sum)) + options.append(tuple(available)) + endpoints.append(noptions - 1) - options = tuple(tuple(option_sums) for option_sums in layouts.values()) supports = tuple( frozenset( atomic - for option in option_sums + for _, option in available for monomial in option for atomic in monomial.atomics) - for option_sums in options) + for available in options) remaining = set(range(len(options))) components = [] @@ -124,14 +175,13 @@ def _factorisation_candidates( for component in components: candidates = OrderedDict() baseline = OrderedDict() - choices = ( - tuple(enumerate(options[i])) for i in component) + choices = (options[i] for i in component) for selected in itertools.product(*choices): positions, monomial_sums = zip(*selected) candidate = MonomialSum.sum(*monomial_sums) key = tuple(candidate) candidates.setdefault(key, candidate) - if all(position in {0, len(options[i]) - 1} + if all(position in {0, endpoints[i]} for i, position in zip(component, positions)): baseline.setdefault(key, candidate) result.append(FactorisationCandidates( @@ -272,20 +322,33 @@ def _select_factorisation_plan( MonomialSum Selected factorisation plan. """ - plans = [] + scores = {} node_budget = 0 for group in candidate_groups: - alternatives = [] - scores = {} - for candidate in group.alternatives: + baseline_scores = [] + for candidate in group.baseline: + key = tuple(candidate) score = _candidate_score(_optimise_candidate( variable, candidate, quadrature_indices, index_replacer)) + scores[key] = score + baseline_scores.append(score) + node_budget += min(baseline_scores)[3] + + plans = [] + for group in candidate_groups: + alternatives = [] + for candidate in group.alternatives: + key = tuple(candidate) + try: + score = scores[key] + except KeyError: + if _factor_dag_size(candidate) > node_budget: + continue + score = _candidate_score(_optimise_candidate( + variable, candidate, quadrature_indices, index_replacer)) + scores[key] = score alternatives.append((score, candidate)) - scores[tuple(candidate)] = score plans.append(alternatives) - node_budget += min( - (scores[tuple(candidate)] for candidate in group.baseline), - key=lambda score: score)[3] # score -> operations, total storage, largest intermediate, DAG nodes states = {(0, 0, 0, 0): ()} From 3e5b364690e19aae8ecbd3f697ab671f4b570495 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 7 Aug 2026 16:39:00 +0100 Subject: [PATCH 08/26] WIP: lower ragged contractions to exact domains --- tsfc/loopy.py | 54 ++++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 45 insertions(+), 9 deletions(-) diff --git a/tsfc/loopy.py b/tsfc/loopy.py index e41b8348fc..6b1cd7d265 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -122,6 +122,7 @@ class LoopyContext(object): def __init__(self, target=None): self.indices = {} # indices for declarations and referencing values, from ImperoC self.active_indices = {} # gem index -> pymbolic variable + self.index_lengths = {} # iname -> (parent inames, tabulated extents) 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 @@ -258,7 +259,8 @@ 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(), ctx.index_parents) + domains = create_domains( + ctx.index_extent.items(), ctx.index_parents, ctx.index_lengths) # Create loopy kernel knl = lp.make_kernel( @@ -277,23 +279,50 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name return knl, event_name -def create_domains(indices, index_parents=None): - """ Create ISL domains from indices +def create_domains(indices, index_parents=None, index_lengths=None): + """Create ISL domains for independent and dependent indices. - :arg indices: iterable of (index_name, extent) pairs - :arg index_parents: optional mapping from index_name to a tuple of parent - index names; the domain of a jagged index is parametrized by its - parents, with upper bound extent minus the sum of the parents. - :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. + index_lengths : mapping, optional + Parent inames and tabulated extents for ragged bounds. + Returns + ------- + list of isl.Set + Iteration domains for Loopy. + """ domains = [] for idx, extent in indices: + if index_lengths and idx in index_lengths: + parents, lengths = index_lengths[idx] + inames = isl.make_zero_and_vars([idx], parents) + domain = None + for point in numpy.ndindex(lengths.shape): + length = int(lengths[point]) + if length == 0: + continue + piece = (inames[0].le_set(inames[idx]) + & inames[idx].lt_set(inames[0] + length)) + for parent, value in zip(parents, point): + piece = piece & inames[parent].eq_set( + inames[0] + value) + domain = piece if domain is None else domain.union(piece) + assert domain is not None + domains.append(domain) + continue + 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)))) + domains.append(inames[0].le_set(inames[idx]) + & inames[idx].lt_set(bound)) if not domains: domains = [isl.BasicSet("[] -> {[]}")] @@ -331,6 +360,13 @@ def statement_for(tree, ctx): # remains correct: jagged expressions are zero-padded. ctx.index_parents[idx] = tuple(ctx.active_indices[parent].name for parent in tree.index.parents) + elif isinstance(tree.index, gem.RaggedIndex) and \ + all(parent in ctx.active_indices for parent in tree.index.parents): + ctx.index_lengths[idx] = ( + tuple(ctx.active_indices[parent].name + for parent in tree.index.parents), + tree.index.lengths, + ) with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active: return statement(tree.children[0], ctx_active) From 66838703e30214f039a07486d3a3eb51e2c02b95 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 7 Aug 2026 22:27:22 +0100 Subject: [PATCH 09/26] Consolidate sum-factorisation plan selection --- tsfc/spectral.py | 199 +++++++++++++---------------------------------- 1 file changed, 56 insertions(+), 143 deletions(-) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index affa60511c..1d744deb4e 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -1,7 +1,6 @@ from collections import OrderedDict, defaultdict, namedtuple from collections.abc import Iterable from functools import partial -import itertools from itertools import chain, zip_longest import math @@ -17,7 +16,7 @@ from gem.refactorise import (ATOMIC, COMPOUND, OTHER, ExpansionLimitExceeded, MonomialSum, collect_monomials) from gem.unconcatenate import unconcatenate -from gem.coffee import sum_factorise_monomial_sum +from gem.coffee import optimise_monomial_sum from gem.utils import groupby @@ -25,7 +24,7 @@ 'quadrature_multiindex', 'argument_indices']) FactorisationCandidates = namedtuple( - 'FactorisationCandidates', ['alternatives', 'baseline']) + 'FactorisationCandidates', ['alternatives', 'fallback']) def _factor_dag_size(monomial_sum): """Count nodes in the irreducible factors of a polynomial. @@ -83,8 +82,8 @@ def _delta_inside(node, self): def _factorisation_candidates( expression, argument_indices, - delta_inside) -> tuple[FactorisationCandidates, ...]: - """Build alternative pre-expansion grouping plans. + delta_inside) -> FactorisationCandidates: + """Build bounded pre-expansion grouping alternatives. Parameters ---------- @@ -97,23 +96,20 @@ def _factorisation_candidates( Returns ------- - tuple of FactorisationCandidates - Distinct monomial representations for each independent group of - choices, together with the all-or-nothing grouping plans. + FactorisationCandidates + Distinct alternatives and the compact fully grouped fallback. """ - terms = factorisation_group_options(expression, argument_indices) layouts = OrderedDict() - for term, layout, options in terms: - records = layouts.setdefault(layout, []) - records.append((term, options)) + for term, layout, options in factorisation_group_options( + expression, argument_indices): + layouts.setdefault(layout, []).append((term, options)) def collect(records, position, max_monomials=None): result = MonomialSum() for term, options in records: - groups = options[position] classifier = partial( classify, argument_indices, - delta_inside=delta_inside, groups=groups) + delta_inside=delta_inside, groups=options[position]) monomial_sum, = collect_monomials( [term], classifier, max_monomials=max_monomials) result = MonomialSum.sum(result, monomial_sum) @@ -122,73 +118,39 @@ def collect(records, position, max_monomials=None): raise ExpansionLimitExceeded return result - grouped = { - layout: collect(records, -1) - for layout, records in layouts.items() - } - expansion_budget = _factor_dag_size( - MonomialSum.sum(*grouped.values())) + compact = tuple(collect(records, -1) + for records in layouts.values()) + fallback = MonomialSum.sum(*compact) + budget = _factor_dag_size(fallback) - options = [] - endpoints = [] - for layout, records in layouts.items(): + candidates = (MonomialSum(),) + for records, grouped in zip(layouts.values(), compact): noptions = len(records[0][1]) - assert all(len(term_options) == noptions - for _, term_options in records) - available = [] + assert all(len(options) == noptions for _, options in records) + alternatives = [] for position in range(noptions): if position == noptions - 1: - monomial_sum = grouped[layout] + candidate = grouped else: try: - monomial_sum = collect(records, position, expansion_budget) + candidate = collect(records, position, budget) except ExpansionLimitExceeded: continue - available.append((position, monomial_sum)) - options.append(tuple(available)) - endpoints.append(noptions - 1) - - supports = tuple( - frozenset( - atomic - for _, option in available - for monomial in option - for atomic in monomial.atomics) - for available in options) - - remaining = set(range(len(options))) - components = [] - while remaining: - component = {remaining.pop()} - support = set().union(*(supports[i] for i in component)) - while True: - neighbours = { - i for i in remaining if supports[i].intersection(support)} - if not neighbours: - break - component.update(neighbours) - remaining.difference_update(neighbours) - support.update(*(supports[i] for i in neighbours)) - components.append(tuple(sorted(component))) - - result = [] - for component in components: - candidates = OrderedDict() - baseline = OrderedDict() - choices = (options[i] for i in component) - for selected in itertools.product(*choices): - positions, monomial_sums = zip(*selected) - candidate = MonomialSum.sum(*monomial_sums) - key = tuple(candidate) - candidates.setdefault(key, candidate) - if all(position in {0, endpoints[i]} - for i, position in zip(component, positions)): - baseline.setdefault(key, candidate) - result.append(FactorisationCandidates( - tuple(candidates.values()), tuple(baseline.values()))) - return tuple(result) - - + alternatives.append(candidate) + + combined = OrderedDict() + for left in candidates: + for right in alternatives: + candidate = MonomialSum.sum(left, right) + if _factor_dag_size(candidate) <= budget: + combined.setdefault(tuple(candidate), candidate) + candidates = tuple(combined.values()) + + alternatives = OrderedDict( + (tuple(candidate), candidate) for candidate in candidates) + alternatives.setdefault(tuple(fallback), fallback) + return FactorisationCandidates( + tuple(alternatives.values()), fallback) def _sum_factorisation_order( indices: Iterable[Index], monomial_sum: MonomialSum) -> tuple[Index, ...]: @@ -297,21 +259,16 @@ def _candidate_score(pairs: tuple[tuple, ...]) -> tuple[int, ...]: def _select_factorisation_plan( - variable, candidate_groups, quadrature_indices, + variable, candidates, quadrature_indices, index_replacer) -> MonomialSum: - """Choose a minimum-work plan under the prior DAG-size budget. - - The all-or-nothing grouping choices define the complexity budget of the - previous factorisation algorithm. Independent groups are combined with a - dynamic program, retaining the least expensive plan for each total DAG - size. This exposes partial groupings without allowing expression growth. + """Choose the least-cost plan no larger than the compact fallback. Parameters ---------- variable : Node Assignment variable. - candidate_groups : tuple of FactorisationCandidates - Independent factorisation choices. + candidates : FactorisationCandidates + Bounded factorisation alternatives. quadrature_indices : tuple of Index Preferred quadrature contraction order. index_replacer : MemoizerArg @@ -322,65 +279,20 @@ def _select_factorisation_plan( MonomialSum Selected factorisation plan. """ - scores = {} - node_budget = 0 - for group in candidate_groups: - baseline_scores = [] - for candidate in group.baseline: - key = tuple(candidate) - score = _candidate_score(_optimise_candidate( - variable, candidate, quadrature_indices, index_replacer)) - scores[key] = score - baseline_scores.append(score) - node_budget += min(baseline_scores)[3] - - plans = [] - for group in candidate_groups: - alternatives = [] - for candidate in group.alternatives: - key = tuple(candidate) - try: - score = scores[key] - except KeyError: - if _factor_dag_size(candidate) > node_budget: - continue - score = _candidate_score(_optimise_candidate( - variable, candidate, quadrature_indices, index_replacer)) - scores[key] = score - alternatives.append((score, candidate)) - plans.append(alternatives) - - # score -> operations, total storage, largest intermediate, DAG nodes - states = {(0, 0, 0, 0): ()} - for alternatives in plans: - updated = {} - for left, selected in states.items(): - for right, candidate in alternatives: - score = ( - left[0] + right[0], - left[1] + right[1], - max(left[2], right[2]), - left[3] + right[3], - ) - if score[3] <= node_budget: - previous = updated.get(score[3]) - if previous is None or score < previous[0]: - updated[score[3]] = (score, selected + (candidate,)) - - # A state with both a larger DAG and a worse lexicographic cost can - # never become optimal as subsequent component costs are additive. - states = {} - best = None - for size in sorted(updated): - score, selected = updated[size] - if best is None or score[:3] < best: - states[score] = selected - best = score[:3] - - score = min(states) - return MonomialSum.sum(*states[score]) - - + fallback = candidates.fallback + fallback_score = _candidate_score(_optimise_candidate( + variable, fallback, quadrature_indices, index_replacer)) + budget = fallback_score[3] + best = fallback_score, fallback + + for candidate in candidates.alternatives: + if candidate is fallback or _factor_dag_size(candidate) > budget: + continue + score = _candidate_score(_optimise_candidate( + variable, candidate, quadrature_indices, index_replacer)) + if score[3] <= budget and score < best[0]: + best = score, candidate + return best[1] def flatten(var_reps, index_cache): quadrature_indices = OrderedDict() @@ -521,5 +433,6 @@ def prune(factors): def sum_factorise(variable, tail_ordering, monomial_sum): - return sum_factorise_monomial_sum( - monomial_sum, tuple(tail_ordering), variable.index_ordering()) + return optimise_monomial_sum( + monomial_sum, variable.index_ordering(), + tuple(tail_ordering)) From 9cfe6623612830c5bb5dedc35de7ade5fe8e27c7 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 7 Aug 2026 22:46:02 +0100 Subject: [PATCH 10/26] Fix sum-factorisation lint --- tsfc/spectral.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 1d744deb4e..4f81d77b03 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -26,6 +26,7 @@ FactorisationCandidates = namedtuple( 'FactorisationCandidates', ['alternatives', 'fallback']) + def _factor_dag_size(monomial_sum): """Count nodes in the irreducible factors of a polynomial. @@ -151,6 +152,8 @@ def collect(records, position, max_monomials=None): alternatives.setdefault(tuple(fallback), fallback) return FactorisationCandidates( tuple(alternatives.values()), fallback) + + def _sum_factorisation_order( indices: Iterable[Index], monomial_sum: MonomialSum) -> tuple[Index, ...]: @@ -293,6 +296,8 @@ def _select_factorisation_plan( if score[3] <= budget and score < best[0]: best = score, candidate return best[1] + + def flatten(var_reps, index_cache): quadrature_indices = OrderedDict() From 875bad47bce819da2798ee4cd7eb1e53b7dc718b Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 13 Aug 2026 09:59:13 +0100 Subject: [PATCH 11/26] Preserve finite element factorisation plans --- tests/tsfc/test_sum_factorisation.py | 27 ++++++ tsfc/spectral.py | 124 ++++----------------------- 2 files changed, 46 insertions(+), 105 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 1c77d8b095..7882cf18db 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -227,6 +227,33 @@ def test_sum_factorisation_order() -> None: assert ordering == (q0, q1) +def test_bernstein_candidate_selection( + monkeypatch: pytest.MonkeyPatch) -> None: + """Select the least-cost algebraic plan for a Bernstein form.""" + mesh = Mesh(VectorElement("CG", triangle, 1)) + space = FunctionSpace(mesh, FiniteElement("Bernstein", triangle, 6)) + u = TrialFunction(space) + v = TestFunction(space) + form = inner(grad(u), grad(v)) * dx + + selected, = compile_form(form, parameters={"mode": "spectral"}) + + def minimum_cost_plan( + variable, candidates, quadrature_indices, index_replacer): + return min( + candidates, + key=lambda candidate: tsfc.spectral._candidate_score( + tsfc.spectral._optimise_candidate( + variable, candidate, quadrature_indices, + index_replacer))) + + monkeypatch.setattr( + tsfc.spectral, "_select_factorisation_plan", minimum_cost_plan) + minimum, = compile_form(form, parameters={"mode": "spectral"}) + + assert selected.flop_count == minimum.flop_count + + if __name__ == "__main__": import os import sys diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 4f81d77b03..98ba5a86ae 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -9,12 +9,11 @@ from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import ( - estimate_cost, factorisation_group_options, hoist_linear_index, - replace_division, + estimate_cost, hoist_linear_index, replace_division, unroll_indexsum, ) -from gem.refactorise import (ATOMIC, COMPOUND, OTHER, ExpansionLimitExceeded, - MonomialSum, collect_monomials) +from gem.refactorise import (ATOMIC, COMPOUND, OTHER, MonomialSum, + collect_factorisation_plans) from gem.unconcatenate import unconcatenate from gem.coffee import optimise_monomial_sum from gem.utils import groupby @@ -23,28 +22,6 @@ Integral = namedtuple('Integral', ['expression', 'quadrature_multiindex', 'argument_indices']) -FactorisationCandidates = namedtuple( - 'FactorisationCandidates', ['alternatives', 'fallback']) - - -def _factor_dag_size(monomial_sum): - """Count nodes in the irreducible factors of a polynomial. - - Parameters - ---------- - monomial_sum : MonomialSum - Polynomial whose atomic and scalar factors are counted. - - Returns - ------- - int - Number of distinct GEM nodes reachable from the factors. - """ - factors = tuple( - factor - for monomial in monomial_sum - for factor in (*monomial.atomics, monomial.rest)) - return len(tuple(traversal(factors))) def Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): @@ -83,8 +60,8 @@ def _delta_inside(node, self): def _factorisation_candidates( expression, argument_indices, - delta_inside) -> FactorisationCandidates: - """Build bounded pre-expansion grouping alternatives. + delta_inside) -> tuple[MonomialSum, ...]: + """Build expanded and map-preserving factorisation plans. Parameters ---------- @@ -97,61 +74,13 @@ def _factorisation_candidates( Returns ------- - FactorisationCandidates - Distinct alternatives and the compact fully grouped fallback. + tuple of MonomialSum + Distinct contraction plans. """ - layouts = OrderedDict() - for term, layout, options in factorisation_group_options( - expression, argument_indices): - layouts.setdefault(layout, []).append((term, options)) - - def collect(records, position, max_monomials=None): - result = MonomialSum() - for term, options in records: - classifier = partial( - classify, argument_indices, - delta_inside=delta_inside, groups=options[position]) - monomial_sum, = collect_monomials( - [term], classifier, max_monomials=max_monomials) - result = MonomialSum.sum(result, monomial_sum) - if (max_monomials is not None - and len(result) > max_monomials): - raise ExpansionLimitExceeded - return result - - compact = tuple(collect(records, -1) - for records in layouts.values()) - fallback = MonomialSum.sum(*compact) - budget = _factor_dag_size(fallback) - - candidates = (MonomialSum(),) - for records, grouped in zip(layouts.values(), compact): - noptions = len(records[0][1]) - assert all(len(options) == noptions for _, options in records) - alternatives = [] - for position in range(noptions): - if position == noptions - 1: - candidate = grouped - else: - try: - candidate = collect(records, position, budget) - except ExpansionLimitExceeded: - continue - alternatives.append(candidate) - - combined = OrderedDict() - for left in candidates: - for right in alternatives: - candidate = MonomialSum.sum(left, right) - if _factor_dag_size(candidate) <= budget: - combined.setdefault(tuple(candidate), candidate) - candidates = tuple(combined.values()) - - alternatives = OrderedDict( - (tuple(candidate), candidate) for candidate in candidates) - alternatives.setdefault(tuple(fallback), fallback) - return FactorisationCandidates( - tuple(alternatives.values()), fallback) + classifier = partial( + classify, argument_indices, delta_inside=delta_inside) + return collect_factorisation_plans( + expression, classifier, argument_indices) def _sum_factorisation_order( @@ -264,14 +193,14 @@ def _candidate_score(pairs: tuple[tuple, ...]) -> tuple[int, ...]: def _select_factorisation_plan( variable, candidates, quadrature_indices, index_replacer) -> MonomialSum: - """Choose the least-cost plan no larger than the compact fallback. + """Choose the least-cost admitted factorisation plan. Parameters ---------- variable : Node Assignment variable. - candidates : FactorisationCandidates - Bounded factorisation alternatives. + candidates : tuple of MonomialSum + Algebraically equivalent factorisation plans. quadrature_indices : tuple of Index Preferred quadrature contraction order. index_replacer : MemoizerArg @@ -282,20 +211,10 @@ def _select_factorisation_plan( MonomialSum Selected factorisation plan. """ - fallback = candidates.fallback - fallback_score = _candidate_score(_optimise_candidate( - variable, fallback, quadrature_indices, index_replacer)) - budget = fallback_score[3] - best = fallback_score, fallback - - for candidate in candidates.alternatives: - if candidate is fallback or _factor_dag_size(candidate) > budget: - continue - score = _candidate_score(_optimise_candidate( - variable, candidate, quadrature_indices, index_replacer)) - if score[3] <= budget and score < best[0]: - best = score, candidate - return best[1] + return min( + candidates, + key=lambda candidate: _candidate_score(_optimise_candidate( + variable, candidate, quadrature_indices, index_replacer))) def flatten(var_reps, index_cache): @@ -369,7 +288,7 @@ def group_key(pair): finalise_options = dict(replace_delta=True, remove_componenttensors=False) -def classify(argument_indices, expression, delta_inside, groups=frozenset()): +def classify(argument_indices, expression, delta_inside): """Classify one expression for multilinear factorization. Parameters @@ -380,16 +299,11 @@ def classify(argument_indices, expression, delta_inside, groups=frozenset()): Expression to classify. delta_inside : callable Predicate detecting delta nodes. - groups : frozenset of Node - Algebraic groups selected by contraction-plan optimization. - Returns ------- str Refactorization label. """ - if expression in groups: - return ATOMIC n = len(argument_indices.intersection(expression.free_indices)) if n == 0: return OTHER From e87705292e67742368606512354ec0d4431c1e71 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 13 Aug 2026 22:12:56 +0100 Subject: [PATCH 12/26] Update factorisation checks for current kernels --- tests/tsfc/test_sum_factorisation.py | 3 --- tsfc/spectral.py | 6 +++--- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 7882cf18db..af192831e4 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -195,8 +195,6 @@ def test_shared_physically_mapped_tabulation( lambda expression, indices: expression) baseline, = compile_form(form, parameters={"mode": "spectral"}) - optimized_source = str(optimized.ast) - baseline_source = str(baseline.ast) optimized_shapes = [ temporary.shape for temporary in optimized.ast.default_entrypoint.temporary_variables.values()] @@ -209,7 +207,6 @@ def test_shared_physically_mapped_tabulation( < sum(not shape for shape in baseline_shapes) assert sum(shape == (15,) for shape in optimized_shapes) \ > sum(shape == (15,) for shape in baseline_shapes) - assert optimized_source.count(" if ") < baseline_source.count(" if ") def test_sum_factorisation_order() -> None: diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 98ba5a86ae..3b0428dd7a 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -91,9 +91,9 @@ def _sum_factorisation_order( A quadrature direction belongs earlier when its minimal dependency frontier retains fewer non-quadrature indices. COFFEE can then isolate that contraction before introducing more strongly coupled factors. The - support sizes are equal for ordinary tensor products, retaining the - established extent-based stable ordering, while nested simplex factors - naturally order themselves from least to most coupled. + Ordinary tensor products have equal support sizes. They retain the + extent-based stable ordering. Nested simplex factors naturally order + themselves from least to most coupled. Parameters ---------- From cafa708f8816f1eea94ab4f63f86f28069a42654 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 13 Aug 2026 22:50:58 +0100 Subject: [PATCH 13/26] Compact products of simplex lattice temporaries --- tsfc/loopy.py | 13 ++++++++++++- tsfc/spectral.py | 6 +++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 6b1cd7d265..7c1bfd9fb1 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -126,6 +126,7 @@ def __init__(self, target=None): 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.compact_indices = {} # temporary -> compact index layout self.name_gen = UniqueNameGenerator() self.target = target self.loop_priorities = set() # used to avoid disadvantageous loop interchanges @@ -175,6 +176,13 @@ 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( + gem.simplex_lattice_rank(item, self.active_indices) + 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: @@ -241,8 +249,11 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name 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 + shape, layout = gem.compact_index_layout( + tuple(ctx.indices[temp])) + shape += temp.shape 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 diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 3b0428dd7a..da828af923 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -90,10 +90,10 @@ def _sum_factorisation_order( A quadrature direction belongs earlier when its minimal dependency frontier retains fewer non-quadrature indices. COFFEE can then isolate - that contraction before introducing more strongly coupled factors. The + that contraction before introducing more strongly coupled factors. Ordinary tensor products have equal support sizes. They retain the - extent-based stable ordering. Nested simplex factors naturally order - themselves from least to most coupled. + extent-based stable ordering. Nested factors naturally order themselves + from least to most coupled. Parameters ---------- From c95a4148570d973037004973f3705e1414cb9a27 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 14 Aug 2026 11:45:58 +0100 Subject: [PATCH 14/26] Share mapped tabulations during factorisation --- tests/tsfc/test_sum_factorisation.py | 7 ++++--- tsfc/spectral.py | 5 +---- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index af192831e4..dfd6274100 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -2,6 +2,7 @@ import pytest import gem +import gem.coffee import tsfc.spectral from gem.gem import one from gem.refactorise import MonomialSum @@ -175,7 +176,7 @@ def test_vector_laplace_action(cell, order): def test_shared_physically_mapped_tabulation( monkeypatch: pytest.MonkeyPatch) -> None: - """Check that a mapped tabulation is shared by both argument axes. + """Share a mapped tabulation between both argument axes. Parameters ---------- @@ -191,8 +192,8 @@ def test_shared_physically_mapped_tabulation( optimized, = compile_form(form, parameters={"mode": "spectral"}) monkeypatch.setattr( - tsfc.spectral, "hoist_linear_index", - lambda expression, indices: expression) + gem.coffee, "_share_linear_maps", + lambda monomial_sum, indices: monomial_sum) baseline, = compile_form(form, parameters={"mode": "spectral"}) optimized_shapes = [ diff --git a/tsfc/spectral.py b/tsfc/spectral.py index da828af923..6dc6a7bc74 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -9,8 +9,7 @@ from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import ( - estimate_cost, hoist_linear_index, replace_division, - unroll_indexsum, + estimate_cost, replace_division, unroll_indexsum, ) from gem.refactorise import (ATOMIC, COMPOUND, OTHER, MonomialSum, collect_factorisation_plans) @@ -280,8 +279,6 @@ def group_key(pair): sum_indices, monomial_sum) # Apply sum factorisation combined with COFFEE technology expression = sum_factorise(variable, sum_indices, monomial_sum) - expression = hoist_linear_index( - expression, variable.free_indices) yield (variable, expression) From e57e7a96e20534becfdcb3703ced74e1553f728a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 14 Aug 2026 11:50:57 +0100 Subject: [PATCH 15/26] Search quadrature contraction orderings --- tests/tsfc/test_sum_factorisation.py | 21 +++++-- tsfc/spectral.py | 83 +++++++++++++--------------- 2 files changed, 52 insertions(+), 52 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index dfd6274100..b96e5d760a 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -13,7 +13,7 @@ TensorProductElement, HCurlElement, HDivElement) from tsfc import compile_form -from tsfc.spectral import _sum_factorisation_order +from tsfc.spectral import _optimise_contraction_order def helmholtz(cell, degree): @@ -211,18 +211,27 @@ def test_shared_physically_mapped_tabulation( def test_sum_factorisation_order() -> None: - """Contract the quadrature direction with least argument support first.""" + """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) - ordering = _sum_factorisation_order( - (q1, q0), monomial_sum) - - assert ordering == (q0, q1) + score, _ = _optimise_contraction_order( + variable, (q1, q0), monomial_sum) + candidates = [ + tsfc.spectral._candidate_score(( + (variable, + tsfc.spectral.sum_factorise( + variable, ordering, monomial_sum)),)) + for ordering in ((q1, q0), (q0, q1)) + ] + + assert score == min(candidates) def test_bernstein_candidate_selection( diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 6dc6a7bc74..8df8f37c25 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -1,11 +1,10 @@ from collections import OrderedDict, defaultdict, namedtuple -from collections.abc import Iterable from functools import partial -from itertools import chain, zip_longest -import math +from itertools import chain, permutations, zip_longest -from gem.gem import Conditional, Delta, Index, Indexed, Sum, index_sum, one -from gem.node import Memoizer, MemoizerArg, traversal +from gem.gem import (Conditional, Delta, Indexed, Node, Sum, + index_sum, one) +from gem.node import Memoizer, MemoizerArg from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination from gem.optimise import ( @@ -82,20 +81,15 @@ def _factorisation_candidates( expression, classifier, argument_indices) -def _sum_factorisation_order( - indices: Iterable[Index], - monomial_sum: MonomialSum) -> tuple[Index, ...]: - """Order quadrature contractions by their retained index support. - - A quadrature direction belongs earlier when its minimal dependency - frontier retains fewer non-quadrature indices. COFFEE can then isolate - that contraction before introducing more strongly coupled factors. - Ordinary tensor products have equal support sizes. They retain the - extent-based stable ordering. Nested factors naturally order themselves - from least to most coupled. +def _optimise_contraction_order( + variable: Node, indices, monomial_sum: MonomialSum +) -> tuple[tuple[int, ...], Node]: + """Choose the least-cost quadrature contraction ordering. Parameters ---------- + variable : Node + Assignment variable whose indices identify multilinear axes. indices : iterable of Index Quadrature indices to contract. monomial_sum : MonomialSum @@ -103,29 +97,28 @@ def _sum_factorisation_order( Returns ------- - tuple of Index - Contraction indices from outermost to innermost stage. + score + Estimated operations, storage, and expression size. + expression + Factorized GEM expression for the selected ordering. + + Notes + ----- + The search is exhaustive in the number of quadrature axes. Finite + element integration supplies one axis per reference-cell direction, so + this space is independent of polynomial degree. + """ indices = tuple(indices) - contraction_indices = frozenset(indices) - factors = tuple( - factor - for monomial in monomial_sum - for factor in (*monomial.atomics, monomial.rest)) - nodes = tuple(traversal(factors)) - - def support_size(index: Index) -> int: - support = set() - for node in nodes: - if (index in node.free_indices - and not any(index in child.free_indices - for child in node.children)): - support.update( - set(node.free_indices) - contraction_indices) - return math.prod(argument.extent for argument in support) - - return tuple(sorted( - indices, key=lambda index: (support_size(index), index.extent))) + plans = ( + sum_factorise(variable, ordering, monomial_sum) + for ordering in permutations(indices) + ) + return min( + ((estimate_cost((expression,)), expression) + for expression in plans), + key=lambda plan: plan[0], + ) def _optimise_candidate( @@ -162,12 +155,11 @@ def _optimise_candidate( candidate = simplified[var] contracted = set(chain.from_iterable( monomial.sum_indices for monomial in candidate)) - ordering = sorted( - (index for index in quadrature_indices - if index in contracted), - key=lambda index: index.extent) - pairs.append(( - var, sum_factorise(var, ordering, candidate))) + ordering = tuple(index for index in quadrature_indices + if index in contracted) + _, expression = _optimise_contraction_order( + var, ordering, candidate) + pairs.append((var, expression)) return tuple(pairs) @@ -275,10 +267,9 @@ def group_key(pair): 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] - sum_indices = _sum_factorisation_order( - sum_indices, monomial_sum) # Apply sum factorisation combined with COFFEE technology - expression = sum_factorise(variable, sum_indices, monomial_sum) + _, expression = _optimise_contraction_order( + variable, sum_indices, monomial_sum) yield (variable, expression) From 5183f2b3ff9b07848f86c98f538d1c6b8bf80b47 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 14 Aug 2026 11:52:41 +0100 Subject: [PATCH 16/26] Leave compact simplex temporaries to simplex lowering --- tsfc/loopy.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 7c1bfd9fb1..6b1cd7d265 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -126,7 +126,6 @@ def __init__(self, target=None): 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.compact_indices = {} # temporary -> compact index layout self.name_gen = UniqueNameGenerator() self.target = target self.loop_priorities = set() # used to avoid disadvantageous loop interchanges @@ -176,13 +175,6 @@ 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( - gem.simplex_lattice_rank(item, self.active_indices) - 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: @@ -249,11 +241,8 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name 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, layout = gem.compact_index_layout( - tuple(ctx.indices[temp])) - shape += temp.shape + 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)) - ctx.compact_indices[temp] = layout ctx.gem_to_pymbolic[temp] = p.Variable(name) # Create instructions From 36b6fdc67b97a88a7c53a513470e02a1599e678d Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 14 Aug 2026 19:11:23 +0100 Subject: [PATCH 17/26] Preserve mapped tabulations through factorisation --- tests/tsfc/test_sum_factorisation.py | 42 ++------- tsfc/spectral.py | 136 ++++----------------------- 2 files changed, 26 insertions(+), 152 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index b96e5d760a..60c6754c46 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -5,6 +5,7 @@ import gem.coffee import tsfc.spectral from gem.gem import one +from gem.optimise import estimate_cost from gem.refactorise import MonomialSum from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, action, interval, triangle, @@ -198,10 +199,14 @@ def test_shared_physically_mapped_tabulation( optimized_shapes = [ temporary.shape for temporary in - optimized.ast.default_entrypoint.temporary_variables.values()] + 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()] + 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) \ @@ -224,43 +229,14 @@ def test_sum_factorisation_order() -> None: score, _ = _optimise_contraction_order( variable, (q1, q0), monomial_sum) candidates = [ - tsfc.spectral._candidate_score(( - (variable, - tsfc.spectral.sum_factorise( - variable, ordering, monomial_sum)),)) + estimate_cost((tsfc.spectral.sum_factorise( + variable, ordering, monomial_sum),)) for ordering in ((q1, q0), (q0, q1)) ] assert score == min(candidates) -def test_bernstein_candidate_selection( - monkeypatch: pytest.MonkeyPatch) -> None: - """Select the least-cost algebraic plan for a Bernstein form.""" - mesh = Mesh(VectorElement("CG", triangle, 1)) - space = FunctionSpace(mesh, FiniteElement("Bernstein", triangle, 6)) - u = TrialFunction(space) - v = TestFunction(space) - form = inner(grad(u), grad(v)) * dx - - selected, = compile_form(form, parameters={"mode": "spectral"}) - - def minimum_cost_plan( - variable, candidates, quadrature_indices, index_replacer): - return min( - candidates, - key=lambda candidate: tsfc.spectral._candidate_score( - tsfc.spectral._optimise_candidate( - variable, candidate, quadrature_indices, - index_replacer))) - - monkeypatch.setattr( - tsfc.spectral, "_select_factorisation_plan", minimum_cost_plan) - minimum, = compile_form(form, parameters={"mode": "spectral"}) - - assert selected.flop_count == minimum.flop_count - - if __name__ == "__main__": import os import sys diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 8df8f37c25..8d2c536b77 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -1,3 +1,14 @@ +"""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, permutations, zip_longest @@ -11,7 +22,7 @@ estimate_cost, replace_division, unroll_indexsum, ) from gem.refactorise import (ATOMIC, COMPOUND, OTHER, MonomialSum, - collect_factorisation_plans) + collect_monomials) from gem.unconcatenate import unconcatenate from gem.coffee import optimise_monomial_sum from gem.utils import groupby @@ -56,31 +67,6 @@ def _delta_inside(node, self): for child in node.children) -def _factorisation_candidates( - expression, argument_indices, - delta_inside) -> tuple[MonomialSum, ...]: - """Build expanded and map-preserving factorisation plans. - - Parameters - ---------- - expression : Node - Multilinear integrand to factorize. - argument_indices : set of Index - Free argument indices. - delta_inside : callable - Memoized predicate detecting delta nodes. - - Returns - ------- - tuple of MonomialSum - Distinct contraction plans. - """ - classifier = partial( - classify, argument_indices, delta_inside=delta_inside) - return collect_factorisation_plans( - expression, classifier, argument_indices) - - def _optimise_contraction_order( variable: Node, indices, monomial_sum: MonomialSum ) -> tuple[tuple[int, ...], Node]: @@ -121,93 +107,6 @@ def _optimise_contraction_order( ) -def _optimise_candidate( - variable, monomial_sum, quadrature_indices, - index_replacer) -> tuple[tuple, ...]: - """Apply delta elimination and contraction optimization to one plan. - - Parameters - ---------- - variable : Node - Assignment variable. - monomial_sum : MonomialSum - Candidate polynomial representation. - quadrature_indices : tuple of Index - Preferred quadrature contraction order. - index_replacer : MemoizerArg - Shared index-substitution mapper. - - Returns - ------- - tuple - Optimized assignment pairs for cost evaluation. - """ - narrow_variables = OrderedDict() - simplified = defaultdict(MonomialSum) - for monomial in monomial_sum: - var, indices, atomics, rest = delta_elimination( - variable, *monomial, index_replacer) - narrow_variables.setdefault(var) - simplified[var].add(indices, atomics, rest) - - pairs = [] - for var in narrow_variables: - candidate = simplified[var] - contracted = set(chain.from_iterable( - monomial.sum_indices for monomial in candidate)) - ordering = tuple(index for index in quadrature_indices - if index in contracted) - _, expression = _optimise_contraction_order( - var, ordering, candidate) - pairs.append((var, expression)) - return tuple(pairs) - - -def _candidate_score(pairs: tuple[tuple, ...]) -> tuple[int, ...]: - """Estimate work and storage in a GEM contraction candidate. - - Parameters - ---------- - pairs - Assignment pairs to schedule. - - Returns - ------- - tuple of int - Operations, total contraction storage, largest contraction, and - expression-node count. Lexicographic ordering prioritizes arithmetic - work. - """ - return estimate_cost(expression for _, expression in pairs) - - -def _select_factorisation_plan( - variable, candidates, quadrature_indices, - index_replacer) -> MonomialSum: - """Choose the least-cost admitted factorisation plan. - - Parameters - ---------- - variable : Node - Assignment variable. - candidates : tuple of MonomialSum - Algebraically equivalent factorisation plans. - quadrature_indices : tuple of Index - Preferred quadrature contraction order. - index_replacer : MemoizerArg - Shared index-substitution mapper. - - Returns - ------- - MonomialSum - Selected factorisation plan. - """ - return min( - candidates, - key=lambda candidate: _candidate_score(_optimise_candidate( - variable, candidate, quadrature_indices, index_replacer))) - - def flatten(var_reps, index_cache): quadrature_indices = OrderedDict() @@ -249,12 +148,11 @@ def group_key(pair): for free_indices, pair_group in groupby(pairs, group_key): variables, expressions = zip(*pair_group) argument_indices = set(free_indices) - for variable, expression in zip(variables, expressions): - candidate_groups = _factorisation_candidates( - expression, argument_indices, delta_inside) - monomial_sum = _select_factorisation_plan( - variable, candidate_groups, quadrature_indices, - index_replacer) + classifier = partial( + classify, argument_indices, delta_inside=delta_inside) + monomial_sums = collect_monomials( + expressions, classifier, argument_indices) + 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) From 3f211adcc88dd1b9c6ed0020031369ed58d4ecfe Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 14 Aug 2026 19:11:33 +0100 Subject: [PATCH 18/26] Distinguish kernel tables from writable storage --- benchmarks/johnson_mercier.py | 72 +++++++++++++++++++++++++---------- 1 file changed, 52 insertions(+), 20 deletions(-) diff --git a/benchmarks/johnson_mercier.py b/benchmarks/johnson_mercier.py index 689362f915..c573084367 100644 --- a/benchmarks/johnson_mercier.py +++ b/benchmarks/johnson_mercier.py @@ -35,8 +35,9 @@ def compile_target(dim: int) -> object: return compile_form(form, parameters={"mode": "spectral"})[0] -def temporary_metrics(kernel: object) -> tuple[int, int, int, int]: - """Measure statically allocated Loopy temporaries. +def temporary_metrics( + kernel: object) -> tuple[int, int, int, int, int, int]: + """Separate writable intermediates from immutable tables. Parameters ---------- @@ -47,21 +48,49 @@ def temporary_metrics(kernel: object) -> tuple[int, int, int, int]: ------- scalar_count Number of scalar temporaries. - array_count - Number of array temporaries. - array_elements - Total entries in array temporaries. - largest_array - Entries in the largest array temporary. + mutable_array_count + Number of writable array temporaries. + mutable_elements + Total entries in writable array temporaries. + largest_mutable_array + Entries in the largest writable array temporary. + table_count + Number of read-only initialized arrays. + table_elements + Total entries in read-only initialized arrays. + + 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 = kernel.ast.default_entrypoint.temporary_variables.values() - shapes = [temporary.shape for temporary in temporaries] - sizes = [numpy.prod(shape, dtype=int) for shape in shapes if shape] + temporaries = tuple(temporaries) + 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 ( - sum(not shape for shape in shapes), - len(sizes), - sum(sizes), - max(sizes, default=0), + sum(not temporary.shape for temporary in temporaries), + len(mutable_sizes), + sum(mutable_sizes), + max(mutable_sizes, default=0), + len(table_sizes), + sum(table_sizes), ) @@ -71,19 +100,22 @@ def main() -> None: parser.add_argument("--dims", nargs="+", type=int, default=(2, 3)) args = parser.parse_args() print("") - print("| dim | compile (s) | flops | scalar temps | array temps | " - "elements | bytes | largest | AST lines |") + print("| dim | compile (s) | flops | scalar temps | mutable arrays | " + "mutable elements | mutable bytes | largest mutable | tables | " + "table elements | AST lines |") print("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | " - "---: |") + "---: | ---: | ---: |") for dim in args.dims: start = time.perf_counter() kernel = compile_target(dim) elapsed = time.perf_counter() - start source = str(kernel.ast) - nscalar, narray, nelem, largest = temporary_metrics(kernel) + (nscalar, nmutable, nmutable_elements, largest_mutable, + ntables, ntable_elements) = temporary_metrics(kernel) print(f"| {dim} | {elapsed:.6f} | {kernel.flop_count:.0f} | " - f"{nscalar} | {narray} | {nelem} | {8 * nelem} | {largest} | " - f"{len(source.splitlines())} |") + f"{nscalar} | {nmutable} | {nmutable_elements} | " + f"{8 * nmutable_elements} | {largest_mutable} | {ntables} | " + f"{ntable_elements} | {len(source.splitlines())} |") if __name__ == "__main__": From 166ce86cb428a075502b6a18b78978ba669efaab Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 15 Aug 2026 08:31:21 +0100 Subject: [PATCH 19/26] Reuse explicit component tensor loops in Loopy --- tests/tsfc/test_sum_factorisation.py | 10 ++++++++-- tsfc/loopy.py | 15 ++++++++++++--- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 60c6754c46..e3c37d3b5e 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -183,6 +183,13 @@ def test_shared_physically_mapped_tabulation( ---------- 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. 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) @@ -211,8 +218,7 @@ def test_shared_physically_mapped_tabulation( 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 sum(shape == (15,) for shape in optimized_shapes) \ - > sum(shape == (15,) for shape in baseline_shapes) + assert [shape for shape in optimized_shapes if shape] == [(15,)] * 7 def test_sum_factorisation_order() -> None: diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 6b1cd7d265..4600124741 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -411,10 +411,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) From 99a2eea7782c5cfc8c4e6cb0058752625eb926b7 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 15 Aug 2026 08:42:55 +0100 Subject: [PATCH 20/26] Use the GEM contraction planner from TSFC --- tests/tsfc/test_sum_factorisation.py | 2 +- tsfc/spectral.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index e3c37d3b5e..1dcbba87f6 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -5,7 +5,7 @@ import gem.coffee import tsfc.spectral from gem.gem import one -from gem.optimise import estimate_cost +from gem.contraction import estimate_cost from gem.refactorise import MonomialSum from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, action, interval, triangle, diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 8d2c536b77..241f7ad386 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -15,12 +15,11 @@ 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 delta_elimination as _delta_elimination -from gem.optimise import ( - estimate_cost, replace_division, unroll_indexsum, -) +from gem.optimise import replace_division, unroll_indexsum from gem.refactorise import (ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials) from gem.unconcatenate import unconcatenate From 8be3ab23e3626e35cda8655cc545a157f8b82f29 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 15 Aug 2026 14:50:40 +0100 Subject: [PATCH 21/26] Cost linear map preservation against expansion --- tests/tsfc/test_sum_factorisation.py | 36 +++++++++ tsfc/spectral.py | 112 ++++++++++++++++++--------- 2 files changed, 110 insertions(+), 38 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 1dcbba87f6..fad42d0a93 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -221,6 +221,42 @@ def test_shared_physically_mapped_tabulation( assert [shape for shape in optimized_shapes if shape] == [(15,)] * 7 +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)) diff --git a/tsfc/spectral.py b/tsfc/spectral.py index 241f7ad386..fc1dbced09 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -106,6 +106,72 @@ def _optimise_contraction_order( ) +def _optimise_plan( + pairs: tuple[tuple[Node, Node], ...], + quadrature_indices: tuple, + preserve_maps: bool) -> tuple[tuple[Node, Node], ...]: + """Optimize one representation of the finite element linear maps. + + Parameters + ---------- + pairs + Output variables and their integral expressions. + quadrature_indices + Quadrature indices in deterministic source order. + preserve_maps + Keep one-axis sums as finite element linear operands when true; + expose their scalar polynomial structure when false. + + Returns + ------- + tuple of tuple + Optimized output variables and GEM expressions. + + Notes + ----- + Preserving a linear map exposes tabulation reuse and map-level code + motion. Expanding it exposes scalar factorization. These + transformations are not composable in general, so plan selection must + compare their optimized contraction trees. + + """ + 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) + + result = [] + for variable in narrow_variables: + monomial_sum = delta_simplified[variable] + contracted = set(chain.from_iterable( + monomial.sum_indices for monomial in monomial_sum)) + indices = tuple( + index for index in quadrature_indices + if index in contracted) + _, expression = _optimise_contraction_order( + variable, indices, monomial_sum) + result.append((variable, expression)) + return tuple(result) + + def flatten(var_reps, index_cache): quadrature_indices = OrderedDict() @@ -129,45 +195,15 @@ 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) quadrature_indices = tuple(quadrature_indices) - # Group assignment pairs by argument indices - 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) - monomial_sums = collect_monomials( - expressions, classifier, argument_indices) - 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] - # Apply sum factorisation combined with COFFEE technology - _, expression = _optimise_contraction_order( - variable, sum_indices, monomial_sum) - yield (variable, expression) + plans = tuple( + _optimise_plan(tuple(pairs), quadrature_indices, preserve_maps) + for preserve_maps in (False, True)) + plan = min( + plans, + key=lambda plan: estimate_cost( + expression for _, expression in plan)) + yield from plan finalise_options = dict(replace_delta=True, remove_componenttensors=False) From 21684f44fd83696ffc768fbbf1300da478b3f828 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 15 Aug 2026 16:13:45 +0100 Subject: [PATCH 22/26] DROP BEFORE MERGE --- .github/actions/install/action.yml | 1 + 1 file changed, 1 insertion(+) 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 From db3ad6fc42b0be2d6e499bc7dae2a5e27ec02957 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 15 Aug 2026 21:31:18 +0100 Subject: [PATCH 23/26] Lower sparse basis maps without a data-dependent loop bound The Loopy lowering built one iteration domain piece per parent value and unioned them. Loopy needs a convex domain, so that union was rejected for most sparsity patterns, and the bound was dropped without warning whenever the parent loop was out of scope. The sparse basis map now arrives as one rectangular contraction, so the lowering needs no such bound. test_coffee_optimise pinned the operand order of a factorised product. That order came from the greedy association, which paired the unit literal with the first factor and appended the folded result. The contraction planner associates the same product without that step, so the expected trees now name the operands in the order it produces. Co-Authored-By: Claude Opus 5 --- tests/tsfc/test_codegen.py | 48 ++++++++++++++++++++++++++++ tests/tsfc/test_coffee_optimise.py | 14 ++++---- tests/tsfc/test_sum_factorisation.py | 5 +-- tsfc/loopy.py | 33 ++----------------- 4 files changed, 60 insertions(+), 40 deletions(-) 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_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index fad42d0a93..68d4939d7c 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -187,7 +187,8 @@ def test_shared_physically_mapped_tabulation( Notes ----- Johnson--Mercier has six mapped basis outputs in two dimensions. The - seventh writable vector holds geometry data. Algebra shared while + 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. """ @@ -218,7 +219,7 @@ def test_shared_physically_mapped_tabulation( 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,)] * 7 + assert [shape for shape in optimized_shapes if shape] == [(15,)] * 8 def test_linear_map_representation_is_costed( diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 4600124741..4304e58422 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -122,7 +122,6 @@ class LoopyContext(object): def __init__(self, target=None): self.indices = {} # indices for declarations and referencing values, from ImperoC self.active_indices = {} # gem index -> pymbolic variable - self.index_lengths = {} # iname -> (parent inames, tabulated extents) 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 @@ -259,8 +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(), ctx.index_parents, ctx.index_lengths) + domains = create_domains(ctx.index_extent.items(), ctx.index_parents) # Create loopy kernel knl = lp.make_kernel( @@ -279,7 +277,7 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name return knl, event_name -def create_domains(indices, index_parents=None, index_lengths=None): +def create_domains(indices, index_parents=None): """Create ISL domains for independent and dependent indices. Parameters @@ -288,8 +286,6 @@ def create_domains(indices, index_parents=None, index_lengths=None): Index names and their static extents. index_parents : mapping, optional Parent inames for simplex-lattice bounds. - index_lengths : mapping, optional - Parent inames and tabulated extents for ragged bounds. Returns ------- @@ -298,24 +294,6 @@ def create_domains(indices, index_parents=None, index_lengths=None): """ domains = [] for idx, extent in indices: - if index_lengths and idx in index_lengths: - parents, lengths = index_lengths[idx] - inames = isl.make_zero_and_vars([idx], parents) - domain = None - for point in numpy.ndindex(lengths.shape): - length = int(lengths[point]) - if length == 0: - continue - piece = (inames[0].le_set(inames[idx]) - & inames[idx].lt_set(inames[0] + length)) - for parent, value in zip(parents, point): - piece = piece & inames[parent].eq_set( - inames[0] + value) - domain = piece if domain is None else domain.union(piece) - assert domain is not None - domains.append(domain) - continue - parents = index_parents.get(idx, ()) if index_parents else () inames = isl.make_zero_and_vars([idx], parents) bound = inames[0] + extent @@ -360,13 +338,6 @@ def statement_for(tree, ctx): # remains correct: jagged expressions are zero-padded. ctx.index_parents[idx] = tuple(ctx.active_indices[parent].name for parent in tree.index.parents) - elif isinstance(tree.index, gem.RaggedIndex) and \ - all(parent in ctx.active_indices for parent in tree.index.parents): - ctx.index_lengths[idx] = ( - tuple(ctx.active_indices[parent].name - for parent in tree.index.parents), - tree.index.lengths, - ) with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active: return statement(tree.children[0], ctx_active) From 52ae07764187c385061db80f29dc0132ffe42c1a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 16 Aug 2026 01:37:27 +0100 Subject: [PATCH 24/26] Measure Johnson--Mercier assembly, not just compilation The benchmark compiled a form through TSFC and reported kernel metrics, so it could show neither what the generated code costs to run nor what a user waits for. It now builds a real mesh, assembles the form, and separates three costs: TSFC compilation, a cold-cache assembly, and the cell loop itself. The on-disk TSFC and PyOP2 caches are redirected to a fresh directory before Firedrake is imported, so compilation is genuinely cold rather than whatever the previous run left behind. Execution is timed by calling the compiled cell loop directly. Going through `assemble` would fold Python and PETSc work into the runtime number; on the tetrahedral form the two agree to within 2%, which is itself the useful result. Assembly overhead is negligible, and a 4.7 s cold assemble surrounds a 0.22 s loop. Co-Authored-By: Claude Opus 5 --- benchmarks/johnson_mercier.py | 222 +++++++++++++++++++++++++++++----- 1 file changed, 194 insertions(+), 28 deletions(-) diff --git a/benchmarks/johnson_mercier.py b/benchmarks/johnson_mercier.py index c573084367..6ba967a747 100644 --- a/benchmarks/johnson_mercier.py +++ b/benchmarks/johnson_mercier.py @@ -1,37 +1,71 @@ #!/usr/bin/env python -"""Measure simplex Johnson--Mercier code generation.""" +"""Measure simplex Johnson--Mercier code generation and assembly.""" import argparse +import cProfile +import os +import pstats +import tempfile import time +import types import numpy -from finat.ufl import FiniteElement, VectorElement -from tsfc import compile_form -from ufl import FunctionSpace, Mesh, TestFunction, TrialFunction, div, dx, inner -from ufl.cell import Cell +def isolate_caches() -> None: + """Point the TSFC and PyOP2 caches at a fresh directory. + + Notes + ----- + Must run before ``import firedrake``, which fills these variables in if + they are unset. Johnson--Mercier assembly is almost entirely code + generation, so a warm disk cache hides the quantity being measured. + """ + cache = tempfile.mkdtemp(prefix="johnson-mercier-") + os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] = os.path.join(cache, "tsfc") + os.environ["PYOP2_CACHE_DIR"] = os.path.join(cache, "pyop2") -def compile_target(dim: int) -> object: - """Compile the JM mass-plus-divergence target on a simplex. + +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 ------- - object - Compiled TSFC kernel. + form + The bilinear form. + space + The Johnson--Mercier function space it is posed on. """ - cell = Cell(("triangle", "tetrahedron")[dim - 2]) - mesh = Mesh(VectorElement("CG", cell, 1)) - element = FiniteElement("Johnson-Mercier", cell, 1) - space = FunctionSpace(mesh, element) + 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) - form = (inner(u, v) + inner(div(u), div(v))) * dx + 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] @@ -94,28 +128,160 @@ def temporary_metrics( ) +def time_kernel(form: object, repeats: int) -> float: + """Time repeated calls to the compiled cell kernel. + + Parameters + ---------- + form + The bilinear form. + 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 firedrake import assemble + from pyop2.global_kernel import GlobalKernel, compile_global_kernel + + 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) + finally: + GlobalKernel.__call__ = original + + kernel, comm, arguments = max( + calls, key=lambda call: call[0].local_kernel.num_flops) + 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 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 + + warm = time_kernel(form, repeats) + + (nscalar, nmutable, nmutable_elements, largest_mutable, + ntables, ntable_elements) = temporary_metrics(kernel) + return types.SimpleNamespace( + dim=dim, + dofs=space.dim(), + cells=space.mesh().num_cells(), + compile_time=compile_time, + cold=cold, + warm=warm, + flops=kernel.flop_count, + nscalar=nscalar, + nmutable=nmutable, + nmutable_elements=nmutable_elements, + largest_mutable=largest_mutable, + ntables=ntables, + ntable_elements=ntable_elements, + ast_lines=len(str(kernel.ast).splitlines()), + ) + + +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 compiler metrics as copyable Markdown.""" + """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() + + if args.profile: + for dim in args.dims: + profile(dim, args.size, args.profile) + return + print("") - print("| dim | compile (s) | flops | scalar temps | mutable arrays | " + 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("| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | " - "---: | ---: | ---: |") + print("| ---: " * 16 + "|") for dim in args.dims: - start = time.perf_counter() - kernel = compile_target(dim) - elapsed = time.perf_counter() - start - source = str(kernel.ast) - (nscalar, nmutable, nmutable_elements, largest_mutable, - ntables, ntable_elements) = temporary_metrics(kernel) - print(f"| {dim} | {elapsed:.6f} | {kernel.flop_count:.0f} | " - f"{nscalar} | {nmutable} | {nmutable_elements} | " - f"{8 * nmutable_elements} | {largest_mutable} | {ntables} | " - f"{ntable_elements} | {len(source.splitlines())} |") + 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__": From 16ecde1f35d0bda3a74f84ad682f0bd8558b08c8 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 16 Aug 2026 23:14:42 +0100 Subject: [PATCH 25/26] Share benchmark metrics across sum-factorisation cases `benchmarks/metrics.py` holds the cache isolation, kernel metrics, kernel timing and source dumping that the benchmarks have in common. `johnson_mercier.py` uses it instead of its own copies. `benchmarks/sum_factorisation.py` measures the two forms of `docs/notebooks/10-sum-factorisation.py` on an extruded hexahedral mesh: a Laplacian on CG and a curl-curl form on NCE. Each is compiled vanilla, on a Gauss-Legendre rule, and on the collocated Gauss-Lobatto-Legendre rule. Forms are measured as operator actions, so the timing sees the local contraction rather than the insertion of an element matrix. Co-Authored-By: Claude Opus 5 --- benchmarks/johnson_mercier.py | 154 +--------------- benchmarks/metrics.py | 188 +++++++++++++++++++ benchmarks/sum_factorisation.py | 316 ++++++++++++++++++++++++++++++++ 3 files changed, 514 insertions(+), 144 deletions(-) create mode 100644 benchmarks/metrics.py create mode 100644 benchmarks/sum_factorisation.py diff --git a/benchmarks/johnson_mercier.py b/benchmarks/johnson_mercier.py index 6ba967a747..7d347616e5 100644 --- a/benchmarks/johnson_mercier.py +++ b/benchmarks/johnson_mercier.py @@ -3,27 +3,11 @@ import argparse import cProfile -import os import pstats -import tempfile import time import types -import numpy - - -def isolate_caches() -> None: - """Point the TSFC and PyOP2 caches at a fresh directory. - - Notes - ----- - Must run before ``import firedrake``, which fills these variables in if - they are unset. Johnson--Mercier assembly is almost entirely code - generation, so a warm disk cache hides the quantity being measured. - """ - cache = tempfile.mkdtemp(prefix="johnson-mercier-") - os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] = os.path.join(cache, "tsfc") - os.environ["PYOP2_CACHE_DIR"] = os.path.join(cache, "pyop2") +from metrics import isolate_caches, kernel_metrics, time_kernel def build_form(dim: int, size: int) -> tuple[object, object]: @@ -69,112 +53,6 @@ def compile_target(form: object) -> object: return compile_form(form, parameters={"mode": "spectral"})[0] -def temporary_metrics( - kernel: object) -> tuple[int, int, int, int, int, int]: - """Separate writable intermediates from immutable tables. - - Parameters - ---------- - kernel - Compiled TSFC kernel. - - Returns - ------- - scalar_count - Number of scalar temporaries. - mutable_array_count - Number of writable array temporaries. - mutable_elements - Total entries in writable array temporaries. - largest_mutable_array - Entries in the largest writable array temporary. - table_count - Number of read-only initialized arrays. - table_elements - Total entries in read-only initialized arrays. - - 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 = kernel.ast.default_entrypoint.temporary_variables.values() - temporaries = tuple(temporaries) - 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 ( - sum(not temporary.shape for temporary in temporaries), - len(mutable_sizes), - sum(mutable_sizes), - max(mutable_sizes, default=0), - len(table_sizes), - sum(table_sizes), - ) - - -def time_kernel(form: object, repeats: int) -> float: - """Time repeated calls to the compiled cell kernel. - - Parameters - ---------- - form - The bilinear form. - 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 firedrake import assemble - from pyop2.global_kernel import GlobalKernel, compile_global_kernel - - 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) - finally: - GlobalKernel.__call__ = original - - kernel, comm, arguments = max( - calls, key=lambda call: call[0].local_kernel.num_flops) - 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 measure(dim: int, size: int, repeats: int) -> types.SimpleNamespace: """Time compilation, cold assembly and the generated cell kernel. @@ -203,26 +81,14 @@ def measure(dim: int, size: int, repeats: int) -> types.SimpleNamespace: assemble(form) cold = time.perf_counter() - start - warm = time_kernel(form, repeats) - - (nscalar, nmutable, nmutable_elements, largest_mutable, - ntables, ntable_elements) = temporary_metrics(kernel) - return types.SimpleNamespace( - dim=dim, - dofs=space.dim(), - cells=space.mesh().num_cells(), - compile_time=compile_time, - cold=cold, - warm=warm, - flops=kernel.flop_count, - nscalar=nscalar, - nmutable=nmutable, - nmutable_elements=nmutable_elements, - largest_mutable=largest_mutable, - ntables=ntables, - ntable_elements=ntable_elements, - ast_lines=len(str(kernel.ast).splitlines()), - ) + 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: @@ -260,7 +126,7 @@ def main() -> None: args = parser.parse_args() if not args.warm_cache: - isolate_caches() + isolate_caches("johnson-mercier-") if args.profile: for dim in args.dims: 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() From 41d6ac0938854f0d63e0515026add54bdf606a2a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 16 Aug 2026 23:14:51 +0100 Subject: [PATCH 26/26] Bound contraction storage when selecting a factorisation plan Contracting each assignment in its own quadrature ordering minimises arithmetic, but assignments that disagree build loop nests Impero cannot merge, so a value they share widens to the whole quadrature grid. Plan selection now also builds a candidate in which every assignment shares one ordering. Storage is a feasibility constraint rather than a term traded against arithmetic: among plans whose temporaries fit `storage_budget`, take the one with least arithmetic, and fall back to the narrowest when none fit. One rule chooses both the representation of the finite element linear maps and the contraction ordering. `_declared_storage` measures a candidate by scheduling it, since a temporary's width follows from the loop nest a value outlives, which an expression DAG does not record. Selection lives in `_select_plan`, over the orderings from `_plans` and the two representations from `_factorise`. `flatten` returns its assignments rather than yielding them, as it computes them eagerly. NCE degree 7 on a Gauss-Legendre rule falls from 96,273 to 7,018 declared scalar entries, and its largest temporary from 3,375 to 1,575, at 1.6% more arithmetic. The collocated rule keeps its per-assignment orderings. Co-Authored-By: Claude Opus 5 --- tests/tsfc/test_sum_factorisation.py | 41 ++++- tsfc/spectral.py | 217 ++++++++++++++++++--------- 2 files changed, 182 insertions(+), 76 deletions(-) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py index 68d4939d7c..d7fc40165c 100644 --- a/tests/tsfc/test_sum_factorisation.py +++ b/tests/tsfc/test_sum_factorisation.py @@ -7,14 +7,14 @@ from gem.gem import one from gem.contraction import estimate_cost from gem.refactorise import MonomialSum -from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, +from ufl import (Coefficient, Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, action, interval, triangle, 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 _optimise_contraction_order +from tsfc.spectral import _plans def helmholtz(cell, degree): @@ -269,15 +269,46 @@ def test_sum_factorisation_order() -> None: monomial_sum = MonomialSum() monomial_sum.add((q0, q1), (inner * outer,), one) - score, _ = _optimise_contraction_order( - variable, (q1, q0), monomial_sum) + 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 score == min(candidates) + 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__": diff --git a/tsfc/spectral.py b/tsfc/spectral.py index fc1dbced09..d503b74873 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -13,15 +13,16 @@ from functools import partial from itertools import chain, permutations, zip_longest -from gem.gem import (Conditional, Delta, Indexed, Node, Sum, - index_sum, one) +import numpy + +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) +from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials from gem.unconcatenate import unconcatenate from gem.coffee import optimise_monomial_sum from gem.utils import groupby @@ -31,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 @@ -66,58 +72,53 @@ def _delta_inside(node, self): for child in node.children) -def _optimise_contraction_order( - variable: Node, indices, monomial_sum: MonomialSum -) -> tuple[tuple[int, ...], Node]: - """Choose the least-cost quadrature contraction ordering. +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 ---------- - variable : Node - Assignment variable whose indices identify multilinear axes. - indices : iterable of Index - Quadrature indices to contract. - monomial_sum : MonomialSum - Factorized integrand. + plan + Output variables and their factorized GEM expressions. + quadrature_indices + Every quadrature index of the integral, in source order. Returns ------- - score - Estimated operations, storage, and expression size. - expression - Factorized GEM expression for the selected ordering. - - Notes - ----- - The search is exhaustive in the number of quadrature axes. Finite - element integration supplies one axis per reference-cell direction, so - this space is independent of polynomial degree. + int + Scalar entries in the temporaries, or zero if the plan schedules + to nothing. """ - indices = tuple(indices) - plans = ( - sum_factorise(variable, ordering, monomial_sum) - for ordering in permutations(indices) - ) - return min( - ((estimate_cost((expression,)), expression) - for expression in plans), - key=lambda plan: plan[0], - ) - - -def _optimise_plan( + 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], ...], - quadrature_indices: tuple, - preserve_maps: bool) -> tuple[tuple[Node, Node], ...]: - """Optimize one representation of the finite element linear maps. + 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. - quadrature_indices - Quadrature indices in deterministic source order. preserve_maps Keep one-axis sums as finite element linear operands when true; expose their scalar polynomial structure when false. @@ -125,14 +126,7 @@ def _optimise_plan( Returns ------- tuple of tuple - Optimized output variables and GEM expressions. - - Notes - ----- - Preserving a linear map exposes tabulation reuse and map-level code - motion. Expanding it exposes scalar factorization. These - transformations are not composable in general, so plan selection must - compare their optimized contraction trees. + Output variables and their delta-cancelled monomial sums. """ index_replacer = MemoizerArg(filtered_replace_indices) @@ -158,18 +152,106 @@ def _optimise_plan( delta_simplified[var].add( indices, atomics, rest) - result = [] - for variable in narrow_variables: - monomial_sum = delta_simplified[variable] - contracted = set(chain.from_iterable( + 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)) - indices = tuple( - index for index in quadrature_indices - if index in contracted) - _, expression = _optimise_contraction_order( - variable, indices, monomial_sum) - result.append((variable, expression)) - return tuple(result) + 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): @@ -195,15 +277,8 @@ def flatten(var_reps, index_cache): # Split Concatenate nodes pairs = unconcatenate(pairs, cache=index_cache) - quadrature_indices = tuple(quadrature_indices) - plans = tuple( - _optimise_plan(tuple(pairs), quadrature_indices, preserve_maps) - for preserve_maps in (False, True)) - plan = min( - plans, - key=lambda plan: estimate_cost( - expression for _, expression in plan)) - yield from plan + + return _select_plan(tuple(pairs), tuple(quadrature_indices)) finalise_options = dict(replace_delta=True, remove_componenttensors=False)